diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7a46725c80c7..b910494a009a 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -249,8 +249,8 @@ if(ENABLE_MULTI_DEVICE) endif() # TRT dependencies -find_package(TensorRT 10 REQUIRED COMPONENTS OnnxParser) -set(TRT_LIB TensorRT::NvInfer) +# TensorRT dependency removed: the C++ tree no longer links the TensorRT library. +set(TRT_LIB "") get_filename_component(TRT_LLM_ROOT_DIR ${CMAKE_CURRENT_SOURCE_DIR} PATH) @@ -291,7 +291,6 @@ include_directories( ${CUDAToolkit_INCLUDE_DIRS} ${CUDAToolkit_INCLUDE_DIRS}/cccl ${CUDNN_ROOT_DIR}/include - $ ${maybe_nvtx_includedir} ${CMAKE_BINARY_DIR}/_deps/cutlass-src/include ${CMAKE_BINARY_DIR}/_deps/cutlass-src/tools/util/include @@ -683,6 +682,6 @@ if(MEASURE_BUILD_TIME) endif() set(BUILD_WHEEL_TARGETS - tensorrt_llm;nvinfer_plugin_tensorrt_llm + tensorrt_llm CACHE STRING "Targets used to build wheel") add_custom_target(build_wheel_targets DEPENDS ${BUILD_WHEEL_TARGETS}) diff --git a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h index 80a06ea5ddf7..56793b9ee24e 100644 --- a/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h +++ b/cpp/include/tensorrt_llm/batch_manager/cacheTransceiver.h @@ -233,7 +233,7 @@ class CacheTransceiver : public BaseCacheTransceiver public: CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, @@ -242,7 +242,7 @@ class CacheTransceiver : public BaseCacheTransceiver CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, std::vector numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType = executor::kv_cache::CacheState::AttentionType::kDEFAULT, std::optional cacheTransceiverConfig = std::nullopt, diff --git a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h index bc619a34bc03..394d13a03f54 100644 --- a/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h +++ b/cpp/include/tensorrt_llm/batch_manager/createNewDecoderRequests.h @@ -66,14 +66,14 @@ class CreateNewDecoderRequests : Algorithm std::vector> operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, + 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, - nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, + 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; diff --git a/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h b/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h index 9a577b61ad51..9d5df24fcc15 100644 --- a/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h +++ b/cpp/include/tensorrt_llm/batch_manager/guidedDecoder.h @@ -39,7 +39,7 @@ class GuidedDecoder using BitmaskT = uint32_t; GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences, - SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, runtime::BufferManager const& runtimeBufferManager); + SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDtype, runtime::BufferManager const& runtimeBufferManager); void build(ScheduledRequests const& scheduledRequests); void execute(DecoderInputBuffers const& decoderInputBuffers, runtime::BufferManager const& runtimeBufferManager); @@ -51,7 +51,7 @@ class GuidedDecoder SizeType32 mMaxNumSequences; SizeType32 mVocabSizePadded; SizeType32 mBitmaskSize; // CeilDiv(vocabSizePadded, 32) - nvinfer1::DataType mLogitsDtype; + tensorrt_llm::DataType mLogitsDtype; TensorPtr mLogitsBitmask; // [mMaxNumRequests, mBitmaskSize] TensorPtr mLogitsBitmaskHost; // [mMaxNumRequests, mBitmaskSize] diff --git a/cpp/include/tensorrt_llm/batch_manager/handleContextLogits.h b/cpp/include/tensorrt_llm/batch_manager/handleContextLogits.h deleted file mode 100644 index cb77545578c8..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/handleContextLogits.h +++ /dev/null @@ -1,53 +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. - */ - -#pragma once - -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime -{ -class BufferManager; -class CudaStream; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -class DecoderInputBuffers; -class MedusaBuffers; - -class HandleContextLogits : Algorithm -{ -public: - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - constexpr static auto name{"HandleContextLogits"}; - - HandleContextLogits() = default; - - runtime::SizeType32 operator()(DecoderInputBuffers& inputBuffers, RequestVector const& contextRequests, - runtime::ITensor::SharedPtr const& logits, std::vector const& numContextLogitsVec, - runtime::ModelConfig const& modelConfig, runtime::BufferManager const& manager, - OptionalRef medusaBuffers) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/handleGenerationLogits.h b/cpp/include/tensorrt_llm/batch_manager/handleGenerationLogits.h deleted file mode 100644 index f9fd58800a6f..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/handleGenerationLogits.h +++ /dev/null @@ -1,53 +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. - */ - -#pragma once - -#include "common.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime -{ -class BufferManager; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -class DecoderInputBuffers; -class RuntimeBuffers; -class MedusaBuffers; - -class HandleGenerationLogits : Algorithm -{ -public: - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - constexpr static auto name{"HandleGenerationLogits"}; - - HandleGenerationLogits() = default; - - void operator()(DecoderInputBuffers& inputBuffers, RequestVector const& generationRequests, - runtime::ITensor::SharedPtr const& logits, runtime::SizeType32 logitsIndex, - runtime::ModelConfig const& modelConfig, runtime::BufferManager const& manager, - OptionalRef genRuntimeBuffers, OptionalRef medusaBuffers) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h index c665f7a8df95..42c820f72498 100644 --- a/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/kvCacheManager.h @@ -32,7 +32,7 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -128,7 +128,7 @@ struct PoolConfiguration { SizeType32 windowSize; SizeType32 sizePerHead; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; }; struct LinearAttentionMetadata @@ -838,7 +838,7 @@ class WindowBlockManager using BlockMap = std::unordered_multimap; using BlockMapIterRange = std::pair; - explicit WindowBlockManager(nvinfer1::DataType dtype, SizeType32 windowSize, + explicit WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 windowSize, std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, @@ -1003,7 +1003,7 @@ class WindowBlockManager //! host pools with mixed precisions when constructed with a per-window //! dtype map. Empty pools or NVFP4-scale pools are routed through the //! per-pool tensor metadata instead. - [[nodiscard]] nvinfer1::DataType getDataType() const noexcept + [[nodiscard]] tensorrt_llm::DataType getDataType() const noexcept { return mDataType; } @@ -1093,7 +1093,7 @@ class WindowBlockManager [[nodiscard]] SizeType32 getNumEltsPerContainer() const { #ifdef ENABLE_FP4 - return mDataType == nvinfer1::DataType::kFP4 ? 2 : 1; + return mDataType == tensorrt_llm::DataType::kFP4 ? 2 : 1; #else return 1; #endif @@ -1158,7 +1158,7 @@ class WindowBlockManager return mLayerToIndexWithinPool.at(layerIdx); } - void setOffsets(kernels::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, SizeType32 beamIdx, + void setOffsets(kernels::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId) const; //! \brief Bring offloaded block from secondary to primary memory. @@ -1319,7 +1319,7 @@ class WindowBlockManager } private: - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; SizeType32 mWindowSize; // Number of blocks in pools @@ -1445,7 +1445,7 @@ class BlockManager explicit BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, std::shared_ptr eventManager = nullptr, bool enablePartialReuse = true, @@ -1527,7 +1527,7 @@ class BlockManager void releaseLastBlock(GenerationRequest& sequence, SizeType32 windowSize); - void setOffsets(kernels::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, SizeType32 beamIdx, + void setOffsets(kernels::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId, SizeType32 windowSize) const; //! \brief Combined prefix reuse analysis — single radix tree walk. @@ -1585,9 +1585,9 @@ class BlockManager //! \brief Convenience: window_size -> dataType, derived from getPoolConfigurations(). //! For one-pool-per-window managers only; multi-pool-per-window will collide. - [[nodiscard]] std::map getDataTypePerWindow() const + [[nodiscard]] std::map getDataTypePerWindow() const { - std::map result; + std::map result; for (auto const& [windowSize, manager] : mWindowBlockManagers) { result[windowSize] = manager.getDataType(); @@ -1600,7 +1600,7 @@ class BlockManager return mWindowBlockManagers.at(windowSize).getSizePerHead(); } - [[nodiscard]] nvinfer1::DataType getDataTypeForWindow(SizeType32 windowSize) const + [[nodiscard]] tensorrt_llm::DataType getDataTypeForWindow(SizeType32 windowSize) const { return mWindowBlockManagers.at(windowSize).getDataType(); } @@ -2153,7 +2153,7 @@ class BaseKVCacheManager /// head_dim=512). Empty vector = uniform @p sizePerHead / @p dtype across all windows. /// @return Map from window size to tuple of (primary blocks, secondary blocks) [[nodiscard]] static BlocksPerWindow calculateMaxNumBlocks(executor::KvCacheConfig const& config, - nvinfer1::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, + tensorrt_llm::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, tensorrt_llm::runtime::WorldConfig const& worldConfig, std::map> const& windowSizeToLayers, uint64_t allottedPrimaryMemBytes, uint64_t allottedSecondaryMemBytes, size_t extraCostMemory, SizeType32 kvFactor, SizeType32 maxBatchSize, @@ -2240,7 +2240,7 @@ class KVCacheManager : public BaseKVCacheManager //! and disagg transfer machinery applies natively. Empty vector = uniform. KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, @@ -2254,7 +2254,7 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, @@ -2268,7 +2268,7 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = true, CacheType cacheType = CacheType::kSELF, std::optional secondaryOffloadMinPriority = std::nullopt, @@ -2282,7 +2282,7 @@ class KVCacheManager : public BaseKVCacheManager KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, SizeType32 sinkTokenLength, + std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse = false, CacheType cacheType = CacheType::kSELF, bool enablePartialReuse = true, bool copyOnpartialReuse = true, bool enableIndexerKCache = false, SizeType32 indexerKCacheQuantBlockSize = 128, @@ -2628,7 +2628,7 @@ class KVCacheManager : public BaseKVCacheManager SizeType32 mMaxNumSequences; // Maximum beam width SizeType32 mMaxBeamWidth; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; // Maximum kv cache length per sequence SizeType32 mMaxAttentionWindow; // Number of tokens per block diff --git a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h index 644813e381b7..d1fa77cc4116 100644 --- a/cpp/include/tensorrt_llm/batch_manager/llmRequest.h +++ b/cpp/include/tensorrt_llm/batch_manager/llmRequest.h @@ -1312,7 +1312,7 @@ class GenericLlmRequest mEncoderOutput = std::move(encoderOutput); } - void allocEncoderOutputHost(SizeType32 encoderHiddenSize, nvinfer1::DataType dataType) + void allocEncoderOutputHost(SizeType32 encoderHiddenSize, tensorrt_llm::DataType dataType) { mEncoderOutputHost = runtime::BufferManager::pinned( runtime::ITensor::makeShape({getEncoderOutputLen(), encoderHiddenSize}), dataType); @@ -1328,13 +1328,13 @@ class GenericLlmRequest return mEncoderHiddenStates; } - void allocEncoderOutput(runtime::BufferManager const& manager, nvinfer1::DataType dataType) + void allocEncoderOutput(runtime::BufferManager const& manager, tensorrt_llm::DataType dataType) { // unique_ptr --> shared_ptr ownership move mEncoderOutput = std::move(manager.emptyTensor(runtime::MemoryType::kGPU, dataType)); } - void allocEncoderHiddenStates(runtime::BufferManager const& manager, nvinfer1::DataType dataType) + void allocEncoderHiddenStates(runtime::BufferManager const& manager, tensorrt_llm::DataType dataType) { // unique_ptr --> shared_ptr ownership move mEncoderHiddenStates = std::move(manager.emptyTensor(runtime::MemoryType::kGPU, dataType)); @@ -1452,7 +1452,7 @@ class GenericLlmRequest mContextLogitsHost = std::move(contextLogitsHost); } - void allocContextLogitsHost(SizeType32 vocabSizePadded, nvinfer1::DataType logitsDataType) + void allocContextLogitsHost(SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDataType) { mContextLogitsHost = runtime::BufferManager::pinnedPool( runtime::ITensor::makeShape({mPromptLen, vocabSizePadded}), logitsDataType); @@ -1471,7 +1471,7 @@ class GenericLlmRequest mGenerationLogitsHost = std::move(generationLogitsHost); } - void allocGenerationLogitsHost(SizeType32 vocabSizePadded, nvinfer1::DataType logitsDataType) + void allocGenerationLogitsHost(SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDataType) { if (mIsStreaming) { @@ -1490,7 +1490,7 @@ class GenericLlmRequest } } - void allocTargetModelAcceptedTokenLogitsHost(SizeType32 vocabSizePadded, nvinfer1::DataType logitsDataType) + void allocTargetModelAcceptedTokenLogitsHost(SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDataType) { mGenerationLogitsHost = runtime::BufferManager::pinnedPool( runtime::ITensor::makeShape({1, getNumDraftTokens() + 1, vocabSizePadded}), logitsDataType); @@ -2344,7 +2344,7 @@ class GenericLlmRequest auto const numWords = static_cast(words.size()); auto const shape = runtime::ITensor::makeShape({2, numWords}); - auto tensor = runtime::BufferManager::pinnedPool(shape, nvinfer1::DataType::kINT32); + auto tensor = runtime::BufferManager::pinnedPool(shape, tensorrt_llm::DataType::kINT32); auto* data = runtime::bufferCast(*tensor); std::memcpy(data, words.data(), numWords * sizeof(int32_t)); std::memcpy(data + numWords, offsets.data(), numWords * sizeof(int32_t)); diff --git a/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h b/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h deleted file mode 100644 index 1916a915e337..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/logitsPostProcessor.h +++ /dev/null @@ -1,53 +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. - */ - -#pragma once - -#include "common.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::runtime -{ -class CudaStream; -} - -namespace tensorrt_llm::batch_manager -{ -class DecoderInputBuffers; - -class LogitsPostProcessor : Algorithm -{ -public: - using CudaStreamPtr = std::shared_ptr; - - using LogitsPostProcessorBatched = std::function const&, - std::vector&, - std::vector> const&, CudaStreamPtr const&, - std::vector> const&)>; - - constexpr static auto name{"LogitsPostProcessor"}; - - LogitsPostProcessor() = default; - - bool operator()(DecoderInputBuffers& inputBuffers, bool replicateLogitsPostProcessor, - runtime::WorldConfig const& worldConfig, CudaStreamPtr const& stream, - std::optional const& logitsPostProcessorBatched = std::nullopt) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h b/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h deleted file mode 100644 index 245f4b4b5286..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.h +++ /dev/null @@ -1,56 +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. - */ - -#pragma once - -#include "common.h" -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/common/optionalRef.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iGptDecoderBatched.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime::decoder -{ -class DecoderState; -} // namespace tensorrt_llm::runtime::decoder - -namespace tensorrt_llm::batch_manager -{ -class DecoderInputBuffers; -class RuntimeBuffers; - -class MakeDecodingBatchInputOutput : Algorithm -{ -public: - constexpr static auto name{"MakeDecodingBatchInputOutput"}; - - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - MakeDecodingBatchInputOutput() = default; - - void operator()(DecoderInputBuffers& inputBuffers, runtime::decoder::DecoderState& decoderState, - runtime::ModelConfig const& modelConfig, OptionalRef fusedRuntimeBuffers) const; - - static void createDecoderBatchInputs(DecoderInputBuffers& inputBuffers, std::vector const& activeSlots, - runtime::decoder::DecoderState const& decoderState); -}; - -} // 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 index ba29be6ede81..5342591840a8 100644 --- a/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h +++ b/cpp/include/tensorrt_llm/batch_manager/medusaBuffers.h @@ -22,7 +22,6 @@ #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/promptTuningParams.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" #include "tensorrt_llm/runtime/worldConfig.h" namespace tensorrt_llm::batch_manager @@ -36,10 +35,6 @@ class MedusaBuffers using TensorPtr = runtime::ITensor::SharedPtr; using TensorMap = runtime::StringPtrMap; - MedusaBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, runtime::TllmRuntime const& runtime); - void reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 tokensPerStep); void insertInputTensors( diff --git a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h index cf65753783e8..ed928e96d811 100644 --- a/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/peftCacheManager.h @@ -25,7 +25,7 @@ #include "tensorrt_llm/runtime/workerPool.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -147,7 +147,7 @@ class PeftCacheManager : public BasePeftCacheManager void updateTaskState(uint64_t taskId, uint64_t reqId, bool terminate = false, bool pause = false); static std::pair getMaxNumSlots(PeftCacheManagerConfig const& config, - nvinfer1::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, + tensorrt_llm::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, runtime::BufferManager const& bufferManager); static std::pair getPageManagerConfig( diff --git a/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h b/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h deleted file mode 100644 index a1d8849a8811..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/promptTuningBuffers.h +++ /dev/null @@ -1,106 +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/common.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 PromptTuningBuffers -{ - -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using ITensor = tensorrt_llm::runtime::ITensor; - using TensorPtr = runtime::ITensor::SharedPtr; - - runtime::PromptTuningParams mPromptTuningParams; - SizeType32 mMaxPromptVocabSize; - - PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - - PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, bool promptTableOffloading); - - void validate(std::optional const& optReqPromptEmbeddingTable, - std::optional const& optReqPromptVocabSize); - - void fill(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::BufferManager const& manager, bool packed); - - /* - * The below functions are specific for Chunked Prefill mode - * Chunk Ptable with Ping-Pong Buffer Implementation - * ----------------------------------------------- - * - * Overview: - * The chunk ptable (prompt tuning table) system uses a ping-pong buffer mechanism to efficiently - * manage large embedding tables when operating in context Prefill mode. This allows - * for processing of large embedding tables by loading them in chunks from CPU to GPU memory, - * enabling support for tables that exceed available GPU memory. - * - * Key Components: - * 1. Ping-Pong Buffers (mChunkPtableBuffers): - * - Two alternating GPU buffers that store chunks of the embedding table - * - While the current buffer is being processed by the model, - * the next chunk can be asynchronously loaded into the other buffer - * - Managed through mChunkPtableCurrentIndex (toggles between 0 and 1) - * 2. Start Positions Tracking (mChunkPtableBufferStartPositions): - * - Mainly used for multi-batch processing - * - Maintains the starting position of each batch's data within each buffer - * - Maintained separately for each ping-pong buffer - * - * Memory Optimization: - * - Only two GPU buffers are maintained regardless of total embedding table size - * - Each buffer size is limited to contextChunkSize * hiddenSize - * - Efficient memory usage through chunk-based processing - */ - - bool mPromptTableOffloading; - - bool mChunkPtableInitialized{false}; - std::optional> mChunkPtableBuffers; - std::optional>> mChunkPtableBufferStartPositions; - size_t mChunkPtableCurrentIndex{0}; - - void initializeChunkPtableBuffers(runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - SizeType32 contextChunkSize, std::shared_ptr const& llmReq); - - void switchChunkPtableBuffer(); - - size_t getChunkPtableCurrentIndex(); - - [[nodiscard]] TensorPtr& getChunkPtableBuffer(size_t index); - - [[nodiscard]] SizeType32 getChunkPtableBufferSliceSize(size_t index, size_t batchIdx); - - [[nodiscard]] SizeType32 getChunkPtableBufferStartPosition(size_t index, size_t batchIdx); - - void updateBufferStartPosition(size_t index, SizeType32 numRows); - - void clearBufferStartPositions(size_t index); -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h b/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h index 5c0bfe136de2..1a816e8ad43a 100644 --- a/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h +++ b/cpp/include/tensorrt_llm/batch_manager/rnnStateManager.h @@ -42,8 +42,8 @@ class RnnStateManager runtime::WorldConfig const& worldConfig, tensorrt_llm::runtime::BufferManager const& bufferManager); RnnStateManager(SizeType32 dState, SizeType32 dConv, SizeType32 numHeads, SizeType32 nGroups, SizeType32 headDim, - SizeType32 maxBatchSize, runtime::WorldConfig const& worldConfig, int64_t stream, nvinfer1::DataType dtype, - nvinfer1::DataType ssmCacheDtype, std::vector const& ppLayers, SizeType32 numLayers); + SizeType32 maxBatchSize, runtime::WorldConfig const& worldConfig, int64_t stream, tensorrt_llm::DataType dtype, + tensorrt_llm::DataType ssmCacheDtype, std::vector const& ppLayers, SizeType32 numLayers); void getPtrBuffers(TensorMap& inputBuffers, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) const; @@ -68,9 +68,9 @@ class RnnStateManager [[nodiscard]] TensorPtr getSsmStates() const; - [[nodiscard]] nvinfer1::DataType getConvStateDataType() const noexcept; + [[nodiscard]] tensorrt_llm::DataType getConvStateDataType() const noexcept; - [[nodiscard]] nvinfer1::DataType getSsmStateDataType() const noexcept; + [[nodiscard]] tensorrt_llm::DataType getSsmStateDataType() const noexcept; [[nodiscard]] executor::kv_cache::CacheState::RnnModelConfig getRnnCacheStateModelConfig() const noexcept; @@ -111,8 +111,8 @@ class RnnStateManager std::vector mFreeBlocks; std::unordered_map mCacheIndex; std::optional mBufferManager; - nvinfer1::DataType mDtype{nvinfer1::DataType::kFLOAT}; - nvinfer1::DataType mSsmCacheDtype{nvinfer1::DataType::kFLOAT}; + tensorrt_llm::DataType mDtype{tensorrt_llm::DataType::kFLOAT}; + tensorrt_llm::DataType mSsmCacheDtype{tensorrt_llm::DataType::kFLOAT}; // RNN model config (global values before TP/PP split) SizeType32 mDState{0}; diff --git a/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h b/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h deleted file mode 100644 index 13bde6d07a5e..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/runtimeBuffers.h +++ /dev/null @@ -1,321 +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 "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/common/optionalRef.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 "tensorrt_llm/runtime/loraManager.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include -#include -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; - -namespace decoder -{ -class DecoderState; -} // namespace decoder -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -namespace kv_cache_manager -{ -class BaseKVCacheManager; -} // namespace kv_cache_manager - -class LlmRequest; - -class EncoderBuffers; -class LoraBuffers; -class MedusaBuffers; -class PromptTuningBuffers; -class RnnStateBuffers; -class TransformerBuffers; - -class RuntimeBuffers -{ -public: - static constexpr auto kLogitsTensorName = "logits"; - static constexpr auto kHiddenStatesOutputTensorName = "hidden_states_output"; - static constexpr auto kHiddenStatesInputTensorName = "hidden_states_input"; - static constexpr auto kInputIdsTensorName = "input_ids"; - static constexpr auto kLastTokenIdsTensorName = "last_token_ids"; - static constexpr auto kHostRequestTypesTensorName = "host_request_types"; - static constexpr auto kContextLengthsTensorName = "context_lengths"; - static constexpr auto kHostContextLengthsTensorName = "host_context_lengths"; - static constexpr auto kSequenceLengthsTensorName = "sequence_length"; - static constexpr auto kPromptEmbeddingTableTensorName = "prompt_embedding_table"; - static constexpr auto kTasksTensorName = "tasks"; - static constexpr auto kPromptVocabSizeTensorName = "prompt_vocab_size"; - static constexpr auto kMRopeRotaryCosSinTensorName = "mrope_rotary_cos_sin"; - static constexpr auto kMRopePositionDeltasTensorName = "mrope_position_deltas"; - - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::ITensor::TensorMap; - using PeftTable = runtime::LoraManager::PeftTable; - template - using OptionalRef = tensorrt_llm::common::OptionalRef; - - [[nodiscard]] SizeType32 constexpr getContextIndex() const noexcept - { - return contextIndex; - }; - - void constexpr setContextIndex(SizeType32 index) noexcept - { - contextIndex = index; - }; - - [[nodiscard]] SizeType32 constexpr getNumContextTokens() const noexcept - { - return numContextTokens; - }; - - [[nodiscard]] BatchState getBatchState() const noexcept - { - return {numContextRequests, numGenRequests, getNumTokens(), maxKvCacheLengthRounded}; - }; - -private: - [[nodiscard]] SizeType32 constexpr getNumRequests() const noexcept - { - return numContextRequests + numGenRequests; - }; - - [[nodiscard]] SizeType32 constexpr getNumSequences() const noexcept - { - return numContextRequests + numGenSequences; - }; - - [[nodiscard]] SizeType32 constexpr getNumTokens() const noexcept - { - return numContextTokens + numGenTokens; - }; - - //! Sizes - SizeType32 numContextRequests{}; - SizeType32 numGenRequests{}; - SizeType32 numGenSequences{}; - SizeType32 numContextTokens{}; - SizeType32 numGenTokens{}; - SizeType32 numLogits{}; - SizeType32 maxKvCacheLengthRounded{}; - - //! General - TensorPtr inputsIds; - - TensorPtr contextLengthsHost; - TensorPtr contextLengthsDevice; - TensorPtr sequenceLengthsHost; - - //! Index of selected runtime context. - SizeType32 contextIndex{}; - SizeType32 maxContextLength{}; - -public: - TensorPtr sequenceLengthsDevice; - bool promptTableOffloading; - - //! Prompt-Tuning - std::unique_ptr promptTuningBuffers; - -private: - //! Runtime - //! Type of host tensor: 0 for context, 1 for generation - TensorPtr requestTypes; - - TensorPtr lastTokenIdsHost; - TensorPtr lastTokenIdsDevice; - TensorPtr logitsIdsHost; - - //! Pipeline-Parallelism - TensorPtr hiddenStates; - - //! Mrope - TensorPtr mropeRotaryCosSin; - TensorPtr mropePositionDeltas; - - //! LoRA - std::unique_ptr loraBuffers; - -public: - //! Additional buffers depending on model type - std::unique_ptr transformerBuffers; - std::unique_ptr rnnStateBuffers; - - //! Encoder-Decoder - std::unique_ptr encoderBuffers; - - //! Medusa - std::unique_ptr mMedusaBuffers; - //! Lookahead decoding - std::unique_ptr mLookaheadBuffers; - //! Explicit draft tokens decoding - std::unique_ptr mExplicitDraftTokensBuffers; - //! Eagle decoding - std::unique_ptr mEagleBuffers; - - //! Language adapter routing information if language adapter is presented, [numTokens, numLanguages] - TensorPtr languageAdapterRoutings; - - TensorPtr cacheIndirDecoderIOBatchedCopySrcOffsets; - TensorPtr cacheIndirDecoderIOBatchedCopyDstOffsets; - TensorPtr cacheIndirDecoderIOBatchedCopySizes; - - //! Logits - std::vector numContextLogits; - TensorPtr logits; - - //! Helper cache for store generation logits - struct GenerationLogitsCache - { - static constexpr auto kCACHE_LENGTH = 8; - - //! Buffer for logits between steps to prevent from being overwritten - //! [kCACHE_LENGTH, maxBatchSize * maxBeamWidth, vocabSizePadded] - TensorPtr logits; - //! Record the usage offset of the cacheGenerationLogits buffer - SizeType32 offset{0}; - - //! Temporarily store the transposed results of multiple fragment logits, [maxBeamWidth, kCACHE_LENGTH] - TensorPtr transposedLogits; - - //! Temporarily store logits buffer address during the transposing, [kCACHE_LENGTH] - TensorPtr fragmentPointerDevice; - - //! Temporarily store logits buffer address during the transposing, [maxBatchSize, kCACHE_LENGTH] - TensorPtr fragmentPointerHost; - - //! Cycling index for workspace - size_t workIdx{0}; - - void cycleWorkIdx() - { - workIdx = (workIdx + 1) % (fragmentPointerHost->getShape().d[0]); - } - - [[nodiscard]] TensorPtr getFragmentPointerHost() - { - TensorPtr slice = runtime::ITensor::slice(fragmentPointerHost, workIdx, 1); - cycleWorkIdx(); - return slice; - }; - }; - - GenerationLogitsCache generationLogitsCache; - - //! Mapping from batch idx to slot id - TensorPtr seqSlots; - TensorPtr seqSlotsDevice; - - //! Explicitly device-copy src offsets to reduce warp stalls in copy batch kernel invocation - //! [mMaxNumRequests], on gpu - TensorPtr mCacheIndirDecoderIOBatchedCopySrcOffsetsSliceDevice; - //! Explicitly device-copy dst offsets to reduce warp stalls in copy batch kernel invocation - //! [mMaxNumRequests], on gpu - TensorPtr mCacheIndirDecoderIOBatchedCopyDstOffsetsSliceDevice; - //! Explicitly device-copy size to reduce warp stalls in copy batch kernel invocation - //! [mMaxNumRequests], on gpu - TensorPtr mCacheIndirDecoderIOBatchedCopyCopySizesDevice; - -private: - //! Re-capture cuda graph when max kv cache len of the batch has changed on kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE. - static SizeType32 constexpr kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE{256}; - - TensorMap mAdditionalOutputTensors; // Tensors storing additional output tensors. - - //! Engine I/O - TensorMap inputMap; - TensorMap outputMap; - -public: - RuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, - bool gatherGenerationLogits, std::optional maxNumTokens = std::nullopt, - std::optional> const& additionalModelOutputs = std::nullopt, - bool promptTableOffloading = false); - - RuntimeBuffers(RuntimeBuffers const& other) = delete; - RuntimeBuffers& operator=(RuntimeBuffers const& other) = delete; - RuntimeBuffers(RuntimeBuffers&& other) = delete; - RuntimeBuffers& operator=(RuntimeBuffers&& other) = delete; - - ~RuntimeBuffers(); - - std::tuple prepareStep(RequestVector const& contextRequests, - RequestVector const& genRequests, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - runtime::decoder::DecoderState const& decoderState, kv_cache_manager::BaseKVCacheManager* kvCacheManager, - kv_cache_manager::BaseKVCacheManager* crossKvCacheManager, rnn_state_manager::RnnStateManager* rnnStateManager, - PeftTable const& peftTable, runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool gatherGenerationLogits, bool trtOverlap, - OptionalRef newOutputTokens = std::nullopt); - - void prepareBuffersForCudaGraph(SizeType32 maxSequenceLength); - - void prepareExplicitDraftTokenBuffers(runtime::ExplicitDraftTokensBuffers::Inputs const& explicitDraftTokensBuffers, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig); - - void prepareEagleBuffers(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::EagleBuffers::Inputs const& eagleBuffers, runtime::TllmRuntime const& runtime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - -private: - void create(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, runtime::TllmRuntime const& runtime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, bool gatherGenerationLogits, - std::optional> const& additionalModelOutputs = std::nullopt); - - //! @brief set max sizes for pre-allocation - void setMaxBufferSizes(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::ModelConfig const& modelConfig, - std::optional maxNumRuntimeTokens); - - //! @brief set sizes depending on scheduled requests - void setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests); - - void reshape(runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool gatherGenerationLogits); - - void setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, SizeType32 maxBeamWidth, - SizeType32 maxAttentionWindow, runtime::decoder::DecoderState const& decoderState, - kv_cache_manager::BaseKVCacheManager* kvCacheManagerPtr, - kv_cache_manager::BaseKVCacheManager* crossKvCacheManagerPtr, - rnn_state_manager::RnnStateManager* rnnStateManagerPtr, PeftTable const& peftTable, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool trtOverlap, OptionalRef newOutputTokens); - - void fillIOMaps(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/transformerBuffers.h b/cpp/include/tensorrt_llm/batch_manager/transformerBuffers.h deleted file mode 100644 index b5254c6357b4..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/transformerBuffers.h +++ /dev/null @@ -1,145 +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 "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/kvCacheType.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -class MulticastTensor; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -namespace kv_cache_manager -{ -class BaseKVCacheManager; -} - -class TransformerBuffers -{ -public: - using SizeType32 = runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - static constexpr auto kCrossAttentionMaskTensorName = "cross_attention_mask"; - static constexpr auto kCrossAttentionPackedMaskTensorName = "cross_attention_packed_mask"; - static constexpr auto kPositionIdsTensorName = "position_ids"; - static constexpr auto kCacheIndirectionsTensorName = "cache_indirection"; - static constexpr auto kHostPastKeyValueLengthsTensorName = "host_past_key_value_lengths"; - static constexpr auto kHostSinkTokenLengthTensorName = "host_sink_token_length"; - static constexpr auto kHostMaxAttentionWindowSizesTensorName = "host_max_attention_window_sizes"; - static constexpr auto kHostContextProgressTensorName = "host_context_progress"; - static constexpr auto kKvCacheBlockOffsetsTensorName = "kv_cache_block_offsets"; - static constexpr auto kHostKvCacheBlockOffsetsTensorName = "host_kv_cache_block_offsets"; - static constexpr auto kCrossKvCacheBlockOffsetsTensorName = "cross_kv_cache_block_offsets"; - static constexpr auto kHostCrossKvCacheBlockOffsetsTensorName = "host_cross_kv_cache_block_offsets"; - static constexpr auto kHostCrossKvCachePoolPointersTensorName = "host_cross_kv_cache_pool_pointers"; - static constexpr auto kHostCrossKvCachePoolMappingTensorName = "host_cross_kv_cache_pool_mapping"; - static constexpr auto kSkipCrossAttentionBlocksTensorName = "skip_cross_attn_blocks"; - - TensorPtr pastKeyValueLengths; // Host tensor - TensorPtr positionIds; - - // max kv cache lengths. - TensorPtr maxAttentionWindows; - // sink token lengths. - TensorPtr sinkTokenLengths; - TensorPtr cacheIndirection; - TensorPtr kvCacheBlockOffsetsHost; // [numPools, maxBatch * maxBeamWidth, 2, maxBlocksPerSeq] - TensorPtr kvCacheBlockOffsetsDevice; // [numPools, maxBatch * maxBeamWidth, 2, maxBlocksPerSeq] - TensorPtr contextProgressHost; - - // Cross attention buffers - TensorPtr crossKvCacheBlockPoolPointers = nullptr; - TensorPtr crossKvCacheBlockPoolMapping = nullptr; - TensorPtr crossKvCacheBlockOffsetsHost = nullptr; - TensorPtr crossKvCacheBlockOffsetsDevice = nullptr; - TensorPtr crossAttentionMaskCopySrcOffsets = nullptr; // [maxNumRequest] pinned memory. - TensorPtr crossAttentionMaskCopyDstOffsets = nullptr; // [maxNumRequest] pinned memory. - TensorPtr crossAttentionMaskCopySizes = nullptr; // [maxNumRequest] pinned memory. - TensorPtr crossAttentionMaskDevice = nullptr; // [maxNumTokens, maxEncoderOutputLen] - // This is created to allow mixed memory types of crossAttentionMask (i.e. CPU and GPU). - TensorPtr crossAttentionMaskPinnedHost = nullptr; // [maxNumTokens, maxEncoderOutputLen] - // See more details in tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaPackedMask.cu. - // The attention packed mask for FMHA where each bit represents one mask. - TensorPtr crossAttentionPackedMaskDevice - = nullptr; // [maxBatchSize, maxInputLengthInBatch, roundUp(maxEncoderOutputLen, 32)] - // The number of cumulative Q sequence lengths in the mask input, which is used to get mask offsets for different - // requests. - TensorPtr crossAttentionCuQSeqLensDevice = nullptr; // [maxBatchSize + 1] - // The number of cumulative Q sequence lengths in the packed mask, which is used to get mask offsets for different - // requests. - TensorPtr crossAttentionPackedMaskCuMaskRowsDevice = nullptr; // [maxBatchSize + 1] - - TensorPtr cacheIndirBatchedCopySrcOffsets; - TensorPtr cacheIndirBatchedCopyDstOffsets; - TensorPtr cacheIndirBatchedCopySizes; - - TensorPtr fillValuesAlt; - TensorPtr fillValuesAltDevice; - TensorPtr seqSlotsAlt; - TensorPtr seqSlotsAltDevice; - TensorPtr skipCrossAttnBlocks; - - std::shared_ptr gemmAllReduceOutput; - - TransformerBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig); - - void reshape(SizeType32 numSequences, SizeType32 numInputTokens); - - void reshapeKvTensors(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxBlocksPerSeq, - kv_cache_manager::CacheType kvCacheType, SizeType32 numPools, runtime::BufferManager const& manager); - - void getBuffers(TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::ModelConfig const& modelConfig) const; - - void copyPositionIds(runtime::TllmRuntime const& runtime, std::vector const& positionIdsHost, - bool isChatGlm, TensorPtr const& decoderPositionIds); - - void copyKvBlockOffsets(RequestVector const& contextRequests, RequestVector const& genRequests, - kv_cache_manager::BaseKVCacheManager const* kvCacheManager, - kv_cache_manager::BaseKVCacheManager const* crossKvCacheManager, runtime::BufferManager const& manager); - - // Copy CacheIndirection from `decoderCacheIndirectionOutput` to `this->cacheIndirection` - void copyCacheIndirection(RequestVector const& genRequests, TensorPtr const& decoderCacheIndirectionOutput, - runtime::CudaStream const& stream); - - void copyCrossAttentionMasks(RequestVector const& contextRequests, RequestVector const& genRequests, - TensorPtr const& decoderContextLengthsDevice, TensorPtr const& encoderInputLengths, - SizeType32 maxDecoderContextLength, SizeType32 maxEncoderInputLengthInBatch, - runtime::TllmRuntime const& runtime); - - void copySkipCrossAttnBlocks(bool const& _skipCrossAttnBlocks, runtime::TllmRuntime const& runtime); - -private: - SizeType32 maxInputLen; - SizeType32 maxEncoderOutputLen; - SizeType32 maxNumTokens; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/batch_manager/updateDecoderBuffers.h b/cpp/include/tensorrt_llm/batch_manager/updateDecoderBuffers.h deleted file mode 100644 index 526a756e5546..000000000000 --- a/cpp/include/tensorrt_llm/batch_manager/updateDecoderBuffers.h +++ /dev/null @@ -1,51 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/algorithm.h" -#include "tensorrt_llm/runtime/modelConfig.h" - -namespace tensorrt_llm::runtime -{ -class BufferManager; -class CudaEvent; - -namespace decoder -{ -class DecoderState; -} // namespace decoder -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -class DecoderOutputBuffers; - -class UpdateDecoderBuffers : Algorithm -{ -public: - constexpr static auto name{"UpdateDecoderBuffers"}; - - UpdateDecoderBuffers() = default; - - runtime::CudaEvent operator()(runtime::ModelConfig const& modelConfig, DecoderOutputBuffers& decoderOutputBuffers, - runtime::BufferManager const& copyBufferManager, runtime::decoder::DecoderState const& decoderState, - bool returnLogProbs, runtime::CudaEvent const& decoderFinishEvent) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/include/tensorrt_llm/common/dataType.h b/cpp/include/tensorrt_llm/common/dataType.h index 2f19404f9c94..9b3bb5fdf0f0 100644 --- a/cpp/include/tensorrt_llm/common/dataType.h +++ b/cpp/include/tensorrt_llm/common/dataType.h @@ -19,7 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/tllmException.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -27,61 +27,61 @@ TRTLLM_NAMESPACE_BEGIN namespace common { -constexpr static size_t getDTypeSize(nvinfer1::DataType type) +constexpr static size_t getDTypeSize(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kINT64: return 8; - case nvinfer1::DataType::kINT32: [[fallthrough]]; - case nvinfer1::DataType::kFLOAT: return 4; - case nvinfer1::DataType::kBF16: [[fallthrough]]; - case nvinfer1::DataType::kHALF: return 2; - case nvinfer1::DataType::kBOOL: [[fallthrough]]; - case nvinfer1::DataType::kUINT8: [[fallthrough]]; - case nvinfer1::DataType::kINT8: [[fallthrough]]; - case nvinfer1::DataType::kFP8: return 1; - case nvinfer1::DataType::kINT4: TLLM_THROW("Cannot determine size of INT4 data type"); - case nvinfer1::DataType::kFP4: TLLM_THROW("Cannot determine size of FP4 data type"); + case tensorrt_llm::DataType::kINT64: return 8; + case tensorrt_llm::DataType::kINT32: [[fallthrough]]; + case tensorrt_llm::DataType::kFLOAT: return 4; + case tensorrt_llm::DataType::kBF16: [[fallthrough]]; + case tensorrt_llm::DataType::kHALF: return 2; + case tensorrt_llm::DataType::kBOOL: [[fallthrough]]; + case tensorrt_llm::DataType::kUINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kFP8: return 1; + case tensorrt_llm::DataType::kINT4: TLLM_THROW("Cannot determine size of INT4 data type"); + case tensorrt_llm::DataType::kFP4: TLLM_THROW("Cannot determine size of FP4 data type"); default: TLLM_THROW("Unknown dtype %d", static_cast(type)); } return 0; } -constexpr static size_t getDTypeSizeInBits(nvinfer1::DataType type) +constexpr static size_t getDTypeSizeInBits(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kINT64: return 64; - case nvinfer1::DataType::kINT32: [[fallthrough]]; - case nvinfer1::DataType::kFLOAT: return 32; - case nvinfer1::DataType::kBF16: [[fallthrough]]; - case nvinfer1::DataType::kHALF: return 16; - case nvinfer1::DataType::kBOOL: [[fallthrough]]; - case nvinfer1::DataType::kUINT8: [[fallthrough]]; - case nvinfer1::DataType::kINT8: [[fallthrough]]; - case nvinfer1::DataType::kFP8: return 8; - case nvinfer1::DataType::kINT4: [[fallthrough]]; - case nvinfer1::DataType::kFP4: return 4; + case tensorrt_llm::DataType::kINT64: return 64; + case tensorrt_llm::DataType::kINT32: [[fallthrough]]; + case tensorrt_llm::DataType::kFLOAT: return 32; + case tensorrt_llm::DataType::kBF16: [[fallthrough]]; + case tensorrt_llm::DataType::kHALF: return 16; + case tensorrt_llm::DataType::kBOOL: [[fallthrough]]; + case tensorrt_llm::DataType::kUINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kINT8: [[fallthrough]]; + case tensorrt_llm::DataType::kFP8: return 8; + case tensorrt_llm::DataType::kINT4: [[fallthrough]]; + case tensorrt_llm::DataType::kFP4: return 4; default: TLLM_THROW("Unknown dtype %d", static_cast(type)); } return 0; } -[[maybe_unused]] static std::string getDtypeString(nvinfer1::DataType type) +[[maybe_unused]] static std::string getDtypeString(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kFLOAT: return "fp32"; break; - case nvinfer1::DataType::kHALF: return "fp16"; break; - case nvinfer1::DataType::kINT8: return "int8"; break; - case nvinfer1::DataType::kINT32: return "int32"; break; - case nvinfer1::DataType::kBOOL: return "bool"; break; - case nvinfer1::DataType::kUINT8: return "uint8"; break; - case nvinfer1::DataType::kFP8: return "fp8"; break; - case nvinfer1::DataType::kBF16: return "bf16"; break; - case nvinfer1::DataType::kINT64: return "int64"; break; - case nvinfer1::DataType::kINT4: return "int4"; break; - case nvinfer1::DataType::kFP4: return "fp4"; break; + case tensorrt_llm::DataType::kFLOAT: return "fp32"; break; + case tensorrt_llm::DataType::kHALF: return "fp16"; break; + case tensorrt_llm::DataType::kINT8: return "int8"; break; + case tensorrt_llm::DataType::kINT32: return "int32"; break; + case tensorrt_llm::DataType::kBOOL: return "bool"; break; + case tensorrt_llm::DataType::kUINT8: return "uint8"; break; + case tensorrt_llm::DataType::kFP8: return "fp8"; break; + case tensorrt_llm::DataType::kBF16: return "bf16"; break; + case tensorrt_llm::DataType::kINT64: return "int64"; break; + case tensorrt_llm::DataType::kINT4: return "int4"; break; + case tensorrt_llm::DataType::kFP4: return "fp4"; break; default: throw std::runtime_error("Unsupported data type"); break; } diff --git a/cpp/include/tensorrt_llm/common/tllmDataType.h b/cpp/include/tensorrt_llm/common/tllmDataType.h new file mode 100644 index 000000000000..a77bd6297e75 --- /dev/null +++ b/cpp/include/tensorrt_llm/common/tllmDataType.h @@ -0,0 +1,96 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2022-2026 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 + +//! \file tllmDataType.h +//! +//! Standalone, TensorRT-free runtime types that replace the few \c nvinfer1 +//! types the shared C++ core historically used as common currency: +//! \c tensorrt_llm::DataType, \c tensorrt_llm::Dims and \c tensorrt_llm::ILogger. These are +//! defined here so the retained tree (runtime, batch manager, executor, kernels +//! and the nanobind bridge) compiles and links without the TensorRT library. +//! +//! The \c DataType enumerator values intentionally mirror the legacy +//! \c tensorrt_llm::DataType integer values so that previously-serialized executor +//! configs and KV-cache metadata remain byte-compatible. The \c Dims layout +//! mirrors \c tensorrt_llm::Dims (\c int32_t \c nbDims followed by \c int64_t +//! \c d[8]) for the same reason. + +namespace tensorrt_llm +{ + +//! \brief Standalone data-type enum. Values mirror the legacy +//! \c tensorrt_llm::DataType for serialization/format compatibility. +enum class DataType : int32_t +{ + kFLOAT = 0, + kHALF = 1, + kINT8 = 2, + kINT32 = 3, + kBOOL = 4, + kUINT8 = 5, + kFP8 = 6, + kBF16 = 7, + kINT64 = 8, + kINT4 = 9, + kFP4 = 10, + kE8M0 = 11, +}; + +//! \brief Standalone dimensions type. Layout mirrors \c tensorrt_llm::Dims +//! (rank plus up to \c MAX_DIMS 64-bit extents) so serialized shapes remain +//! compatible. +class Dims +{ +public: + //! The maximum rank (number of dimensions) supported for a tensor. + static constexpr int32_t MAX_DIMS{8}; + + //! The rank (number of dimensions). + int32_t nbDims; + + //! The extent of each dimension. + int64_t d[MAX_DIMS]; +}; + +//! \brief Severity levels for the internal logger, mirroring the legacy +//! \c tensorrt_llm::ILogger::Severity values. +enum class ILoggerSeverity : int32_t +{ + kINTERNAL_ERROR = 0, + kERROR = 1, + kWARNING = 2, + kINFO = 3, + kVERBOSE = 4, +}; + +//! \brief Minimal, TensorRT-free logger interface replacing +//! \c tensorrt_llm::ILogger for the retained C++ tree. +class ILogger +{ +public: + using Severity = ILoggerSeverity; + + virtual void log(Severity severity, char const* msg) noexcept = 0; + + virtual ~ILogger() = default; +}; + +} // namespace tensorrt_llm diff --git a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h index 5067ae61dc83..77ebcfff4f66 100644 --- a/cpp/include/tensorrt_llm/executor/dataTransceiverState.h +++ b/cpp/include/tensorrt_llm/executor/dataTransceiverState.h @@ -50,7 +50,7 @@ class CacheState final }; CacheState(ModelConfig modelConfig, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableBlockReuse = false, bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128, bool indexerKCacheUseFp4 = false) @@ -71,7 +71,7 @@ class CacheState final CacheState(std::vector nbKvHeadPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, SizeType32 tensorParallelism, SizeType32 pipelineParallelism, SizeType32 contextParallelism, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false, int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128, @@ -92,7 +92,7 @@ class CacheState final CacheState(SizeType32 nbAttentionLayers, SizeType32 nbKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, SizeType32 tensorParallelism, SizeType32 pipelineParallelism, SizeType32 contextParallelism, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, AttentionType attentionType = AttentionType::kDEFAULT, int kvFactor = 2, bool enableAttentionDP = false, int DPrank = 0, int DPsize = 0, bool enableBlockReuse = false, bool enablePartialReuse = false, bool hasIndexerKCache = false, SizeType32 indexerDimPerHead = 0, SizeType32 indexerKCacheQuantBlockSize = 128, @@ -238,8 +238,8 @@ class CacheState final RnnModelConfig mModelConfig; /// Number of RNN layers per pipeline parallelism rank. std::vector mLayerNumPerPP; - nvinfer1::DataType mConvStateDataType; - nvinfer1::DataType mSsmStateDataType; + tensorrt_llm::DataType mConvStateDataType; + tensorrt_llm::DataType mSsmStateDataType; [[nodiscard]] bool operator==(RnnCacheState const& other) const noexcept { @@ -263,7 +263,7 @@ class CacheState final return mAttentionConfig; } - [[nodiscard]] nvinfer1::DataType const& getDataType() const + [[nodiscard]] tensorrt_llm::DataType const& getDataType() const { return mDataType; } @@ -308,7 +308,7 @@ class CacheState final } void setRnnConfig(RnnModelConfig rnnModelConfig, std::vector rnnLayerNumPerPP, - nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + tensorrt_llm::DataType convStateDataType, tensorrt_llm::DataType ssmStateDataType) { mRnnCacheState = RnnCacheState{ std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType}; @@ -325,12 +325,12 @@ class CacheState final return getRnnCacheState().mModelConfig; } - [[nodiscard]] nvinfer1::DataType getConvStateDataType() const + [[nodiscard]] tensorrt_llm::DataType getConvStateDataType() const { return getRnnCacheState().mConvStateDataType; } - [[nodiscard]] nvinfer1::DataType getSsmStateDataType() const + [[nodiscard]] tensorrt_llm::DataType getSsmStateDataType() const { return getRnnCacheState().mSsmStateDataType; } @@ -395,7 +395,7 @@ class CacheState final friend class tensorrt_llm::executor::Serialization; ModelConfig mModelConfig; ParallelConfig mParallelConfig; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; AttentionConfig mAttentionConfig; bool mEnableBlockReuse{false}; bool mEnablePartialReuse{false}; diff --git a/cpp/include/tensorrt_llm/executor/disaggServerUtil.h b/cpp/include/tensorrt_llm/executor/disaggServerUtil.h deleted file mode 100644 index b68dce78738a..000000000000 --- a/cpp/include/tensorrt_llm/executor/disaggServerUtil.h +++ /dev/null @@ -1,158 +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. - */ -#pragma once - -#include "tensorrt_llm/executor/executor.h" - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::executor::disagg_executor -{ - -namespace texec = tensorrt_llm::executor; - -struct ResponseWithId -{ - - tensorrt_llm::executor::Response response; - IdType gid; - - ResponseWithId(tensorrt_llm::executor::Response&& response, IdType gid) - : response(std::move(response)) - , gid(gid) - { - } - - ResponseWithId(tensorrt_llm::executor::Response const& response, IdType gid) - : response(response) - , gid(gid) - { - } - - ResponseWithId(ResponseWithId&& other) noexcept - : response(std::move(other.response)) - , gid(other.gid) - { - other.gid = {}; - } - - ResponseWithId(ResponseWithId const& other) = default; - - ResponseWithId& operator=(ResponseWithId&& other) noexcept - { - if (this != &other) - { - response = std::move(other.response); - gid = other.gid; - other.gid = {}; - } - return *this; - } - - ResponseWithId& operator=(ResponseWithId const& other) - { - - if (this != &other) - { - response = other.response; - gid = other.gid; - } - return *this; - } - - ~ResponseWithId() = default; -}; - -class DisaggExecutorOrchestrator -{ -public: - /// @brief Constructs a DisaggExecutorOrchestrator object. - /// - /// @param ctxEnginePaths A vector of file paths to context engine files. - /// @param genEnginePaths A vector of file paths to generation engine files. - /// @param ctxExecutorConfigs A vector of ExecutorConfig for context executors. - /// @param genExecutorConfigs A vector of ExecutorConfig for generation executors. - /// @param hasContextAwaitThreads Whether or not there are threads that receive response for each generation - /// executor. - /// @param hasGenAwaitThreads Whether or not there are threads that receive response for each generation executor. - - DisaggExecutorOrchestrator(std::vector const& ctxEnginePaths, - std::vector const& genEnginePaths, - std::vector const& ctxExecutorConfigs, - std::vector const& genExecutorConfigs, bool hasContextAwaitThreads, - bool hasGenAwaitThreads); - - /// @brief Enqueue context-only requests to context executors. - /// @param requests A vector of context-only requests. - /// @param selectContextId The index of the context executor to use. If `std::nullopt`, the executor that has the - /// smallest number of inflight requests will be used. - /// @param batch If true,enqueue requests in same context executor.If false, will try to use a different executor - /// for each request. - /// @return A vector of global request ids, corresponding to the order of the requests in `requests`, the id - /// returned may be different from the request id in each executor. - [[nodiscard]] std::vector enqueueContext(std::vector const& requests, - std::optional selectContextId = std::nullopt, bool batch = false); - - /// @brief Enqueue generation-only requests to generation executors. - /// @param requests A vector of generation-only requests. - /// @param globalRequestIds A vector of global request ids, corresponding to the order of the requests,and must be - /// the ids returned by the enqueueContext function. - /// @param selectGenIdx The index of the generation executor to use. If `std::nullopt`, the executor that has the - /// smallest number of inflight requests will be used. - /// @param batch If true,enqueue requests in same generation executor.If false, will try to use a different executor - /// for each request. - - void enqueueGeneration(std::vector const& requests, std::vector const& globalRequestIds, - std::optional selectGenIdx = std::nullopt, bool batch = false); - - /// @brief Await for context responses - /// @param timeout The maximum time to wait for new responses - /// @param contextIdx The index of the context executor to use. If `std::nullopt`, return ready responses in all - /// context executors,if `hasContextAwaitThreads` is true, then this parameter must be std::nullopt. - /// @return A vector of responses with corresponding global request ids - - [[nodiscard]] std::vector awaitContextResponses( - std::optional const& timeout, std::optional contextIdx = std::nullopt); - - /// @brief Await for generation responses - /// @param timeout The maximum time to wait for new responses. - /// @param genIdx The index of the generation executor to use. If `std::nullopt`, return ready responses in all - /// generation executors,if `hasGenAwaitThreads` is true, then this parameter must be std::nullopt. - /// @return A vector of responses with corresponding global request ids. - [[nodiscard]] std::vector awaitGenerationResponses( - std::optional const& timeout, std::optional genIdx = std::nullopt); - - /// @brief Indicates if the current process is allowed to enqueueRequests - [[nodiscard]] bool canEnqueue() const; - - /// @brief Get context executors - [[nodiscard]] std::vector> const& getContextExecutors() const; - - /// @brief Get generation executors - [[nodiscard]] std::vector> const& getGenExecutors() const; - - ~DisaggExecutorOrchestrator(); - -private: - class Impl; - std::unique_ptr mImpl; -}; -} // namespace tensorrt_llm::executor::disagg_executor diff --git a/cpp/include/tensorrt_llm/plugins/api/tllmPlugin.h b/cpp/include/tensorrt_llm/plugins/api/tllmPlugin.h deleted file mode 100644 index e3d4613e3d00..000000000000 --- a/cpp/include/tensorrt_llm/plugins/api/tllmPlugin.h +++ /dev/null @@ -1,74 +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 - -// Forward declarations -namespace nvinfer1 -{ -class ILoggerFinder; -class ILogger; - -namespace v_1_0 -{ -class IPluginCreator; -class IPluginCreatorV3One; -class IPluginCreatorInterface; -} // namespace v_1_0 - -} // namespace nvinfer1 - -namespace tensorrt_llm::plugins::api -{ - -auto constexpr kDefaultNamespace = "tensorrt_llm"; - -class LoggerManager -{ -public: - //! Set the logger finder. - void setLoggerFinder(nvinfer1::ILoggerFinder* finder); - - //! Get the logger. - [[maybe_unused]] nvinfer1::ILogger* logger(); - - static LoggerManager& getInstance() noexcept; - - static nvinfer1::ILogger* defaultLogger() noexcept; - -private: - LoggerManager() = default; - - nvinfer1::ILoggerFinder* mLoggerFinder{nullptr}; - std::mutex mMutex; -}; -} // namespace tensorrt_llm::plugins::api - -extern "C" -{ - // This function is used for explicitly registering the TRT-LLM plugins and the default logger. - bool initTrtLlmPlugins(void* logger = tensorrt_llm::plugins::api::LoggerManager::defaultLogger(), - char const* libNamespace = tensorrt_llm::plugins::api::kDefaultNamespace); - - // The functions below are used by TensorRT to when loading a shared plugin library with automatic registering. - // see https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#generating-plugin-library - [[maybe_unused]] void setLoggerFinder([[maybe_unused]] nvinfer1::ILoggerFinder* finder); - [[maybe_unused]] nvinfer1::v_1_0::IPluginCreator* const* getPluginCreators(std::int32_t& nbCreators); - [[maybe_unused]] nvinfer1::v_1_0::IPluginCreatorInterface* const* getCreators(std::int32_t& nbCreators); -} diff --git a/cpp/include/tensorrt_llm/runtime/bufferManager.h b/cpp/include/tensorrt_llm/runtime/bufferManager.h index 8357443dc5ea..8d7c94583cd0 100644 --- a/cpp/include/tensorrt_llm/runtime/bufferManager.h +++ b/cpp/include/tensorrt_llm/runtime/bufferManager.h @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/cudaStream.h" #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -62,63 +62,63 @@ class BufferManager } } - static auto constexpr kBYTE_TYPE = nvinfer1::DataType::kUINT8; + static auto constexpr kBYTE_TYPE = tensorrt_llm::DataType::kUINT8; //! \brief Allocates an `IBuffer` of the given size on the GPU, using cudaMallocAsync. - [[nodiscard]] IBufferPtr gpu(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE) const; + [[nodiscard]] IBufferPtr gpu(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Allocates an `ITensor` of the given dimensions on the GPU, using cudaMallocAsync. - [[nodiscard]] ITensorPtr gpu(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE) const; + [[nodiscard]] ITensorPtr gpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Allocates an `IBuffer` of the given size on the GPU, using cudaMalloc. - [[nodiscard]] static IBufferPtr gpuSync(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr gpuSync(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions on the GPU, using cudaMalloc. - [[nodiscard]] static ITensorPtr gpuSync(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr gpuSync(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `IBuffer` of the given size on the CPU. - [[nodiscard]] static IBufferPtr cpu(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr cpu(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions on the CPU. - [[nodiscard]] static ITensorPtr cpu(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr cpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `IBuffer` of the given size on the CPU. - [[nodiscard]] static IBufferPtr pinned(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr pinned(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `ITensor` of the given dimensions on the CPU. - [[nodiscard]] static ITensorPtr pinned(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr pinned(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `IBuffer` of the given size on the CPU in the default memory pool. - [[nodiscard]] static IBufferPtr pinnedPool(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr pinnedPool(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates a pinned `ITensor` of the given dimensions on the CPU in the default memory pool. - [[nodiscard]] static ITensorPtr pinnedPool(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr pinnedPool(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `IBuffer` of the given size in UVM. - [[nodiscard]] static IBufferPtr managed(std::size_t size, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static IBufferPtr managed(std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions in UVM. - [[nodiscard]] static ITensorPtr managed(nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE); + [[nodiscard]] static ITensorPtr managed(tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE); //! \brief Allocates an `ITensor` of the given dimensions for NVLS - [[nodiscard]] static ITensorPtr ipcNvls(std::set ranks, nvinfer1::Dims dims, nvinfer1::DataType type); + [[nodiscard]] static ITensorPtr ipcNvls(std::set ranks, tensorrt_llm::Dims dims, tensorrt_llm::DataType type); //! \brief Allocates an `IBuffer` of the given size and memory type. [[nodiscard]] IBufferPtr allocate( - MemoryType memoryType, std::size_t size, nvinfer1::DataType type = kBYTE_TYPE) const; + MemoryType memoryType, std::size_t size, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Allocates an `ITensor` of the given dimensions and memory type. [[nodiscard]] ITensorPtr allocate( - MemoryType memoryType, nvinfer1::Dims dims, nvinfer1::DataType type = kBYTE_TYPE) const; + MemoryType memoryType, tensorrt_llm::Dims dims, tensorrt_llm::DataType type = kBYTE_TYPE) const; //! \brief Create an empty `IBuffer` of the given memory type. It may be resized later. - [[nodiscard]] IBufferPtr emptyBuffer(MemoryType memoryType, nvinfer1::DataType type = kBYTE_TYPE) const + [[nodiscard]] IBufferPtr emptyBuffer(MemoryType memoryType, tensorrt_llm::DataType type = kBYTE_TYPE) const { return allocate(memoryType, 0, type); } //! \brief Create an empty `ITensor` of the given memory type. It may be reshaped later. - [[nodiscard]] ITensorPtr emptyTensor(MemoryType memoryType, nvinfer1::DataType type = kBYTE_TYPE) const + [[nodiscard]] ITensorPtr emptyTensor(MemoryType memoryType, tensorrt_llm::DataType type = kBYTE_TYPE) const { return allocate(memoryType, ITensor::makeShape({}), type); } @@ -167,7 +167,7 @@ class BufferManager //! \brief Copy `src` into a new `ITensor` with a potentially different memory type. template - [[nodiscard]] ITensorPtr copyFrom(T* src, nvinfer1::Dims dims, MemoryType memoryType) const + [[nodiscard]] ITensorPtr copyFrom(T* src, tensorrt_llm::Dims dims, MemoryType memoryType) const { auto buffer = allocate(memoryType, dims, TRTDataType>::value); copy(src, *buffer); @@ -176,7 +176,7 @@ class BufferManager //! \brief Copy `src` into a new `ITensor` with a potentially different memory type. template - [[nodiscard]] ITensorPtr copyFrom(std::vector const& src, nvinfer1::Dims dims, MemoryType memoryType) const + [[nodiscard]] ITensorPtr copyFrom(std::vector const& src, tensorrt_llm::Dims dims, MemoryType memoryType) const { TLLM_CHECK_WITH_INFO(src.size() == ITensor::volumeNonNegative(dims), common::fmtstr("[TensorRT-LLM][ERROR] Incompatible size %lu and dims %s", src.size(), diff --git a/cpp/include/tensorrt_llm/runtime/decoderState.h b/cpp/include/tensorrt_llm/runtime/decoderState.h index 95d7ff0ffac9..73d52b682e4c 100644 --- a/cpp/include/tensorrt_llm/runtime/decoderState.h +++ b/cpp/include/tensorrt_llm/runtime/decoderState.h @@ -52,7 +52,7 @@ class DecoderState //! @brief Setup buffers for the decoder excluding speculative decoding. void setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, nvinfer1::DataType dtype, + SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); //! @brief Setup buffers for the cache indirection. @@ -62,7 +62,7 @@ class DecoderState //! @brief Setup buffers for speculative decoding. void setupSpeculativeDecoding(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, nvinfer1::DataType dtype, ModelConfig const& modelConfig, + SizeType32 maxTokensPerEngineStep, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager); //! @brief Disable lookahead decoding. @@ -199,7 +199,7 @@ class DecoderState [[nodiscard]] DecodingOutput& getJointDecodingOutput() const; private: - void setupBuffers(nvinfer1::DataType dtype, BufferManager const& bufferManager); + 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); @@ -209,7 +209,7 @@ class DecoderState SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow); void setupSpeculativeDecodingBuffers( - SpeculativeDecodingMode speculativeDecodingMode, nvinfer1::DataType dtype, BufferManager const& bufferManager); + 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); diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoder.h b/cpp/include/tensorrt_llm/runtime/gptDecoder.h index 7e0cc1bb56d2..876418597e13 100644 --- a/cpp/include/tensorrt_llm/runtime/gptDecoder.h +++ b/cpp/include/tensorrt_llm/runtime/gptDecoder.h @@ -22,7 +22,7 @@ #include "tensorrt_llm/runtime/decodingOutput.h" #include "tensorrt_llm/runtime/samplingConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -55,7 +55,7 @@ class IGptDecoder /// @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 explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) = 0; @@ -70,7 +70,7 @@ class IGptDecoder std::optional const& samplingConfig, SizeType32 batchSize, TensorConstPtr batchSlots) = 0; - static std::unique_ptr create(executor::DecodingMode const& mode, nvinfer1::DataType dtype, + 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); @@ -90,7 +90,7 @@ class GptDecoder : public virtual IGptDecoder void setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, std::optional const& output = std::nullopt, - std::optional explicitDraftTokensDType = std::nullopt, + std::optional explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) override; @@ -121,17 +121,17 @@ class GptDecoder : public virtual IGptDecoder executor::DecodingMode mDecodingMode; }; -inline std::unique_ptr IGptDecoder::create(executor::DecodingMode const& mode, nvinfer1::DataType dtype, +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 nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: return std::make_unique>( mode, maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, stream, speculativeDecodingModule); - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: return std::make_unique>( mode, maxNumSequences, maxBeamWidth, vocabSize, vocabSizePadded, stream, speculativeDecodingModule); default: diff --git a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h index 9fcd3262c8ca..ac1f49c31dd8 100644 --- a/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h +++ b/cpp/include/tensorrt_llm/runtime/gptDecoderBatched.h @@ -48,7 +48,7 @@ class GptDecoderBatched : public IGptDecoderBatched explicit GptDecoderBatched(CudaStreamPtr stream); void setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) override; + tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) override; void disableLookahead(RequestVector const& genRequests, TensorPtr const& batchSlots) override; diff --git a/cpp/include/tensorrt_llm/runtime/iBuffer.h b/cpp/include/tensorrt_llm/runtime/iBuffer.h index 91d5cd739f32..bf63d3a0da7b 100644 --- a/cpp/include/tensorrt_llm/runtime/iBuffer.h +++ b/cpp/include/tensorrt_llm/runtime/iBuffer.h @@ -22,7 +22,7 @@ #include "tensorrt_llm/kernels/kvCacheIndex.h" #include "tensorrt_llm/runtime/common.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #ifdef ENABLE_FP8 @@ -88,13 +88,13 @@ struct MemoryTypeString }; //! \brief For converting a TensorRT data type to a C++ data type. -template +template struct DataTypeTraits { }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = float; static char constexpr name[] = "float"; @@ -102,7 +102,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = half; static char constexpr name[] = "half"; @@ -110,7 +110,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::int8_t; static char constexpr name[] = "int8"; @@ -118,7 +118,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::int32_t; static char constexpr name[] = "int32"; @@ -126,7 +126,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::int64_t; static char constexpr name[] = "int64"; @@ -134,7 +134,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::uint32_t; static char constexpr name[] = "uint32"; @@ -142,7 +142,7 @@ struct DataTypeTraits }; template <> -struct DataTypeTraits +struct DataTypeTraits { using type = std::uint64_t; static char constexpr name[] = "uint64"; @@ -150,7 +150,7 @@ struct DataTypeTraits }; template -struct DataTypeTraits +struct DataTypeTraits { using type = bool; static char constexpr name[] = "bool"; @@ -158,7 +158,7 @@ struct DataTypeTraits }; template -struct DataTypeTraits +struct DataTypeTraits { using type = std::uint8_t; static char constexpr name[] = "uint8"; @@ -167,7 +167,7 @@ struct DataTypeTraits #ifdef ENABLE_BF16 template <> -struct DataTypeTraits +struct DataTypeTraits { using type = __nv_bfloat16; static char constexpr name[] = "bfloat16"; @@ -177,7 +177,7 @@ struct DataTypeTraits #ifdef ENABLE_FP8 template <> -struct DataTypeTraits +struct DataTypeTraits { using type = __nv_fp8_e4m3; static char constexpr name[] = "fp8"; @@ -185,7 +185,7 @@ struct DataTypeTraits }; #endif -template +template struct DataTypeTraits { using type = typename DataTypeTraits::type*; @@ -193,26 +193,26 @@ struct DataTypeTraits static auto constexpr size = sizeof(type); }; -//! \brief A wrapper around `nvinfer1::DataType` that provides a support for pointer types. +//! \brief A wrapper around `tensorrt_llm::DataType` that provides a support for pointer types. class BufferDataType { public: constexpr BufferDataType( // NOLINT(*-explicit-constructor) - nvinfer1::DataType dataType, bool _unsigned = false, bool pointer = false) + tensorrt_llm::DataType dataType, bool _unsigned = false, bool pointer = false) : mDataType{dataType} , mUnsigned{_unsigned} , mPointer{pointer} { } - static auto constexpr kTrtPointerType = nvinfer1::DataType::kINT64; + static auto constexpr kTrtPointerType = tensorrt_llm::DataType::kINT64; - constexpr operator nvinfer1::DataType() const noexcept // NOLINT(*-explicit-constructor) + constexpr operator tensorrt_llm::DataType() const noexcept // NOLINT(*-explicit-constructor) { return mPointer ? kTrtPointerType : mDataType; } - [[nodiscard]] constexpr nvinfer1::DataType getDataType() const noexcept + [[nodiscard]] constexpr tensorrt_llm::DataType getDataType() const noexcept { return mDataType; } @@ -226,24 +226,24 @@ class BufferDataType { switch (mDataType) { - case nvinfer1::DataType::kBOOL: [[fallthrough]]; - case nvinfer1::DataType::kUINT8: return true; + case tensorrt_llm::DataType::kBOOL: [[fallthrough]]; + case tensorrt_llm::DataType::kUINT8: return true; default: return mUnsigned; } } [[nodiscard]] constexpr std::size_t getSize() const noexcept { - return tensorrt_llm::common::getDTypeSize(static_cast(*this)); + return tensorrt_llm::common::getDTypeSize(static_cast(*this)); } [[nodiscard]] constexpr std::size_t getSizeInBits() const noexcept { - return tensorrt_llm::common::getDTypeSizeInBits(static_cast(*this)); + return tensorrt_llm::common::getDTypeSizeInBits(static_cast(*this)); } private: - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; bool mUnsigned; bool mPointer; }; @@ -257,62 +257,62 @@ struct TRTDataType template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kFLOAT; + static constexpr auto value = tensorrt_llm::DataType::kFLOAT; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kHALF; + static constexpr auto value = tensorrt_llm::DataType::kHALF; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kINT8; + static constexpr auto value = tensorrt_llm::DataType::kINT8; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kINT32; + static constexpr auto value = tensorrt_llm::DataType::kINT32; }; template <> struct TRTDataType { - static constexpr auto value = BufferDataType{nvinfer1::DataType::kINT32, true}; + static constexpr auto value = BufferDataType{tensorrt_llm::DataType::kINT32, true}; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kINT64; + static constexpr auto value = tensorrt_llm::DataType::kINT64; }; template <> struct TRTDataType { - static constexpr auto value = BufferDataType{nvinfer1::DataType::kINT64, true}; + static constexpr auto value = BufferDataType{tensorrt_llm::DataType::kINT64, true}; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kBOOL; + static constexpr auto value = tensorrt_llm::DataType::kBOOL; }; template <> struct TRTDataType { - static constexpr auto value = nvinfer1::DataType::kUINT8; + static constexpr auto value = tensorrt_llm::DataType::kUINT8; }; #ifdef ENABLE_BF16 template <> struct TRTDataType<__nv_bfloat16> { - static constexpr auto value = nvinfer1::DataType::kBF16; + static constexpr auto value = tensorrt_llm::DataType::kBF16; }; #endif @@ -320,7 +320,7 @@ struct TRTDataType<__nv_bfloat16> template <> struct TRTDataType<__nv_fp8_e4m3> { - static constexpr auto value = nvinfer1::DataType::kFP8; + static constexpr auto value = tensorrt_llm::DataType::kFP8; }; #endif @@ -380,7 +380,7 @@ class IBuffer using SharedPtr = std::shared_ptr; using UniqueConstPtr = std::unique_ptr; using SharedConstPtr = std::shared_ptr; - using DataType = nvinfer1::DataType; + using DataType = tensorrt_llm::DataType; //! //! \brief Returns a pointer to underlying array. diff --git a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h index ab55b754f9be..d6df7aba0a64 100644 --- a/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h +++ b/cpp/include/tensorrt_llm/runtime/iGptDecoderBatched.h @@ -51,7 +51,7 @@ class IGptDecoderBatched //! @brief Setup the decoder before calling `forward()` virtual void setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) + tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) = 0; //! @brief Disable Lookahead decoding. diff --git a/cpp/include/tensorrt_llm/runtime/iTensor.h b/cpp/include/tensorrt_llm/runtime/iTensor.h index eb5c10eeb691..9f51847ee65d 100644 --- a/cpp/include/tensorrt_llm/runtime/iTensor.h +++ b/cpp/include/tensorrt_llm/runtime/iTensor.h @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iBuffer.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -33,10 +33,6 @@ #include #include -namespace nvinfer1 -{ -class IExecutionContext; -} namespace tensorrt_llm::runtime { @@ -50,7 +46,7 @@ class ITensor : virtual public IBuffer using SharedPtr = std::shared_ptr; using UniqueConstPtr = std::unique_ptr; using SharedConstPtr = std::shared_ptr; - using Shape = nvinfer1::Dims; + using Shape = tensorrt_llm::Dims; using DimType64 = std::remove_reference_t; using TensorMap = runtime::StringPtrMap; @@ -352,9 +348,9 @@ class ITensor : virtual public IBuffer //! \param shape The shape of the tensor. //! \param capacity The capacity of the buffer. //! \return An `ITensor`. - static UniquePtr wrap(void* data, nvinfer1::DataType type, Shape const& shape, std::size_t capacity); + static UniquePtr wrap(void* data, tensorrt_llm::DataType type, Shape const& shape, std::size_t capacity); - static UniquePtr wrap(void* data, nvinfer1::DataType type, Shape const& shape) + static UniquePtr wrap(void* data, tensorrt_llm::DataType type, Shape const& shape) { return wrap(data, type, shape, volumeNonNegative(shape)); } diff --git a/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h b/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h index ecaa439f2d52..26c6e3886be4 100644 --- a/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h +++ b/cpp/include/tensorrt_llm/runtime/lookaheadBuffers.h @@ -17,9 +17,9 @@ #pragma once #include "tensorrt_llm/executor/executor.h" +#include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" #include "tensorrt_llm/runtime/worldConfig.h" namespace tensorrt_llm::runtime @@ -37,47 +37,4 @@ class LookaheadDecodingBuffers TensorPtr positionIds; }; -class LookaheadRuntimeBuffers -{ -public: - using TensorPtr = ITensor::SharedPtr; - using TensorMap = StringPtrMap; - - LookaheadRuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, BufferManager const& manager, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, - TllmRuntime const& runtime); - - void setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, ITensor const& requestTypes, - ITensor const& seqSlots, LookaheadDecodingBuffers const& decoderLookaheadBuffers, TllmRuntime const& runtime, - ModelConfig const& modelConfig, WorldConfig const& worldConfig) const; - - void reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 tokensPerStep); - - void insertInputTensors(TensorMap& inputBuffers, TensorMap& outputBuffers, WorldConfig const& worldConfig) const; - - void enableLookaheadDecoding(SizeType32 maxBatchSize, SizeType32 tokensPerStep); - - void disableLookaheadDecoding(); - -public: - TensorPtr cumSumLength; // [1] the cumulative sum of generation length, on pinned - TensorPtr packedMasksDevice; // [forwardBatchSize, tokensPerStep, numPackedMasks], on gpu - TensorPtr generationLengthsDevice; // [forwardBatchSize], on gpu - TensorPtr positionOffsetsDevice; // [forwardBatchSize, tokensPerStep], on gpu - TensorPtr positionIdsDevice; // [forwardBatchSize, tokensPerStep], on gpu - - TensorPtr packedMaskHost; - TensorPtr generationLengthsHost; - TensorPtr positionOffsetsHost; - TensorPtr positionIdsHost; - - TensorPtr packedMaskHostCopy; - TensorPtr generationLengthsHostCopy; - TensorPtr positionOffsetsHostCopy; - TensorPtr positionIdsHostCopy; - TensorPtr useSpecDecoding; - - TensorPtr batchSlotsHostCopy; -}; - } // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/loraCache.h b/cpp/include/tensorrt_llm/runtime/loraCache.h index 1d242cdc80c5..d293d4f7b490 100644 --- a/cpp/include/tensorrt_llm/runtime/loraCache.h +++ b/cpp/include/tensorrt_llm/runtime/loraCache.h @@ -25,7 +25,7 @@ #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include diff --git a/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h b/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h index a995304e94ec..cf1b1a6aac18 100644 --- a/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h +++ b/cpp/include/tensorrt_llm/runtime/loraCachePageManagerConfig.h @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/iBuffer.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -36,7 +36,7 @@ namespace tensorrt_llm::runtime class LoraCachePageManagerConfig { public: - explicit constexpr LoraCachePageManagerConfig(runtime::MemoryType memType, nvinfer1::DataType dType, + explicit constexpr LoraCachePageManagerConfig(runtime::MemoryType memType, tensorrt_llm::DataType dType, SizeType32 totalNumPages, SizeType32 maxPagesPerBlock, SizeType32 slotsPerPage, SizeType32 pageWidth, SizeType32 numCopyStreams) : mMemoryType(memType) @@ -59,12 +59,12 @@ class LoraCachePageManagerConfig mMemoryType = memoryType; } - [[nodiscard]] nvinfer1::DataType constexpr getDataType() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getDataType() const noexcept { return mDataType; } - void constexpr setDataType(nvinfer1::DataType const& dtype) noexcept + void constexpr setDataType(tensorrt_llm::DataType const& dtype) noexcept { mDataType = dtype; } @@ -131,7 +131,7 @@ class LoraCachePageManagerConfig private: runtime::MemoryType mMemoryType; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; /* * Number cache pages in the cache. @@ -154,7 +154,7 @@ inline std::ostream& operator<<(std::ostream& os, LoraCachePageManagerConfig con { os << "{" << "memoryType=" << static_cast::type>(c.getMemoryType()) - << " dataType=" << static_cast::type>(c.getDataType()) + << " dataType=" << static_cast::type>(c.getDataType()) << " totalNumPages=" << c.getTotalNumPages() << " maxPagesPerBlock=" << c.getMaxPagesPerBlock() << " slotsPerPage=" << c.getSlotsPerPage() << " pageWidth=" << c.getPageWidth() << " initToZero=" << c.getInitToZero() << "}"; diff --git a/cpp/include/tensorrt_llm/runtime/modelConfig.h b/cpp/include/tensorrt_llm/runtime/modelConfig.h index 5bfe7bce9d58..b5f18da07f3b 100644 --- a/cpp/include/tensorrt_llm/runtime/modelConfig.h +++ b/cpp/include/tensorrt_llm/runtime/modelConfig.h @@ -23,7 +23,7 @@ #include "tensorrt_llm/runtime/speculativeDecodingMode.h" #include "tensorrt_llm/runtime/speculativeDecodingModule.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include namespace tensorrt_llm::runtime @@ -101,7 +101,7 @@ class ModelConfig }; explicit ModelConfig(SizeType32 vocabSize, SizeType32 nbLayers, SizeType32 nbAttentionLayers, - SizeType32 nbRnnLayers, SizeType32 nbHeads, SizeType32 hiddenSize, nvinfer1::DataType dtype) + SizeType32 nbRnnLayers, SizeType32 nbHeads, SizeType32 hiddenSize, tensorrt_llm::DataType dtype) : mVocabSize(vocabSize) , mNbLayers(nbLayers) , mNbAttentionLayers(nbAttentionLayers) @@ -137,7 +137,7 @@ class ModelConfig , mUsePositionEmbedding(false) , mUseTokenTypeEmbedding(false) , mSpeculativeDecodingMode(SpeculativeDecodingMode::None()) - , mLogitsDtype(nvinfer1::DataType::kFLOAT) + , mLogitsDtype(tensorrt_llm::DataType::kFLOAT) , mUseShapeInference(true) , mManageWeightsType(ManageWeightsType::kDisabled) , mSkipCrossAttnBlocks(false) @@ -331,7 +331,7 @@ class ModelConfig mSizePerHead = sizePerHead; } - [[nodiscard]] nvinfer1::DataType constexpr getDataType() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getDataType() const noexcept { return mDataType; } @@ -735,20 +735,20 @@ class ModelConfig resetSpeculativeDecodingModule(); } - [[nodiscard]] nvinfer1::DataType getKvDataType() const + [[nodiscard]] tensorrt_llm::DataType getKvDataType() const { if (getQuantMode().hasFp8KvCache()) { - return nvinfer1::DataType::kFP8; + return tensorrt_llm::DataType::kFP8; } if (getQuantMode().hasInt8KvCache()) { - return nvinfer1::DataType::kINT8; + return tensorrt_llm::DataType::kINT8; } else if (getQuantMode().hasFp4KvCache()) { #ifdef ENABLE_FP4 - return nvinfer1::DataType::kFP4; + return tensorrt_llm::DataType::kFP4; #else throw std::runtime_error("Model has FP4 KV cache, but TRT-LLM was not compiled with FP4 enabled."); #endif @@ -800,22 +800,22 @@ class ModelConfig return mSpeculativeDecodingMode; } - void setLogitsDtype(nvinfer1::DataType inputDtype) noexcept + void setLogitsDtype(tensorrt_llm::DataType inputDtype) noexcept { mLogitsDtype = inputDtype; } - [[nodiscard]] nvinfer1::DataType constexpr getLogitsDtype() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getLogitsDtype() const noexcept { return mLogitsDtype; } - void setGemmAllReduceDtype(nvinfer1::DataType inputDtype) noexcept + void setGemmAllReduceDtype(tensorrt_llm::DataType inputDtype) noexcept { mGemmAllReduceDtype = inputDtype; } - [[nodiscard]] nvinfer1::DataType constexpr getGemmAllReduceDtype() const noexcept + [[nodiscard]] tensorrt_llm::DataType constexpr getGemmAllReduceDtype() const noexcept { return mGemmAllReduceDtype; } @@ -945,10 +945,10 @@ class ModelConfig SizeType32 mNbHeads; SizeType32 mHiddenSize; SizeType32 mSizePerHead; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; bool mUseGptAttentionPlugin; bool mUseGemmAllReducePlugin; - nvinfer1::DataType mGemmAllReduceDtype; + tensorrt_llm::DataType mGemmAllReduceDtype; bool mUseMambaConv1dPlugin; bool mInputPacked; bool mPagedState; @@ -998,7 +998,7 @@ class ModelConfig SpeculativeDecodingMode mSpeculativeDecodingMode; // Logits datatype - nvinfer1::DataType mLogitsDtype; + tensorrt_llm::DataType mLogitsDtype; bool mUseShapeInference; ManageWeightsType mManageWeightsType; std::string mModelName; diff --git a/cpp/include/tensorrt_llm/runtime/rawEngine.h b/cpp/include/tensorrt_llm/runtime/rawEngine.h deleted file mode 100644 index b219cbe03382..000000000000 --- a/cpp/include/tensorrt_llm/runtime/rawEngine.h +++ /dev/null @@ -1,123 +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/executor/tensor.h" - -#include -#include -#include -#include - -namespace tensorrt_llm::runtime -{ - -class RawEngine -{ -public: - enum Type - { - FilePath, - AddressWithSize, - HostMemory - }; - - explicit RawEngine(std::filesystem::path enginePath) noexcept - : mType(FilePath) - , mEnginePath(std::move(enginePath)) - { - } - - explicit RawEngine(void const* engineAddr, std::size_t engineSize) noexcept - : mType(AddressWithSize) - , mEngineAddr(engineAddr) - , mEngineSize(engineSize) - { - } - - explicit RawEngine(nvinfer1::IHostMemory const* engineBuffer) noexcept - : mType(HostMemory) - , mEngineBuffer(engineBuffer) - { - } - - [[nodiscard]] Type getType() const - { - return mType; - } - - [[nodiscard]] std::filesystem::path getPath() const - { - TLLM_CHECK(mEnginePath.has_value()); - return mEnginePath.value(); - } - - [[nodiscard]] std::optional getPathOpt() const - { - return mEnginePath; - } - - void setPath(std::filesystem::path enginePath) - { - mEnginePath = std::move(enginePath); - } - - [[nodiscard]] std::optional> const& - getManagedWeightsMapOpt() const - { - return mManagedWeightsMap; - } - - void setManagedWeightsMap(std::map managedWeightsMap) - { - mManagedWeightsMap = std::move(managedWeightsMap); - } - - [[nodiscard]] void const* getAddress() const - { - TLLM_CHECK(mType == AddressWithSize); - return mEngineAddr; - } - - [[nodiscard]] std::size_t getSize() const - { - TLLM_CHECK(mType == AddressWithSize); - return mEngineSize; - } - - [[nodiscard]] nvinfer1::IHostMemory const* getHostMemory() const - { - TLLM_CHECK(mType == HostMemory); - return mEngineBuffer; - } - -private: - Type mType; - std::optional mEnginePath; - - struct - { - void const* mEngineAddr{}; - std::size_t mEngineSize{}; - }; - - nvinfer1::IHostMemory const* mEngineBuffer{}; - std::optional> mManagedWeightsMap; -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/tllmLogger.h b/cpp/include/tensorrt_llm/runtime/tllmLogger.h deleted file mode 100644 index dd3806ec5242..000000000000 --- a/cpp/include/tensorrt_llm/runtime/tllmLogger.h +++ /dev/null @@ -1,34 +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 - -namespace tensorrt_llm::runtime -{ - -class TllmLogger : public nvinfer1::ILogger -{ -public: - void log(Severity severity, nvinfer1::AsciiChar const* msg) noexcept override; - - Severity getLevel(); - - void setLevel(Severity level); -}; - -} // namespace tensorrt_llm::runtime diff --git a/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h b/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h index 68064d74c7e2..c86b23a9c740 100644 --- a/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h +++ b/cpp/include/tensorrt_llm/runtime/utils/debugUtils.h @@ -24,7 +24,7 @@ template bool tensorHasInvalid(ITensor const& tensor, BufferManager const& manager, std::string const& infoStr); bool tensorHasInvalid( - size_t M, size_t K, nvinfer1::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr); + size_t M, size_t K, tensorrt_llm::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr); int stallStream( char const* name, std::optional stream = std::nullopt, std::optional delay = std::nullopt); diff --git a/cpp/include/tensorrt_llm/runtime/worldConfig.h b/cpp/include/tensorrt_llm/runtime/worldConfig.h index 9ff2d0970df7..96703a8ef1c2 100644 --- a/cpp/include/tensorrt_llm/runtime/worldConfig.h +++ b/cpp/include/tensorrt_llm/runtime/worldConfig.h @@ -18,7 +18,7 @@ #include "tensorrt_llm/runtime/common.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include diff --git a/cpp/tensorrt_llm/CMakeLists.txt b/cpp/tensorrt_llm/CMakeLists.txt index 06aefb3886ba..df1fb27b0116 100644 --- a/cpp/tensorrt_llm/CMakeLists.txt +++ b/cpp/tensorrt_llm/CMakeLists.txt @@ -145,7 +145,7 @@ add_subdirectory(kernels) add_subdirectory(layers) add_subdirectory(runtime) add_subdirectory(testing) -add_subdirectory(executor_worker) +# add_subdirectory(executor_worker) # TensorRT-engine worker removed set(BATCH_MANAGER_TARGET tensorrt_llm_batch_manager_static) set(BATCH_MANAGER_TARGET_ARCH ${TARGET_ARCH}) @@ -309,4 +309,4 @@ if(BUILD_FLASH_MLA) add_subdirectory(flash_mla) endif() -add_subdirectory(plugins) +# add_subdirectory(plugins) # TensorRT plugin layer removed diff --git a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt index 88e2484cc6f1..6290acaccde9 100644 --- a/cpp/tensorrt_llm/batch_manager/CMakeLists.txt +++ b/cpp/tensorrt_llm/batch_manager/CMakeLists.txt @@ -32,34 +32,21 @@ set(SRCS contextProgress.cpp dataTransceiver.cpp decoderBuffers.cpp - encoderBuffers.cpp guidedDecoder.cpp - handleContextLogits.cpp - handleGenerationLogits.cpp kvCacheManager.cpp kvCacheEventManager.cpp kvCacheTransferManager.cpp kvCacheManagerV2Utils.cpp kvCacheManagerV2Utils.cu llmRequest.cpp - logitsPostProcessor.cpp - loraBuffers.cpp - makeDecodingBatchInputOutput.cpp medusaBuffers.cpp microBatchScheduler.cpp pauseRequests.cpp peftCacheManager.cpp - promptTuningBuffers.cpp - rnnStateBuffers.cpp rnnStateManager.cpp rnnCacheFormatter.cpp rnnCacheTransBuffer.cpp - runtimeBuffers.cpp sequenceSlotManager.cpp - transformerBuffers.cpp - trtEncoderModel.cpp - trtGptModelInflightBatching.cpp - updateDecoderBuffers.cpp utils/debugUtils.cpp utils/inflightBatchingUtils.cpp utils/logitsThread.cpp diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp index 3be8ee22bc1f..ea77d3eb9a8a 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.cpp @@ -51,7 +51,7 @@ void BufferIndexHolder::release() noexcept } BaseTransBufferManager::BaseTransBufferManager( - size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens) + size_t transferBufferSize, tensorrt_llm::DataType dataType, std::optional maxNumTokens) : mDataType{dataType} , mBufferManager{std::make_shared()} , mMaxNumTokens{maxNumTokens} diff --git a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h index 2cbf9f514bd7..214375cccd43 100644 --- a/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h +++ b/cpp/tensorrt_llm/batch_manager/baseTransBuffer.h @@ -197,7 +197,7 @@ class BaseTransBufferManager /// @param dataType Data type for the buffers. /// @param maxNumTokens Optional max tokens for sizing. BaseTransBufferManager( - size_t transferBufferSize, nvinfer1::DataType dataType, std::optional maxNumTokens = std::nullopt); + size_t transferBufferSize, tensorrt_llm::DataType dataType, std::optional maxNumTokens = std::nullopt); struct ConcurrenceResource { @@ -224,7 +224,7 @@ class BaseTransBufferManager bool mOnlyUseDynamicBuffer; bool mUseFabricMemory; size_t mNumberOfElements; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; ConcurrenceResource mConcurrenceSendResource; ConcurrenceResource mConcurrenceRecvResource; runtime::BufferManager mBufferManager; diff --git a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h index 458cac8d4382..066898a7db9b 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheFormatter.h +++ b/cpp/tensorrt_llm/batch_manager/cacheFormatter.h @@ -28,7 +28,7 @@ #include "tensorrt_llm/executor/dataTransceiverState.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp index b6f1e4e8df15..a396244eeafd 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransBuffer.cpp @@ -21,7 +21,7 @@ #include "tensorrt_llm/common/opUtils.h" #include "tensorrt_llm/executor/executor.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include namespace tensorrt_llm::batch_manager::kv_cache_manager @@ -194,7 +194,7 @@ bool FabricMemory::supportFbaricMemory() size_t CacheTransBufferManager::computeTransferBufferSize( KVCacheManager::BaseKVCacheManager* cacheManager, std::optional maxNumTokens, bool transferIndexerKCache) { - nvinfer1::DataType dataType; + tensorrt_llm::DataType dataType; if (transferIndexerKCache) { dataType = cacheManager->getIndexerKCachePool()->getDataType(); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp index 18c2b19e9eeb..7f50e824fab4 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransceiver.cpp @@ -272,7 +272,7 @@ std::unique_ptr CacheTransceiverFactory::createCacheTransc CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheManager, executor::kv_cache::CacheState::ModelConfig const& cacheStateModelCfg, runtime::WorldConfig const& worldConfig, - std::vector const& attentionLayerNumPerPP, nvinfer1::DataType dataType, + std::vector const& attentionLayerNumPerPP, tensorrt_llm::DataType dataType, executor::kv_cache::CacheState::AttentionType attentionType, std::optional cacheTransceiverConfig, rnn_state_manager::RnnStateManager* rnnStateManager, std::vector const& rnnLayerNumPerPP) @@ -387,20 +387,20 @@ CacheTransceiver::CacheTransceiver(kv_cache_manager::BaseKVCacheManager* cacheMa // Pool dtype is UINT8 (raw byte storage), so we cannot use pool->getDataType(). // Only the byte size matters for split/concat kernel stride calculations — the actual // dtype enum is not interpreted numerically, just used for getDTypeSize() dispatch. - auto dtypeFromSize = [](SizeType32 size) -> nvinfer1::DataType + auto dtypeFromSize = [](SizeType32 size) -> tensorrt_llm::DataType { switch (size) { - case 4: return nvinfer1::DataType::kFLOAT; - case 2: return nvinfer1::DataType::kBF16; - case 1: return nvinfer1::DataType::kFP8; + case 4: return tensorrt_llm::DataType::kFLOAT; + case 2: return tensorrt_llm::DataType::kBF16; + case 1: return tensorrt_llm::DataType::kFP8; default: TLLM_THROW("Unsupported RNN state dtype size: %d", size); } }; TLLM_CHECK_WITH_INFO(linearMeta->rnnSsmDtypeSize > 0, "rnnSsmDtypeSize not set in LinearAttentionMetadata"); TLLM_CHECK_WITH_INFO(linearMeta->rnnConvDtypeSize > 0, "rnnConvDtypeSize not set in LinearAttentionMetadata"); - nvinfer1::DataType ssmDtype = dtypeFromSize(linearMeta->rnnSsmDtypeSize); - nvinfer1::DataType convDtype = dtypeFromSize(linearMeta->rnnConvDtypeSize); + tensorrt_llm::DataType ssmDtype = dtypeFromSize(linearMeta->rnnSsmDtypeSize); + tensorrt_llm::DataType convDtype = dtypeFromSize(linearMeta->rnnConvDtypeSize); mCacheState->setRnnConfig(rnnModelCfg, rnnLayerNumPerPP, convDtype, ssmDtype); // Create RnnCacheTransBufferManager for unified pool path. diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp index c013fd75c6e4..d8a0a5c27a76 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.cpp @@ -122,7 +122,7 @@ void CacheTransferLayer::unformat(TransferSession& session) const } void CacheTransferLayer::setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, tensorrt_llm::DataType ssmStateDataType) { mCacheState.setRnnConfig( std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType); diff --git a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h index 0506e98197e6..c82202591f40 100644 --- a/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h +++ b/cpp/tensorrt_llm/batch_manager/cacheTransferLayer.h @@ -79,8 +79,8 @@ class CacheTransferLayer /// @brief Update the RNN config on the internal CacheState. /// Used by CppMambaHybridCacheManager path where RNN config is set after construction. void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType); + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType); [[nodiscard]] kv_cache_manager::BaseKVCacheManager* getCacheManager() const noexcept; diff --git a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp index 5c5d3e11a01c..5d22d1c9cecb 100644 --- a/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp +++ b/cpp/tensorrt_llm/batch_manager/createNewDecoderRequests.cpp @@ -33,7 +33,7 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/speculativeChoicesUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" using namespace tensorrt_llm::runtime; @@ -93,7 +93,7 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe /// @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(nvinfer1::DataType logitsType, TensorPtr const& tensor) +[[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) @@ -102,7 +102,7 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe } // Support FP32 input for FP16 embedding bias (in the case of FP8 models) - if (tensor->getDataType() == nvinfer1::DataType::kFLOAT && logitsType == nvinfer1::DataType::kHALF) + 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( @@ -133,7 +133,7 @@ void copySequenceLengths(RequestVector const& contextRequests, DecoderInputBuffe std::tuple, std::vector, std::vector> CreateNewDecoderRequests::operator()(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, nvinfer1::DataType logitsType, + 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 @@ -235,7 +235,7 @@ void initializeBeamSearch(DecodingInput& dJointInput, DecodingOutput& dJointOutp } void initializeEmbeddingBias(DecodingInput& dJointInput, SizeType32 batchSlot, - std::optional const& embeddingBias, nvinfer1::DataType logitsType, + 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); @@ -631,7 +631,7 @@ void newRequestSpeculativeDecoding(DecodingInput& jointDecodingInput, DecodingOu std::tuple, std::vector> CreateNewDecoderRequests::createDecoderRequests(RequestVector const& finishedContextRequests, TensorPtr const& inputIds, executor::DecodingConfig const& decodingConfig, runtime::decoder::DecoderState& decoderState, - nvinfer1::DataType logitsType, runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, + 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 { diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp index 853866687d53..c05e5c34e65f 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.cpp @@ -20,7 +20,6 @@ #include "tensorrt_llm/batch_manager/cacheFormatter.h" #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/kvCacheUtils.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" #include "tensorrt_llm/common/envUtils.h" #include "tensorrt_llm/common/logger.h" #include "tensorrt_llm/common/tllmException.h" @@ -735,8 +734,8 @@ class CacheSender::Impl public: void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType) { mCacheTransferLayer.setRnnConfig(rnnModelConfig, rnnLayerNumPerPP, convStateDataType, ssmStateDataType); mSelfState.setCacheState(mCacheTransferLayer.getCacheState()); @@ -1225,8 +1224,8 @@ class CacheReceiver::Impl public: void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType) { mCacheTransferLayer.setRnnConfig(rnnModelConfig, rnnLayerNumPerPP, convStateDataType, ssmStateDataType); mSelfState.setCacheState(mCacheTransferLayer.getCacheState()); @@ -1302,7 +1301,7 @@ void CacheSender::sendReadySignal(LlmRequest::RequestIdType requestId, bool isRe } void CacheSender::setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, tensorrt_llm::DataType ssmStateDataType) { mImpl->setRnnConfig(std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType); } @@ -1341,7 +1340,7 @@ bool CacheReceiver::receiveReadySignal(TransferSession& session) } void CacheReceiver::setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, nvinfer1::DataType ssmStateDataType) + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, tensorrt_llm::DataType ssmStateDataType) { mImpl->setRnnConfig(std::move(rnnModelConfig), std::move(rnnLayerNumPerPP), convStateDataType, ssmStateDataType); } diff --git a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h index 3362574da902..f1aa33e780de 100644 --- a/cpp/tensorrt_llm/batch_manager/dataTransceiver.h +++ b/cpp/tensorrt_llm/batch_manager/dataTransceiver.h @@ -290,8 +290,8 @@ class CacheSender /// @brief Update the RNN config on the internal CacheState copies. /// Used by CppMambaHybridCacheManager path where RNN config is set after construction. void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType); + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType); /// @brief Destructor. virtual ~CacheSender(); @@ -341,8 +341,8 @@ class CacheReceiver /// @brief Update the RNN config on the internal CacheState copies. /// Used by CppMambaHybridCacheManager path where RNN config is set after construction. void setRnnConfig(executor::kv_cache::CacheState::RnnModelConfig rnnModelConfig, - std::vector rnnLayerNumPerPP, nvinfer1::DataType convStateDataType, - nvinfer1::DataType ssmStateDataType); + std::vector rnnLayerNumPerPP, tensorrt_llm::DataType convStateDataType, + tensorrt_llm::DataType ssmStateDataType); /// @brief Destructor. virtual ~CacheReceiver(); diff --git a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp index fd67bb55e89d..a7e01ce270a6 100644 --- a/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp +++ b/cpp/tensorrt_llm/batch_manager/decoderBuffers.cpp @@ -70,21 +70,21 @@ DecoderOutputBuffers::DecoderOutputBuffers(SizeType32 maxNumSequences, SizeType3 auto constexpr TRTTokenIdType = runtime::TRTDataType::value; sequenceLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kINT32); - finishedSumHost = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::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}), nvinfer1::DataType::kFLOAT); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kFLOAT); logProbsHost = BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kFLOAT); + ITensor::makeShape({maxNumSequences, maxBeamWidth, maxSeqLen}), tensorrt_llm::DataType::kFLOAT); finishReasonsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), nvinfer1::DataType::kUINT8); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences, maxBeamWidth}), tensorrt_llm::DataType::kUINT8); } void DecoderOutputBuffers::enableLookaheadDecoding(SizeType32 maxNumSequences, SizeType32 maxTokensPerStep) @@ -115,9 +115,9 @@ void DecoderOutputBuffers::setupSpeculativeDecoding( if (speculativeDecodingMode.variableDraftLength()) { nextDraftTokensLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); prevDraftTokensLengthsHost - = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinned(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); } } } @@ -307,17 +307,17 @@ DecoderSlotAsyncSend::~DecoderSlotAsyncSend() SlotDecoderBuffers::SlotDecoderBuffers(SizeType32 maxBeamWidth, SizeType32 maxSeqLen, BufferManager const& manager) { - outputIds = manager.gpu(ITensor::makeShape({maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kINT32); - outputIdsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), nvinfer1::DataType::kINT32); + 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}), nvinfer1::DataType::kINT32); - sequenceLengthsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), nvinfer1::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}), nvinfer1::DataType::kFLOAT); - cumLogProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth}), nvinfer1::DataType::kFLOAT); + 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}), nvinfer1::DataType::kFLOAT); - logProbsHost = BufferManager::pinned(ITensor::makeShape({maxBeamWidth, maxSeqLen}), nvinfer1::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/encoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/encoderBuffers.cpp deleted file mode 100644 index 56fd393c68d7..000000000000 --- a/cpp/tensorrt_llm/batch_manager/encoderBuffers.cpp +++ /dev/null @@ -1,560 +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 "encoderBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/nvtxUtils.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 - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -EncoderBuffers::EncoderBuffers( - SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - // init empty buffers on cpu/gpu/pinned - init(maxBatchSize, modelConfig, worldConfig, runtime); - - // pre-allocate based on max buffer sizes - // Note: pre-allocation can be done directly instead of empty-->reshape, but it is ok extract the common reshape() - // utility because the buffer shapes can be dynamically set during runtime as well - initBufferSizes(maxBatchSize, modelConfig, worldConfig, runtime); -} - -void EncoderBuffers::init( - SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - - auto hiddenStatesType = modelConfig.getDataType(); - - inputFeatures = manager.emptyTensor(MemoryType::kGPU, hiddenStatesType); - inputIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - // in PP, only rank 0 needs the following input fields - if (modelConfig.usePositionEmbedding() && worldConfig.isFirstPipelineParallelRank()) - { - positionIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - positionIdsReserved.resize(maxBatchSize * modelConfig.getMaxInputLen()); - std::iota(positionIdsReserved.begin(), positionIdsReserved.end(), 0); - } - if (modelConfig.useTokenTypeEmbedding() && worldConfig.isFirstPipelineParallelRank()) - { - tokenTypeIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - tokenTypeIdsReserved.resize(maxBatchSize * modelConfig.getMaxInputLen()); - std::fill(tokenTypeIdsReserved.begin(), tokenTypeIdsReserved.end(), 0); - } - - inputLengths = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - maxInputLength = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - if (worldConfig.isPipelineParallel()) - { - hiddenStates = manager.emptyTensor(MemoryType::kGPU, hiddenStatesType); - } - if (worldConfig.isLastPipelineParallelRank()) - { - encoderOutput = manager.emptyTensor(MemoryType::kGPU, hiddenStatesType); - } - - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings = manager.emptyTensor(MemoryType::kGPU, TRTDataType::value); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::initBufferSizes( - SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // get buffer shape based on max values - numRequests = maxBatchSize; - encoderInputLen = maxBatchSize * modelConfig.getMaxInputLen(); - encoderOutputLen = maxBatchSize * modelConfig.getMaxInputLen(); // assume output length <= input length - maxInputLengthInBatch = modelConfig.getMaxInputLen(); - - // update buffer shapes - reshape(runtime, modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::updateBufferSizes(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - numRequests = requests.size(); - encoderInputLen = 0; - encoderOutputLen = 0; - maxInputLengthInBatch = 0; - - // get buffer shape based on actual batched requests - for (auto const& req : requests) - { - encoderInputLen += req->getEncoderInputLen(); - encoderOutputLen += req->getEncoderOutputLen(); - maxInputLengthInBatch - = std::max(maxInputLengthInBatch, req->getEncoderInputLen()); // Decoder input is encoder output - } - - // update buffer shapes - reshape(runtime, modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::reshape(TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (modelConfig.isMultiModal()) - { - return; // multimodal models do not need to set position id, etc. or any output tensors - } - - inputIds->reshape(ITensor::makeShape({encoderInputLen})); - if (positionIds) - { - if (modelConfig.isWhisper()) - { - positionIds->reshape(ITensor::makeShape({encoderOutputLen})); - } - else - { - positionIds->reshape(ITensor::makeShape({encoderInputLen})); - } - } - if (tokenTypeIds) - { - tokenTypeIds->reshape(ITensor::makeShape({encoderInputLen})); - } - - inputLengths->reshape(ITensor::makeShape({numRequests})); - maxInputLength->reshape(ITensor::makeShape({maxInputLengthInBatch})); - - if (worldConfig.isPipelineParallel()) - { - hiddenStates->reshape( - ITensor::makeShape({encoderOutputLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - if (worldConfig.isLastPipelineParallelRank()) - { - encoderOutput->reshape( - ITensor::makeShape({encoderOutputLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings->reshape(ITensor::makeShape({encoderInputLen, 1})); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::setFromInputs(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(encoderBuffersSetFromInputs); - - if (!worldConfig.isFirstPipelineParallelRank()) - { - return; - } - - auto const& manager = runtime.getBufferManager(); - - std::vector inputIdsAll; - std::vector positionIdsAll; - std::vector tokenTypeIdsAll; - std::vector inputLengthsAll; - std::vector languageAdapterRoutingAll; - // use shape to indicates max input length, content is not important - // TODO: change to a scalar value for this from engine side - std::vector maxInputLengthAll(maxInputLengthInBatch, 0); - - if (requests.front()->getEncoderInputFeatures()) - { - if (modelConfig.isMultiModal()) - { - auto batchedInputShape = requests.front()->getEncoderInputFeatures()->getShape(); // [1, 3, H, W] - batchedInputShape.d[0] = encoderInputLen; // [batch_size, 3, H, W] - inputFeatures->reshape(batchedInputShape); - } - else - { - SizeType32 const featureDim = requests.front()->getEncoderInputFeatures()->getShape().d[1]; - TLLM_LOG_DEBUG("EncoderBuffers::setFromInputs - featureDim = %d", featureDim); - inputFeatures->reshape(ITensor::makeShape({encoderInputLen, featureDim})); - } - } - - SizeType32 offset = 0; - - for (auto const& llmReq : requests) - { - SizeType32 const inputLength = llmReq->getEncoderInputLen(); - SizeType32 const outputLength = llmReq->getEncoderOutputLen(); - if (llmReq->getEncoderInputFeatures()) - { - auto const& reqFeatures - = llmReq - ->getEncoderInputFeatures(); // whisper: [length, featureDim]; Vision: [batch_size, channel, W, H] - TLLM_LOG_DEBUG("EncoderBuffers::setFromInputs - request id = %d, input features length = %d", - llmReq->mRequestId, inputLength); - manager.copy(*reqFeatures, *ITensor::slice(inputFeatures, offset, inputLength)); - offset += inputLength; - } - else - { - auto const& reqTokens = *llmReq->getEncoderTokens().value(); - inputIdsAll.insert(inputIdsAll.end(), reqTokens.begin(), reqTokens.end()); - if (tokenTypeIds) - { - tokenTypeIdsAll.insert( - tokenTypeIdsAll.end(), tokenTypeIdsReserved.begin(), tokenTypeIdsReserved.begin() + inputLength); - } - } - if (positionIds) - { - SizeType32 const length = modelConfig.isWhisper() ? outputLength : inputLength; - positionIdsAll.insert( - positionIdsAll.end(), positionIdsReserved.begin(), positionIdsReserved.begin() + length); - } - if (modelConfig.useLanguageAdapter()) - { - auto const languageAdapterRouting - = llmReq->getLanguageAdapterRouting(modelConfig.getNumLanguages().value(), inputLength); - languageAdapterRoutingAll.insert( - languageAdapterRoutingAll.end(), std::begin(languageAdapterRouting), std::end(languageAdapterRouting)); - } - inputLengthsAll.push_back(inputLength); - } - - // copy inputs from host to device - { - NVTX3_SCOPED_RANGE(bufferCopies); - if (requests.front()->getEncoderTokens()) - { - manager.copy(inputIdsAll.data(), *inputIds); - if (tokenTypeIds) - { - manager.copy(tokenTypeIdsAll.data(), *tokenTypeIds); - } - manager.copy(maxInputLengthAll.data(), *maxInputLength); - } - if (positionIds) - { - manager.copy(positionIdsAll.data(), *positionIds); - } - manager.copy(inputLengthsAll.data(), *inputLengths); - if (modelConfig.useLanguageAdapter()) - { - manager.copy(languageAdapterRoutingAll.data(), *languageAdapterRoutings); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::fillIOMaps(ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersFillIOMaps); - - inputMap.clear(); - outputMap.clear(); - - // inputs - if (modelConfig.isMultiModal()) - { - inputMap.insert_or_assign("input", inputFeatures); - } - else if (modelConfig.isWhisper()) - { - inputMap.insert_or_assign("input_features", inputFeatures); - inputMap.insert_or_assign("input_lengths", inputLengths); - inputMap.insert_or_assign("position_ids", positionIds); - } - else - { - if (worldConfig.isFirstPipelineParallelRank()) - { - inputMap.insert_or_assign("input_ids", inputIds); - if (positionIds) - { - inputMap.insert_or_assign("position_ids", positionIds); - } - if (tokenTypeIds) - { - inputMap.insert_or_assign("token_type_ids", tokenTypeIds); - } - } - else - { - inputMap.insert_or_assign("hidden_states_input", hiddenStates); - } - inputMap.insert_or_assign("input_lengths", inputLengths); - inputMap.insert_or_assign("max_input_length", maxInputLength); - if (modelConfig.useLanguageAdapter()) - { - inputMap.insert_or_assign("language_adapter_routings", languageAdapterRoutings); - } - } - - // outputs - if (worldConfig.isLastPipelineParallelRank()) - { - outputMap.insert_or_assign("encoder_output", encoderOutput); - } - else - { - outputMap.insert_or_assign("hidden_states_output", hiddenStates); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::pair EncoderBuffers::prepareIO( - RequestVector const& requests, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - updateBufferSizes(requests, modelConfig, worldConfig, runtime); - - setFromInputs(requests, modelConfig, worldConfig, runtime); - - fillIOMaps(modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - - return {inputMap, outputMap}; -} - -void EncoderBuffers::rearrangeOutputs(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(encoderBuffersRearrangeOutput); - - auto const& manager = runtime.getBufferManager(); - - SizeType32 offset = 0, size = 0; - - updateReqOutputShape(requests, runtime, worldConfig, modelConfig); - - for (auto const& req : requests) - { - // copy from internal buffer to request-owned external buffers - size = req->getEncoderOutputLen(); - TLLM_LOG_DEBUG("EncoderBuffers::rearrangeOutputs - req: %d, encoderOutput shape = (%d, %d)", req->mClientId, - req->getEncoderOutput()->getShape().d[0], req->getEncoderOutput()->getShape().d[1]); - TLLM_LOG_DEBUG("EncoderBuffers::rearrangeOutputs - req: %d, enc output size = %d", req->mClientId, size); - - if (worldConfig.isPipelineParallel()) - { - manager.copy(*ITensor::slice(hiddenStates, offset, size), *req->getEncoderHiddenStates()); - } - if (worldConfig.isLastPipelineParallelRank()) - { - if (modelConfig.isMultiModal()) - { - manager.copy( - *ITensor::slice(encoderOutput, offset, size), *(req->getPromptEmbeddingTableMutable().value())); - } - else - { - manager.copy(*ITensor::slice(encoderOutput, offset, size), *req->getEncoderOutput()); - } - } - offset += size; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::updateReqOutputShape(RequestVector const& requests, TllmRuntime const& runtime, - WorldConfig const& worldConfig, ModelConfig const& modelConfig) -{ - auto const& manager = runtime.getBufferManager(); - - for (auto const& req : requests) - { - if (modelConfig.isMultiModal()) - { - auto shape = encoderOutput->getShape(); // [batch_size, prompt_vocab_size, feature_dim] - shape.d[0] = req->getEncoderOutputLen(); - req->getPromptEmbeddingTableMutable() = manager.emptyTensor(MemoryType::kGPU, encoderOutput->getDataType()); - req->getPromptEmbeddingTableMutable().value()->reshape(shape); - req->setPromptVocabSize(shape.d[1]); - // TODO: extra ids for kv cache reuse - } - else - { - auto encOutLen = req->getEncoderOutputLen(); - // update request-owned external buffer for each request - if (worldConfig.isPipelineParallel()) - { - req->getEncoderHiddenStates()->reshape( - ITensor::makeShape({encOutLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - if (worldConfig.isLastPipelineParallelRank()) - { - req->getEncoderOutput()->reshape( - ITensor::makeShape({encOutLen, modelConfig.getHiddenSize() * worldConfig.getTensorParallelism()})); - } - } - } -} - -void EncoderBuffers::create(SizeType32 maxBatchSize, ModelConfig const& modelConfig, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - - inputLengths = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - maxInputLength = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - hiddenSize = modelConfig.getEncoderHiddenSize(); // full hidden size - // assume encoder & decoder use the same data type - encoderOutput = manager.emptyTensor(MemoryType::kGPU, modelConfig.getDataType()); - encoderOutputReserved = manager.gpu(ITensor::makeShape({1, hiddenSize}), modelConfig.getDataType()); - - crossKvCacheGen = manager.gpu(ITensor::makeShape({1}), nvinfer1::DataType::kBOOL); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::setMaxBufferSizes(SizeType32 maxBatchSize, runtime::ModelConfig const& modelConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - numRequests = maxBatchSize; - encoderInputLen = maxBatchSize * modelConfig.getMaxEncoderLen(); - encoderOutputLen = maxBatchSize * modelConfig.getMaxEncoderLen(); - maxInputLengthInBatch = modelConfig.getMaxEncoderLen(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - numRequests = 0; /// total number of requests that need encoder information (context requests + - /// generation requests * beam width) - encoderInputLen = 0; - encoderOutputLen = 0; - maxInputLengthInBatch = 1; /// maximum encoder length in a batch - - for (auto const& llmReq : contextRequests) - { - numRequests += 1; - encoderInputLen += llmReq->getEncoderInputLen(); - encoderOutputLen += llmReq->getEncoderOutputLen(); - maxInputLengthInBatch = std::max(maxInputLengthInBatch, llmReq->getEncoderInputLen()); - } - - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - numRequests += reqBeamWidth; // tile by beam width - maxInputLengthInBatch = std::max(maxInputLengthInBatch, llmReq->getEncoderInputLen()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::reshape() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - inputLengths->reshape(ITensor::makeShape({numRequests})); - maxInputLength->reshape(ITensor::makeShape({maxInputLengthInBatch})); - encoderOutput->reshape(ITensor::makeShape({encoderOutputLen, hiddenSize})); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::fill( - RequestVector const& ctxRequests, RequestVector const& genRequests, runtime::BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(encoderBufferCopies); - - std::vector inputLengthsAll; - std::vector maxInputLengthAll(maxInputLength->getShape().d[0], 0); - - SizeType32 offset = 0, size = 0; - for (auto const& requests : {ctxRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - // 1. only ctx requests should gather the encoder output - // 2. only gen requests should tile encoder input lengths info by beam width - bool isCtx = llmReq->isContextInitState(); - if (isCtx) - { - size = llmReq->getEncoderOutputLen(); - auto const encoderOutputSlice = runtime::ITensor::slice(encoderOutput, offset, size); - manager.copy(*llmReq->getEncoderOutput(), *encoderOutputSlice); - offset += size; - - inputLengthsAll.emplace_back(size); - } - else - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - std::fill_n(std::back_inserter(inputLengthsAll), reqBeamWidth, - llmReq->getEncoderOutputLen()); // although encoder output is not needed, gen phase still needs the - // encoder length info for cross kv cache. Also tile by beam width - } - } - } - manager.copy(inputLengthsAll.data(), *inputLengths); - manager.copy(maxInputLengthAll.data(), *maxInputLength); - // crossKvCacheGen unused in engine for now, use default tensor - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EncoderBuffers::insertInputTensors(TensorMap& inputMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - inputMap.insert_or_assign("encoder_output", encoderOutput); - inputMap.insert_or_assign("encoder_input_lengths", inputLengths); - inputMap.insert_or_assign("encoder_max_input_length", maxInputLength); - inputMap.insert_or_assign("cross_kv_cache_gen", crossKvCacheGen); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/encoderBuffers.h b/cpp/tensorrt_llm/batch_manager/encoderBuffers.h deleted file mode 100644 index 64d416280f21..000000000000 --- a/cpp/tensorrt_llm/batch_manager/encoderBuffers.h +++ /dev/null @@ -1,140 +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/common.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::batch_manager -{ - -class EncoderBuffers -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using ITensor = tensorrt_llm::runtime::ITensor; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - using ModelConfig = runtime::ModelConfig; - using WorldConfig = runtime::WorldConfig; - using TllmRuntime = runtime::TllmRuntime; - - TensorPtr inputIds; - TensorPtr positionIds = nullptr; - TensorPtr tokenTypeIds = nullptr; - - TensorPtr inputLengths; // [numEncoderRequests] - TensorPtr maxInputLength; // [maxInputLengthInBatch] - - // intermediate states in pipeline parallelism - TensorPtr hiddenStates; // [numTokens, hiddenSize] - - // features for multimodal encoders (audio, image, etc.) - TensorPtr - inputFeatures; // [totalNumOfFeatures, featureDim] if remove_padding else [batchSize, featureDim, featureLength] - - // language adapter routing information for encoders if language adapter is presented. - TensorPtr languageAdapterRoutings; // [numTokens, numLanguages] - - // encoder output - TensorPtr encoderOutput; // [numEncoderTokens, hiddenSize] - - // output buffer owned by llmRequest, such that it's per-request output buffer - // encoderBuffers class can init and reshape each buffer, without maintaining a list/set of inflight buffers - // TODO in progress: to support BS>1 encoder, need (1) internal scratch space tensors to save the contiguous - // batched output (2) copy from CONTIGUOUS scratch tensor to individual request's DISCRETE output tensor after - // execution To standardize the implementation, for both BS=1 and BS>1, we use internal buffer to store BS=1/BS>1 - // results, and copy to request's external buffers. For BS=1, this introduces a redundancy copy, but ok for now. - - EncoderBuffers() = default; - EncoderBuffers(SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - std::pair prepareIO(RequestVector const& requests, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, TllmRuntime const& runtime); - - void rearrangeOutputs(RequestVector const& requests, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - //! @brief set shape of individual request's encoder output (Ptuning embedding table if multimodal) - void updateReqOutputShape(RequestVector const& requests, TllmRuntime const& runtime, WorldConfig const& worldConfig, - ModelConfig const& modelConfig); - -private: - SizeType32 numRequests{}; - SizeType32 encoderInputLen{}; - SizeType32 encoderOutputLen{}; - SizeType32 maxInputLengthInBatch{}; // max input length in a batch - - // prefilled with deterministic values to avoid runtime creation - std::vector positionIdsReserved; - std::vector tokenTypeIdsReserved; - - // engine I/O - TensorMap inputMap; - TensorMap outputMap; - - void init(SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - //! @brief pre-allocate max buffer sizes during init - void initBufferSizes(SizeType32 maxBatchSize, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - //! @brief update actual buffer usage of requests during runtime - void updateBufferSizes(RequestVector const& requests, ModelConfig const& modelConfig, - WorldConfig const& worldConfig, TllmRuntime const& runtime); - - void reshape(TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig); - - void setFromInputs(RequestVector const& requests, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - TllmRuntime const& runtime); - - void fillIOMaps(ModelConfig const& modelConfig, WorldConfig const& worldConfig); - - // additional members that are Encoder-Decoder specific -private: - TensorPtr encoderOutputReserved; // [1, hiddenSize], dummy tensor for gen phase - TensorPtr crossKvCacheGen; // [1] - SizeType32 hiddenSize; // full hidden size (after multiplying tensor parallelism) - -public: - void create(SizeType32 maxBatchSize, ModelConfig const& modelConfig, TllmRuntime const& runtime); - - SizeType32 getMaxInputLengthInBatch() const - { - return maxInputLengthInBatch; - }; - - void setMaxBufferSizes(SizeType32 maxBatchSize, runtime::ModelConfig const& modelConfig); - - void setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests); - - void reshape(); - - void fill( - RequestVector const& ctxRequests, RequestVector const& genRequests, runtime::BufferManager const& manager); - - void insertInputTensors(TensorMap& inputMap); -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp b/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp index cb2264ec8003..8fb5fe0be818 100644 --- a/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp +++ b/cpp/tensorrt_llm/batch_manager/guidedDecoder.cpp @@ -30,7 +30,7 @@ namespace tensorrt_llm::batch_manager { GuidedDecoder::GuidedDecoder(executor::GuidedDecodingConfig const& guidedDecodingConfig, SizeType32 maxNumSequences, - SizeType32 vocabSizePadded, nvinfer1::DataType logitsDtype, BufferManager const& runtimeBufferManager) + SizeType32 vocabSizePadded, tensorrt_llm::DataType logitsDtype, BufferManager const& runtimeBufferManager) : mGuidedDecodingBackend{guidedDecodingConfig.getBackend()} , mMaxNumSequences{maxNumSequences} , mVocabSizePadded{vocabSizePadded} @@ -201,13 +201,13 @@ void GuidedDecoder::execute(DecoderInputBuffers const& decoderInputBuffers, Buff *ITensor::slice(mLogitsBitmaskPtrVec, 0, batchIdx)); auto logitsBitmaskPtrVec = bufferCast(*mLogitsBitmaskPtrVec); - if (mLogitsDtype == nvinfer1::DataType::kFLOAT) + if (mLogitsDtype == tensorrt_llm::DataType::kFLOAT) { auto logitsPtrVec = bufferCast(*mLogitsPtrVec); tensorrt_llm::kernels::invokeLogitsBitmask( logitsPtrVec, logitsBitmaskPtrVec, batchIdx, mVocabSizePadded, stream.get()); } - else if (mLogitsDtype == nvinfer1::DataType::kHALF) + else if (mLogitsDtype == tensorrt_llm::DataType::kHALF) { auto logitsPtrVec = bufferCast(*mLogitsPtrVec); tensorrt_llm::kernels::invokeLogitsBitmask( diff --git a/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp b/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp deleted file mode 100644 index 6f4a541ffcbb..000000000000 --- a/cpp/tensorrt_llm/batch_manager/handleContextLogits.cpp +++ /dev/null @@ -1,176 +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/handleContextLogits.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/runtimeBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -namespace tr = tensorrt_llm::runtime; -namespace tru = tensorrt_llm::runtime::utils; - -namespace tensorrt_llm::batch_manager -{ - -using BufferManager = tensorrt_llm::runtime::BufferManager; -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -namespace -{ - -//! @brief Copy logits from context phase to beginning of generation logits. -//! @details Usually, this concerns logits of 1 token. In speculative decoding this concerns draftLen + 1 tokens. -void copyLastContextLogits(TensorPtr const& contextLogits, LlmRequest& llmReq, BufferManager const& bufferManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const numLogits = contextLogits->getShape().d[0]; - for (int beam = 0; beam < llmReq.getBeamWidthByIter(); beam++) - { - // [beamWidth, mMaxNewTokens, vocabSizePadded] -> [numLogits, vocabSizePadded] - auto beamHostTensorPtr = ITensor::slice(llmReq.getGenerationLogitsHost(), {beam, 0}, numLogits); - bufferManager.copy(*contextLogits, *beamHostTensorPtr); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void setupMedusaLogits(std::vector& medusaLogitsHeads, TensorPtr const& medusaLogitsDevice, - SizeType32 medusaHeads, SizeType32 logitsIndex, SizeType32 numLogits) -{ - for (SizeType32 hi = 0; hi < medusaHeads; ++hi) - { - TensorPtr logitsHead = ITensor::slice(medusaLogitsDevice, hi, 1); - logitsHead->squeeze(0); - medusaLogitsHeads[hi] = ITensor::slice(logitsHead, logitsIndex, numLogits); - } -} - -} // namespace - -SizeType32 HandleContextLogits::operator()(DecoderInputBuffers& inputBuffers, RequestVector const& contextRequests, - tr::ITensor::SharedPtr const& logits, std::vector const& numContextLogitsVec, - tr::ModelConfig const& modelConfig, tr::BufferManager const& manager, - OptionalRef medusaBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(HandleContextLogits); - - auto& decoderRequests = inputBuffers.decoderRequests; - decoderRequests.clear(); - decoderRequests.reserve(contextRequests.size()); - auto& allDecoderLogits = inputBuffers.decoderLogits; - allDecoderLogits.clear(); - allDecoderLogits.reserve(contextRequests.size()); - - SizeType32 batchIndex{0}; - SizeType32 logitsIndex{0}; - // Copy logits into decoderBuffers.logits - for (auto const& llmReq : contextRequests) - { - auto const numContextLogits = numContextLogitsVec.at(batchIndex); - auto const draftLength = llmReq->isLastContextChunk() ? llmReq->getNumDraftTokens() : 0; - - TLLM_LOG_DEBUG("logitsIndex: %d", logitsIndex); - TLLM_LOG_DEBUG("numContextLogits %d", numContextLogits); - TLLM_LOG_DEBUG("draftLength: %d", draftLength); - - if (modelConfig.computeContextLogits()) - { - // Since the computational graph has been modified, only the last token is needed. - TLLM_CHECK_WITH_INFO(!modelConfig.getSpeculativeDecodingMode().isMedusa() - && !modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding(), - "Return context logits is not supported with Medusa and Lookahead decoding"); - - if (llmReq->getReturnContextLogits()) - { - if (llmReq->getPrepopulatedPromptLen() > 0) - { - TLLM_LOG_WARNING( - "Because of KV cache reuse, not all context logits could be produced for request %lu.", - llmReq->mRequestId); - } - TensorPtr contextLogitsDeviceView = ITensor::slice(logits, logitsIndex, numContextLogits); - TensorPtr contextLogitsHostView = ITensor::slice( - llmReq->getContextLogitsHost(), llmReq->getContextCurrentPosition(), numContextLogits); - // Copy to host directly - manager.copy(*contextLogitsDeviceView, *contextLogitsHostView); - } - } - logitsIndex += numContextLogits + draftLength; - - // Get the logits from the last context token and draft tokens - auto const numDecoderLogits = 1 + draftLength; - auto const seqSlot = llmReq->mSeqSlot.value(); - TensorPtr logitsView = ITensor::slice(logits, logitsIndex - numDecoderLogits, numDecoderLogits); - - if (modelConfig.getSpeculativeDecodingMode().hasDraftLogits()) - { - auto& medusaLogitsHeads = inputBuffers.predictedDraftLogits.at(seqSlot); - TLLM_CHECK(medusaBuffers); - setupMedusaLogits(medusaLogitsHeads, medusaBuffers->medusaLogitsDevice, - modelConfig.getSpeculativeDecodingModule().getMaxDraftPathLen(), logitsIndex - numDecoderLogits, - numDecoderLogits); - } - - // Save the last token logits of context into generation logits or - // save the accepted token logits from target model - if (llmReq->getReturnGenerationLogits()) - { - copyLastContextLogits(logitsView, *llmReq, manager); - } - - TLLM_CHECK_DEBUG_WITH_INFO(tru::tensorHasInvalid(*logitsView, manager, "logits") == false, - "Found invalid number (NaN or Inf) in logits"); - - if (llmReq->isLastContextChunk()) - { - TensorPtr decoderLogits; - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - if (reqBeamWidth > 1) - { - // Tile logits of context requests - auto const& logitsShape = logitsView->getShape(); - auto const logitsType = logitsView->getDataType(); - decoderLogits = manager.gpu(ITensor::makeShape({reqBeamWidth, logitsShape.d[1]}), logitsType); - tensorrt_llm::runtime::kernels::tileTensor( - *decoderLogits, *logitsView, reqBeamWidth, manager.getStream()); - decoderLogits->unsqueeze(0); - } - else - { - decoderLogits = logitsView; - decoderLogits->unsqueeze(1); - } - decoderRequests.push_back(llmReq); - allDecoderLogits.emplace_back(std::move(decoderLogits)); - } - - ++batchIndex; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return logitsIndex; -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp b/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp deleted file mode 100644 index e2a7486b050a..000000000000 --- a/cpp/tensorrt_llm/batch_manager/handleGenerationLogits.cpp +++ /dev/null @@ -1,161 +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/handleGenerationLogits.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/runtimeBuffers.h" -#include "tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -namespace tr = tensorrt_llm::runtime; -namespace tru = tensorrt_llm::runtime::utils; - -namespace tensorrt_llm::batch_manager -{ - -using BufferManager = tensorrt_llm::runtime::BufferManager; -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -namespace -{ - -//! @brief Copy logits from generation phase under streaming mode. -void copyStreamingGenerationLogits(BufferManager const& bufferManager, LlmRequest& llmReq) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - // If llmRequest is streaming, directly copy to host. - // Only one token's logits needs to be copied each time. - TLLM_CHECK(llmReq.getGenerationLogitsFragmentsSize() == 1); - - SizeType32 numGenerationToken = llmReq.getMaxBeamNumTokens() - llmReq.mPromptLen; - TensorPtr const& generationLogitsHost - = llmReq.getGenerationLogitsHost(); // [mMaxNewTokens (or 1), beamWidth, vocabSizePadded] - - TensorPtr hostTensorPtr - = ITensor::slice(generationLogitsHost, numGenerationToken, 1); // [1, beamWidth, vocabSizePadded] - TensorPtr deviceTensorPtr = *(llmReq.getGenerationLogitsFragments().begin()); - - bufferManager.copy(*deviceTensorPtr, *hostTensorPtr); - llmReq.clearGenerationLogitsFragments(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void setupMedusaLogits(std::vector& medusaLogitsHeads, TensorPtr const& medusaLogitsDevice, - SizeType32 medusaHeads, SizeType32 logitsIndex, SizeType32 numLogits) -{ - for (SizeType32 hi = 0; hi < medusaHeads; ++hi) - { - TensorPtr logitsHead = ITensor::slice(medusaLogitsDevice, hi, 1); - logitsHead->squeeze(0); - medusaLogitsHeads[hi] = ITensor::slice(logitsHead, logitsIndex, numLogits); - } -} - -} // namespace - -void HandleGenerationLogits::operator()(DecoderInputBuffers& inputBuffers, RequestVector const& generationRequests, - tr::ITensor::SharedPtr const& logits, tr::SizeType32 logitsIndex, tr::ModelConfig const& modelConfig, - tr::BufferManager const& manager, OptionalRef genRuntimeBuffers, - OptionalRef medusaBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(HandleGenerationLogits); - - auto& decoderRequests = inputBuffers.decoderRequests; - decoderRequests.reserve(decoderRequests.size() + generationRequests.size()); - auto& allDecoderLogits = inputBuffers.decoderLogits; - allDecoderLogits.reserve(allDecoderLogits.size() + generationRequests.size()); - - for (auto const& llmReq : generationRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const seqSlot = llmReq->mSeqSlot.value(); - - auto const draftLength = llmReq->getNumDraftTokens(); - auto const numLogits = draftLength + reqBeamWidth; - - TLLM_CHECK(draftLength == 0 || reqBeamWidth == 1); - - TLLM_LOG_DEBUG("logitsIndex: %d", logitsIndex); - TLLM_LOG_DEBUG("draftLength: %d", draftLength); - TLLM_LOG_DEBUG("reqBeamWidth: %d", reqBeamWidth); - - // genRuntimeBuffers.logits shape: [numGen*reqBeamWidth, vocabSize] - // logitsView shape: [numLogits, vocabSize] - TensorPtr logitsView = ITensor::slice(logits, logitsIndex, numLogits); - TLLM_CHECK_DEBUG_WITH_INFO(tru::tensorHasInvalid(*logitsView, manager, "logits") == false, - "Found invalid number (NaN or Inf) in logits"); - - TLLM_CHECK(llmReq->isGenerationInProgressState()); - TensorPtr decoderLogits; - if (reqBeamWidth > 1) - { - decoderLogits = logitsView; - decoderLogits->unsqueeze(0); - } - else - { - decoderLogits = logitsView; - decoderLogits->unsqueeze(1); - } - decoderRequests.push_back(llmReq); - allDecoderLogits.emplace_back(std::move(decoderLogits)); - - if (llmReq->getReturnGenerationLogits()) - { - TLLM_CHECK_WITH_INFO(modelConfig.getSpeculativeDecodingMode().isNone() - || modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal(), - "Only speculative decoding with external draft tokens supports returning generation logits"); - - // Push into fragments vector - llmReq->addGenerationLogitsFragment(logitsView); - TLLM_CHECK( - llmReq->getGenerationLogitsFragmentsSize() <= RuntimeBuffers::GenerationLogitsCache::kCACHE_LENGTH); - if (llmReq->isStreaming()) - { - copyStreamingGenerationLogits(manager, *llmReq); - } - // Copy back to host for every kCACHE_LENGTH steps to mitigate GPU memory pressure - else if (llmReq->getGenerationLogitsFragmentsSize() == RuntimeBuffers::GenerationLogitsCache::kCACHE_LENGTH) - { - TLLM_CHECK(genRuntimeBuffers); - auto constexpr beforeDecoder = true; - utils::copyGenerationLogits(genRuntimeBuffers->generationLogitsCache, manager, *llmReq, beforeDecoder); - } - } - if (modelConfig.getSpeculativeDecodingMode().hasDraftLogits()) - { - auto& medusaLogitsHeads = inputBuffers.predictedDraftLogits.at(seqSlot); - TLLM_CHECK(medusaBuffers); - setupMedusaLogits(medusaLogitsHeads, medusaBuffers->medusaLogitsDevice, - modelConfig.getSpeculativeDecodingModule().getMaxDraftPathLen(), logitsIndex, draftLength); - } - logitsIndex += numLogits; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp index 0fb8af1527ae..ccd56b121435 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheManager.cpp @@ -576,7 +576,7 @@ std::map BlockManager::calculateWindowSizeToShare( BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, CudaStreamPtr stream, SizeType32 maxSequenceLength, SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, - nvinfer1::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, + tensorrt_llm::DataType dtype, SizeType32 sinkBubbleLength, SizeType32 chunkSize, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, std::shared_ptr kvCacheConnectorManager, @@ -730,7 +730,7 @@ BlockManager::BlockManager(std::vector const& numKvHeadsPerLayer, Si "Maybe you tried changing either of them to an std::unordered_map?"); } -WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 windowSize, +WindowBlockManager::WindowBlockManager(tensorrt_llm::DataType dtype, SizeType32 windowSize, std::vector const& managedLayers, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, bool isSWA, SizeType32 blocksInPrimaryPool, SizeType32 blocksInSecondaryPool, SizeType32 maxNumSequences, std::shared_ptr stream, @@ -830,7 +830,7 @@ WindowBlockManager::WindowBlockManager(nvinfer1::DataType dtype, SizeType32 wind // to specify FP4 related parameters (scale dtypes, etc)? This can also be passed // in the constructor. constexpr SizeType32 kQuantBlockSizeNVFP4 = 16; - if (dtype == nvinfer1::DataType::kFP4) + if (dtype == tensorrt_llm::DataType::kFP4) { createBlockScalePools(kQuantBlockSizeNVFP4); } @@ -1082,7 +1082,7 @@ void BlockManager::allocatePools(bool useUvm) void WindowBlockManager::allocatePools(bool useUvm) { - constexpr nvinfer1::DataType kScaleDtypeNVFP4 = nvinfer1::DataType::kFP8; + constexpr tensorrt_llm::DataType kScaleDtypeNVFP4 = tensorrt_llm::DataType::kFP8; // Allocate a memory pool backing the blocks for each numKvHeads // TODO(oargov): allocate pools in a single buffer and split it, to avoid fragmentation @@ -1091,21 +1091,21 @@ void WindowBlockManager::allocatePools(bool useUvm) auto blockSize = pool.blockSize; auto poolDtype = pool.containsBlockScales ? kScaleDtypeNVFP4 : mDataType; #ifdef ENABLE_FP4 - auto const poolIsFP4 = poolDtype == nvinfer1::DataType::kFP4; + auto const poolIsFP4 = poolDtype == tensorrt_llm::DataType::kFP4; #else auto const poolIsFP4 = false; #endif if (poolIsFP4) { - poolDtype = nvinfer1::DataType::kINT8; + poolDtype = tensorrt_llm::DataType::kINT8; } if (pool.containsIndexerKCache) { - poolDtype = nvinfer1::DataType::kUINT8; + poolDtype = tensorrt_llm::DataType::kUINT8; } - nvinfer1::Dims cacheShape = isRecurrentState() + tensorrt_llm::Dims cacheShape = isRecurrentState() ? ITensor::makeShape({pool.numLayers, mNumPrimaryBlocks, mKVFactor, blockSize}) : ITensor::makeShape({mNumPrimaryBlocks, pool.numLayers, mKVFactor, blockSize}); pool.layerFirstLayout = isRecurrentState(); @@ -1121,7 +1121,7 @@ void WindowBlockManager::allocatePools(bool useUvm) pool.primaryPtr = mBufferManager.gpuSync(cacheShape, poolDtype); if (mNumSecondaryBlocks > 0) { - nvinfer1::Dims cacheShapeOffload = isRecurrentState() + tensorrt_llm::Dims cacheShapeOffload = isRecurrentState() ? ITensor::makeShape({pool.numLayers, mNumSecondaryBlocks, mKVFactor, blockSize}) : ITensor::makeShape({mNumSecondaryBlocks, pool.numLayers, mKVFactor, blockSize}); TLLM_LOG_DEBUG("[%s] Allocating secondary pool with %d blocks for %d layers with %d kv heads", @@ -1292,7 +1292,7 @@ BlockPtr WindowBlockManager::getFreeBlock(GenerationRequest& sequence, executor: return block; } -void WindowBlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, +void WindowBlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId) const { auto constexpr kIdx = 0; @@ -1330,7 +1330,7 @@ void WindowBlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, nvinfer1::Dims } } -void BlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, nvinfer1::Dims const& offsetsShape, SizeType32 beamIdx, +void BlockManager::setOffsets(tk::KVCacheIndex* offsetsPtr, tensorrt_llm::Dims const& offsetsShape, SizeType32 beamIdx, SizeType32 blockIdx, KVCacheBlock::IdType blockId, SizeType32 windowSize) const { mWindowBlockManagers.at(windowSize).setOffsets(offsetsPtr, offsetsShape, beamIdx, blockIdx, blockId); @@ -3168,7 +3168,7 @@ void WindowBlockManager::schedulingReleaseBlocks(RequestIdType requestId) KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, bool enablePartialReuse, bool copyOnPartialReuse, bool enableIndexerKCache, SizeType32 indexerKCacheQuantBlockSize, SizeType32 indexerKCacheIndexHeadDim, @@ -3185,7 +3185,7 @@ KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, Size KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, int64_t stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, @@ -3204,7 +3204,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, @@ -3246,7 +3246,7 @@ KVCacheManager::KVCacheManager(std::vector const& numKvHeadsPerLayer KVCacheManager::KVCacheManager(SizeType32 numLayers, SizeType32 numKvHeads, SizeType32 sizePerHead, SizeType32 tokensPerBlock, BlocksPerWindow const& blocksPerWindow, SizeType32 maxNumSequences, - SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, nvinfer1::DataType dtype, + SizeType32 maxBeamWidth, std::vector const& maxAttentionWindowVec, tensorrt_llm::DataType dtype, SizeType32 sinkTokenLength, CudaStreamPtr stream, runtime::SizeType32 maxSequenceLength, SizeType32 chunkSize, bool enableBlockReuse, CacheType cacheType, std::optional secondaryOffloadMinPriority, std::shared_ptr eventManager, bool enablePartialReuse, bool copyOnPartialReuse, @@ -3279,7 +3279,7 @@ void KVCacheManager::allocatePools(bool useUvm) // a future per-window override map can mix precisions inside a single manager. auto const poolDataType = primaryPool->getDataType(); #ifdef ENABLE_FP4 - auto const isFp4 = poolDataType == nvinfer1::DataType::kFP4; + auto const isFp4 = poolDataType == tensorrt_llm::DataType::kFP4; #else auto const isFp4 = false; #endif @@ -4200,7 +4200,7 @@ std::map computeWindowSizeShares( } // namespace BlocksPerWindow BaseKVCacheManager::calculateMaxNumBlocks(executor::KvCacheConfig const& config, - nvinfer1::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, + tensorrt_llm::DataType dtype, std::vector const& numKvHeadsPerLayer, SizeType32 sizePerHead, SizeType32 tokensPerBlock, WorldConfig const& worldConfig, std::map> const& windowSizeToLayers, uint64_t allottedPrimaryMemBytes, uint64_t allottedSecondaryMemBytes, size_t extraCostMemory, SizeType32 kvFactor, SizeType32 maxBatchSize, diff --git a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp index c28d1e476137..a9d14261e724 100644 --- a/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/kvCacheTransferManager.cpp @@ -163,8 +163,8 @@ void KVCacheTransferManager::copyBlock(BlockPtr const& src, BlockPtr const& dst, // If no partial tokens or if the dataType is not supported for partial copy, copy entire block. // Note that nvfp4 kv cache SFs use an interleaved layout, so we need to copy the entire block. - if (numTokensToCopy <= 0 || srcPtr->getDataType() == nvinfer1::DataType::kINT4 - || srcPtr->getDataType() == nvinfer1::DataType::kFP4 || containsBlockScales) + if (numTokensToCopy <= 0 || srcPtr->getDataType() == tensorrt_llm::DataType::kINT4 + || srcPtr->getDataType() == tensorrt_llm::DataType::kFP4 || containsBlockScales) { // For partial copy not implemented with these data types, // just do a full copy. @@ -461,8 +461,8 @@ std::size_t KVCacheTransferManager::computeBlockTransferBytes( // Mirror the logic in copyBlock: a partial copy only happens when numTokensToCopy > 0, // the data type supports it (not kINT4/kFP4), not block scales, and numTokensToCopy < tokensPerBlock. - bool const isPartialCopy = numTokensToCopy > 0 && dataType != nvinfer1::DataType::kINT4 - && dataType != nvinfer1::DataType::kFP4 && !pool.containsBlockScales + bool const isPartialCopy = numTokensToCopy > 0 && dataType != tensorrt_llm::DataType::kINT4 + && dataType != tensorrt_llm::DataType::kFP4 && !pool.containsBlockScales && numTokensToCopy < pool.tokensPerBlock; if (isPartialCopy) diff --git a/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp b/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp deleted file mode 100644 index 95b324f0f2ec..000000000000 --- a/cpp/tensorrt_llm/batch_manager/logitsPostProcessor.cpp +++ /dev/null @@ -1,88 +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/logitsPostProcessor.h" - -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tr = tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -bool LogitsPostProcessor::operator()(DecoderInputBuffers& inputBuffers, bool replicateLogitsPostProcessor, - tr::WorldConfig const& worldConfig, CudaStreamPtr const& stream, - std::optional const& logitsPostProcessorBatched) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(LogitsPostProcessor); - - // Arguments for batched processor - std::vector reqIdsVec; - std::vector logitsVec; - std::vector> beamTokensVec; - std::vector> clientIdsVec; - - bool logitsPostProcessorIsApplied = false; - for (size_t batchIdx = 0; batchIdx < inputBuffers.decoderRequests.size(); ++batchIdx) - { - auto const& llmReq = inputBuffers.decoderRequests.at(batchIdx); - auto& logits = inputBuffers.decoderLogits.at(batchIdx); - - // Invoke non-batched processor or collect arguments for batched processor - if (llmReq->mLogitsPostProcessor) - { - logitsPostProcessorIsApplied = true; - if (replicateLogitsPostProcessor || worldConfig.isFirstTensorParallelRank()) - { - (*llmReq->mLogitsPostProcessor)( - llmReq->mRequestId, logits, llmReq->getTokens(), stream, llmReq->mClientId); - } - } - else if (llmReq->mApplyLogitsPostProcessorBatched) - { - reqIdsVec.push_back(llmReq->mRequestId); - logitsVec.push_back(logits); - beamTokensVec.emplace_back(llmReq->getTokens()); - clientIdsVec.push_back(llmReq->mClientId); - } - } - - // Invoke batched processor - if (!reqIdsVec.empty()) - { - logitsPostProcessorIsApplied = true; - if (replicateLogitsPostProcessor || worldConfig.isFirstTensorParallelRank()) - { - (*logitsPostProcessorBatched)(reqIdsVec, logitsVec, beamTokensVec, stream, clientIdsVec); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - - return logitsPostProcessorIsApplied; -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/loraBuffers.cpp b/cpp/tensorrt_llm/batch_manager/loraBuffers.cpp deleted file mode 100644 index b67b72f6c49a..000000000000 --- a/cpp/tensorrt_llm/batch_manager/loraBuffers.cpp +++ /dev/null @@ -1,109 +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 "loraBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/runtime/loraUtils.h" - -namespace tensorrt_llm::batch_manager -{ - -LoraBuffers::LoraBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::TllmRuntime const& tllmRuntime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) -{ - auto const localNbLayers - = modelConfig.getNbAttentionLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - auto const firstLayerId = worldConfig.getPipelineParallelRank() * localNbLayers; - - auto nbModelConfigs = static_cast(modelConfig.getLoraModules().size()); - - // there are 3 pointers: LoRA A, LoRA B, and a DoRA magnitude (null if not DoRA) - auto loraWeightsPtrsShape - = runtime::ITensor::makeShape({nbModelConfigs, localNbLayers, maxBatchSize * maxBeamWidth, 3}); - auto loraAdapterSizesShape - = runtime::ITensor::makeShape({nbModelConfigs, localNbLayers, maxBatchSize * maxBeamWidth}); - - auto firstModuleName = std::string(modelConfig.getLoraModules().front().name()); - auto ptrsFieldName = firstModuleName + "_lora_weights_pointers_" + std::to_string(firstLayerId); - auto rankFieldName = firstModuleName + "_lora_ranks_" + std::to_string(firstLayerId); - auto weightsPtrDtype = tllmRuntime.getEngine().getTensorDataType(ptrsFieldName.c_str()); - auto ranksDtype = tllmRuntime.getEngine().getTensorDataType(rankFieldName.c_str()); - - mLoraManager.create(modelConfig); - - mLoraWeightsPointersHost = runtime::BufferManager::pinned(loraWeightsPtrsShape, weightsPtrDtype); - mLoraAdapterSizesHost = runtime::BufferManager::pinned(loraAdapterSizesShape, ranksDtype); -} - -void LoraBuffers::fill(RequestVector const& contextRequests, RequestVector const& genRequests, - PeftTable const& peftTable, runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) -{ - manager.setZero(*mLoraWeightsPointersHost); - manager.setZero(*mLoraAdapterSizesHost); - - SizeType32 batchIdx{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - auto const optReqLoraWeights = llmReq->getLoraWeights(); - auto const optReqLoraConfig = llmReq->getLoraConfig(); - - auto const isContextRequest = llmReq->isContextInitState(); - auto const beamWidth = isContextRequest ? 1 : llmReq->mSamplingConfig.beamWidth; - auto const peftIt = peftTable.find(llmReq->mRequestId); - if (peftIt != peftTable.end()) - { - auto const& peftValues = peftIt->second; - if (!peftValues.empty()) - { - mLoraManager.fillInputTensors(mLoraWeightsPointersHost, mLoraAdapterSizesHost, peftIt->second, - batchIdx, beamWidth, modelConfig, worldConfig); - } - } - ++batchIdx; - } - } -} - -void LoraBuffers::validate(std::optional const& optTaskId, - std::optional const& optReqLoraWeights, std::optional const& optReqLoraConfig, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) -{ - runtime::lora::loraValidateRequestTensors(optTaskId, optReqLoraWeights, optReqLoraConfig, modelConfig, worldConfig); -} - -void LoraBuffers::insertInputTensors(TensorMap& inputTensors, TensorPtr weightsPtrs, TensorPtr adapterSizes, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) const -{ - mLoraManager.insertInputTensors(inputTensors, weightsPtrs, adapterSizes, modelConfig, worldConfig); -} - -void LoraBuffers::reshape(SizeType32 numSequences) -{ - auto weightsPtrsShape = mLoraWeightsPointersHost->getShape(); - weightsPtrsShape.d[2] = numSequences; - mLoraWeightsPointersHost->reshape(weightsPtrsShape); - - auto adapterSizesShape = mLoraAdapterSizesHost->getShape(); - adapterSizesShape.d[2] = numSequences; - mLoraAdapterSizesHost->reshape(adapterSizesShape); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/loraBuffers.h b/cpp/tensorrt_llm/batch_manager/loraBuffers.h deleted file mode 100644 index 3ba68995518f..000000000000 --- a/cpp/tensorrt_llm/batch_manager/loraBuffers.h +++ /dev/null @@ -1,61 +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/common.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/loraManager.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -namespace tensorrt_llm::batch_manager -{ - -class LoraBuffers -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using PeftTable = runtime::LoraManager::PeftTable; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - TensorPtr mLoraWeightsPointersHost; - TensorPtr mLoraAdapterSizesHost; - - runtime::LoraManager mLoraManager; - - LoraBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::TllmRuntime const& tllmRuntime, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - - static void validate(std::optional const& optTaskId, - std::optional const& optReqLoraWeights, std::optional const& optReqLoraConfig, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig); - - void fill(RequestVector const& contextRequests, RequestVector const& genRequests, PeftTable const& peftTable, - runtime::BufferManager const& manager, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig); - - void insertInputTensors(TensorMap& inputTensors, TensorPtr weightsPtrs, TensorPtr adapterSizes, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) const; - - void reshape(SizeType32 numSequences); -}; -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp b/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp deleted file mode 100644 index 3e494a6383ec..000000000000 --- a/cpp/tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.cpp +++ /dev/null @@ -1,198 +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/makeDecodingBatchInputOutput.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/iGptDecoderBatched.h" - -namespace tr = tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ -using SizeType32 = MakeDecodingBatchInputOutput::SizeType32; -using TensorPtr = MakeDecodingBatchInputOutput::TensorPtr; - -void MakeDecodingBatchInputOutput::createDecoderBatchInputs(DecoderInputBuffers& inputBuffers, - std::vector const& activeSlots, runtime::decoder::DecoderState const& decoderState) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& numDecodingEngineTokens = decoderState.getNumDecodingEngineTokens(); - auto const& maxDecodingEngineTokens = decoderState.getMaxDecodingEngineTokens(); - auto const& maxDecodingDecoderTokens = decoderState.getMaxDecodingDecoderTokens(); - auto const maxDecoderSteps = common::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 = common::ceilDiv(numDecodingEngineTokens.at(slot), maxDecodingDecoderTokens); - maxActiveDecoderSteps = std::max(maxActiveDecoderSteps, numDecoderSteps); - for (SizeType32 step = 0; step < numDecoderSteps; ++step) - { - auto batchSlotsRange = tr::BufferRange(*batchSlots.at(step)); - batchSlotsRange[batchSizes[step]] = slot; - batchSizes[step]++; - auto logitsSlice = tr::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; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ - -std::pair, std::vector> getActiveSlots(RequestVector const& decoderRequests) -{ - std::vector activeSlots; - std::vector generationSteps; - for (auto const& llmReq : decoderRequests) - { - activeSlots.push_back(llmReq->mSeqSlot.value()); - generationSteps.push_back(llmReq->getDecodingIter()); - } - - return {activeSlots, generationSteps}; -} - -//! @brief Sets inputs for explicit draft tokens. -void setExplicitDraftTokensInputs(tr::DecodingInput& dInput, RuntimeBuffers const& fusedRuntimeBuffers) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(fusedRuntimeBuffers.mExplicitDraftTokensBuffers); - auto const& explicitDraftTokensInputs = fusedRuntimeBuffers.mExplicitDraftTokensBuffers->engineOutputs; - auto const& explicitDraftTokensLastInputs = fusedRuntimeBuffers.mExplicitDraftTokensBuffers->engineInputs; - - dInput.explicitDraftTokensInputs = tr::DecodingInput::ExplicitDraftTokensInputs(); - dInput.explicitDraftTokensInputs->nextDraftTokens = explicitDraftTokensInputs.nextDraftTokens; - dInput.explicitDraftTokensInputs->nextFlatTokens = explicitDraftTokensInputs.nextFlatTokens; - dInput.explicitDraftTokensInputs->nextDraftIndices = explicitDraftTokensInputs.nextDraftIndices; - dInput.explicitDraftTokensInputs->nextDraftProbs = explicitDraftTokensInputs.nextDraftProbs; - dInput.explicitDraftTokensInputs->lastDraftTokens = explicitDraftTokensLastInputs.draftTokens; - dInput.explicitDraftTokensInputs->lastDraftIndices = explicitDraftTokensLastInputs.draftIndices; - dInput.explicitDraftTokensInputs->lastPositionIdsBase = explicitDraftTokensLastInputs.positionIdsBase; - dInput.explicitDraftTokensInputs->masks = explicitDraftTokensInputs.masks; - dInput.explicitDraftTokensInputs->packedPositionIds = explicitDraftTokensInputs.packedPositionIds; - dInput.explicitDraftTokensInputs->bestPathLengths = explicitDraftTokensInputs.bestPathLengths; - dInput.explicitDraftTokensInputs->bestPathIndices = explicitDraftTokensInputs.bestPathIndices; - dInput.explicitDraftTokensInputs->nextGenerationLengths = explicitDraftTokensInputs.nextGenerationLengths; - dInput.explicitDraftTokensInputs->lastGenerationLengths = explicitDraftTokensLastInputs.generationLengths; - dInput.explicitDraftTokensInputs->maxGenLengthDevice = explicitDraftTokensInputs.maxGenToken; - // Slots in request order - dInput.explicitDraftTokensInputs->seqSlots = fusedRuntimeBuffers.seqSlots; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -//! @brief Sets inputs for eagle decoding. -void setEagleInputs(tr::DecodingInput& dInput, RuntimeBuffers const& fusedRuntimeBuffers) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(fusedRuntimeBuffers.mEagleBuffers); - auto const& eagleInputs = fusedRuntimeBuffers.mEagleBuffers->engineOutputs; - auto const& eagleLastInputs = fusedRuntimeBuffers.mEagleBuffers->engineInputs; - - dInput.eagleInputs = tr::DecodingInput::EagleInputs(); - dInput.eagleInputs->nextDraftTokens = eagleInputs.nextDraftTokens; - dInput.eagleInputs->nextDraftLens = eagleInputs.nextDraftLens; - dInput.eagleInputs->nextDraftPaths = eagleInputs.nextDraftPaths; - dInput.eagleInputs->lastDraftTokens = eagleLastInputs.draftTokens; - dInput.eagleInputs->lastDraftLens = eagleLastInputs.draftLens; - dInput.eagleInputs->lastDraftPaths = eagleLastInputs.draftPaths; - dInput.eagleInputs->acceptedTokens = eagleInputs.acceptedTokens; - dInput.eagleInputs->acceptedLens = eagleInputs.acceptedLens; - dInput.eagleInputs->acceptedPathIds = eagleInputs.acceptedPaths; - dInput.eagleInputs->chunkedContextNextTokens = eagleInputs.chunkedContextNextTokens; - // Slots in request order - dInput.eagleInputs->seqSlots = fusedRuntimeBuffers.seqSlots; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace - -void MakeDecodingBatchInputOutput::operator()(DecoderInputBuffers& inputBuffers, - runtime::decoder::DecoderState& decoderState, runtime::ModelConfig const& modelConfig, - OptionalRef fusedRuntimeBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto [activeSlots, generationSteps] = getActiveSlots(inputBuffers.decoderRequests); - - createDecoderBatchInputs(inputBuffers, activeSlots, decoderState); - - auto const maxBeamWidth = decoderState.getMaxBeamWidth(); - if (maxBeamWidth > 1) - { - // For Variable-Beam-Width-Search - decoderState.getJointDecodingInput().generationSteps = generationSteps; - } - - if (modelConfig.getSpeculativeDecodingMode().hasDraftLogits()) - { - decoderState.getJointDecodingInput().medusaInputs->medusaLogits = inputBuffers.predictedDraftLogits; - } - - if (modelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens()) - { - TLLM_CHECK(fusedRuntimeBuffers); - // requires mCtxGenFusion == true - setExplicitDraftTokensInputs(decoderState.getJointDecodingInput(), *fusedRuntimeBuffers); - } - else if (modelConfig.getSpeculativeDecodingMode().isEagle()) - { - TLLM_CHECK(fusedRuntimeBuffers); - // requires mCtxGenFusion == true - setEagleInputs(decoderState.getJointDecodingInput(), *fusedRuntimeBuffers); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp b/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp index eb40208739cf..32935e683b83 100644 --- a/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp +++ b/cpp/tensorrt_llm/batch_manager/medusaBuffers.cpp @@ -17,99 +17,10 @@ #include "tensorrt_llm/batch_manager/medusaBuffers.h" #include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/medusaModule.h" -#include "tensorrt_llm/runtime/utils/speculativeChoicesUtils.h" namespace tensorrt_llm::batch_manager { -MedusaBuffers::MedusaBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, runtime::TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK_WITH_INFO(maxBeamWidth == 1, "Medusa does not support beam search"); - - auto const& engine = runtime.getEngine(); - - auto const maxNumSequences = maxBatchSize; - - auto const medusaModule = std::dynamic_pointer_cast( - modelConfig.getSpeculativeDecodingModulePtr()); - - auto const medusaHeads = medusaModule->getMaxDraftPathLen(); - auto const maxPathLen = medusaModule->getMaxPathLen(); // medusaHeads + 1 - auto const maxMedusaTokens = medusaModule->getMaxDecodingDraftTokens(); - auto const maxDecodingTokens = medusaModule->getMaxDecodingTokens(); // maxMedusaTokens + 1 - auto const numPackedMasks = medusaModule->getNumPackedMasks(); - - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - if (worldConfig.isLastPipelineParallelRank()) - { - auto logitsType = engine.getTensorDataType("medusa_logits"); - medusaLogitsDevice = manager.gpu( - ITensor::makeShape({medusaHeads, maxBatchSize, maxDecodingTokens, vocabSizePadded}), logitsType); - } - - // Note: reserved for variable sequence length support. - medusaGenerationLengthsHost - = runtime::BufferManager::pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - // TODO: pack batch and tokensPerStep into one dim to support variable sequence length without padddings. - attentionPackedMaskHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxDecodingTokens, numPackedMasks}), nvinfer1::DataType::kINT32); - medusaPositionOffsetsHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxDecodingTokens}), nvinfer1::DataType::kINT32); - medusaTreeIdsHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxMedusaTokens}), nvinfer1::DataType::kINT32); - medusaPathsHost = runtime::BufferManager::pinned( - ITensor::makeShape({maxNumSequences, maxDecodingTokens, maxPathLen}), nvinfer1::DataType::kINT32); - - TensorPtr medusaPositionOffsetsHostSlice = ITensor::slice(medusaPositionOffsetsHost, 0, 1); - medusaPositionOffsetsHostSlice->squeeze(0); - TensorPtr medusaTreeIdsHostSlice = ITensor::slice(medusaTreeIdsHost, 0, 1); - medusaTreeIdsHostSlice->squeeze(0); - TensorPtr medusaPathsHostSlice = ITensor::slice(medusaPathsHost, 0, 1); - medusaPathsHostSlice->squeeze(0); - TensorPtr attentionPackedMaskHostSlice = ITensor::slice(attentionPackedMaskHost, 0, 1); - attentionPackedMaskHostSlice->squeeze(0); - - // Init buffers for 1 request - auto const& choices = decodingConfig.getMedusaChoices().value_or(medusaModule->getMedusaChoices()); - runtime::utils::initTensorsFromChoices(*medusaModule, choices, mTopKs, medusaGenerationLengthsHost, - medusaPositionOffsetsHostSlice, medusaTreeIdsHostSlice, medusaPathsHostSlice, attentionPackedMaskHostSlice); - - auto scatterToBatch = [maxBatchSize, &manager](TensorPtr& data) - { - auto srcSlice = ITensor::slice(data, 0, 1); - // Populate data from the 1st request to the other requests in the batch - for (SizeType32 bi = 1; bi < maxBatchSize; ++bi) - { - auto dstSlice = ITensor::slice(data, bi, 1); - manager.copy(*srcSlice, *dstSlice); - } - }; - - scatterToBatch(medusaPositionOffsetsHost); - scatterToBatch(medusaTreeIdsHost); - scatterToBatch(medusaPathsHost); - scatterToBatch(attentionPackedMaskHost); - - // Copy buffers to device - // 1st dimension of packed mask is num_total_generation_tokens now (packed without paddings). - attentionPackedMaskHost->reshape(ITensor::makeShape({maxNumSequences * maxDecodingTokens, numPackedMasks})); - attentionPackedMaskDevice = manager.copyFrom(*attentionPackedMaskHost, runtime::MemoryType::kGPU); - medusaGenerationLengthsDevice = manager.copyFrom(*medusaGenerationLengthsHost, runtime::MemoryType::kGPU); - medusaPositionOffsetsDevice = manager.copyFrom(*medusaPositionOffsetsHost, runtime::MemoryType::kGPU); - medusaTreeIdsDevice = manager.copyFrom(*medusaTreeIdsHost, runtime::MemoryType::kGPU); - medusaPathsDevice = manager.copyFrom(*medusaPathsHost, runtime::MemoryType::kGPU); - - // use speculative decoding buffer - medusaUseSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - runtime::bufferCast(*medusaUseSpecDecoding)[0] = 1; - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - void MedusaBuffers::reshape(SizeType32 /* numCtxSequences */, SizeType32 numGenSequences, SizeType32 tokensPerStep) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); diff --git a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp index 0bf9a989fd65..dffa83a292e3 100644 --- a/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/peftCacheManager.cpp @@ -30,7 +30,7 @@ #include "tensorrt_llm/runtime/workerPool.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -52,7 +52,7 @@ PeftTaskNotCachedException::PeftTaskNotCachedException(std::string const& msg) PeftTaskNotCachedException::~PeftTaskNotCachedException() noexcept = default; std::pair PeftCacheManager::getMaxNumSlots(PeftCacheManagerConfig const& config, - nvinfer1::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, runtime::BufferManager const& bufferManager) + tensorrt_llm::DataType dataType, uint64_t pageWidth, uint64_t max1dModSize, runtime::BufferManager const& bufferManager) { TLLM_LOG_DEBUG("max1dModeSize=%llu", max1dModSize); TLLM_LOG_DEBUG("pageWidth=%llu", pageWidth); diff --git a/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp b/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp deleted file mode 100644 index 1cf73a2c0d21..000000000000 --- a/cpp/tensorrt_llm/batch_manager/promptTuningBuffers.cpp +++ /dev/null @@ -1,323 +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/promptTuningBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/common/nvtxUtils.h" - -namespace tensorrt_llm::batch_manager -{ -using SizeType32 = tensorrt_llm::runtime::SizeType32; -using TensorPtr = runtime::ITensor::SharedPtr; - -PromptTuningBuffers::PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig) -{ - auto maxPromptEmbeddingTableSize = modelConfig.getMaxPromptEmbeddingTableSize(); - auto const hiddenSize = modelConfig.getHiddenSize() * worldConfig.getTensorParallelism(); - - // vocabSize and mMaxPromptVocabSize - mPromptTuningParams.vocabSize = manager.gpu(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - mMaxPromptVocabSize = maxPromptEmbeddingTableSize / maxBatchSize; - - auto promptVocabSizeHost - = runtime::BufferManager::pinned(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - auto promptVocabSizeHostData = runtime::bufferCast(*promptVocabSizeHost); - promptVocabSizeHostData[0] = mMaxPromptVocabSize; - manager.copy(*promptVocabSizeHost, *mPromptTuningParams.vocabSize); - - // embeddingTable - mPromptTuningParams.embeddingTable = manager.gpu( - runtime::ITensor::makeShape({maxPromptEmbeddingTableSize, hiddenSize}), modelConfig.getDataType()); - - // tasks - mPromptTuningParams.tasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); -} - -PromptTuningBuffers::PromptTuningBuffers(SizeType32 maxBatchSize, runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, bool promptTableOffloading) -{ - auto maxPromptEmbeddingTableSize = modelConfig.getMaxPromptEmbeddingTableSize(); - auto const hiddenSize = modelConfig.getHiddenSize() * worldConfig.getTensorParallelism(); - - // vocabSize and mMaxPromptVocabSize - mPromptTuningParams.vocabSize = manager.gpu(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - mMaxPromptVocabSize = maxPromptEmbeddingTableSize / maxBatchSize; - mPromptTableOffloading = promptTableOffloading; - - auto promptVocabSizeHost - = runtime::BufferManager::pinned(runtime::ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - auto promptVocabSizeHostData = runtime::bufferCast(*promptVocabSizeHost); - promptVocabSizeHostData[0] = mMaxPromptVocabSize; - manager.copy(*promptVocabSizeHost, *mPromptTuningParams.vocabSize); - - // embeddingTable - mPromptTuningParams.embeddingTable = manager.gpu( - runtime::ITensor::makeShape({maxPromptEmbeddingTableSize, hiddenSize}), modelConfig.getDataType()); - - // tasks - mPromptTuningParams.tasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); -} - -void PromptTuningBuffers::validate( - std::optional const& optReqPromptEmbeddingTable, std::optional const& optReqPromptVocabSize) -{ - // Need to copy request embeddingTable to promptEmbeddingTable - if (optReqPromptEmbeddingTable.has_value()) - { - - auto reqPromptEmbeddingTable = optReqPromptEmbeddingTable.value(); - auto reqPromptVocabSize = optReqPromptVocabSize.value(); - - if (reqPromptVocabSize > mMaxPromptVocabSize) - { - std::string errStr = "Prompt vocab size" + std::to_string(reqPromptVocabSize) - + " is larger than max prompt vocab size of " + std::to_string(mMaxPromptVocabSize) - + ". Max prompt vocab size is computed from max_prompt_embedding_table_size / max_batch_size. "; - TLLM_LOG_ERROR(errStr); - throw std::runtime_error(errStr); - } - else - { - // Check that type matches model weights - if (reqPromptEmbeddingTable->getDataType() != mPromptTuningParams.embeddingTable->getDataType()) - { - std::string errStr = "Request embedding table data type doesn't match model weight data type."; - TLLM_LOG_ERROR(errStr); - throw std::runtime_error(errStr); - } - - if (reqPromptEmbeddingTable->getShape().d[1] != reqPromptVocabSize) - { - std::string errStr - = "First dimension of request embedding table is expected to be equal to prompt vocab size"; - TLLM_LOG_ERROR(errStr); - throw std::runtime_error(errStr); - } - } - } -} - -void PromptTuningBuffers::fill(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::BufferManager const& manager, bool packed) -{ - NVTX3_SCOPED_RANGE_WITH_NAME(range, "PromptTuningBuffers::fill"); - - auto const numContextRequests = static_cast(contextRequests.size()); - - std::vector reqBeamWidths; - std::vector reqPromptLengths; - mPromptTuningParams.promptTuningEnabled.clear(); - - SizeType32 batchIdx{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - reqBeamWidths.push_back(llmReq->mSamplingConfig.beamWidth); - if (batchIdx < numContextRequests) - { - SizeType32 numContextTokens = 0; - auto const draftLength = llmReq->isLastContextChunk() ? llmReq->getNumDraftTokens() : 0; - auto const contextChunkSize = llmReq->getContextChunkSize(); - numContextTokens += contextChunkSize + draftLength; - reqPromptLengths.push_back(numContextTokens); - } - - std::optional optReqPromptEmbeddingTable = std::nullopt; - std::optional optReqPromptVocabSize = std::nullopt; - - if (mPromptTableOffloading) - { - optReqPromptEmbeddingTable = getChunkPtableBuffer(getChunkPtableCurrentIndex()); - optReqPromptVocabSize = getChunkPtableBufferSliceSize(getChunkPtableCurrentIndex(), batchIdx); - } - else - { - optReqPromptEmbeddingTable = llmReq->getPromptEmbeddingTable(); - optReqPromptVocabSize = llmReq->getPromptVocabSize(); - } - - mPromptTuningParams.promptTuningEnabled.push_back(optReqPromptEmbeddingTable.has_value()); - - // If context request & has embedding table, validate it - if (optReqPromptEmbeddingTable.has_value()) - { - // If a context request, validate prompt tensors and move to GPU - if (batchIdx < numContextRequests) - { - if (mPromptTableOffloading) - { - // Need to slice the ptable since we don't need the entire buffer - // The size depends on optReqPromptVocabSize which stores how many fake prompts are in the chunk - auto slicedPtable = runtime::ITensor::slice( - optReqPromptEmbeddingTable.value(), 0, optReqPromptVocabSize.value()); - slicedPtable->unsqueeze(0); - optReqPromptEmbeddingTable = std::move(slicedPtable); - } - else - { - // Move to GPU - llmReq->movePromptEmbeddingTableToGpu(manager); - optReqPromptEmbeddingTable = llmReq->getPromptEmbeddingTable(); - } - - // Validate the table, prompt_vocab_size - validate(optReqPromptEmbeddingTable, optReqPromptVocabSize); - } - - auto const reqPromptEmbeddingTable = optReqPromptEmbeddingTable.value(); - auto const reqPromptVocabSize = optReqPromptVocabSize.value(); - - // TODO: Use invokeCopyBatch to avoid multiple bs1 copies - // Copy into large prompt embedding table - TensorPtr reqPromptEmbeddingTableView = runtime::ITensor::view(reqPromptEmbeddingTable); - reqPromptEmbeddingTableView->squeeze(0); - auto const promptEmbeddingTableSlice = runtime::ITensor::slice( - mPromptTuningParams.embeddingTable, batchIdx * mMaxPromptVocabSize, reqPromptVocabSize); - manager.copy(*reqPromptEmbeddingTable, *promptEmbeddingTableSlice); - // TODO: src: 2007040 (llmReq->getPromptEmbeddingTable()) != dst: 1003520 (reqPromptVocabSize) - // (original shape passed from - // python == 196 * 5120, fp16) - // VILA mode 1 , 2 images in one request - } - ++batchIdx; - } - } - - auto const batchSize = batchIdx; - std::vector tasksHostVec(batchSize); - std::iota(tasksHostVec.begin(), tasksHostVec.end(), 0); - - // Create a tensor that wraps the vector and convert unique_ptr to shared_ptr - auto tasksHost = std::shared_ptr( - runtime::ITensor::wrap(tasksHostVec, runtime::ITensor::makeShape({batchSize})).release()); - - mPromptTuningParams.fillTasksTensor( - tasksHost, batchSize, numContextRequests, reqBeamWidths, reqPromptLengths, manager, packed); -} - -void PromptTuningBuffers::initializeChunkPtableBuffers(runtime::BufferManager const& manager, - runtime::ModelConfig const& modelConfig, SizeType32 contextChunkSize, std::shared_ptr const& llmReq) -{ - if (mChunkPtableInitialized) - { - return; - } - - std::array buffers; - std::vector> startPositions(2); - for (int i = 0; i < 2; i++) - { - startPositions[i].emplace_back(0); - auto memType = llmReq->getPromptEmbeddingTable().value()->getDataType(); - buffers[i] = manager.gpu(runtime::ITensor::makeShape({contextChunkSize, modelConfig.getHiddenSize()}), memType); - } - - mChunkPtableBuffers = std::move(buffers); - mChunkPtableBufferStartPositions = std::move(startPositions); - - mChunkPtableCurrentIndex = 0; - mChunkPtableInitialized = true; -} - -void PromptTuningBuffers::switchChunkPtableBuffer() -{ - mChunkPtableCurrentIndex = 1 - mChunkPtableCurrentIndex; - clearBufferStartPositions(mChunkPtableCurrentIndex); -} - -size_t PromptTuningBuffers::getChunkPtableCurrentIndex() -{ - return mChunkPtableCurrentIndex; -} - -TensorPtr& PromptTuningBuffers::getChunkPtableBuffer(size_t index) -{ - if (!mChunkPtableBuffers.has_value()) - { - TLLM_THROW("Chunk ptable buffers not initialized"); - } - if (!mChunkPtableBuffers.value()[index]) - { - TLLM_THROW("Chunk ptable buffer at index %zu is null", index); - } - return mChunkPtableBuffers.value()[index]; -} - -SizeType32 PromptTuningBuffers::getChunkPtableBufferSliceSize(size_t index, size_t batchIdx) -{ - if (!mChunkPtableBufferStartPositions.has_value()) - { - return 0; - } - - if (batchIdx + 1 >= mChunkPtableBufferStartPositions.value()[index].size()) - { - TLLM_THROW("Batch index %zu + 1 out of bounds for buffer %zu (size: %zu)", batchIdx, index, - mChunkPtableBufferStartPositions.value()[index].size()); - } - - return mChunkPtableBufferStartPositions.value()[index][batchIdx + 1] - - mChunkPtableBufferStartPositions.value()[index][batchIdx]; -} - -SizeType32 PromptTuningBuffers::getChunkPtableBufferStartPosition(size_t index, size_t batchIdx) -{ - if (!mChunkPtableBufferStartPositions.has_value()) - { - return 0; - } - - if (batchIdx >= mChunkPtableBufferStartPositions.value()[index].size()) - { - TLLM_THROW("Batch index %zu out of bounds for buffer %zu (size: %zu)", batchIdx, index, - mChunkPtableBufferStartPositions.value()[index].size()); - } - - // For first batch, return the value directly - if (batchIdx == 0) - { - return mChunkPtableBufferStartPositions.value()[index][0]; - } - - // For other batches, return difference from previous position - return mChunkPtableBufferStartPositions.value()[index][batchIdx] - - mChunkPtableBufferStartPositions.value()[index][batchIdx - 1]; -} - -void PromptTuningBuffers::updateBufferStartPosition(size_t index, SizeType32 numRows) -{ - if (!mChunkPtableBufferStartPositions.has_value()) - { - return; - } - auto& positions = mChunkPtableBufferStartPositions.value()[index]; - positions.push_back(positions.back() + numRows); -} - -void PromptTuningBuffers::clearBufferStartPositions(size_t index) -{ - if (mChunkPtableBufferStartPositions.has_value()) - { - mChunkPtableBufferStartPositions.value()[index].clear(); - mChunkPtableBufferStartPositions.value()[index].emplace_back(0); - } -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp index 1e52a04b9580..4696c97c97fd 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnCacheTransBuffer.cpp @@ -65,7 +65,7 @@ size_t RnnCacheTransBufferManager::computeTransferBufferSize( RnnCacheTransBufferManager::RnnCacheTransBufferManager( RnnStateManager* rnnStateManager, std::optional maxNumTokens) : BaseTransBufferManager(computeTransferBufferSize(rnnStateManager, maxNumTokens), - nvinfer1::DataType::kUINT8, // Use byte buffer for mixed dtypes + tensorrt_llm::DataType::kUINT8, // Use byte buffer for mixed dtypes maxNumTokens) , mRnnStateManager{rnnStateManager} { @@ -141,7 +141,7 @@ size_t RnnCacheTransBufferManager::computeTransferBufferSizeFromPool( RnnCacheTransBufferManager::RnnCacheTransBufferManager(kv_cache_manager::BaseKVCacheManager* kvCacheManager, executor::kv_cache::CacheState const& cacheState, std::optional maxNumTokens) : BaseTransBufferManager(computeTransferBufferSizeFromPool(kvCacheManager, cacheState, maxNumTokens), - nvinfer1::DataType::kUINT8, maxNumTokens) + tensorrt_llm::DataType::kUINT8, maxNumTokens) , mRnnStateManager{nullptr} { TLLM_CHECK(kvCacheManager != nullptr); diff --git a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.cpp b/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.cpp deleted file mode 100644 index 6fc7977ef8f1..000000000000 --- a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.cpp +++ /dev/null @@ -1,78 +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 "rnnStateBuffers.h" - -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -RnnStateBuffers::RnnStateBuffers(SizeType32 maxBatchSize, runtime::TllmRuntime const& runtime) -{ - auto const& manager = runtime.getBufferManager(); - - slotMappingHost = BufferManager::cpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - slotMappingDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); -} - -void RnnStateBuffers::reshape(SizeType32 numSequences) -{ - slotMappingHost->reshape(ITensor::makeShape({numSequences})); - slotMappingDevice->reshape(ITensor::makeShape({numSequences})); -} - -void RnnStateBuffers::fillSlotMappings( - RequestVector const& contextRequests, rnn_state_manager::RnnStateManager* rnnStateManager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(rnnStateBuffersFillSlotMappings); - - SizeType32 batchIdx{0}; - for (auto const& llmReq : contextRequests) - { - auto const seqSlot = llmReq->mSeqSlot.value(); - auto const reqBeamWidth = llmReq->mSamplingConfig.beamWidth; - rnnStateManager->fillSlotMapping(*slotMappingHost, batchIdx, seqSlot, reqBeamWidth); - ++batchIdx; - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RnnStateBuffers::copySlotMappingH2D(runtime::TllmRuntime const& runtime) -{ - auto const& manager = runtime.getBufferManager(); - manager.copy(*slotMappingHost, *slotMappingDevice); -} - -void RnnStateBuffers::getBuffers(TensorMap& inputBuffers) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(rnnStateBuffersGetBuffers); - - inputBuffers.insert_or_assign("slot_mapping", slotMappingDevice); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.h b/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.h deleted file mode 100644 index e25df47382a1..000000000000 --- a/cpp/tensorrt_llm/batch_manager/rnnStateBuffers.h +++ /dev/null @@ -1,58 +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/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ - -namespace rnn_state_manager -{ -class RnnStateManager; -} - -class RnnStateBuffers -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TensorPtr = runtime::ITensor::SharedPtr; - using TensorMap = runtime::StringPtrMap; - - // others should be in rnnStateManager, we only need slotMapping here. - TensorPtr slotMappingHost; // [batch_size] - TensorPtr slotMappingDevice; // [batch_size] - - RnnStateBuffers(SizeType32 maxBatchSize, runtime::TllmRuntime const& runtime); - - void reshape(SizeType32 numSequences); - - void fillSlotMappings(RequestVector const& contextRequests, rnn_state_manager::RnnStateManager* rnnStateManager); - - void copySlotMappingH2D(runtime::TllmRuntime const& runtime); - - void getBuffers(TensorMap& inputBuffers) const; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp b/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp index 7608079fb396..dd5bb8abe999 100644 --- a/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp +++ b/cpp/tensorrt_llm/batch_manager/rnnStateManager.cpp @@ -80,7 +80,7 @@ RnnStateManager::RnnStateManager(SizeType32 maxNumSequences, tensorrt_llm::runti {localNbLayers, mMaxNumSequences * mBeamSlotsPerSequence, convKernel - 1, rnnConvDimSize}); mDtype = dataType; - mSsmCacheDtype = nvinfer1::DataType::kFLOAT; + mSsmCacheDtype = tensorrt_llm::DataType::kFLOAT; // Store RNN model config for CacheTransceiver mDState = stateSize; @@ -117,7 +117,7 @@ RnnStateManager::RnnStateManager(SizeType32 maxNumSequences, tensorrt_llm::runti RnnStateManager::RnnStateManager(SizeType32 dState, SizeType32 dConv, SizeType32 numHeads, SizeType32 nGroups, SizeType32 headDim, SizeType32 maxBatchSize, WorldConfig const& worldConfig, int64_t stream, - nvinfer1::DataType dtype, nvinfer1::DataType ssmCacheDtype, std::vector const& ppLayers, + tensorrt_llm::DataType dtype, tensorrt_llm::DataType ssmCacheDtype, std::vector const& ppLayers, SizeType32 numLayers) : mMaxNumSequences(maxBatchSize) , mMaxBeamWidth{1} @@ -297,12 +297,12 @@ RnnStateManager::TensorPtr RnnStateManager::getSsmStates() const return pagedRnnStates; } -nvinfer1::DataType RnnStateManager::getConvStateDataType() const noexcept +tensorrt_llm::DataType RnnStateManager::getConvStateDataType() const noexcept { return mDtype; } -nvinfer1::DataType RnnStateManager::getSsmStateDataType() const noexcept +tensorrt_llm::DataType RnnStateManager::getSsmStateDataType() const noexcept { return mSsmCacheDtype; } diff --git a/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp b/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp deleted file mode 100644 index 691fb9c7efda..000000000000 --- a/cpp/tensorrt_llm/batch_manager/runtimeBuffers.cpp +++ /dev/null @@ -1,1029 +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/runtimeBuffers.h" - -#include "tensorrt_llm/batch_manager/encoderBuffers.h" -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/loraBuffers.h" -#include "tensorrt_llm/batch_manager/medusaBuffers.h" -#include "tensorrt_llm/batch_manager/promptTuningBuffers.h" -#include "tensorrt_llm/batch_manager/rnnStateBuffers.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/batch_manager/transformerBuffers.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/stlUtils.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/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" - -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; - -namespace tensorrt_llm::batch_manager -{ - -RuntimeBuffers::RuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, bool gatherGenerationLogits, std::optional maxNumTokens, - std::optional> const& additionalModelOutputs, - bool promptTableOffloadingParam) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - promptTableOffloading = promptTableOffloadingParam; - - create(maxBatchSize, maxBeamWidth, maxAttentionWindowVec, maxAttentionWindow, sinkTokenLen, runtime, modelConfig, - worldConfig, decodingConfig, gatherGenerationLogits, additionalModelOutputs); - - // pre-allocate - setMaxBufferSizes(maxBatchSize, maxBeamWidth, modelConfig, maxNumTokens); - reshape(runtime, modelConfig, worldConfig, gatherGenerationLogits); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -RuntimeBuffers::~RuntimeBuffers() = default; - -void RuntimeBuffers::create(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - executor::DecodingConfig const& decodingConfig, bool gatherGenerationLogits, - std::optional> const& additionalModelOutputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - auto const& engine = runtime.getEngine(); - - if (modelConfig.isTransformerBased()) - { - transformerBuffers = std::make_unique(maxBatchSize, maxBeamWidth, maxAttentionWindowVec, - maxAttentionWindow, sinkTokenLen, runtime, modelConfig, worldConfig); - } - if (modelConfig.isRnnBased()) - { - rnnStateBuffers = std::make_unique(maxBatchSize, runtime); - } - - auto constexpr nvTokenIdType = TRTDataType::value; - inputsIds = manager.emptyTensor(MemoryType::kGPU, nvTokenIdType); - - mropeRotaryCosSin = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - mropePositionDeltas = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - if (worldConfig.isLastPipelineParallelRank()) - { - auto const logitsType = engine.getTensorDataType(batch_manager::RuntimeBuffers::kLogitsTensorName); - logits = manager.emptyTensor(MemoryType::kGPU, logitsType); - } - - // TODO: check which tensors can be allocated as pinned for max size - requestTypes = manager.emptyTensor(MemoryType::kCPU, TRTDataType::value); - - contextLengthsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - contextLengthsDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - sequenceLengthsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - sequenceLengthsDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - lastTokenIdsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - lastTokenIdsDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - logitsIdsHost = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - - inputsIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - if (worldConfig.isPipelineParallel()) - { - hiddenStates = manager.emptyTensor(MemoryType::kGPU, modelConfig.getDataType()); - } - - auto const maxBatchSizeShape = ITensor::makeShape({maxBatchSize}); - seqSlots = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT32); - seqSlotsDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT32); - - cacheIndirDecoderIOBatchedCopySrcOffsets - = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT64); - cacheIndirDecoderIOBatchedCopyDstOffsets - = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT64); - cacheIndirDecoderIOBatchedCopySizes - = tensorrt_llm::runtime::BufferManager::pinnedPool(maxBatchSizeShape, nvinfer1::DataType::kINT64); - mCacheIndirDecoderIOBatchedCopySrcOffsetsSliceDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT64); - mCacheIndirDecoderIOBatchedCopyDstOffsetsSliceDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT64); - mCacheIndirDecoderIOBatchedCopyCopySizesDevice = manager.gpu(maxBatchSizeShape, nvinfer1::DataType::kINT64); - - // Pre-allocate buffer for saving generation logits for model w/o draft tokens - if (gatherGenerationLogits - && (modelConfig.getSpeculativeDecodingMode().isDraftTokensExternal() - || modelConfig.getSpeculativeDecodingMode().isNone()) - && worldConfig.isLastPipelineParallelRank()) - { - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - auto const logitsType = engine.getTensorDataType(batch_manager::RuntimeBuffers::kLogitsTensorName); - - generationLogitsCache.transposedLogits = manager.gpu( - ITensor::makeShape({maxBeamWidth, GenerationLogitsCache::kCACHE_LENGTH, vocabSizePadded}), logitsType); - generationLogitsCache.logits = manager.gpu( - ITensor::makeShape({GenerationLogitsCache::kCACHE_LENGTH, maxBatchSize * maxBeamWidth, vocabSizePadded}), - logitsType); - - generationLogitsCache.fragmentPointerDevice - = manager.gpu(ITensor::makeShape({GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64); - generationLogitsCache.fragmentPointerHost = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize, GenerationLogitsCache::kCACHE_LENGTH}), nvinfer1::DataType::kINT64); - } - - if (modelConfig.useCrossAttention()) - { - encoderBuffers = std::make_unique(); - encoderBuffers->create(maxBatchSize, modelConfig, runtime); - } - - if (modelConfig.usePromptTuning()) - { - promptTuningBuffers = std::make_unique( - maxBatchSize, manager, modelConfig, worldConfig, promptTableOffloading); - } - - if (modelConfig.useLoraPlugin()) - { - loraBuffers = std::make_unique(maxBatchSize, maxBeamWidth, runtime, modelConfig, worldConfig); - } - - if (modelConfig.getSpeculativeDecodingMode().isMedusa()) - { - mMedusaBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig, decodingConfig, runtime); - } - else if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - mLookaheadBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig, decodingConfig, runtime); - } - else if (modelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens()) - { - mExplicitDraftTokensBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig); - } - else if (modelConfig.getSpeculativeDecodingMode().isEagle()) - { - mEagleBuffers = std::make_unique( - maxBatchSize, maxBeamWidth, manager, modelConfig, worldConfig, decodingConfig); - } - - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings = manager.emptyTensor(MemoryType::kGPU, TRTDataType::value); - } - - for (auto const& output : additionalModelOutputs.value_or(std::vector{})) - { - auto const& engine = runtime.getEngine(); - auto const dataType = engine.getTensorDataType(output.name.c_str()); - mAdditionalOutputTensors.emplace(output.name, manager.emptyTensor(runtime::MemoryType::kGPU, dataType)); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::setMaxBufferSizes(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - runtime::ModelConfig const& modelConfig, std::optional maxNumRuntimeTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // `maxNumSequences` is reached when all requests are in generation - numContextRequests = 0; - numGenRequests = maxBatchSize; - numGenSequences = maxBatchSize * maxBeamWidth; - - auto const maxDraftTokens = modelConfig.getMaxDecodingDraftTokens(); - // Draft-Tokens and Beam-Search are mutually exclusive - numLogits = maxBatchSize * std::max(1 + maxDraftTokens, maxBeamWidth); - auto const maxNumModelTokens = modelConfig.getMaxNumTokens(); - auto const maxNumContextTokens = maxBatchSize * modelConfig.getMaxInputLen(); - auto const maxNumGenTokens = numLogits; - // For pre-allocation - numContextTokens = 0; // Set in `setBufferSizes` rather than here for `computeContextLogits` - numGenTokens - = maxNumRuntimeTokens.value_or(maxNumModelTokens.value_or(std::max(maxNumContextTokens, maxNumGenTokens))); - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->setMaxBufferSizes(maxBatchSize, modelConfig); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::setBufferSizes(RequestVector const& contextRequests, RequestVector const& genRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersSetBufferSizes); - - // set context sizes - numContextRequests = static_cast(contextRequests.size()); - auto numContextLogits = numContextRequests; - numContextTokens = 0; - maxContextLength = 0; - for (auto const& llmReq : contextRequests) - { - auto const draftLength = llmReq->isLastContextChunk() ? llmReq->getNumDraftTokens() : 0; - numContextLogits += draftLength; - - auto const contextChunkSize = llmReq->getContextChunkSize(); - numContextTokens += contextChunkSize + draftLength; - if (maxContextLength < llmReq->mPromptLen) - { - maxContextLength = llmReq->mPromptLen; - } - } - - // set generation sizes - numGenRequests = static_cast(genRequests.size()); - numGenSequences = 0; - numGenTokens = 0; - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - numGenSequences += reqBeamWidth; - auto const draftLen = llmReq->getNumDraftTokens(); - numGenTokens += draftLen + reqBeamWidth; - } - - numLogits = numContextLogits + numGenTokens; - - if (encoderBuffers) - { - encoderBuffers->setBufferSizes(contextRequests, genRequests); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::reshape(TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - bool gatherGenerationLogits) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersReshape); - - if (worldConfig.isLastPipelineParallelRank()) - { - auto const vocabSizePadded = modelConfig.getVocabSizePadded(worldConfig.getSize()); - - if (modelConfig.computeContextLogits() && (numContextRequests > 0)) - { - // Only when need to return context logits, and there are new requests will execute context phase, - // logits buffer need to be re-allocated with size of [numContextTokens + numGenSequences, vocabSizePadded] - auto const& engine = runtime.getEngine(); - auto const& manager = runtime.getBufferManager(); - auto const logitsType = engine.getTensorDataType(kLogitsTensorName); - logits = manager.gpu(ITensor::makeShape({numContextTokens + numGenSequences, vocabSizePadded}), logitsType); - } - else if (gatherGenerationLogits && modelConfig.getSpeculativeDecodingMode().isNone()) - { - // If need to return generation logits, re-point the logit buffer to avoid overwrite, - // so we could write back GenerationLogitsCache::kCACHE_LENGTH steps' logits together - // logits shape: [1, maxBatchSize * maxBeamWidth, vocabSizePadded] - // which is large enough to cover both numContextRequests and numGenSequences - logits = ITensor::slice(generationLogitsCache.logits, generationLogitsCache.offset, 1); - generationLogitsCache.offset = (generationLogitsCache.offset + 1) % GenerationLogitsCache::kCACHE_LENGTH; - logits->squeeze(0); - } - else - { - logits->reshape(ITensor::makeShape({numLogits, vocabSizePadded})); - } - } - - auto const numSequences = getNumSequences(); - auto const numSequencesShape = ITensor::makeShape({numSequences}); - requestTypes->reshape(numSequencesShape); - contextLengthsHost->reshape(numSequencesShape); - contextLengthsDevice->reshape(numSequencesShape); - sequenceLengthsHost->reshape(numSequencesShape); - sequenceLengthsDevice->reshape(numSequencesShape); - - auto const numLogitsShape = ITensor::makeShape({numLogits}); - lastTokenIdsHost->reshape(numLogitsShape); - lastTokenIdsDevice->reshape(numLogitsShape); - logitsIdsHost->reshape(numLogitsShape); - - if (transformerBuffers) - { - transformerBuffers->reshape(numSequences, numContextTokens + numGenTokens); - } - - if (rnnStateBuffers) - { - rnnStateBuffers->reshape(numSequences); - } - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->reshape(); - } - - if (modelConfig.useLoraPlugin()) - { - loraBuffers->reshape(numSequences); - } - - if (mMedusaBuffers) - { - mMedusaBuffers->reshape( - numContextRequests, numGenRequests, modelConfig.getSpeculativeDecodingModulePtr()->getMaxDecodingTokens()); - } - - if (mLookaheadBuffers && modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - mLookaheadBuffers->reshape( - numContextRequests, numGenRequests, modelConfig.getSpeculativeDecodingModulePtr()->getMaxDecodingTokens()); - } - - if (mExplicitDraftTokensBuffers) - { - mExplicitDraftTokensBuffers->reshape(numContextRequests, numGenRequests, modelConfig); - } - - if (mEagleBuffers) - { - mEagleBuffers->reshape(numContextRequests, numGenRequests, modelConfig); - } - - auto const numRequests = getNumRequests(); - auto const numRequestsShape = ITensor::makeShape({numRequests}); - seqSlots->reshape(numRequestsShape); - seqSlotsDevice->reshape(numRequestsShape); - - auto const numTokens = getNumTokens(); - inputsIds->reshape(ITensor::makeShape({numTokens})); - - if (modelConfig.useMrope()) - { - auto const mropeRotaryCosSinSize = modelConfig.getMaxPositionEmbeddings() * modelConfig.getRotaryEmbeddingDim(); - mropeRotaryCosSin->reshape(ITensor::makeShape({numSequences, mropeRotaryCosSinSize})); - mropePositionDeltas->reshape(ITensor::makeShape({numSequences, 1})); - } - - if (worldConfig.isPipelineParallel()) - { - auto const hiddenSize = (!modelConfig.getPpReduceScatter() || worldConfig.isFirstPipelineParallelRank()) - ? modelConfig.getHiddenSize() * worldConfig.getTensorParallelism() - : modelConfig.getHiddenSize(); - - auto const hiddenStatesShape = ITensor::makeShape({numTokens, hiddenSize}); - hiddenStates->reshape(hiddenStatesShape); - } - - if (modelConfig.useLanguageAdapter()) - { - languageAdapterRoutings->reshape(ITensor::makeShape({numTokens, 1})); - } - - for (auto const& outputTensor : mAdditionalOutputTensors) - { - auto const& [name, tensor] = outputTensor; - auto const& engine = runtime.getEngine(); - auto shape = engine.getTensorShape(name.c_str()); - TLLM_CHECK_WITH_INFO( - shape.d[0] == -1, "First dimension of additional output tensor '%s' must be dynamic", name.c_str()); - shape.d[0] = numTokens; - tensor->reshape(shape); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::prepareBuffersForCudaGraph(SizeType32 maxSequenceLength) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(prepareBuffersForCudaGraph); - - TLLM_CHECK(numContextRequests == 0); - - if (transformerBuffers) - { - // Set pastKeyValueLength for graph capturing. This way we will capture graph with - // maxKvCacheLengthRounded rounded to the next kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE. - // MMHA will launch excessive amount of blocks and some of them will exit early during the actual launch. - // We can reuse the same graph for the next kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE iterations. - - // make sure the size does not overflow the max allowed pastKvCacheLength - auto const pastKvCacheLength = std::min(maxSequenceLength - 1, maxKvCacheLengthRounded); - - auto* pastKeyValueLengthsPtr = bufferCast(*transformerBuffers->pastKeyValueLengths); - std::fill_n(pastKeyValueLengthsPtr, getNumSequences(), pastKvCacheLength); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::setFromInputs(RequestVector const& contextRequests, RequestVector const& genRequests, - SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, runtime::decoder::DecoderState const& decoderState, - kv_cache_manager::BaseKVCacheManager* kvCacheManagerPtr, - kv_cache_manager::BaseKVCacheManager* crossKvCacheManagerPtr, - rnn_state_manager::RnnStateManager* rnnStateManagerPtr, PeftTable const& peftTable, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, bool trtOverlap, OptionalRef newOutputTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersSetFromInputs); - - auto const& manager = runtime.getBufferManager(); - auto const& stream = runtime.getStream(); - - // Fill requestTypes - { - auto* hostRequestTypes = bufferCast(*requestTypes); - std::fill_n(hostRequestTypes, numContextRequests, runtime::RequestType::kCONTEXT); - std::fill_n(hostRequestTypes + numContextRequests, numGenSequences, runtime::RequestType::kGENERATION); - } - - SizeType32 totalInputSize = 0; - std::vector inputHost; - std::vector positionIdsHost; - std::vector positionIdsHostRow2; - std::vector mropePositionDeltasHost; - std::vector languageAdapterRoutingsHost; - - auto* contextLengthsHostPtr = bufferCast(*contextLengthsHost); - auto* sequenceLengthsHostPtr = bufferCast(*sequenceLengthsHost); - auto* pastKeyValueLengthsPtr - = transformerBuffers ? bufferCast(*transformerBuffers->pastKeyValueLengths) : nullptr; - SizeType32 totalNumLogits{0}; - auto* logitsIdsHostPtr = bufferCast(*logitsIdsHost); - bool const isChatGlm = modelConfig.getModelVariant() == ModelConfig::ModelVariant::kChatGlm; - bool const isGlm = modelConfig.getModelVariant() == ModelConfig::ModelVariant::kGlm; - auto const mropeRotaryCosSinSize = modelConfig.getMaxPositionEmbeddings() * modelConfig.getRotaryEmbeddingDim(); - - { - NVTX3_SCOPED_RANGE(seqSlotsLoop); - auto* seqSlotIndices = bufferCast(*seqSlots); - - SizeType32 batchIdx{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - // Get position of the current sequence in the decoder - auto const seqSlot = llmReq->mSeqSlot.value(); - seqSlotIndices[batchIdx] = seqSlot; - ++batchIdx; - } - } - - TLLM_CHECK(seqSlots->getSize() == static_cast(batchIdx)); - manager.copy(*seqSlots, *seqSlotsDevice); - } - - // context preparation loop - if (!contextRequests.empty()) - { - NVTX3_SCOPED_RANGE(contextPrepareLoop); - numContextLogits.resize(contextRequests.size()); - - SizeType32 batchIdx{0}; - for (auto const& llmReq : contextRequests) - { - TLLM_CHECK_WITH_INFO(llmReq->isContextInitState() || llmReq->isDisaggGenerationTransmissionComplete(), - "The request should be in context phase or disaggregated generation tranmissionComplete phase."); - TLLM_CHECK_WITH_INFO( - llmReq->getMaxNumGeneratedTokens() == 0, "Context request should not have generated tokens."); - - auto const& reqTokens = llmReq->getTokens(0); - auto const& draftTokens = llmReq->getDraftTokens(); - auto const draftLength = llmReq->getNumDraftTokens(); - auto const& positionIds = llmReq->getPositionIds(); - - auto const contextChunkSize = llmReq->getContextChunkSize(); - auto const beginCompute = llmReq->getContextCurrentPosition(); - auto const endCompute = beginCompute + contextChunkSize; - inputHost.insert(inputHost.end(), reqTokens.begin() + beginCompute, reqTokens.begin() + endCompute); - - logitsIdsHostPtr[totalNumLogits++] = contextChunkSize; - numContextLogits.at(batchIdx) = modelConfig.computeContextLogits() ? contextChunkSize : 1; - - if (llmReq->isLastContextChunk()) - { - inputHost.insert(inputHost.end(), draftTokens->begin(), draftTokens->end()); - std::fill_n(logitsIdsHostPtr + totalNumLogits, draftLength, 1); - totalNumLogits += draftLength; - } - auto const inputLength = contextChunkSize + (llmReq->isLastContextChunk() ? draftLength : 0); - contextLengthsHostPtr[batchIdx] = inputLength; - auto const sequenceLen = inputLength + llmReq->getContextCurrentPosition(); - sequenceLengthsHostPtr[batchIdx] = sequenceLen; - - if (static_cast(pastKeyValueLengthsPtr)) - { - pastKeyValueLengthsPtr[batchIdx] = beginCompute + inputLength; - } - - if (positionIds.has_value()) - { - TLLM_CHECK_WITH_INFO(!(isChatGlm || isGlm), "ChatGLM-6B and Glm only use the default initialization"); - positionIdsHost.insert(positionIdsHost.end(), positionIds.value()->begin() + beginCompute, - positionIds.value()->begin() + endCompute); - } - else - { - if (isChatGlm) - { - // Specialize for ChatGLM-6B with 2D-Position-Embedding - positionIdsHost.resize(totalInputSize + inputLength); - std::iota(std::begin(positionIdsHost) + totalInputSize, std::end(positionIdsHost), 0); - positionIdsHost.back() = positionIdsHost.back() - 1; - - positionIdsHostRow2.resize(totalInputSize + inputLength); - positionIdsHostRow2.back() = 1; - } - else if (isGlm) - { - // Specialize for GLM-10B with 2D-Position-Embedding and special value of the mask id position - auto start = inputHost.begin() + totalInputSize; - auto end = start + inputLength; - auto it = std::find_if( - start, end, [](SizeType32 id) { return id == 50260 || id == 50263 || id == 50264; }); - llmReq->mMaskPosition = (it != end) ? std::distance(start, it) : maxContextLength; - - positionIdsHost.resize(totalInputSize + inputLength); - std::iota(std::begin(positionIdsHost) + totalInputSize, std::end(positionIdsHost), 0); - positionIdsHost.back() = llmReq->mMaskPosition; - - positionIdsHostRow2.resize(totalInputSize + inputLength); - positionIdsHostRow2.back() = 1; - } - else - { - // Other models - positionIdsHost.resize(totalInputSize + inputLength); - std::iota(std::begin(positionIdsHost) + totalInputSize, - std::begin(positionIdsHost) + totalInputSize + inputLength, beginCompute); - } - } - if (modelConfig.useMrope()) - { - auto optMropeRotaryCosSin = llmReq->getMropeRotaryCosSin().value(); - TLLM_CHECK_WITH_INFO(optMropeRotaryCosSin->getShape().d[0] == mropeRotaryCosSinSize, - "Provided MropeRotarySinCos is %ld and expected is %d.\n", optMropeRotaryCosSin->getShape().d[0], - int(mropeRotaryCosSinSize)); - - auto const mropeRotaryCosSinCtx = ITensor::slice(mropeRotaryCosSin, batchIdx, 1); - manager.copy(*optMropeRotaryCosSin, *mropeRotaryCosSinCtx); - } - - if (modelConfig.useLanguageAdapter()) - { - auto const languageAdapterRouting = llmReq->getLanguageAdapterRouting( - modelConfig.getNumLanguages().value(), endCompute - beginCompute); - languageAdapterRoutingsHost.insert(languageAdapterRoutingsHost.end(), - std::begin(languageAdapterRouting), std::end(languageAdapterRouting)); - } - totalInputSize += inputLength; - ++batchIdx; - } - - if (rnnStateBuffers) - { - rnnStateBuffers->fillSlotMappings(contextRequests, rnnStateManagerPtr); - } - } - - // generation preparation loop - if (!genRequests.empty()) - { - NVTX3_SCOPED_RANGE(genPrepareLoop); - - auto const numContextRequests = static_cast(contextRequests.size()); - auto numSequences = numContextRequests; - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const draftLength = llmReq->getNumDraftTokens(); - auto const& draftTokens = llmReq->getDraftTokens(); - auto const numLogits = draftLength + reqBeamWidth; - TLLM_CHECK(draftLength == 0 || reqBeamWidth == 1); - - auto const promptLen = llmReq->mPromptLen; - auto const sequenceLen - = promptLen + llmReq->getMaxNumGeneratedTokens() + static_cast(trtOverlap); - auto const& positionIds = llmReq->getPositionIds(); - for (int beam = 0; beam < reqBeamWidth; ++beam) - { - auto const numTokens = llmReq->getNumTokens(beam) + static_cast(trtOverlap); - // TODO: can this be removed completely? - if (!trtOverlap) - { - auto const lastToken = llmReq->getLastTokens(beam); - inputHost.push_back(lastToken); - if (draftLength > 0) - { - inputHost.insert(inputHost.end(), draftTokens->begin(), draftTokens->end()); - } - } - - // If model updates generation position ids do not append them here. - if (!modelConfig.getSpeculativeDecodingMode().updatesPositionIds()) - { - if (positionIds.has_value()) - { - TLLM_CHECK_WITH_INFO( - !(isChatGlm || isGlm), "ChatGLM-6B and Glm only use the default initialization"); - auto last_context_position_id = positionIds.value()->back(); - positionIdsHost.push_back( - static_cast(last_context_position_id + sequenceLen - promptLen)); - } - else - { - if (isChatGlm) // ChatGLM-6B - { - positionIdsHost.push_back(static_cast(promptLen - 2)); - positionIdsHostRow2.push_back(static_cast(sequenceLen - promptLen + 1)); - } - else if (isGlm) - { - positionIdsHost.push_back(llmReq->mMaskPosition); - positionIdsHostRow2.push_back(static_cast(sequenceLen - promptLen + 1)); - } - else // GPT / ChatGLM2-6B / ChatGLM3-6B / BART - { - // positionIds is just the size of tokens -1 - positionIdsHost.push_back(numTokens - 1); - } - } - } - - if (modelConfig.useMrope()) - { - auto optMropePositionDeltas = llmReq->getMropePositionDeltas().value(); - mropePositionDeltasHost.push_back(optMropePositionDeltas); - } - - if (modelConfig.useLanguageAdapter()) - { - // Generation requests only have one token per sequence - auto const languageAdapterRouting - = llmReq->getLanguageAdapterRouting(modelConfig.getNumLanguages().value(), 1); - languageAdapterRoutingsHost.insert(languageAdapterRoutingsHost.end(), - std::begin(languageAdapterRouting), std::end(languageAdapterRouting)); - } - } - - if (static_cast(pastKeyValueLengthsPtr)) - { - SizeType32 pastKeyValueLength = sequenceLen - 1; - std::fill_n(pastKeyValueLengthsPtr + numSequences, reqBeamWidth, pastKeyValueLength); - } - totalInputSize += numLogits; - - std::fill_n(logitsIdsHostPtr + totalNumLogits, numLogits, 1); - - totalNumLogits += numLogits; - - if (rnnStateBuffers) - { - auto const seqSlot = llmReq->mSeqSlot.value(); - auto& rnnStateManager = *rnnStateManagerPtr; - rnnStateManager.fillSlotMapping(*rnnStateBuffers->slotMappingHost, numSequences, seqSlot, reqBeamWidth); - } - numSequences += reqBeamWidth; - } - - if (transformerBuffers && maxBeamWidth > 1) - { - transformerBuffers->copyCacheIndirection(genRequests, decoderState.getCacheIndirectionOutput(), stream); - } - - numSequences = numContextRequests; - for (auto const& llmReq : genRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const draftLength = llmReq->getNumDraftTokens(); - - auto const contextQLength = llmReq->mPromptLen + draftLength; - auto const sequenceLen - = contextQLength + llmReq->getMaxNumGeneratedTokens() + static_cast(trtOverlap); - - std::fill_n(contextLengthsHostPtr + numSequences, reqBeamWidth, contextQLength); - std::fill_n(sequenceLengthsHostPtr + numSequences, reqBeamWidth, sequenceLen); - numSequences += reqBeamWidth; - } - if (modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - // copy from lookahead decoding buffer - mLookaheadBuffers->setFromInputs(numContextRequests, numGenRequests, *requestTypes, *seqSlots, - decoderState.getLookaheadBuffers(), runtime, modelConfig, worldConfig); - } - } - - // check skipCrossAttnBlocks - if (transformerBuffers && modelConfig.skipCrossAttnBlocks()) - { - bool isSkipCrossAttn = true; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - bool tmpValue = false; - if (llmReq->getSkipCrossAttnBlocks() != nullptr) - { - manager.copy(*llmReq->getSkipCrossAttnBlocks(), &tmpValue); - } - isSkipCrossAttn &= tmpValue; - } - } - transformerBuffers->copySkipCrossAttnBlocks(isSkipCrossAttn, runtime); - } - - if (isChatGlm || isGlm) - { - positionIdsHost.reserve(totalInputSize * 2); - positionIdsHost.insert(positionIdsHost.end(), positionIdsHostRow2.begin(), positionIdsHostRow2.end()); - } - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->fill(contextRequests, genRequests, manager); - } - if (modelConfig.usePromptTuning()) - { - promptTuningBuffers->fill(contextRequests, genRequests, manager, modelConfig.usePackedInput()); - } - if (modelConfig.useLoraPlugin()) - { - loraBuffers->fill(contextRequests, genRequests, peftTable, manager, modelConfig, worldConfig); - } - if (modelConfig.useMrope()) - { - if (!mropePositionDeltasHost.empty()) - { - auto mropePositionDeltasGen = ITensor::slice(mropePositionDeltas, 0, numGenSequences); - manager.copy(mropePositionDeltasHost.data(), *mropePositionDeltasGen); - } - } - - { - NVTX3_SCOPED_RANGE(bufferCopies); - if (trtOverlap) - { - auto contextInputsIds = ITensor::slice(inputsIds, 0, numContextTokens); - manager.copy(inputHost.data(), *contextInputsIds); - - if (!genRequests.empty()) - { - auto generationInputsIds = ITensor::slice(inputsIds, numContextTokens); - auto seqSlotsDeviceSlice = ITensor::slice(seqSlotsDevice, numContextRequests); - runtime::kernels::invokeGatherBatch( - *generationInputsIds, *newOutputTokens, *seqSlotsDeviceSlice, maxBeamWidth, stream); - } - } - else - { - manager.copy(inputHost.data(), *inputsIds); - } - // In generation phase, device ptr of context lengths need to be tiled. - manager.copy(*contextLengthsHost, *contextLengthsDevice); - manager.copy(*sequenceLengthsHost, *sequenceLengthsDevice); - auto const logitsIdsHostRange = BufferRange(*logitsIdsHost); - auto lastTokenIdsHostRange = BufferRange(*lastTokenIdsHost); - common::stl_utils::inclusiveScan( - logitsIdsHostRange.begin(), logitsIdsHostRange.end(), lastTokenIdsHostRange.begin()); - manager.copy(*lastTokenIdsHost, *lastTokenIdsDevice); - if (transformerBuffers) - { - TensorPtr decoderPositionIds = modelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() - ? mLookaheadBuffers->positionIdsDevice - : nullptr; - transformerBuffers->copyPositionIds(runtime, positionIdsHost, isChatGlm || isGlm, decoderPositionIds); - } - if (rnnStateBuffers) - { - rnnStateBuffers->copySlotMappingH2D(runtime); - } - if (modelConfig.useLanguageAdapter()) - { - manager.copy(languageAdapterRoutingsHost.data(), *languageAdapterRoutings); - } - } - - if (transformerBuffers && static_cast(kvCacheManagerPtr)) - { - transformerBuffers->copyKvBlockOffsets( - contextRequests, genRequests, kvCacheManagerPtr, crossKvCacheManagerPtr, manager); - } - - if (modelConfig.useCrossAttention()) - { - transformerBuffers->copyCrossAttentionMasks(contextRequests, genRequests, contextLengthsDevice, - encoderBuffers->inputLengths, maxContextLength, encoderBuffers->getMaxInputLengthInBatch(), runtime); - } - - maxKvCacheLengthRounded = 0; - if (static_cast(pastKeyValueLengthsPtr)) - { - auto const maxKvCacheLength - = *std::max_element(pastKeyValueLengthsPtr, pastKeyValueLengthsPtr + getNumSequences()); - // Round up kv cache length - maxKvCacheLengthRounded = common::ceilDiv(maxKvCacheLength, kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE) - * kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE; - } - - if (modelConfig.getSpeculativeDecodingMode().needsDecoderPrologue()) - { - if (modelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens()) - { - prepareExplicitDraftTokenBuffers( - decoderState.getExplicitDraftTokensBuffers(), runtime, modelConfig, worldConfig); - } - if (modelConfig.getSpeculativeDecodingMode().isEagle()) - { - prepareEagleBuffers( - contextRequests, genRequests, decoderState.getEagleBuffers(), runtime, modelConfig, worldConfig); - } - } - - sync_check_cuda_error(stream.get()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::prepareExplicitDraftTokenBuffers( - runtime::ExplicitDraftTokensBuffers::Inputs const& explicitDraftTokensBuffers, TllmRuntime const& runtime, - ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(mExplicitDraftTokensBuffers); - - mExplicitDraftTokensBuffers->setFromInputs(numContextRequests, numGenRequests, *requestTypes, *seqSlots, - explicitDraftTokensBuffers, *transformerBuffers->positionIds, modelConfig, worldConfig, - runtime.getBufferManager(), runtime.getStream()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void RuntimeBuffers::prepareEagleBuffers(RequestVector const& contextRequests, RequestVector const& genRequests, - runtime::EagleBuffers::Inputs const& eagleBuffers, TllmRuntime const& runtime, ModelConfig const& modelConfig, - WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(mEagleBuffers); - - mEagleBuffers->setFromInputs(contextRequests, genRequests, *requestTypes, *seqSlots, eagleBuffers, - runtime.getBufferManager(), modelConfig, worldConfig); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::tuple RuntimeBuffers::prepareStep( - RequestVector const& contextRequests, RequestVector const& genRequests, SizeType32 maxBeamWidth, - SizeType32 maxAttentionWindow, runtime::decoder::DecoderState const& decoderState, - kv_cache_manager::BaseKVCacheManager* kvCacheManager, kv_cache_manager::BaseKVCacheManager* crossKvCacheManager, - rnn_state_manager::RnnStateManager* rnnStateManager, PeftTable const& peftTable, TllmRuntime const& runtime, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, bool gatherGenerationLogits, bool trtOverlap, - OptionalRef newOutputTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersPrepareStep); - - setBufferSizes(contextRequests, genRequests); - reshape(runtime, modelConfig, worldConfig, gatherGenerationLogits); - - setFromInputs(contextRequests, genRequests, maxBeamWidth, maxAttentionWindow, decoderState, kvCacheManager, - crossKvCacheManager, rnnStateManager, peftTable, runtime, modelConfig, worldConfig, trtOverlap, - newOutputTokens); - - fillIOMaps(modelConfig, worldConfig); - - auto const numTokens = getNumTokens(); - auto const optProfileId = runtime.getOptProfileId(numTokens, ModelConfig::getOptProfilesSplitPoints()); - setContextIndex(optProfileId); - TLLM_LOG_DEBUG("numTokens: %d, optProfileId: %d", numTokens, optProfileId); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {optProfileId, inputMap, outputMap}; -} - -void RuntimeBuffers::fillIOMaps(ModelConfig const& modelConfig, WorldConfig const& worldConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(runtimeBuffersFillIOMaps); - - inputMap.clear(); - outputMap.clear(); - - if (transformerBuffers) - { - transformerBuffers->getBuffers(inputMap, outputMap, modelConfig); - } - if (rnnStateBuffers) - { - rnnStateBuffers->getBuffers(inputMap); - } - - if (worldConfig.isLastPipelineParallelRank()) - { - // feed a view to TensorRT runtime so reshaping does not change logits buffer - outputMap.insert_or_assign(kLogitsTensorName, ITensor::view(logits)); - } - else - { - outputMap.insert_or_assign(kHiddenStatesOutputTensorName, hiddenStates); - } - - if (worldConfig.isFirstPipelineParallelRank()) - { - inputMap.insert_or_assign(kInputIdsTensorName, inputsIds); - } - else - { - inputMap.insert_or_assign(kHiddenStatesInputTensorName, hiddenStates); - } - - inputMap.insert_or_assign(kLastTokenIdsTensorName, lastTokenIdsDevice); - - inputMap.insert_or_assign(kHostRequestTypesTensorName, requestTypes); - // In the generation phase, we still pass context lengths. - inputMap.insert_or_assign(kContextLengthsTensorName, contextLengthsDevice); - inputMap.insert_or_assign(kHostContextLengthsTensorName, contextLengthsHost); - inputMap.insert_or_assign(kSequenceLengthsTensorName, sequenceLengthsDevice); - - if (modelConfig.useCrossAttention()) - { - encoderBuffers->insertInputTensors(inputMap); - } - if (modelConfig.usePromptTuning()) - { - auto const& promptTuningParams = promptTuningBuffers->mPromptTuningParams; - inputMap.insert_or_assign(kPromptEmbeddingTableTensorName, promptTuningParams.embeddingTable); - inputMap.insert_or_assign(kTasksTensorName, promptTuningParams.tasks); - inputMap.insert_or_assign(kPromptVocabSizeTensorName, promptTuningParams.vocabSize); - } - if (modelConfig.useMrope()) - { - - inputMap.insert_or_assign(kMRopeRotaryCosSinTensorName, mropeRotaryCosSin); - inputMap.insert_or_assign(kMRopePositionDeltasTensorName, mropePositionDeltas); - } - if (modelConfig.useLoraPlugin()) - { - loraBuffers->insertInputTensors(inputMap, loraBuffers->mLoraWeightsPointersHost, - loraBuffers->mLoraAdapterSizesHost, modelConfig, worldConfig); - } - if (modelConfig.useLanguageAdapter()) - { - inputMap.insert_or_assign("language_adapter_routings", languageAdapterRoutings); - } - - if (mMedusaBuffers) - { - mMedusaBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - if (mLookaheadBuffers) - { - mLookaheadBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - if (mExplicitDraftTokensBuffers) - { - mExplicitDraftTokensBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - if (mEagleBuffers) - { - mEagleBuffers->insertInputTensors(inputMap, outputMap, worldConfig); - } - - for (auto const& outputTensor : mAdditionalOutputTensors) - { - outputMap.insert_or_assign(outputTensor.first, outputTensor.second); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/transformerBuffers.cpp b/cpp/tensorrt_llm/batch_manager/transformerBuffers.cpp deleted file mode 100644 index 4f81c8926682..000000000000 --- a/cpp/tensorrt_llm/batch_manager/transformerBuffers.cpp +++ /dev/null @@ -1,679 +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/transformerBuffers.h" - -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/kernels/attentionMask.h" -#include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaPackedMask.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmBuffers.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include - -using namespace tensorrt_llm::runtime; -namespace tk = tensorrt_llm::kernels; - -namespace tensorrt_llm::batch_manager -{ - -TransformerBuffers::TransformerBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - std::vector const& maxAttentionWindowVec, SizeType32 maxAttentionWindow, SizeType32 sinkTokenLen, - runtime::TllmRuntime const& runtime, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig) - : maxInputLen(modelConfig.getMaxInputLen()) - , maxEncoderOutputLen(modelConfig.getMaxEncoderLen()) -{ - auto const& manager = runtime.getBufferManager(); - auto const& engine = runtime.getEngine(); - - positionIds = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - auto const localNbAttnLayers - = modelConfig.getNbAttentionLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - // find the index of the first attention layer in the current rank - auto const firstLayerId = modelConfig.countLowerRankLayers(runtime::ModelConfig::LayerType::kATTENTION, - worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - - cacheIndirection - = manager.gpu(ITensor::makeShape({maxBatchSize, maxBeamWidth, maxAttentionWindow}), nvinfer1::DataType::kINT32); - - if (!modelConfig.getMaxNumTokens().has_value()) - { - TLLM_THROW("Model must configure a max number of tokens."); - } - maxNumTokens = modelConfig.getMaxNumTokens().value(); - - if (modelConfig.isKVCacheEnabled()) - { - auto const kvCacheBlockOffsetsType = engine.getTensorDataType("kv_cache_block_offsets"); - kvCacheBlockOffsetsHost = manager.emptyTensor(MemoryType::kPINNEDPOOL, kvCacheBlockOffsetsType); - kvCacheBlockOffsetsDevice = manager.emptyTensor(MemoryType::kGPU, kvCacheBlockOffsetsType); - - if (modelConfig.useCrossAttention()) - { - crossKvCacheBlockOffsetsHost = manager.emptyTensor(MemoryType::kPINNEDPOOL, kvCacheBlockOffsetsType); - crossKvCacheBlockOffsetsDevice = manager.emptyTensor(MemoryType::kGPU, kvCacheBlockOffsetsType); - crossAttentionMaskDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kBOOL); - crossAttentionMaskPinnedHost = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens, maxEncoderOutputLen}), nvinfer1::DataType::kBOOL); - crossAttentionPackedMaskDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - crossAttentionCuQSeqLensDevice = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - crossAttentionPackedMaskCuMaskRowsDevice - = manager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); - - // Pinned memory for batch copy of attention masks. - // There will be paddings in the dim1, so copy it by tokens. - crossAttentionMaskCopySrcOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens}), nvinfer1::DataType::kINT64); - crossAttentionMaskCopyDstOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens}), nvinfer1::DataType::kINT64); - crossAttentionMaskCopySizes = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxNumTokens}), nvinfer1::DataType::kINT64); - } - } - - fillValuesAlt = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - fillValuesAltDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - seqSlotsAlt = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - seqSlotsAltDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - - cacheIndirBatchedCopySrcOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT64); - cacheIndirBatchedCopyDstOffsets = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT64); - cacheIndirBatchedCopySizes = tensorrt_llm::runtime::BufferManager::pinnedPool( - ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT64); - skipCrossAttnBlocks - = tensorrt_llm::runtime::BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kBOOL); - - pastKeyValueLengths = manager.emptyTensor(MemoryType::kCPU, nvinfer1::DataType::kINT32); - - maxAttentionWindows = BufferManager::cpu(ITensor::makeShape({localNbAttnLayers}), nvinfer1::DataType::kINT32); - auto* maxAttentionWindowsPtr = bufferCast(*maxAttentionWindows); - auto const attentionWindowLength = maxAttentionWindowVec.size(); - for (SizeType32 i = 0; i < localNbAttnLayers; ++i) - { - maxAttentionWindowsPtr[i] = maxAttentionWindowVec[(firstLayerId + i) % attentionWindowLength]; - } - - sinkTokenLengths = BufferManager::cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - bufferCast(*sinkTokenLengths)[0] = sinkTokenLen; - - contextProgressHost = BufferManager::cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT64); - bufferCast(*contextProgressHost)[0] = 0; - - if (modelConfig.useGemmAllReducePlugin() && worldConfig.isTensorParallel()) - { - nvinfer1::DataType ARType = modelConfig.getGemmAllReduceDtype(); - - auto hiddenSize = modelConfig.getHiddenSize() * worldConfig.getTensorParallelism(); - - auto tpGroup = worldConfig.getTensorParallelGroup(); - std::set tpGroupSet(tpGroup.begin(), tpGroup.end()); - - auto outputDims = ITensor::makeShape({modelConfig.getMaxNumTokens().value() * hiddenSize}); - - gemmAllReduceOutput = std::make_shared(outputDims, ARType, tpGroupSet); - } -} - -void TransformerBuffers::reshape(SizeType32 numSequences, SizeType32 numInputTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - pastKeyValueLengths->reshape(ITensor::makeShape({numSequences})); - - if (kvCacheBlockOffsetsHost) - { - auto cacheBlockOffsetsShape = kvCacheBlockOffsetsHost->getShape(); - if (cacheBlockOffsetsShape.nbDims > 0) - { - cacheBlockOffsetsShape.d[1] = numSequences; - kvCacheBlockOffsetsHost->reshape(cacheBlockOffsetsShape); - kvCacheBlockOffsetsDevice->reshape(cacheBlockOffsetsShape); - } - else - { - TLLM_LOG_DEBUG("kvCacheBlockOffsets not allocated yet"); - } - } - - if (crossKvCacheBlockOffsetsHost) - { - TLLM_CHECK_WITH_INFO( - crossKvCacheBlockOffsetsDevice, "crossKvCacheBlockOffsetsDevice is empty for model with cross attention!"); - auto crossCacheBlockOffsetsShape = crossKvCacheBlockOffsetsHost->getShape(); - if (crossCacheBlockOffsetsShape.nbDims > 0) - { - crossCacheBlockOffsetsShape.d[1] = numSequences; - crossKvCacheBlockOffsetsHost->reshape(crossCacheBlockOffsetsShape); - crossKvCacheBlockOffsetsDevice->reshape(crossCacheBlockOffsetsShape); - } - else - { - TLLM_LOG_DEBUG("crossKvCacheBlockOffsets not allocated yet"); - } - } - - if (crossAttentionMaskDevice) - { - auto crossAttentionMaskShape = crossAttentionMaskDevice->getShape(); - if (crossAttentionMaskShape.nbDims > 0) - { - crossAttentionMaskShape.d[0] = numInputTokens; - crossAttentionMaskDevice->reshape(crossAttentionMaskShape); - crossAttentionMaskPinnedHost->reshape(crossAttentionMaskShape); - crossAttentionMaskCopySrcOffsets->reshape(ITensor::makeShape({numInputTokens})); - crossAttentionMaskCopyDstOffsets->reshape(ITensor::makeShape({numInputTokens})); - crossAttentionMaskCopySizes->reshape(ITensor::makeShape({numInputTokens})); - } - else - { - TLLM_LOG_DEBUG("crossAttentionMaskDevice not allocated yet"); - } - } - - if (crossAttentionPackedMaskDevice) - { - auto crossAttentionMaskPackedShape = crossAttentionPackedMaskDevice->getShape(); - if (crossAttentionMaskPackedShape.nbDims > 0) - { - crossAttentionMaskPackedShape.d[0] = numInputTokens; - crossAttentionPackedMaskDevice->reshape(crossAttentionMaskPackedShape); - } - else - { - TLLM_LOG_DEBUG("crossAttentionPackedMaskDevice not allocated yet"); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::reshapeKvTensors(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, SizeType32 maxBlocksPerSeq, - kv_cache_manager::CacheType kvCacheType, SizeType32 numPools, BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // allocate with max shape during init - if (kvCacheType == kv_cache_manager::CacheType::kSELF) - { - auto const cacheBlockOffsetsShape - = ITensor::makeShape({numPools, maxBatchSize * maxBeamWidth, 2, maxBlocksPerSeq}); - - kvCacheBlockOffsetsHost->reshape(cacheBlockOffsetsShape); - manager.setZero(*kvCacheBlockOffsetsHost); - - kvCacheBlockOffsetsDevice->reshape(cacheBlockOffsetsShape); - manager.setZero(*kvCacheBlockOffsetsDevice); - } - else if (kvCacheType == kv_cache_manager::CacheType::kCROSS) - { - auto const crossCacheBlockOffsetsShape - = ITensor::makeShape({numPools, maxBatchSize * maxBeamWidth, 2, maxBlocksPerSeq}); - - crossKvCacheBlockOffsetsHost->reshape(crossCacheBlockOffsetsShape); - manager.setZero(*crossKvCacheBlockOffsetsHost); - - crossKvCacheBlockOffsetsDevice->reshape(crossCacheBlockOffsetsShape); - manager.setZero(*crossKvCacheBlockOffsetsDevice); - - crossAttentionMaskDevice->reshape(ITensor::makeShape({maxNumTokens, maxEncoderOutputLen})); - manager.setZero(*crossAttentionMaskDevice); - manager.setZero(*crossAttentionMaskPinnedHost); - - // Only context attention needs this, so allocate it by shape [maxBatchSize, maxInputLen, maxEncoderOutputLen]. - auto [packedMaskM, packedMaskN] = tk::roundUpPackedMaskMNDims(maxInputLen, maxEncoderOutputLen); - crossAttentionPackedMaskDevice->reshape(ITensor::makeShape({maxBatchSize * packedMaskM, packedMaskN})); - manager.setZero(*crossAttentionPackedMaskDevice); - - crossAttentionCuQSeqLensDevice->reshape(ITensor::makeShape({maxBatchSize + 1})); - manager.setZero(*crossAttentionCuQSeqLensDevice); - - crossAttentionPackedMaskCuMaskRowsDevice->reshape(ITensor::makeShape({maxBatchSize + 1})); - manager.setZero(*crossAttentionPackedMaskCuMaskRowsDevice); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::getBuffers( - TensorMap& inputBuffers, TensorMap& outputBuffers, runtime::ModelConfig const& modelConfig) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(transformerBuffersGetBuffers); - - inputBuffers.insert_or_assign(kPositionIdsTensorName, positionIds); - inputBuffers.insert_or_assign(kHostPastKeyValueLengthsTensorName, pastKeyValueLengths); - inputBuffers.insert_or_assign(kCacheIndirectionsTensorName, cacheIndirection); - inputBuffers.insert_or_assign(kHostSinkTokenLengthTensorName, sinkTokenLengths); - - inputBuffers.insert_or_assign(kHostMaxAttentionWindowSizesTensorName, maxAttentionWindows); - inputBuffers.insert_or_assign(kKvCacheBlockOffsetsTensorName, kvCacheBlockOffsetsDevice); - inputBuffers.insert_or_assign(kHostKvCacheBlockOffsetsTensorName, kvCacheBlockOffsetsHost); - inputBuffers.insert_or_assign(kHostContextProgressTensorName, contextProgressHost); - - if (crossKvCacheBlockOffsetsHost) - { - inputBuffers.insert_or_assign(kCrossKvCacheBlockOffsetsTensorName, crossKvCacheBlockOffsetsDevice); - inputBuffers.insert_or_assign(kHostCrossKvCacheBlockOffsetsTensorName, crossKvCacheBlockOffsetsHost); - inputBuffers.insert_or_assign(kHostCrossKvCachePoolPointersTensorName, crossKvCacheBlockPoolPointers); - inputBuffers.insert_or_assign(kHostCrossKvCachePoolMappingTensorName, crossKvCacheBlockPoolMapping); - inputBuffers.insert_or_assign(kCrossAttentionMaskTensorName, crossAttentionMaskDevice); - inputBuffers.insert_or_assign(kCrossAttentionPackedMaskTensorName, crossAttentionPackedMaskDevice); - } - - if (skipCrossAttnBlocks) - { - inputBuffers.insert_or_assign(kSkipCrossAttentionBlocksTensorName, skipCrossAttnBlocks); - } - - if (modelConfig.useGemmAllReducePlugin()) - { - for (int idx = 0; idx < modelConfig.getNbAttentionLayers() * 2; ++idx) - { - // XXX (xsimmons): this is a bit hacky as it assumes - // 2x RowLinear layers per attention block. - // This will be fixed soon when I remove coupling between model - // and runtime. - auto gemmARViewUC = gemmAllReduceOutput->getTensorView(MulticastTensorView::ViewType::kUNICAST); - auto gemmARViewMC = gemmAllReduceOutput->getTensorView(MulticastTensorView::ViewType::kMULTICAST); - auto gemmARViewIpc = gemmAllReduceOutput->getTensorView(MulticastTensorView::ViewType::kIPC_LIST); - - outputBuffers.insert_or_assign("gemm_allreduce_uc_out_" + std::to_string(idx), gemmARViewUC); - outputBuffers.insert_or_assign("gemm_allreduce_mc_out_" + std::to_string(idx), gemmARViewMC); - outputBuffers.insert_or_assign("gemm_allreduce_ipc_out_" + std::to_string(idx), gemmARViewIpc); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copyPositionIds(runtime::TllmRuntime const& runtime, - std::vector const& positionIdsHost, bool isChatGlm, TensorPtr const& decoderPositionIds) -{ - auto const& manager = runtime.getBufferManager(); - if (isChatGlm) - { - positionIds->reshape(ITensor::makeShape({2, static_cast(positionIdsHost.size()) / 2})); - manager.copy(positionIdsHost.data(), *positionIds); - } - else if (decoderPositionIds == nullptr) - { - positionIds->reshape(ITensor::makeShape({static_cast(positionIdsHost.size())})); - manager.copy(positionIdsHost.data(), *positionIds); - } - else - { - // concat context phase and generation phase positionIds. - auto const contextPositionIdsLen = static_cast(positionIdsHost.size()); - auto const generationPositionIdsLen = ITensor::volume(decoderPositionIds->getShape()); - positionIds->reshape(ITensor::makeShape({contextPositionIdsLen + generationPositionIdsLen})); - manager.copy(positionIdsHost.data(), *ITensor::slice(positionIds, 0, contextPositionIdsLen)); - manager.copy(*decoderPositionIds, *ITensor::slice(positionIds, contextPositionIdsLen)); - } -} - -void TransformerBuffers::copyKvBlockOffsets(RequestVector const& contextRequests, RequestVector const& genRequests, - kv_cache_manager::BaseKVCacheManager const* kvCacheManager, - kv_cache_manager::BaseKVCacheManager const* crossKvCacheManager, BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(copyKvBlockOffsets); - - auto const& cudaStream = manager.getStream(); - - SizeType32 constexpr contextBeamWidth{1}; - SizeType32 numSequences{0}; - SizeType32 maxBlockCount{0}; - SizeType32 maxCrossBlockCount{0}; - for (auto const& requests : {contextRequests, genRequests}) - { - for (auto const& llmReq : requests) - { - auto const requestId = llmReq->mRequestId; - auto const isContextRequest = llmReq->isContextInitState(); - auto const beamWidth = isContextRequest ? contextBeamWidth : llmReq->getBeamWidthByIter(); - auto const maxBeamBlockCount - = kvCacheManager->copyBlockOffsets(*kvCacheBlockOffsetsHost, numSequences, requestId); - maxBlockCount = std::max(maxBlockCount, maxBeamBlockCount); - if (crossKvCacheBlockOffsetsHost) - { - auto const maxCrossBeamBlockCount - = crossKvCacheManager->copyBlockOffsets(*crossKvCacheBlockOffsetsHost, numSequences, requestId); - maxCrossBlockCount = std::max(maxCrossBlockCount, maxCrossBeamBlockCount); - } - numSequences += beamWidth; - } - } - - // requests' block offsets collected as [totalNumSequences, 2, maxBlocksPerSeq], copy to device - auto copyOffsetsToDevice = [&cudaStream](TensorPtr& offsetsHost, TensorPtr& offsetsDevice, SizeType32 maxBlockCount) - { - // shape should be [totalNumSequences, 2, maxBlocksPerSeq] - auto const& offsetsShape = offsetsHost->getShape(); - auto const maxBlocksPerSeq = offsetsShape.d[3]; - auto const offsetsTypeSize = tensorrt_llm::common::getDTypeSize(offsetsHost->getDataType()); - auto const copyPitch = maxBlocksPerSeq * offsetsTypeSize; - auto const copyHeight = offsetsShape.d[0] * offsetsShape.d[1] * offsetsShape.d[2]; - auto const copyWidth = maxBlockCount * offsetsTypeSize; - auto* srcPtr = bufferCast(*offsetsHost); - auto* dstPtr = bufferCast(*offsetsDevice); - - TLLM_CUDA_CHECK(cudaMemcpy2DAsync( - dstPtr, copyPitch, srcPtr, copyPitch, copyWidth, copyHeight, cudaMemcpyHostToDevice, cudaStream.get())); - }; - - copyOffsetsToDevice(kvCacheBlockOffsetsHost, kvCacheBlockOffsetsDevice, maxBlockCount); - if (crossKvCacheBlockOffsetsHost) - { - copyOffsetsToDevice(crossKvCacheBlockOffsetsHost, crossKvCacheBlockOffsetsDevice, maxCrossBlockCount); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copyCacheIndirection( - RequestVector const& genRequests, TensorPtr const& decoderCacheIndirectionOutput, CudaStream const& stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(copyCacheIndirection); - - auto const numGenerationRequests = genRequests.size(); - - auto batchedCopySrcOffsets = BufferRange(*cacheIndirBatchedCopySrcOffsets); - auto batchedCopyDstOffsets = BufferRange(*cacheIndirBatchedCopyDstOffsets); - auto batchedCopySizes = BufferRange(*cacheIndirBatchedCopySizes); - - auto cacheIndirShape = decoderCacheIndirectionOutput->getShape(); - - // 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. - auto reqBeamWidth = genRequests[0]->getBeamWidthByIter(); - - // Get size of copying from shape of `CacheIndirectionOutput` - cacheIndirShape.d[0] = 1; - cacheIndirShape.d[1] = reqBeamWidth; // Use beam width of current step rather than max beam width as dst offset - auto const copySize = static_cast(ITensor::volume(cacheIndirShape)); - - std::transform(genRequests.begin(), genRequests.end(), batchedCopySrcOffsets.begin(), - [copySize](auto const& llmReq) { return llmReq->mSeqSlot.value() * copySize; }); - std::generate_n( - batchedCopyDstOffsets.begin(), numGenerationRequests, [copySize, i = 0]() mutable { return (i++) * copySize; }); - std::fill_n(batchedCopySizes.begin(), numGenerationRequests, copySize); - - auto const batchedCopySrcOffsetsSlice = ITensor::slice(cacheIndirBatchedCopySrcOffsets, 0, numGenerationRequests); - auto const batchedCopyDstOffsetsSlice = ITensor::slice(cacheIndirBatchedCopyDstOffsets, 0, numGenerationRequests); - auto const batchedCopySizesSlice = ITensor::slice(cacheIndirBatchedCopySizes, 0, numGenerationRequests); - runtime::kernels::invokeCopyBatch(*decoderCacheIndirectionOutput, *cacheIndirection, *batchedCopySrcOffsetsSlice, - *batchedCopyDstOffsetsSlice, *batchedCopySizesSlice, copySize, stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copyCrossAttentionMasks(RequestVector const& contextRequests, RequestVector const& genRequests, - TensorPtr const& decoderContextLengthsDevice, TensorPtr const& encoderInputLengths, - SizeType32 maxDecoderContextLength, SizeType32 maxEncoderInputLengthInBatch, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const& manager = runtime.getBufferManager(); - - // Reshape the tensor to make sure the dim1 matches maxEncoderInputLengthInBatch. - auto crossAttentionMaskShape = crossAttentionMaskDevice->getShape(); - crossAttentionMaskShape.d[1] = maxEncoderInputLengthInBatch; - crossAttentionMaskDevice->reshape(crossAttentionMaskShape); - // Set crossAttentionMask to true by default if it is not provided. - manager.setMem(*crossAttentionMaskDevice, 1); - - // Check if all context requests have cross attention mask. - bool allContextCrossAttentionMaskProvided = true; - for (auto const& llmReq : contextRequests) - { - auto const& crossAttentionMaskRequest = llmReq->getCrossAttentionMask(); - if (bufferCastOrNull(crossAttentionMaskRequest) == nullptr) - { - allContextCrossAttentionMaskProvided = false; - break; - } - } - // If not all requests have cross attention mask, let us create the default ones. - auto const& stream = runtime.getStream(); - if (!allContextCrossAttentionMaskProvided) - { - TLLM_LOG_WARNING("Default padding attention mask will be used as not all requests have cross attention mask."); - tk::AttentionMaskParams attentionMaskParams; - memset((void*) &attentionMaskParams, 0, sizeof(attentionMaskParams)); - // Set parameters. - attentionMaskParams.mask = bufferCastOrNull(crossAttentionMaskDevice); - attentionMaskParams.cuQSeqLens = bufferCastOrNull(crossAttentionCuQSeqLensDevice); - attentionMaskParams.actualQSeqLens = bufferCastOrNull(decoderContextLengthsDevice); - attentionMaskParams.actualKvSeqLens = bufferCastOrNull(encoderInputLengths); - attentionMaskParams.attentionMaskType = tk::AttentionMaskType::PADDING; - attentionMaskParams.batchSize = static_cast(contextRequests.size()); - attentionMaskParams.maxQSeqLen = maxDecoderContextLength; - attentionMaskParams.maxKvSeqLen = maxEncoderInputLengthInBatch; - // Launch the kernel. - tk::invokeBuildAttentionMask(attentionMaskParams, stream.get()); - sync_check_cuda_error(stream.get()); - } - // Use the first request's cross attention mask tensor's pointer address as the primary source pointer. - auto const& attentionMaskSrc = !contextRequests.empty() ? contextRequests[0]->getCrossAttentionMask() - : genRequests[0]->getCrossAttentionMask(); - bool const* primarySrcPtr = bufferCastOrNull(attentionMaskSrc); - - // Pinned-memory buffer preparation for batch copy. - auto batchedCopySrcOffsets = BufferRange(*crossAttentionMaskCopySrcOffsets); - auto batchedCopyDstOffsets = BufferRange(*crossAttentionMaskCopyDstOffsets); - auto batchedCopySizes = BufferRange(*crossAttentionMaskCopySizes); - // Requests with cross-attention-mask don't need to copy. - manager.setZero(*crossAttentionMaskCopySizes); - sync_check_cuda_error(stream.get()); - - SizeType32 numTokens = 0; - SizeType32 numCopiedTokens = 0; - bool* pinnedMemPtr = bufferCastOrNull(crossAttentionMaskPinnedHost); - for (auto const& llmReq : contextRequests) - { - auto const& crossAttentionMaskRequest = llmReq->getCrossAttentionMask(); - auto const position = llmReq->getContextCurrentPosition(); - auto const size = llmReq->getContextChunkSize(); - if (bufferCastOrNull(crossAttentionMaskRequest) != nullptr) - { - auto memType = crossAttentionMaskRequest->getMemoryType(); - auto const crossAttentionMaskRequestDim0 - = static_cast(crossAttentionMaskRequest->getShape().d[0]); - auto const crossAttentionMaskRequestDim1 - = static_cast(crossAttentionMaskRequest->getShape().d[1]); - TLLM_LOG_DEBUG("copyCrossAttentionMasks (shape [%d, %d]) from contextRequests position %d chunkSize %d", - crossAttentionMaskRequestDim0, crossAttentionMaskRequestDim1, position, size); - if ((position + size - 1) >= crossAttentionMaskRequestDim0) - { - TLLM_LOG_WARNING( - "The provided crossAttentionMask input is not complete for context phases, the last row " - "will be " - "used by default."); - } - // copy it to pinned memory if it is a cpu tensor. - if (memType == MemoryType::kCPU) - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on CPU."); - auto const copiedPosition - = std::min(crossAttentionMaskRequestDim0 - 1, static_cast(position)); - auto const copiedSize - = std::min(crossAttentionMaskRequestDim0 - copiedPosition, static_cast(size)); - SizeType64 inputMaskOffset = (copiedPosition * crossAttentionMaskRequestDim1); - SizeType64 inputMaskSize = (copiedSize * crossAttentionMaskRequestDim1); - std::memcpy( - pinnedMemPtr, bufferCastOrNull(crossAttentionMaskRequest) + inputMaskOffset, inputMaskSize); - pinnedMemPtr += inputMaskSize; - for (SizeType32 tokenId = position; tokenId < position + size; tokenId++) - { - SizeType64 tokenIdInPinnedMem - = std::min(copiedSize - 1, static_cast(tokenId - position)); - batchedCopySrcOffsets.begin()[numCopiedTokens] - = (pinnedMemPtr - primarySrcPtr) + tokenIdInPinnedMem * crossAttentionMaskRequestDim1; - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - numCopiedTokens++; - numTokens++; - } - } - else - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on GPU."); - for (SizeType32 tokenId = position; tokenId < position + size; tokenId++) - { - batchedCopySrcOffsets.begin()[numCopiedTokens] - = static_cast(bufferCastOrNull(crossAttentionMaskRequest) - primarySrcPtr) - + std::min(crossAttentionMaskRequestDim0 - 1, static_cast(tokenId)) - * crossAttentionMaskRequestDim1; - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - numCopiedTokens++; - numTokens++; - } - } - } - else - { - numTokens += size; - TLLM_LOG_WARNING( - "CrossAttentionMask is not provided for the request. Default padding attention mask will be " - "created."); - } - } - sync_check_cuda_error(stream.get()); - - for (auto const& llmReq : genRequests) - { - auto const promptLen = llmReq->mPromptLen; - auto const decodingIter = llmReq->getDecodingIter(); - auto const& crossAttentionMaskRequest = llmReq->getCrossAttentionMask(); - if (bufferCastOrNull(crossAttentionMaskRequest) != nullptr) - { - auto const memType = crossAttentionMaskRequest->getMemoryType(); - auto const crossAttentionMaskRequestDim0 - = static_cast(crossAttentionMaskRequest->getShape().d[0]); - auto const crossAttentionMaskRequestDim1 - = static_cast(crossAttentionMaskRequest->getShape().d[1]); - TLLM_LOG_DEBUG("copyCrossAttentionMasks (shape [%d, %d]) from genRequests decodingIter %d", - crossAttentionMaskRequestDim0, crossAttentionMaskRequestDim1, decodingIter); - if (promptLen + decodingIter - 1 >= crossAttentionMaskRequestDim0) - { - TLLM_LOG_WARNING( - "The provided crossAttentionMask input is not complete for generation phases, the last row " - "will be " - "used by default."); - } - // copy it to pinned memory if it is a cpu tensor. - if (memType == MemoryType::kCPU) - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on CPU."); - SizeType64 copiedPosition = std::min( - crossAttentionMaskRequestDim0 - 1, static_cast(promptLen + decodingIter - 1)); - SizeType64 inputMaskOffset = (copiedPosition * crossAttentionMaskRequestDim1); - SizeType64 inputMaskSize = crossAttentionMaskRequestDim1; - std::memcpy( - pinnedMemPtr, bufferCastOrNull(crossAttentionMaskRequest) + inputMaskOffset, inputMaskSize); - pinnedMemPtr += inputMaskSize; - batchedCopySrcOffsets.begin()[numCopiedTokens] = static_cast(pinnedMemPtr - primarySrcPtr); - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - } - else - { - TLLM_LOG_DEBUG("CrossAttentionMask tensor is on GPU."); - batchedCopySrcOffsets.begin()[numCopiedTokens] - = static_cast(bufferCastOrNull(crossAttentionMaskRequest) - primarySrcPtr) - + std::min(crossAttentionMaskRequestDim0 - 1, static_cast(promptLen + decodingIter - 1)) - * crossAttentionMaskRequestDim1; - batchedCopyDstOffsets.begin()[numCopiedTokens] - = numTokens * static_cast(maxEncoderInputLengthInBatch); - batchedCopySizes.begin()[numCopiedTokens] = crossAttentionMaskRequestDim1; - } - numCopiedTokens++; - numTokens++; - } - else - { - numTokens++; - TLLM_LOG_WARNING( - "CrossAttentionMask is not provided for the generation request. Full valid attentionMask will " - "be used " - "by default."); - } - } - sync_check_cuda_error(stream.get()); - - // Copy all requests' attention mask in one kernel. - if (attentionMaskSrc != nullptr) - { - crossAttentionMaskCopySrcOffsets->reshape(ITensor::makeShape({numCopiedTokens})); - crossAttentionMaskCopyDstOffsets->reshape(ITensor::makeShape({numCopiedTokens})); - crossAttentionMaskCopySizes->reshape(ITensor::makeShape({numCopiedTokens})); - runtime::kernels::invokeCopyBatch(*attentionMaskSrc, *crossAttentionMaskDevice, - *crossAttentionMaskCopySrcOffsets, *crossAttentionMaskCopyDstOffsets, *crossAttentionMaskCopySizes, - maxEncoderInputLengthInBatch, stream); - } - sync_check_cuda_error(stream.get()); - - // The packed mask is only needed by context requests now. - if (!contextRequests.empty()) - { - // Set the parameters for creating packed mask for context FMHA. - tk::PackedMaskParams maskParams{}; - maskParams.maskInput = bufferCastOrNull(crossAttentionMaskDevice); - maskParams.cuQSeqLens = bufferCastOrNull(crossAttentionCuQSeqLensDevice); - maskParams.packedMask = bufferCastOrNull(crossAttentionPackedMaskDevice); - maskParams.cuMaskRows = bufferCastOrNull(crossAttentionPackedMaskCuMaskRowsDevice); - maskParams.actualQSeqLens = bufferCastOrNull(decoderContextLengthsDevice); - maskParams.actualKvSeqLens = bufferCastOrNull(encoderInputLengths); - maskParams.batchSize = contextRequests.size(); - maskParams.maxQSeqLen = maxDecoderContextLength; - maskParams.maxKvSeqLen = maxEncoderInputLengthInBatch; - maskParams.attentionMaskType = tk::ContextAttentionMaskType::CUSTOM_MASK; - maskParams.validPosVal = true; - - // Launch the pack mask kernel. - tk::invokeBuildPackedMask(maskParams, stream.get()); - sync_check_cuda_error(stream.get()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TransformerBuffers::copySkipCrossAttnBlocks(bool const& _skipCrossAttnBlocks, runtime::TllmRuntime const& runtime) -{ - auto const& manager = runtime.getBufferManager(); - manager.copy(&_skipCrossAttnBlocks, *skipCrossAttnBlocks); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp b/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp deleted file mode 100644 index de1525b07730..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.cpp +++ /dev/null @@ -1,615 +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 "trtEncoderModel.h" -#include "encoderBuffers.h" -#include "tensorrt_llm/batch_manager/capacityScheduler.h" -#include "tensorrt_llm/batch_manager/microBatchScheduler.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/utils/runtimeUtils.h" - -#include -#include -#include - -using namespace tensorrt_llm::runtime; -using namespace tensorrt_llm::mpi; - -namespace tensorrt_llm::batch_manager -{ - -TrtEncoderModel::TrtEncoderModel(runtime::ModelConfig const& modelConfig, WorldConfig const& worldConfig, - runtime::RawEngine const& rawEngine, std::shared_ptr logger, - executor::ExecutorConfig const& executorConfig) - : TrtGptModel(modelConfig, worldConfig, executorConfig) - , mModelConfig{modelConfig} - , mWorldConfig{worldConfig} - , mDevice{runtime::utils::initDevice(worldConfig)} - , mLogger{logger ? std::move(logger) : std::make_shared()} - , mRuntime{std::make_shared( - rawEngine, mLogger.get(), executorConfig.getUseGpuDirectStorage(), executorConfig.getGpuWeightsPercent())} - , mNumMicroBatches{1} - , mNumBuffers{mNumMicroBatches} - , mCopyBufferManager{std::make_shared()} -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO( - !mWorldConfig.isPipelineParallel(), "Pipeline parallelism is currently not supported for encoder models."); - - createRuntimeContexts(); - - createBuffers(); - - if (mWorldConfig.isPipelineParallel()) - { - auto const& commSession = COMM_SESSION; - mMpiCommPipelinePara = std::make_shared( - commSession.split(mWorldConfig.getTensorParallelRank(), mWorldConfig.getPipelineParallelRank())); - } - - mMicroBatchScheduledRequests.resize(mNumMicroBatches); - // mEncoderWaitEvents.resize(mNumMicroBatches); - - // set noScheduleUntilState to LlmRequestState::kENCODER_INIT for encoder model - // when null kv cache manager is given, request scheduler will use MaxRequests as capacity scheduler, i.e. no - // handling of maximizing utilization or pause/evict - // TODO: finer control on encoder requests scheduling - mCapacityScheduler = std::make_unique( - getMaxBatchSize() * mNumMicroBatches, executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy(), false, - false, LlmRequestState::kENCODER_INIT, LlmRequestState::kCONTEXT_INIT); - - mMicroBatchScheduler = std::make_unique( - std::nullopt, mModelConfig.getMaxInputLen(), LlmRequestState::kENCODER_INIT, LlmRequestState::kCONTEXT_INIT); - - mHiddenSize = modelConfig.getHiddenSize(); - - mMaxInputLen = mModelConfig.getMaxInputLen(); - TLLM_LOG_INFO("TRTEncoderModel mMaxInputLen: reset to %d from build config.", mMaxInputLen); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -BufferManager const& TrtEncoderModel::getBufferManager() const -{ - return mRuntime->getBufferManager(); -} - -BufferManager::CudaStreamPtr TrtEncoderModel::getRuntimeStreamPtr() const -{ - return mRuntime->getStreamPtr(); -} - -nvinfer1::DataType TrtEncoderModel::getTensorDataType(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorDataType(name.c_str()); -} - -nvinfer1::Dims TrtEncoderModel::getTensorShape(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorShape(name.c_str()); -} - -void TrtEncoderModel::getCurrentIterationStats(executor::IterationStats& stats) const -{ - stats.iter = mIterCounter; -} - -void TrtEncoderModel::getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const -{ - stats.iter = mIterCounter; -} - -executor::DebugTensorsPerIteration TrtEncoderModel::getCurrentDebugTensors() const -{ - executor::DebugTensorsPerIteration debugTensors; - debugTensors.iter = mIterCounter; - - TLLM_LOG_WARNING("TrtEncoderModel doesn't support getting debug tensors."); - - return debugTensors; -} - -void TrtEncoderModel::setLayerProfiler() -{ - TLLM_CHECK(mRuntime); - mRuntime->setLayerProfiler(); -} - -std::string TrtEncoderModel::getLayerProfileInfo() const -{ - TLLM_CHECK(mRuntime); - return mRuntime->getLayerProfileInfo(); -} - -void TrtEncoderModel::createRuntimeContexts() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRuntime->clearContexts(); - auto const numProfiles = mRuntime->getNbProfiles(); - TLLM_CHECK_WITH_INFO(numProfiles == 1, "Encoder only expects one optimization profile"); - for (auto i = 0; i < numProfiles; ++i) - { - mRuntime->addContext(i); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::executeContext(SizeType32 runtimeContextId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeContext); - auto enqueueSuccessful = mRuntime->executeContext(runtimeContextId); - if (!enqueueSuccessful) - { - throw std::runtime_error("Executing TRT engine failed!"); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::createBuffers() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - for (SizeType32 i = 0; i < mNumBuffers; ++i) - { - mBuffers.emplace_back( - std::make_shared(getMaxBatchSize(), mModelConfig, mWorldConfig, *mRuntime)); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::executeBatch(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeBatch); - - // encoder model only have one optimization profile for now, so no optimization profile switch - SizeType32 optProfileIndex = 0; - auto const bufferId = getBufferId(); - if (!scheduledRequests.contextRequests.empty()) - { - // engine I/O - auto [inputMap, outputMap] - = mBuffers[bufferId]->prepareIO(scheduledRequests.contextRequests, mModelConfig, mWorldConfig, *mRuntime); - mRuntime->setInputTensors(optProfileIndex, inputMap); - mRuntime->setOutputTensors(optProfileIndex, outputMap); - - // engine run - executeContext(optProfileIndex); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::rearrangeOutputs(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(rearrangeOutputs); - - auto const bufferId = getBufferId(); - if (!scheduledRequests.contextRequests.empty()) - { - mBuffers[bufferId]->rearrangeOutputs(scheduledRequests.contextRequests, mModelConfig, mWorldConfig, *mRuntime); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::forwardSync() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtEncoderModel::forwardSync"); - - auto const device = mWorldConfig.getDevice(); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - // auto& encoderWaitEvent = mEncoderWaitEvents.at(mMicroBatchId); - - if (!currRequests.empty()) - { - if (!mWorldConfig.isPipelineParallel() || !mWorldConfig.isLastPipelineParallelRank()) - { - // TLLM_CHECK_WITH_INFO(mEncStepAsyncSndHdl.get() == nullptr, "encoderSync handle must be nullptr."); - // // Wait for encoding for requests in flight for the current micro batch - // mEncStepAsyncSndHdl = encoderSync(currRequests, encoderWaitEvent); - } - else - { - } - - NVTX3_SCOPED_RANGE(pauseFlaggedCurrRequests); - for (auto const& requests : {currRequests.contextRequests}) - { - for (auto const& llmReq : requests) - { - auto const reqId = llmReq->mRequestId; - mInflightReqIds.erase(reqId); - TLLM_LOG_DEBUG("request ID %u removed from ENCODER inflight set", reqId); - - // If a request in encoder phase had been flagged to be paused, pause it right away - if (mReqIdsToPause.find(reqId) != mReqIdsToPause.end()) - { - terminateRequest(llmReq, true); - mReqIdsToPause.erase(reqId); - } - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::forwardAsync(RequestList const& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtEncoderModel::ForwardAsync"); - auto const device = mWorldConfig.getDevice(); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - try - { - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - // auto& encoderWaitEvent = mEncoderWaitEvents.at(mMicroBatchId); - - // Get a new set of requests for encoder - // The scheduler will not include any requests that are already in flight for encoder models - // TODO: add pause handling logic - TLLM_LOG_DEBUG("Running ENCODER request scheduler"); - - auto [fittingRequests, fittingDisaggeGenInitReuqests, requestsToPause] = (*mCapacityScheduler)(activeRequests); - - TLLM_CHECK_WITH_INFO( - fittingDisaggeGenInitReuqests.empty(), "Disaggregated servering is not support by encoder model."); - - std::tie(currRequests.contextRequests, std::ignore) = (*mMicroBatchScheduler)( - fittingRequests, mInflightReqIds, getMaxBatchSize(), mModelConfig.getMaxNumTokens()); - - { - NVTX3_SCOPED_RANGE(pauseRequestsFlaggedByScheduler); - // Loop over requests flagged to be paused, and if not in flight pause it right away - for (auto const& llmReq : requestsToPause) - { - auto const reqId = llmReq->mRequestId; - if (mInflightReqIds.find(reqId) == mInflightReqIds.end()) - { - // Not in flight, can terminate right away - terminateRequest(llmReq, true); - } - else - { - // In flight, add to set for pausing later - mReqIdsToPause.insert(reqId); - } - } - } - - TLLM_CHECK(currRequests.size() <= static_cast(getMaxBatchSize())); - - if (!currRequests.empty()) - { - TLLM_LOG_DEBUG("Running ENCODER model with batch size: %u", currRequests.size()); - { - NVTX3_SCOPED_RANGE(updateInflightReqIds); - // Add to set of requests in flight - for (auto const& requests : {currRequests.contextRequests}) - { - for (auto const& llmReq : requests) - { - TLLM_LOG_DEBUG("request ID %u added to ENCODER inflight set", llmReq->mRequestId); - mInflightReqIds.insert(llmReq->mRequestId); - } - } - } - - executeBatch(currRequests); - - sync_check_cuda_error(mRuntime->getStream().get()); - - rearrangeOutputs(currRequests); - - sync_check_cuda_error(mRuntime->getStream().get()); - - // encoderWaitEvent = encoderStepAsync(currRequests); - - for (auto const& requests : {currRequests.contextRequests}) - { - for (auto const& llmReq : requests) - { - if (llmReq->isEncoderInitState()) - { - llmReq->setState(LlmRequestState::kCONTEXT_INIT); - TLLM_LOG_DEBUG("request ID: %u finishes encoder phase", llmReq->mRequestId); - } - } - } - } - - // TODO: PP handling - if (!currRequests.empty()) - { - if (mWorldConfig.isPipelineParallel() && mWorldConfig.isLastPipelineParallelRank()) - { - // TLLM_CHECK_WITH_INFO(mEncStepAsyncSndHdl.get() == nullptr, "decoderSync handle must be nullptr."); - // Wait for encoding for requests in flight for the current micro batch - // mEncStepAsyncSndHdl = encoderSync(currRequests, encoderWaitEvent); - } - } - - // Update the micro batch ID - mMicroBatchId = (mMicroBatchId + 1) % mNumMicroBatches; - } - // In case of error, we need to free the batch slot associated with those requests - catch (std::exception const& e) - { - for (auto const& llmReq : activeRequests) - { - terminateRequest(llmReq); - } - throw; - } - - ++mIterCounter; - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::terminateRequest(std::shared_ptr const& llmReq, bool pause) -{ - // For encoder-only models, just change req state here. might need to do more when using an asynced forward - // For enc-dec models, only remove cross kv cache after decoder - // genenration has finished - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - if (llmReq->isEncoderInitState()) - { - llmReq->setState(LlmRequestState::kCONTEXT_INIT); - } - else - { - TLLM_LOG_DEBUG("Non-encoder request terminated in encoder model: id %lu", llmReq->mRequestId); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::terminateRequestSync( - std::shared_ptr const& llmReq, executor::FinishReason finishReason) -{ - terminateRequest(llmReq, false); - llmReq->finishByReason(finishReason); - llmReq->clearGeneratedTokens(); -} - -void TrtEncoderModel::fillEncoderOutputSync(RequestVector const& requestList, TensorMap outputTensors) -{ - auto const totalTokensNb = outputTensors["encoder_output"]->getShape().d[0]; - auto const encoderOutputDtype = mRuntime->getEngine().getTensorDataType("encoder_output"); - SizeType32 const bytesPerValue = (encoderOutputDtype == nvinfer1::DataType::kFLOAT) ? 4 : 2; - std::vector encoderOutputHost( - totalTokensNb * mHiddenSize * bytesPerValue * mWorldConfig.getTensorParallelism()); - TLLM_CHECK_WITH_INFO(encoderOutputHost.size() > 0, "Encoder output size is 0!"); - getBufferManager().copy(*(outputTensors["encoder_output"]), reinterpret_cast(encoderOutputHost.data())); - getBufferManager().getStream().synchronize(); // TODO: change engine call to async to improve perf. Also - // need to store output buffers, cuda events, etc. - - auto encoderOutputHostPtr = encoderOutputHost.data(); - for (auto const& llmReq : requestList) - { - SizeType32 const seqLen = llmReq->getEncoderOutputLen(); - TensorPtr currentEncoderOutput - = mCopyBufferManager.copyFrom(reinterpret_cast(encoderOutputHostPtr), - ITensor::makeShape({seqLen, mHiddenSize * mWorldConfig.getTensorParallelism()}), MemoryType::kCPU); - llmReq->setEncoderOutputHost(currentEncoderOutput); - encoderOutputHostPtr += seqLen * mHiddenSize * bytesPerValue * mWorldConfig.getTensorParallelism(); - - if (llmReq->isEncoderInitState()) - { - llmReq->setState(LlmRequestState::kCONTEXT_INIT); - } - else - { - TLLM_LOG_DEBUG("Non-encoder request terminated in encoder model: id %lu", llmReq->mRequestId); - } - } -} - -void TrtEncoderModel::executeBatch(RequestVector const& requestList) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeBatch); - - auto const modelName = mModelConfig.getModelName(); - TLLM_CHECK_WITH_INFO(modelName == "EncoderModel" || modelName == "WhisperEncoder", "Model not supported."); - TensorMap inputTensors; - TensorMap outputTensors; - TensorPtr rankOutput; - - std::vector inputIdsHost; - std::vector positionIdsHost; - SizeType32 totalOutputLength = 0; - SizeType32 totalInputLength = 0; - std::vector inputLengthsHost; - std::vector inputFeaturesHost; - - inputLengthsHost.reserve(requestList.size()); - SizeType32 maxInputLengthHost = 0; - - for (auto const& llmReq : requestList) - { - SizeType32 length = 0; - if (mModelConfig.getModelName() == "EncoderModel") - { - auto const& reqTokens = *(llmReq->getEncoderTokens().value()); - length = reqTokens.size(); - - inputIdsHost.insert(inputIdsHost.end(), reqTokens.begin(), reqTokens.end()); - maxInputLengthHost = std::max(maxInputLengthHost, static_cast(length)); - } - else if (mModelConfig.getModelName() == "WhisperEncoder") - { - auto const& reqFeatures = llmReq->getEncoderInputFeatures(); // [length, featureDim] - length = reqFeatures->getShape().d[0]; - - auto const curFeatureBytes = reqFeatures->getSizeInBytes(); - auto const srcPtr = reinterpret_cast(reqFeatures->data()); - inputFeaturesHost.insert(inputFeaturesHost.end(), srcPtr, srcPtr + curFeatureBytes); - } - positionIdsHost.reserve(positionIdsHost.size() + length); - auto const newReqPosBegin = positionIdsHost.end(); - positionIdsHost.resize(positionIdsHost.size() + length); - std::iota(newReqPosBegin, positionIdsHost.end(), 0); - - totalOutputLength += llmReq->getEncoderOutputLen(); - totalInputLength += length; - inputLengthsHost.push_back(length); - } - - TensorPtr hiddenStatesInput; - TensorPtr inputLengths = getBufferManager().copyFrom( - inputLengthsHost, ITensor::makeShape({static_cast(inputLengthsHost.size())}), MemoryType::kGPU); - inputTensors.emplace("input_lengths", inputLengths); - - if (mModelConfig.getModelName() == "EncoderModel") - { - // use shape of maxInputLength to indicates max length, content is not important - TensorPtr maxInputLength - = getBufferManager().gpu(ITensor::makeShape({maxInputLengthHost}), nvinfer1::DataType::kINT32); - inputTensors.emplace("max_input_length", maxInputLength); - } - - // engine outputs - rankOutput = getBufferManager().gpu( - ITensor::makeShape({totalOutputLength, mHiddenSize * mWorldConfig.getTensorParallelism()}), - mModelConfig.getDataType()); - - if (mWorldConfig.isFirstPipelineParallelRank()) - { - if (mModelConfig.getModelName() == "EncoderModel") - { - // Engine inputs - TensorPtr inputIds - = getBufferManager().copyFrom(inputIdsHost, ITensor::makeShape({totalInputLength}), MemoryType::kGPU); - TensorPtr positionIds = getBufferManager().copyFrom( - positionIdsHost, ITensor::makeShape({totalInputLength}), MemoryType::kGPU); - inputTensors.emplace("input_ids", inputIds); - inputTensors.emplace("position_ids", positionIds); - } - else if (mModelConfig.getModelName() == "WhisperEncoder") - { - auto inputFeaturesHostPtr = inputFeaturesHost.data(); - auto const featureDim = requestList.front()->getEncoderInputFeatures()->getShape().d[1]; - auto const dtype = requestList.front()->getEncoderInputFeatures()->getDataType(); - TensorPtr inputFeatures = getBufferManager().gpu(ITensor::makeShape({totalInputLength, featureDim}), dtype); - getBufferManager().copy( - reinterpret_cast(inputFeaturesHostPtr), *inputFeatures, runtime::MemoryType::kCPU); - TensorPtr positionIds = getBufferManager().copyFrom( - positionIdsHost, ITensor::makeShape({totalOutputLength}), MemoryType::kGPU); - inputTensors.emplace("input_features", inputFeatures); - inputTensors.emplace("position_ids", positionIds); - } - } - else - { - SizeType32 length = mModelConfig.getModelName() == "WhisperEncoder" ? totalOutputLength : totalInputLength; - hiddenStatesInput - = getBufferManager().gpu(ITensor::makeShape({length, mHiddenSize * mWorldConfig.getTensorParallelism()}), - mModelConfig.getDataType()); - - inputTensors.emplace("hidden_states_input", hiddenStatesInput); - } - - auto const outputName = mWorldConfig.isLastPipelineParallelRank() ? "encoder_output" : "hidden_states_output"; - outputTensors.emplace(outputName, rankOutput); - - // Set input / output tensors to context, encoder model only have one context - mRuntime->setInputTensors(0, inputTensors); - mRuntime->setOutputTensors(0, outputTensors); - - executeContext(0); - - // copy encoder output to llmRequest, if last PP rank - // dispatch result to each llmReq, only needed by the last PP rank - // TODO: more dtypes support - if (mWorldConfig.isLastPipelineParallelRank()) - { - fillEncoderOutputSync(requestList, outputTensors); - } - else - { - getBufferManager().getStream().synchronize(); - } - - // Update the micro batch ID for next microbatches - mMicroBatchId = (mMicroBatchId + 1) % mWorldConfig.getPipelineParallelism(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::forward(RequestVector& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const device = mWorldConfig.getDevice(); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - try - { - if (activeRequests.empty()) - { - return; - } - - executeBatch(activeRequests); - } - catch (std::exception const& e) - { - for (auto& req : activeRequests) - { - terminateRequest(req); - } - throw; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtEncoderModel::setLogitsPostProcessorBatched( - std::optional logitsPostProcessorBatched) -{ - TLLM_CHECK_WITH_INFO(!logitsPostProcessorBatched.has_value(), "TrtEncoderModel does not use logits processor."); -} - -void TrtEncoderModel::setReplicateLogitsPostProcessor(bool replicateLogitsPostProcessor) -{ - TLLM_THROW("TrtEncoderModel does not use logits processor."); -} - -bool TrtEncoderModel::getReplicateLogitsPostProcessor() const -{ - TLLM_THROW("TrtEncoderModel does not use logits processor."); -} - -TrtEncoderModel::~TrtEncoderModel() = default; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.h b/cpp/tensorrt_llm/batch_manager/trtEncoderModel.h deleted file mode 100644 index 31f7d3d0c89b..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtEncoderModel.h +++ /dev/null @@ -1,205 +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/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "trtGptModel.h" - -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -class NcclCommunicator; -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::batch_manager -{ -class CapacityScheduler; -class MicroBatchScheduler; -class EncoderBuffers; - -class TrtEncoderModel : public TrtGptModel -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TokenIdType = tensorrt_llm::runtime::TokenIdType; - using BufferManager = tensorrt_llm::runtime::BufferManager; - using TensorMap = runtime::StringPtrMap; - using TensorPtr = runtime::ITensor::SharedPtr; - - TrtEncoderModel(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - runtime::RawEngine const& rawEngine, std::shared_ptr logger, - executor::ExecutorConfig const& executorConfig); - - ~TrtEncoderModel() override; - - void terminateRequest(std::shared_ptr const& llmRequest, bool pause = false) override; - void terminateRequestSync( - std::shared_ptr const& llmRequest, executor::FinishReason finishReason) override; - - void forward(RequestVector& activeRequests); - - void forwardSync() override; - - void forwardAsync(RequestList const& activeRequests) override; - - [[nodiscard]] runtime::BufferManager const& getBufferManager() const override; - [[nodiscard]] runtime::BufferManager::CudaStreamPtr getRuntimeStreamPtr() const override; - - runtime::ModelConfig const& getModelConfig() const override - { - return mModelConfig; - } - - [[nodiscard]] bool getGatherGenerationLogits() const override - { - return getModelConfig().computeGenerationLogits(); - } - - runtime::WorldConfig const& getWorldConfig() const override - { - return mWorldConfig; - } - - [[nodiscard]] SizeType32 getHiddenSize() const override - { - return mHiddenSize; - } - - [[nodiscard]] SizeType32 getMaxInputLen() const override - { - return mMaxInputLen; - } - - [[nodiscard]] SizeType32 getNumMicroBatches() const override - { - return mNumMicroBatches; - } - - [[nodiscard]] nvinfer1::DataType getLogitDataType() const override - { - return getModelConfig().getDataType(); - } - - nvinfer1::DataType getTensorDataType(std::string const& name) const override; - nvinfer1::Dims getTensorShape(std::string const& name) const override; - - [[nodiscard]] TrtGptModelType getModelType() const override - { - throw std::runtime_error("TrtEncoderModel does not have model type."); // FIXME: - } - - [[nodiscard]] executor::IterationType getIterCounter() const noexcept override - { - return mIterCounter; - } - - void updatePeftCache(std::shared_ptr const& /*llmRequest*/) override - { - throw std::runtime_error("TrtEncoderModel does not have Peft Cache."); - } - - void getCurrentIterationStats(executor::IterationStats& stats) const override; - void getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const override; - [[nodiscard]] executor::DebugTensorsPerIteration getCurrentDebugTensors() const override; - - void setLayerProfiler() override; - std::string getLayerProfileInfo() const override; - - void setLogitsPostProcessorBatched(std::optional logitsPostProcessorBatched) override; - void setReplicateLogitsPostProcessor(bool replicateLogitsPostProcessor) override; - [[nodiscard]] bool getReplicateLogitsPostProcessor() const override; - - void resetIterationStats() override {} - - [[nodiscard]] SizeType32 getMaxCapacityBatchSize(SizeType32 inputLength, SizeType32 outputLength) const override - { - return 0; - }; - -protected: - std::shared_ptr getKVCacheManager() override - { - throw std::runtime_error("TrtEncoderModel does not have KVCache."); - } - - [[nodiscard]] std::shared_ptr getKVCacheManager() const override - { - throw std::runtime_error("TrtEncoderModel does not have KVCache."); - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() override - { - throw std::runtime_error("TrtEncoderModel does not use PEFT."); - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() const override - { - throw std::runtime_error("TrtEncoderModel does not use PEFT."); - } - -private: - [[nodiscard]] SizeType32 getBufferId() const - { - return mMicroBatchId; - } - - void createRuntimeContexts(); - void executeContext(SizeType32 runtimeContextId); - void createBuffers(); - void executeBatch(RequestVector const& requestList); - void executeBatch(ScheduledRequests const& scheduledRequests); - void rearrangeOutputs(ScheduledRequests const& scheduledRequests); - void createCustomAllReduceWorkspace(); - void fillEncoderOutputSync(RequestVector const& requestList, TensorMap outputTensors); - - runtime::ModelConfig const mModelConfig; - runtime::WorldConfig const mWorldConfig; - int mDevice{-1}; - std::shared_ptr mMpiCommPipelinePara; - - std::shared_ptr mLogger; - std::shared_ptr mRuntime; - - SizeType32 mMicroBatchId{0}; - - // TODO: Add runtime buffers for async PP - std::vector> mBuffers; - - SizeType32 mNumMicroBatches; - SizeType32 mNumBuffers; - - std::vector mMicroBatchScheduledRequests; - ReqIdsSet mInflightReqIds; - ReqIdsSet mReqIdsToPause; - - std::unique_ptr mCapacityScheduler; - std::unique_ptr mMicroBatchScheduler; - - SizeType32 mHiddenSize; // already divided by Tensor Parallelism - SizeType32 mMaxInputLen; // WAR for max_input_len == max_seq_len at all circumstances - - runtime::BufferManager mCopyBufferManager; - - // Iteration counter used to distinguish debug output - executor::IterationType mIterCounter{0}; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModel.h b/cpp/tensorrt_llm/batch_manager/trtGptModel.h deleted file mode 100644 index 54ad36b13895..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModel.h +++ /dev/null @@ -1,339 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2026 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/peftCacheManager.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/stlUtils.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/model.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -namespace tc = tensorrt_llm::common; - -namespace tensorrt_llm::batch_manager -{ -enum class TrtGptModelType -{ - InflightBatching, - InflightFusedBatching -}; - -class LlmRequest; - -namespace kv_cache_manager -{ -class BaseKVCacheManager; -} // namespace kv_cache_manager - -class TrtGptModel : public executor::Model -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - - TrtGptModel(runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - executor::ExecutorConfig const& executorConfig) - : mMaxBatchSize{executorConfig.getMaxBatchSize().value_or(modelConfig.getMaxBatchSize())} - , mMaxBeamWidth{executorConfig.getMaxBeamWidth()} - , mMaxSequenceLen{modelConfig.getMaxSequenceLen()} - , mMaxDraftLen{modelConfig.getMaxDecodingDraftTokens()} - , mVocabSizePadded{modelConfig.getVocabSizePadded(worldConfig.getSize())} - , mNormalizeLogProbs{executorConfig.getNormalizeLogProbs()} - , mEnableTrtOverlap{executorConfig.getEnableTrtOverlap()} - , mCudaGraphMode{executorConfig.getExtendedRuntimePerfKnobConfig().getCudaGraphMode()} - { - TLLM_CHECK_WITH_INFO(mMaxBeamWidth <= modelConfig.getMaxBeamWidth(), - "Runtime configured max beam width (%d) must not exceed engine max beam width (%d)", mMaxBeamWidth, - modelConfig.getMaxBeamWidth()); - TLLM_CHECK_WITH_INFO(mMaxBatchSize <= modelConfig.getMaxBatchSize(), - "Runtime configured max batch size (%d) must not exceed engine max batch size (%d)", mMaxBatchSize, - modelConfig.getMaxBatchSize()); - if (executorConfig.getEnableTrtOverlap()) - { - if (mMaxBeamWidth > 1) - { - mEnableTrtOverlap = false; - TLLM_LOG_WARNING( - "TRT overlap is not supported with beam search (maxBeamWidth is set to %d) and will be disabled.", - mMaxBeamWidth); - } - if (!modelConfig.getSpeculativeDecodingMode().isNone()) - { - mEnableTrtOverlap = false; - TLLM_LOG_WARNING("TRT overlap is not supported with speculative decoding and will be disabled."); - } - } - - mMaxAttentionWindow = 0; - if (executorConfig.getKvCacheConfig().getMaxAttentionWindowVec().has_value()) - { - bool warning = false; - auto const& maxAttentionWindowVec = executorConfig.getKvCacheConfig().getMaxAttentionWindowVec(); - for (int maxAttenWin : maxAttentionWindowVec.value()) - { - mMaxAttentionWindowVec.push_back(std::min(maxAttenWin, mMaxSequenceLen)); - mMaxAttentionWindow = std::max(mMaxAttentionWindow, mMaxAttentionWindowVec.back()); - if (maxAttenWin > mMaxSequenceLen) - { - warning = true; - } - TLLM_CHECK_WITH_INFO(mMaxAttentionWindowVec.back() > 0, - "Attention window sizes (elements in maxAttentionWindowVec) must be > 0"); - } - if (warning) - { - TLLM_LOG_WARNING( - "The value of maxAttentionWindow cannot exceed mMaxSequenceLen. " - "Therefore, it has been adjusted to match the value of mMaxSequenceLen."); - } - } - else - { - mMaxAttentionWindowVec.push_back(mMaxSequenceLen); - mMaxAttentionWindow = mMaxSequenceLen; - } - - mSinkTokenLen = executorConfig.getKvCacheConfig().getSinkTokenLength().has_value() - ? executorConfig.getKvCacheConfig().getSinkTokenLength().value() - : 0; - - mMaxNumSequences = mMaxBatchSize * worldConfig.getPipelineParallelism(); - - auto const numTotalAttenLayers = modelConfig.getNbAttentionLayers(); - auto const numRepeatsAttenWindow = numTotalAttenLayers / mMaxAttentionWindowVec.size(); - auto const numRemainsAttenWindow = numTotalAttenLayers % mMaxAttentionWindowVec.size(); - std::string attenWindowRemainInfo = numRemainsAttenWindow > 0 - ? " + " + tc::arr2str(mMaxAttentionWindowVec.data(), numRemainsAttenWindow) - : ""; - - TLLM_LOG_INFO("TRTGptModel maxNumSequences: %d", mMaxNumSequences); - TLLM_LOG_INFO("TRTGptModel maxBatchSize: %d", mMaxBatchSize); - TLLM_LOG_INFO("TRTGptModel maxBeamWidth: %d", mMaxBeamWidth); - TLLM_LOG_INFO("TRTGptModel maxSequenceLen: %d", mMaxSequenceLen); - TLLM_LOG_INFO("TRTGptModel maxDraftLen: %d", mMaxDraftLen); - TLLM_LOG_INFO("TRTGptModel mMaxAttentionWindowSize: %s * %d%s", tc::vec2str(mMaxAttentionWindowVec).c_str(), - numRepeatsAttenWindow, attenWindowRemainInfo.c_str()); - TLLM_LOG_INFO("TRTGptModel enableTrtOverlap: %d", mEnableTrtOverlap); - TLLM_LOG_INFO("TRTGptModel normalizeLogProbs: %d", mNormalizeLogProbs); - - mMaxNumTokens = modelConfig.getMaxNumTokens(); - if (executorConfig.getMaxNumTokens().has_value() && mMaxNumTokens) - { - if (executorConfig.getMaxNumTokens().value() > mMaxNumTokens.value()) - { - TLLM_LOG_WARNING( - "Runtime configured max num tokens (%d) is larger than model max num tokens (%d) and will be " - "ignored.", - executorConfig.getMaxNumTokens().value(), mMaxNumTokens.value()); - } - else - { - mMaxNumTokens = executorConfig.getMaxNumTokens(); - } - } - if (mMaxNumTokens) - { - TLLM_LOG_INFO("TRTGptModel maxNumTokens: %d", mMaxNumTokens.value()); - } - - if (executorConfig.getEnableChunkedContext()) - { - mMaxInputLen = mMaxSequenceLen - 1; - TLLM_LOG_INFO( - "TRTGptModel maxInputLen: %d = maxSequenceLen - 1 since chunked context is enabled", mMaxInputLen); - TLLM_LOG_INFO( - "TRTGptModel If model type is encoder, maxInputLen would be reset in trtEncoderModel to maxInputLen: " - "%d = maxSequenceLen.", - mMaxSequenceLen); - } - else if (modelConfig.getContextFMHA() && modelConfig.usePackedInput()) - { - TLLM_CHECK_WITH_INFO( - mMaxNumTokens, "Max number of tokens has to be set for context FMHA and usePackedInput case."); - mMaxInputLen = std::min(mMaxSequenceLen - 1, mMaxNumTokens.value()); - TLLM_LOG_INFO( - "TRTGptModel maxInputLen: %d = min(maxSequenceLen - 1, maxNumTokens) since context FMHA " - "and usePackedInput are enabled", - mMaxInputLen); - TLLM_LOG_INFO( - "TRTGptModel If model type is encoder, maxInputLen would be reset in trtEncoderModel to maxInputLen: " - "min(maxSequenceLen, maxNumTokens)."); - } - else - { - mMaxInputLen = modelConfig.getMaxInputLen(); - TLLM_LOG_INFO("TRTGptModel maxInputLen: %d = max_input_len (in trtllm-build args)", mMaxInputLen); - } - - using tensorrt_llm::common::stl_utils::toString; - - TLLM_LOG_INFO("Capacity Scheduler Policy: %s", - toString(executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy()).c_str()); - TLLM_LOG_INFO("Context Chunking Scheduler Policy: %s", - toString(executorConfig.getSchedulerConfig().getContextChunkingPolicy()).c_str()); - } - - [[nodiscard]] std::optional getMaxNumTokens() const - { - return mMaxNumTokens; - } - - [[nodiscard]] SizeType32 getMaxNumSequences() const override - { - return mMaxNumSequences; - } - - [[nodiscard]] SizeType32 getMaxBatchSize() const - { - return mMaxBatchSize; - } - - [[nodiscard]] SizeType32 getMaxInputLen() const override - { - return mMaxInputLen; - } - - [[nodiscard]] SizeType32 getHiddenSize() const override - { - return getModelConfig().getHiddenSize(); - }; - - [[nodiscard]] SizeType32 getMaxSequenceLen() const override - { - return mMaxSequenceLen; - } - - [[nodiscard]] virtual TrtGptModelType getModelType() const = 0; - - [[nodiscard]] SizeType32 getVocabSizePadded() const override - { - return mVocabSizePadded; - } - - [[nodiscard]] SizeType32 getMaxDraftLen() const override - { - return mMaxDraftLen; - } - - [[nodiscard]] SizeType32 getOperatingBeamWidth() const override - { - return mMaxBeamWidth; - } - - [[nodiscard]] bool hasSpeculativeDecodingFastLogits() const noexcept override - { - return false; - } - - [[nodiscard]] bool hasGuidedDecoder() const noexcept override - { - return false; - } - - virtual void setLayerProfiler() = 0; - [[nodiscard]] virtual std::string getLayerProfileInfo() const = 0; - - [[nodiscard]] bool hasKVCacheManager() const - { - return getKVCacheManager() != nullptr; - } - -protected: - [[nodiscard]] SizeType32 getMaxBeamWidth() const - { - return mMaxBeamWidth; - } - - [[nodiscard]] std::vector getMaxAttentionWindowVec() const - { - return mMaxAttentionWindowVec; - } - - [[nodiscard]] SizeType32 getMaxAttentionWindow() const - { - return mMaxAttentionWindow; - } - - [[nodiscard]] SizeType32 getSinkTokenLen() const - { - return mSinkTokenLen; - } - - [[nodiscard]] bool isNormalizeLogProbs() const - { - return mNormalizeLogProbs; - } - - [[nodiscard]] bool isTrtOverlap() const - { - return mEnableTrtOverlap; - } - - [[nodiscard]] bool isCudaGraphMode() const - { - return mCudaGraphMode; - } - - void setMaxAttentionWindowVec(std::vector const& maxAttentionWindowVec) - { - TLLM_CHECK_WITH_INFO(maxAttentionWindowVec.size() == mMaxAttentionWindowVec.size(), - "The size of maxAttentionWindowVec must match the size of mMaxAttentionWindowVec"); - mMaxAttentionWindowVec = maxAttentionWindowVec; - mMaxAttentionWindow = *std::max_element(std::begin(mMaxAttentionWindowVec), std::end(mMaxAttentionWindowVec)); - } - - void setMaxSequenceLen(SizeType32 maxSequenceLen) - { - mMaxSequenceLen = maxSequenceLen; - } - - void setMaxInputLen(SizeType32 maxInputLen) - { - mMaxInputLen = maxInputLen; - } - - [[nodiscard]] std::shared_ptr getKVCacheManager() override = 0; - [[nodiscard]] std::shared_ptr getKVCacheManager() const override = 0; - - [[nodiscard]] virtual std::shared_ptr getPeftCacheManager() = 0; - [[nodiscard]] virtual std::shared_ptr getPeftCacheManager() const = 0; - -private: - std::optional mMaxNumTokens; - SizeType32 mMaxNumSequences; - SizeType32 mMaxBatchSize; - SizeType32 mMaxBeamWidth; - SizeType32 mMaxInputLen; - SizeType32 mMaxSequenceLen; - SizeType32 mMaxDraftLen; - - SizeType32 mVocabSizePadded; - std::vector mMaxAttentionWindowVec; - SizeType32 mMaxAttentionWindow; - SizeType32 mSinkTokenLen; - - bool mNormalizeLogProbs; - bool mEnableTrtOverlap; - bool mCudaGraphMode; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelFactory.h b/cpp/tensorrt_llm/batch_manager/trtGptModelFactory.h deleted file mode 100644 index bd4d7c767378..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelFactory.h +++ /dev/null @@ -1,98 +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/trtGptModel.h" -#include "tensorrt_llm/batch_manager/trtGptModelInflightBatching.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include - -#include -#include - -namespace tensorrt_llm::batch_manager -{ - -class TrtGptModelFactory -{ -public: - using SizeType32 = tensorrt_llm::runtime::SizeType32; - - static std::shared_ptr create(std::filesystem::path const& trtEnginePath, TrtGptModelType modelType, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - { - auto const jsonConfig = runtime::GptJsonConfig::parse(trtEnginePath / "config.json"); - auto const& deviceIds = executorConfig.getParallelConfig().value_or(executor::ParallelConfig()).getDeviceIds(); - auto const worldConfig = getWorldConfig(jsonConfig, deviceIds); - auto const enginePath = trtEnginePath / jsonConfig.engineFilename(worldConfig); - - auto const& modelConfig = jsonConfig.getModelConfig(); - return create( - runtime::RawEngine(enginePath), modelConfig, worldConfig, modelType, executorConfig, isLeaderInOrchMode); - } - - static std::shared_ptr create(std::filesystem::path const& trtEnginePath, TrtGptModelType modelType, - runtime::GptJsonConfig const& jsonConfig, runtime::WorldConfig const& worldConfig, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - { - auto const enginePath = trtEnginePath / jsonConfig.engineFilename(worldConfig); - auto const& modelConfig = jsonConfig.getModelConfig(); - return create( - runtime::RawEngine(enginePath), modelConfig, worldConfig, modelType, executorConfig, isLeaderInOrchMode); - } - - static std::shared_ptr create(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, TrtGptModelType modelType, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - { - auto logger = std::make_shared(); - auto const device = worldConfig.getDevice(); - auto const rank = worldConfig.getRank(); - TLLM_LOG_INFO("Rank %d is using GPU %d", rank, device); - TLLM_CUDA_CHECK(cudaSetDevice(device)); - - if ((modelType == TrtGptModelType::InflightBatching) || (modelType == TrtGptModelType::InflightFusedBatching)) - { - executor::ExecutorConfig const& fixedExecutorConfig - = TrtGptModelInflightBatching::executorConfigIsValid(modelConfig, executorConfig) - ? executorConfig - : TrtGptModelInflightBatching::fixExecutorConfig(modelConfig, executorConfig); - bool const ctxGenFusion = modelType == TrtGptModelType::InflightFusedBatching; - return std::make_shared( - logger, modelConfig, worldConfig, rawEngine, ctxGenFusion, fixedExecutorConfig, isLeaderInOrchMode); - } - - throw std::runtime_error("Invalid modelType in trtGptModelFactory"); - } - -private: - static runtime::WorldConfig getWorldConfig( - runtime::GptJsonConfig const& json, std::optional> const& deviceIds) - { - return runtime::WorldConfig::mpi(json.getGpusPerNode(), json.getTensorParallelism(), - json.getPipelineParallelism(), json.getContextParallelism(), deviceIds); - } -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp deleted file mode 100644 index 05ed827a9511..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.cpp +++ /dev/null @@ -1,3098 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 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 "trtGptModelInflightBatching.h" - -#include "tensorrt_llm/batch_manager/allocateKvCache.h" -#include "tensorrt_llm/batch_manager/assignReqSeqSlots.h" -#include "tensorrt_llm/batch_manager/cacheTransceiver.h" -#include "tensorrt_llm/batch_manager/capacityScheduler.h" -#include "tensorrt_llm/batch_manager/common.h" -#include "tensorrt_llm/batch_manager/contextProgress.h" -#include "tensorrt_llm/batch_manager/createNewDecoderRequests.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/batch_manager/guidedDecoder.h" -#include "tensorrt_llm/batch_manager/handleContextLogits.h" -#include "tensorrt_llm/batch_manager/handleGenerationLogits.h" -#include "tensorrt_llm/batch_manager/kvCacheEventManager.h" -#include "tensorrt_llm/batch_manager/kvCacheManager.h" -#include "tensorrt_llm/batch_manager/llmRequest.h" -#include "tensorrt_llm/batch_manager/logitsPostProcessor.h" -#include "tensorrt_llm/batch_manager/makeDecodingBatchInputOutput.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/batch_manager/promptTuningBuffers.h" -#include "tensorrt_llm/batch_manager/rnnStateManager.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" -#include "tensorrt_llm/batch_manager/sequenceSlotManager.h" -#include "tensorrt_llm/batch_manager/transformerBuffers.h" -#include "tensorrt_llm/batch_manager/updateDecoderBuffers.h" -#include "tensorrt_llm/batch_manager/utils/debugUtils.h" -#include "tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h" -#include "tensorrt_llm/batch_manager/utils/logitsThread.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/timestampUtils.h" -#include "tensorrt_llm/kernels/decodingCommon.h" -#include "tensorrt_llm/layers/defaultDecodingParams.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/gptDecoderBatched.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/ipcUtils.h" -#include "tensorrt_llm/runtime/lookaheadModule.h" -#include "tensorrt_llm/runtime/memoryCounters.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" -#include "tensorrt_llm/runtime/tllmLogger.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/utils/runtimeUtils.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; -namespace tk = tensorrt_llm::kernels; - -using tensorrt_llm::batch_manager::CacheTransceiverFactory; - -namespace tensorrt_llm::batch_manager -{ - -std::map TrtGptModelInflightBatching::calculateCacheSizePerTokenForDisagg( - ModelConfig const& modelConfig, WorldConfig const& worldConfig, - std::vector const& maxAttentionWindowVec, bool isCrossAttention, SizeType32 kvFactor) -{ - // These are the number of attention layers on this PP rank. - auto const numLocalAttnLayers - = modelConfig.getNbAttentionLayers(worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - // These are the number of attention layers on all previous PP ranks. - auto const numLowerRankAttnLayers = modelConfig.countLowerRankLayers(ModelConfig::LayerType::kATTENTION, - worldConfig.getPipelineParallelism(), worldConfig.getPipelineParallelRank()); - // Use global ranks of attention layers to lookup from maxAttentionWindowVec. - auto const startAttnLayerId = numLowerRankAttnLayers; - auto const endAttnLayerId = numLowerRankAttnLayers + numLocalAttnLayers; - auto const numNonUniqueWindowSizes = static_cast(maxAttentionWindowVec.size()); - std::map> uniqueWindowSizeToLayers; - for (SizeType32 layerIdx = startAttnLayerId; layerIdx < endAttnLayerId; layerIdx++) - { - // maxAttentionWindowVec may or may not be stretched to the length of numLayers yet. - // If not stretched yet, we cycle through the window sizes. - auto const windowSize = maxAttentionWindowVec.at(layerIdx % numNonUniqueWindowSizes); - uniqueWindowSizeToLayers[windowSize].push_back(layerIdx); - } - std::map cacheSizeBytesPerTokenPerWindow; - for (auto const& [windowSize, globalLayerIds] : uniqueWindowSizeToLayers) - { - auto const nkvh = modelConfig.getNumKvHeadsForGivenLayers(globalLayerIds, isCrossAttention); - auto const sumLocalHeads = std::reduce(nkvh.cbegin(), nkvh.cend()); - auto const cacheSizePerToken = sumLocalHeads * kvFactor * modelConfig.getSizePerHead(); - auto const cacheSizeBytesPerToken = cacheSizePerToken * BufferDataType(modelConfig.getKvDataType()).getSize(); - cacheSizeBytesPerTokenPerWindow[windowSize] = cacheSizeBytesPerToken; - } - - return cacheSizeBytesPerTokenPerWindow; -}; - -bool TrtGptModelInflightBatching::executorConfigIsValid( - ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig) -{ - // Make sure logic in this function matches fixExecutorConfig - if (executorConfig.getKvCacheConfig().getEnableBlockReuse()) - { - if (!modelConfig.getPagedContextFMHA()) - { - return false; - } - // Context logits cannot be returned for reused tokens, so disable reuse - if (modelConfig.computeContextLogits()) - { - return false; - } - } - return true; -} - -executor::ExecutorConfig TrtGptModelInflightBatching::fixExecutorConfig( - ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig) -{ - // Make sure logic in this function matches executorConfigIsValid - if (executorConfig.getKvCacheConfig().getEnableBlockReuse()) - { - auto kvCacheConfig = executorConfig.getKvCacheConfig(); - - if (!modelConfig.getPagedContextFMHA()) - { - TLLM_LOG_WARNING( - "Fixing executorConfig: KV cache reuse disabled because model was not built with paged context FMHA " - "support"); - kvCacheConfig.setEnableBlockReuse(false); - } - if (modelConfig.computeContextLogits()) - { - TLLM_LOG_WARNING( - "Fixing executorConfig: KV cache reuse disabled because model was built to return context logits"); - kvCacheConfig.setEnableBlockReuse(false); - } - - auto fixedExecutorConfig = executor::ExecutorConfig(executorConfig); - fixedExecutorConfig.setKvCacheConfig(kvCacheConfig); - return fixedExecutorConfig; - } - return executorConfig; -} - -TrtGptModelInflightBatching::TrtGptModelInflightBatching(std::shared_ptr logger, - ModelConfig const& modelConfig, WorldConfig const& worldConfig, RawEngine const& rawEngine, bool ctxGenFusion, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode) - : TrtGptModel(modelConfig, worldConfig, executorConfig) - , mModelConfig(modelConfig) - , mWorldConfig(worldConfig) - , mDevice{runtime::utils::initDevice(worldConfig)} - , mDecodingConfig{executorConfig.getDecodingConfig().value_or(executor::DecodingConfig{})} - , mExtendedRuntimePerfKnobConfig{executorConfig.getExtendedRuntimePerfKnobConfig()} - , mDebugConfig{executorConfig.getDebugConfig()} - , mAdditionalModelOutputs{worldConfig.isLastPipelineParallelRank() ? executorConfig.getAdditionalModelOutputs() - : std::nullopt} - , mLogger{logger ? std::move(logger) : std::make_shared()} - , mRuntime{std::make_unique(rawEngine, mLogger.get(), executorConfig.getUseGpuDirectStorage(), - executorConfig.getGpuWeightsPercent(), modelConfig.useShapeInference())} - , mCopyBufferManager{std::make_shared()} - , mCtxGenFusion(ctxGenFusion) - , mOperatingBeamWidth{getMaxBeamWidth()} - , mGatherGenerationLogits{executorConfig.getGatherGenerationLogits()} - , mPromptTableOffloading{executorConfig.getPromptTableOffloading()} - , mIsLeaderInOrchMode{isLeaderInOrchMode} -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_LOG_INFO("gatherContextLogits: %d", mModelConfig.computeContextLogits()); - TLLM_LOG_INFO("gatherGenerationLogits: %d", getGatherGenerationLogits()); - - if (!(mModelConfig.supportsInflightBatching())) - { - throw std::runtime_error( - "TrtGptModelInflightBatching requires GPT attention/Mamba Conv 1d plugin with " - "packed input and paged KV cache."); - } - if (mWorldConfig.isTensorParallel()) - { - mRuntime->initializeUserBuffer(mWorldConfig, mModelConfig.getMaxBatchSize(), mModelConfig.getMaxBeamWidth(), - mModelConfig.getMaxSequenceLen(), mModelConfig.getHiddenSize(), getMaxNumTokens()); - } - if (mWorldConfig.isPipelineParallel()) - { - mNumMicroBatches = mWorldConfig.getPipelineParallelism(); - } - else - { - mNumMicroBatches = isTrtOverlap() ? 2 : 1; - } - - mNumBuffers = (mCtxGenFusion ? 1 : 2) * mNumMicroBatches; - - auto const& kvCacheConfig = executorConfig.getKvCacheConfig(); - - if (mModelConfig.getSpeculativeDecodingMode().isDraftTokensExternal()) - { - TLLM_CHECK_WITH_INFO(kvCacheConfig.getEnableBlockReuse(), - "KV cache block reuse must be enabled for speculative decoding target model"); - } - - if (mCtxGenFusion) - { - TLLM_CHECK_WITH_INFO(!mModelConfig.isRnnBased(), "RNN based model doesn't support context generation fusion."); - TLLM_CHECK_WITH_INFO( - mModelConfig.isTransformerBased(), "Only transformer based model support context generation fusion now."); - } - - if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - mSeamlessLADMaxDraftLen = modelConfig.getMaxDecodingDraftTokens(); - // TODO: enable it when speculativeDecodingMode is None and run with '--lookahead_config' - mUseSeamlessLookahead = false; - } - - setupSpeculativeDecodingModule(mDecodingConfig); - - if (mWorldConfig.isLastPipelineParallelRank() && executorConfig.getGuidedDecodingConfig()) - { - mGuidedDecoder = std::make_unique(executorConfig.getGuidedDecodingConfig().value(), - getMaxNumSequences(), mModelConfig.getVocabSizePadded(mWorldConfig.getSize()), - mModelConfig.getLogitsDtype(), mRuntime->getBufferManager()); - } - - createRuntimeContexts(); - - if (mWorldConfig.isTensorParallel()) - { - createCustomAllReduceWorkspace(); - } - - if (mModelConfig.isTransformerBased()) - { - createRuntimePerfKnobsTensor(mExtendedRuntimePerfKnobConfig); - } - - auto& memCounter = MemoryCounters::getInstance(); - auto const gpuUsage1 = memCounter.getGpu(); - createBuffers(mDecodingConfig, mAdditionalModelOutputs); - auto const gpuUsage2 = memCounter.getGpu(); - TLLM_LOG_INFO("[MemUsageChange] Allocated %s GPU memory for runtime buffers.", - memCounter.bytesToString(gpuUsage2 - gpuUsage1).c_str()); - - createDecoder(mDecodingConfig.getDecodingMode()); - auto const gpuUsage3 = memCounter.getGpu(); - TLLM_LOG_INFO("[MemUsageChange] Allocated %s GPU memory for decoder.", - memCounter.bytesToString(gpuUsage3 - gpuUsage2).c_str()); - - if (modelConfig.getManageWeightsType() != ModelConfig::ManageWeightsType::kDisabled) - { - mRuntime->loadManagedWeights(rawEngine, worldConfig.getLocalRank()); - } - - if (mModelConfig.useLoraPlugin()) - { - auto const peftCacheManagerConfig - = PeftCacheManagerConfig(executorConfig.getPeftCacheConfig().value_or(executor::PeftCacheConfig())); - mPeftCacheManager = std::make_shared( - peftCacheManagerConfig, mModelConfig, mWorldConfig, mRuntime->getBufferManager()); - } - else - { - mPeftCacheManager = std::make_shared(); - } - - if (mModelConfig.isRnnBased()) - { - createRnnStateManager(); - } - if (mModelConfig.isTransformerBased() && modelConfig.isKVCacheEnabled()) - { - auto cacheTransceiverConfig - = executorConfig.getCacheTransceiverConfig().value_or(executor::CacheTransceiverConfig()); - - auto const cacheSizeBytesPerTokenPerWindow = calculateCacheSizePerTokenForDisagg( - mModelConfig, mWorldConfig, getMaxAttentionWindowVec(), mModelConfig.useCrossAttention(), 2); - auto cacheTransPreAllocaSize = kv_cache_manager::CacheTransBufferManager::preAllocBufferSize( - cacheSizeBytesPerTokenPerWindow, mModelConfig.getTokensPerBlock(), cacheTransceiverConfig); - - auto const [freePrimaryMemBytes, freeSecondaryMemBytes] - = BaseKVCacheManager::calculateFreeMemBytes(mRuntime->getBufferManager(), kvCacheConfig); - if (mModelConfig.useCrossAttention()) - { - TLLM_CHECK_WITH_INFO(kvCacheConfig.getCrossKvCacheFraction().has_value(), - "Must set crossKvCacheFraction for encoder-decoder model"); - auto const crossKvCacheFraction = kvCacheConfig.getCrossKvCacheFraction().value(); - mKvCacheManager = createKvCacheManager(kvCacheConfig, KvCacheType::kSELF, - freePrimaryMemBytes * (1.0f - crossKvCacheFraction), - freeSecondaryMemBytes * (1.0f - crossKvCacheFraction), cacheTransPreAllocaSize, - executorConfig.getFailFastOnAttentionWindowTooLarge()); - mCrossKvCacheManager = createKvCacheManager(kvCacheConfig, KvCacheType::kCROSS, - freePrimaryMemBytes * crossKvCacheFraction, freeSecondaryMemBytes * crossKvCacheFraction, - cacheTransPreAllocaSize, executorConfig.getFailFastOnAttentionWindowTooLarge()); - TLLM_LOG_INFO("This is an Encoder-Decoder model, set %0.1f cross KV cache fraction based on the config.", - crossKvCacheFraction); - } - else - { - TLLM_CHECK_WITH_INFO(!kvCacheConfig.getCrossKvCacheFraction().has_value(), - "Do not set crossKvCacheFraction for decoder-only model"); - mKvCacheManager = createKvCacheManager(kvCacheConfig, KvCacheType::kSELF, freePrimaryMemBytes, - freeSecondaryMemBytes, cacheTransPreAllocaSize, executorConfig.getFailFastOnAttentionWindowTooLarge()); - } - - mCacheTransceiver - = CacheTransceiverFactory::createCacheTransceiver(mKvCacheManager.get(), mModelConfig, mWorldConfig, - executor::kv_cache::CacheState::AttentionType::kDEFAULT, executorConfig.getCacheTransceiverConfig()); - } - - if (mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind()) - { - TLLM_CHECK_WITH_INFO( - mModelConfig.isKVCacheEnabled(), "When needsKVCacheRewind() returns true, KV cache needs to be enabled."); - auto const& blockManager = mKvCacheManager->getBlockManager(); - - TLLM_CHECK_WITH_INFO(blockManager.getNumPools() == 1, - "Rewinding KV cache blocks for models with multiple pools is not supported"); - - // Two "redundant" checks given the pool size check above, but those below don't rely on an implementation - // detail I guess. - TLLM_CHECK_WITH_INFO( - !blockManager.isVariableWindow(), "Rewinding KV cache blocks for variable SWA models isn't supported"); - auto const maxBlocksPerSeq = blockManager.getMaxBlockPerSeqWhenSingleWindowSize(); - - // TODO(oargov): VGQA is not supported, assume all layers have the same num_kv_heads - TLLM_CHECK_WITH_INFO( - !blockManager.isVariableGQA(), "Rewinding KV cache blocks for variable GQA models isn't supported"); - auto const numKvHeads = mModelConfig.getNbKvHeads(0); - - mRewindInputs = RewindInputs{maxBlocksPerSeq, /*isUseOneMoreBlock*/ false, numKvHeads}; - } - - if (mWorldConfig.isPipelineParallel()) - { - mAsyncSendWaitThread = std::make_unique( - "asyncSendWaitThread", - [this]() - { - mDecStepAsyncSndHdls.clear(); - mDecSlotAsyncSndHdls.clear(); - }, - [this]() { TLLM_CUDA_CHECK(cudaSetDevice(mWorldConfig.getDevice())); }); - - auto const& commSession = COMM_SESSION; - mMpiCommPipelinePara = std::make_unique( - commSession.split(mWorldConfig.getTensorParallelRank(), mWorldConfig.getPipelineParallelRank())); - mDecSlotAsyncSndHdls.reserve(getMaxBatchSize()); - } - if (mWorldConfig.isTensorParallel()) - { - auto const& commSession = COMM_SESSION; - mMpiCommTensorPara = std::make_unique( - commSession.split(mWorldConfig.getPipelineParallelRank(), mWorldConfig.getTensorParallelRank())); - } - - mSeqSlotManager - = std::make_shared(getMaxNumSequences(), executorConfig.getMaxSeqIdleMicroseconds()); - - mMicroBatchScheduledRequests.resize(mNumMicroBatches); - mDecoderFinishedEvents.resize(mNumMicroBatches); - mPeftTables.resize(mNumMicroBatches); - - if (modelConfig.isRnnBased()) - { - TLLM_CHECK_WITH_INFO(modelConfig.getMaxBeamWidth() == 1, "RNN based model doesn't support beam search now."); - TLLM_CHECK_WITH_INFO( - !executorConfig.getEnableChunkedContext(), "RNN based model doesn't support Chunked Context now."); - TLLM_CHECK_WITH_INFO( - modelConfig.getSpeculativeDecodingMode().isNone(), "RNN based model doesn't support speculative decoding."); - } - - std::optional ctxChunkConfig; - if (executorConfig.getEnableChunkedContext()) - { - TLLM_CHECK_WITH_INFO(modelConfig.isKVCacheEnabled() && mModelConfig.getPagedContextFMHA(), - "Chunked context requires context FMHA, paged kv_cache and paged context FMHA all enabled at the same " - "time."); - SizeType32 chunkUnitSize = mKvCacheManager->getTokensPerBlock(); - // If sliding window attention is used, then make sure the unit size aligns with the paged context fmha's kv - // step size. - if (getMaxInputLen() > getMaxAttentionWindow()) // TODO(nhaber): minAttentionWindow - { - chunkUnitSize = std::max(/* maxKvStepSizeInFmha */ 256, chunkUnitSize); - TLLM_LOG_INFO("ChunkUnitSize is set to %d as sliding window attention is used.", chunkUnitSize); - } - ctxChunkConfig = batch_scheduler::ContextChunkingConfig{ - executorConfig.getSchedulerConfig().getContextChunkingPolicy().value_or( - executor::ContextChunkingPolicy::kFIRST_COME_FIRST_SERVED), - chunkUnitSize}; - } - - auto maxNumTokens = getMaxNumTokens(); - TLLM_CHECK_WITH_INFO(maxNumTokens, "Max number of tokens is not set in model config."); - - // Max context size is limited by `max_num_tokens` for chunked-context or context-FMHA, - // or by `max_input_len` of the model. - auto const maxContextLength = (executorConfig.getEnableChunkedContext() || mModelConfig.getContextFMHA()) - ? maxNumTokens - : std::make_optional(mModelConfig.getMaxInputLen()); - - mMaxBatchSizeTunerRecommended = 0; - mMaxBatchSizeRuntime = getMaxBatchSize(); - mMaxNumTokensStatic = maxNumTokens; - mMaxNumTokensTunerRecommended = 0; - mMaxNumTokensRuntime = maxNumTokens; - - if (mKvCacheManager && ctxChunkConfig) - { - TLLM_CHECK_WITH_INFO(ctxChunkConfig.value().chunkUnitSize % mKvCacheManager->getTokensPerBlock() == 0, - "To prevent cache fragmentation, the context chunk unit size (%d) should be divisible by the number of " - "tokens per kv-cache block (%d).", - ctxChunkConfig.value().chunkUnitSize, mKvCacheManager->getTokensPerBlock()); - } - - mCapacityScheduler = std::make_unique(getMaxNumSequences(), - executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy(), mKvCacheManager != nullptr, - mWorldConfig.isPipelineParallel()); - - mMicroBatchScheduler = std::make_unique(ctxChunkConfig, maxContextLength); - - if (ctxChunkConfig) - { - if (maxContextLength) - { - ctxChunkConfig.value().chunkUnitSize - = std::min(ctxChunkConfig.value().chunkUnitSize, maxContextLength.value()); - } - TLLM_CHECK_WITH_INFO(ctxChunkConfig.value().chunkUnitSize > 0, - "Context chunk size (%d) must be a positive integer.", maxContextLength.value()); - } - else - { - if (maxContextLength && maxNumTokens) - { - TLLM_CHECK_WITH_INFO(maxContextLength.value() <= maxNumTokens.value(), - "Without enabling chunked context, the max context length (%d) needs to be less than or equal to the " - "max number of tokens (%d).", - maxContextLength.value(), maxNumTokens.value()); - } - } - - mPauseRequests = std::make_unique(getMaxInputLen()); - mAssignReqSeqSlots = std::make_unique(); - mAllocateKvCache = std::make_unique(); - - if (isCudaGraphMode()) - { - // Limit cuda graph cache size. Depending on the model one graph is 4-10MB of GPU memory. - SizeType32 cudaGraphCacheSize - = std::min(getMaxBatchSize(), std::max(mExtendedRuntimePerfKnobConfig.getCudaGraphCacheSize(), 1)); - // We can't have common cache for all microbatches as cuda graph is tied to the memory pointers of the runtime - // buffers. - mCudaGraphExecutorCaches.resize(mNumBuffers, utils::CudaGraphExecutorCache(cudaGraphCacheSize)); - } - - mSpeculativeDecodingFastLogits - = executorConfig.getSpecDecConfig().has_value() && executorConfig.getSpecDecConfig()->fastLogits; - if (mSpeculativeDecodingFastLogits && modelConfig.getSpeculativeDecodingMode().isNone() && mIsLeaderInOrchMode) - { - mDraftModelSendLogitsThread - = std::make_unique(&utils::draftModelSendLogitsThread, mDevice, &mDraftModelThreadShouldExit, - &mDraftRequestsWaitingToSendLogits, &mDraftRequestsDoneSendingLogits, &mDraftRequestsMtx); - } - - mCreateNewDecoderRequests = std::make_unique( - mSpeculativeDecodingFastLogits, mIsLeaderInOrchMode, isNormalizeLogProbs()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -TrtGptModelInflightBatching::~TrtGptModelInflightBatching() -{ - if (mCacheTransceiver) - { - mCacheTransceiver->checkContextTransferStatus(1, true); - TLLM_CHECK_WITH_INFO(mCacheTransceiver->checkGenTransferComplete(), "Generation transfer not complete"); - } - if (mAsyncSendWaitThread) - { - mAsyncSendWaitThread.reset(nullptr); - } - if (mDraftModelSendLogitsThread) - { - mDraftModelThreadShouldExit = true; - mDraftModelSendLogitsThread->join(); - mDraftModelSendLogitsThread.reset(nullptr); - } -} - -void TrtGptModelInflightBatching::setupSpeculativeDecodingModule(executor::DecodingConfig const& decodingConfig) -{ - if (mModelConfig.getSpeculativeDecodingMode().isExplicitDraftTokens() - || mModelConfig.getSpeculativeDecodingMode().isEagle()) - { - TLLM_CHECK_WITH_INFO(mCtxGenFusion, "Current speculative decoding mode requires context-gen fusion IFB"); - } - - if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() && decodingConfig.getLookaheadDecodingConfig()) - { - // FIXME choose defaults - auto maxLookaheadConfig = decodingConfig.getLookaheadDecodingConfig().value(); - - SizeType32 maxDraftTokens{0}; - SizeType32 maxDraftPathLen{0}; - std::tie(std::ignore, std::ignore, maxDraftTokens, maxDraftPathLen) - = maxLookaheadConfig.calculateSpeculativeResource(); - TLLM_CHECK(maxDraftTokens <= mModelConfig.getMaxDecodingDraftTokens()); - mModelConfig.getSpeculativeDecodingModulePtr()->setMaxDraftTokens(maxDraftTokens); - mModelConfig.getSpeculativeDecodingModulePtr()->setMaxDraftPathLen(maxDraftPathLen); - - auto lookaheadModulePtr - = std::dynamic_pointer_cast(mModelConfig.getSpeculativeDecodingModulePtr()); - lookaheadModulePtr->setExecutionConfig(maxLookaheadConfig); - } -} - -void TrtGptModelInflightBatching::reshapeKvTensors(OffsetTableDimensions const& dims) -{ - TLLM_CHECK(mBuffers.size() == static_cast(mNumBuffers)); - auto const& manager = mRuntime->getBufferManager(); - for (auto& buffers : mBuffers) - { - TLLM_CHECK(buffers->transformerBuffers); - // any method that operates on transformerBuffers must distinguish between self and cross cache, because - // transformerBuffers is not managed by KVCacheManager same rule applies to kv pool pointers below - buffers->transformerBuffers->reshapeKvTensors( - getMaxBatchSize(), mOperatingBeamWidth, dims.maxBlocksPerSeq, dims.cacheType, dims.numPools, manager); - } -} - -using BlocksPerWindow = std::map>; - -std::pair> -TrtGptModelInflightBatching::clampWindowSizesToFitAtLeastOneSequence( - BlocksPerWindow const& blocksPerWindow, bool const failFastOnAttentionWindowTooLarge) -{ - // At this point, we can only validate that the cheapest sequence in terms of kv-cache resources still fits. More - // validation is needed on a per-request basis, once the prompt / output lengths and the actual beam width are - // known. - auto const promptLength = getMaxInputLen(); - auto const outputLength - = getMaxSequenceLen() - promptLength; // This makes it the best case scenario, as context tokens are 'cheaper' - // in terms of kv-cache resources on average. - auto const sinkTokenLength = getSinkTokenLen(); - auto const maxBeamWidth = getMaxBeamWidth(); - auto const tokensPerBlock = mModelConfig.getTokensPerBlock(); - auto const& oldMaxAttentionWindowVec = getMaxAttentionWindowVec(); - std::vector newMaxAttentionWindowVec; - BlocksPerWindow newBlocksPerWindow; - - newMaxAttentionWindowVec.reserve(oldMaxAttentionWindowVec.size()); - for (auto const windowSize : oldMaxAttentionWindowVec) - { - auto const bestCaseBlockRequirements = kv_cache_manager::KVCacheManager::calculateMaxBlockRequirements( - promptLength, outputLength, sinkTokenLength, windowSize, maxBeamWidth, tokensPerBlock); - auto const [numPrimaryBlocks, numSecondaryBlocks] = blocksPerWindow.at(windowSize); - if (bestCaseBlockRequirements > numPrimaryBlocks) - { - auto const newMaxAttentionWindow = KVCacheManager::calculateMaxAttentionWindow( - promptLength, outputLength, sinkTokenLength, numPrimaryBlocks, maxBeamWidth, tokensPerBlock); - newMaxAttentionWindowVec.push_back(newMaxAttentionWindow); - newBlocksPerWindow[newMaxAttentionWindow] = std::make_tuple(numPrimaryBlocks, numSecondaryBlocks); - } - else - { - newMaxAttentionWindowVec.push_back(windowSize); - newBlocksPerWindow[windowSize] = std::make_tuple(numPrimaryBlocks, numSecondaryBlocks); - } - } - if (newMaxAttentionWindowVec == getMaxAttentionWindowVec()) - { - return {blocksPerWindow, newMaxAttentionWindowVec}; - } - TLLM_LOG_WARNING("maxAttentionWindowVec too large to fit at least one sequence in kvCache. Old: %s, New: %s", - common::vec2str(getMaxAttentionWindowVec()).c_str(), common::vec2str(newMaxAttentionWindowVec).c_str()); - - if (failFastOnAttentionWindowTooLarge) - { - throw std::runtime_error( - "Attention window too large to fit even a single sequence in the KV cache. Failing fast rather than " - "attempting an adjustment of the window sizes. " - "Old: " - + common::vec2str(getMaxAttentionWindowVec()) + ", New: " + common::vec2str(newMaxAttentionWindowVec)); - } - - setMaxAttentionWindowVec(newMaxAttentionWindowVec); - if (getMaxSequenceLen() > getMaxAttentionWindow()) - { - TLLM_LOG_WARNING("maxSequenceLen is reduced to maxAttentionWindow: %d", getMaxAttentionWindow()); - setMaxSequenceLen(getMaxAttentionWindow()); - if (getMaxInputLen() > getMaxSequenceLen() - 1) - { - setMaxInputLen(getMaxSequenceLen() - 1); - TLLM_LOG_WARNING("maxInputLen is reduced to %d", getMaxInputLen()); - } - } - // createBuffers depends on: - // maxAttentionWindow; maxAttentionWindowVec; maxSequenceLen; - // TODO: This is problematic, as createBuffers edits the state of trtGptModelInflightBatching, but - // what if there are different window values for cross+self etc. in encoder+decoder scenario... - createBuffers(mDecodingConfig, mAdditionalModelOutputs); - createDecoder(mDecodingConfig.getDecodingMode()); - return {newBlocksPerWindow, newMaxAttentionWindowVec}; -} - -std::unique_ptr TrtGptModelInflightBatching::createKvCacheManager( - KvCacheConfig const& kvCacheConfig, KvCacheType kvCacheType, uint64_t freePrimaryMemBytes, - uint64_t freeSecondaryMemBytes, size_t extraCostMemory, bool const failFastOnAttentionWindowTooLarge) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - bool isCrossAttention = kvCacheType == KvCacheType::kCROSS; - TLLM_CHECK_WITH_INFO( - mModelConfig.isTransformerBased(), "KvCacheManager is only needed by transformer based model."); - - auto const tokensPerBlock = mModelConfig.getTokensPerBlock(); - auto const kvDtype = mModelConfig.getKvDataType(); - - // init KV cache block manager - auto [numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd] = mModelConfig.getNumKvHeadsPerLayerLocalRange( - mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank(), isCrossAttention); - auto numKvHeadsPerLayer = std::vector(numKvHeadsPerLayerBegin, numKvHeadsPerLayerEnd); - - auto maxAttentionWindowVec = getMaxAttentionWindowVec(); - if (kvCacheType != KvCacheType::kSELF) // TODO(nhaber): more foolproof way of initing cross-kvcache-manager - { - maxAttentionWindowVec = std::vector{mModelConfig.getMaxEncoderLen()}; - } - - auto const numLayers = static_cast(numKvHeadsPerLayer.size()); - auto const windowSizeToLayers = KVCacheManager::groupLayersByWindowSize(maxAttentionWindowVec, numLayers); - auto const sizePerHead = mModelConfig.getSizePerHead(); - auto blocksPerWindow = KVCacheManager::calculateMaxNumBlocks(kvCacheConfig, kvDtype, numKvHeadsPerLayer, - sizePerHead, tokensPerBlock, mWorldConfig, windowSizeToLayers, freePrimaryMemBytes, freeSecondaryMemBytes, - extraCostMemory, 2, getMaxBatchSize()); - - // now we check if any of the window sizes is too large for at least one sequence to fit in kvCache - // this can happen if e.g. maxSeqLen is deduced from the model and is too large - // and user also didn't provide maxAttentionWindow, which leads it to be equal to maxSeqLen - if (kvCacheType == KvCacheType::kSELF) - { - std::tie(blocksPerWindow, maxAttentionWindowVec) - = clampWindowSizesToFitAtLeastOneSequence(blocksPerWindow, failFastOnAttentionWindowTooLarge); - } - - if (kvCacheType == KvCacheType::kCROSS && kvCacheConfig.getEnableBlockReuse()) - { - TLLM_LOG_INFO( - "Cross KV cache does not support reuse because cross attention depends on encoder and decoder input ids. " - "Thus, KV cache reuse is disabled for cross KV cache."); - } - auto const enableBlockReuse = kvCacheType == KvCacheType::kSELF ? kvCacheConfig.getEnableBlockReuse() : false; - - auto kvCacheManager = std::make_unique(numKvHeadsPerLayer, sizePerHead, tokensPerBlock, - blocksPerWindow, getMaxNumSequences(), getMaxBeamWidth(), maxAttentionWindowVec, kvDtype, getSinkTokenLen(), - mRuntime->getStreamPtr(), - kvCacheType == KvCacheType::kCROSS ? mModelConfig.getMaxEncoderLen() : getMaxSequenceLen(), - getMaxNumTokens().value(), enableBlockReuse, kvCacheType, kvCacheConfig.getSecondaryOffloadMinPriority(), - kvCacheConfig.getEventBufferMaxSize() > 0 - ? std::make_unique(kvCacheConfig.getEventBufferMaxSize()) - : nullptr, - kvCacheConfig.getEnablePartialReuse(), kvCacheConfig.getCopyOnPartialReuse()); - - reshapeKvTensors(kvCacheManager->getOffsetTableDimensions()); - - kvCacheManager->allocatePools(kvCacheConfig.getUseUvm()); - - TensorMap inputBuffers; - TensorPtr poolPointers = kvCacheManager->getBlockPoolPointers(); - TensorPtr poolMapping = kvCacheManager->getLayerToPoolMapping(); - - if (kvCacheType == KvCacheType::kSELF) - { - inputBuffers.insert_or_assign("host_kv_cache_pool_pointers", std::move(poolPointers)); - inputBuffers.insert_or_assign("host_kv_cache_pool_mapping", std::move(poolMapping)); - } - else - { - inputBuffers.insert_or_assign("host_cross_kv_cache_pool_pointers", std::move(poolPointers)); - inputBuffers.insert_or_assign("host_cross_kv_cache_pool_mapping", std::move(poolMapping)); - } - mRuntime->setStaticInputTensors(inputBuffers); - - // Emit the `created` event - kvCacheManager->flushIterationEvents(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return kvCacheManager; -} - -void TrtGptModelInflightBatching::createRnnStateManager() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO(mModelConfig.isRnnBased(), "RnnStateManager is only needed by RNN based model."); - - mRnnStateManager = std::make_unique( - getMaxNumSequences(), mModelConfig, mWorldConfig, mRuntime->getBufferManager()); - - TensorMap inputBuffers; - mRnnStateManager->getPtrBuffers(inputBuffers, mModelConfig, mWorldConfig); - mRuntime->setStaticInputTensors(inputBuffers); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createCustomAllReduceWorkspace() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(mWorldConfig.isTensorParallel()); - - auto const& manager = mRuntime->getBufferManager(); - auto const hiddenSize = mModelConfig.getHiddenSize(); - - mAllReduceBuffers = std::make_unique(getMaxBatchSize(), getMaxBeamWidth(), getMaxSequenceLen(), - hiddenSize, manager, mWorldConfig, mRuntime->isUserBufferEnabled()); - - TensorMap inputBuffers; - inputBuffers.insert_or_assign("all_reduce_workspace", mAllReduceBuffers->mAllReduceCommPtrs); - mRuntime->setStaticInputTensors(inputBuffers); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createRuntimePerfKnobsTensor( - executor::ExtendedRuntimePerfKnobConfig const& extendedRuntimePerfKnobConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - SizeType32 constexpr perfKnobSize{16}; - mExtendedRuntimePerfKnobsHost = BufferManager::cpu(ITensor::makeShape({perfKnobSize}), nvinfer1::DataType::kINT64); - auto* runtimePerfKnobsHostPtr = bufferCast(*mExtendedRuntimePerfKnobsHost); - std::fill_n(runtimePerfKnobsHostPtr, perfKnobSize, -1); - SizeType32 multiBlockModeVal = extendedRuntimePerfKnobConfig.getMultiBlockMode() ? 1 : 0; - SizeType32 enableContextFMHAFP32AccVal = extendedRuntimePerfKnobConfig.getEnableContextFMHAFP32Acc() ? 1 : 0; - runtimePerfKnobsHostPtr[0] = multiBlockModeVal; - runtimePerfKnobsHostPtr[1] = enableContextFMHAFP32AccVal; - - TensorMap inputBuffers; - inputBuffers.insert_or_assign("host_runtime_perf_knobs", mExtendedRuntimePerfKnobsHost); - mRuntime->setStaticInputTensors(inputBuffers); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::terminateRequest(LlmRequestPtr const& llmReq, bool pause) -{ - utils::terminateRequest( - *mSeqSlotManager, *llmReq, getMaxInputLen(), mKvCacheManager, mCrossKvCacheManager, mPeftCacheManager, pause); -} - -void TrtGptModelInflightBatching::terminateRequestSync( - LlmRequestPtr const& llmRequest, executor::FinishReason finishReason) -{ - TLLM_LOG_DEBUG("Registering termination for request %lu with finish reason %d", llmRequest->mRequestId, - static_cast(finishReason)); - mReqIdsToTerminate.try_emplace(llmRequest->mRequestId, finishReason); -} - -TrtGptModelInflightBatching::IterationStatsIFB TrtGptModelInflightBatching::fillIterationStats( - ScheduledRequests const& scheduledRequests, RequestVector const& requestsToPause) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(fillIterationStats); - - IterationStatsIFB iterationStatsIfb{mMicroBatchId}; - iterationStatsIfb.numCtxRequests = scheduledRequests.contextRequests.size(); - iterationStatsIfb.numGenRequests = scheduledRequests.generationRequests.size(); - iterationStatsIfb.avgNumDecodedTokensPerIter = 0; - - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - auto const& buffers = mBuffers.at(contextBufferId); - iterationStatsIfb.numCtxTokens = buffers->getNumContextTokens(); - - for (auto const& llmReq : scheduledRequests.contextRequests) - { - iterationStatsIfb.scheduledRequests.insert(llmReq->mRequestId); - } - for (auto const& llmReq : scheduledRequests.generationRequests) - { - iterationStatsIfb.scheduledRequests.insert(llmReq->mRequestId); - iterationStatsIfb.avgNumDecodedTokensPerIter += llmReq->getAvgDecodedTokensPerIter(); - } - if (iterationStatsIfb.numGenRequests > 0) - { - iterationStatsIfb.avgNumDecodedTokensPerIter /= iterationStatsIfb.numGenRequests; - TLLM_LOG_DEBUG( - "iterationStatsIfb.avgNumDecodedTokensPerIter = %.2f", iterationStatsIfb.avgNumDecodedTokensPerIter); - } - for (auto const& llmReq : requestsToPause) - { - iterationStatsIfb.pausedRequests.insert(llmReq->mRequestId); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return iterationStatsIfb; -} - -void TrtGptModelInflightBatching::forwardSync() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtGptModelInflightBatching::forwardSync"); - - TLLM_CUDA_CHECK(cudaSetDevice(mWorldConfig.getDevice())); - - if (!mWorldConfig.isLastPipelineParallelRank()) - { - mAsyncSendWaitThread->waitStop(); - } - - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - - if (!currRequests.empty()) - { - if (!mWorldConfig.isPipelineParallel() || !mWorldConfig.isLastPipelineParallelRank()) - { - for (auto& hdl : mDecStepAsyncSndHdls) - { - TLLM_CHECK_WITH_INFO(hdl.get() == nullptr, "decoderSync handle must be nullptr."); - } - // Wait for decoding for requests in flight for the current micro batch - auto& decoderWaitEvent = mDecoderFinishedEvents.at(mMicroBatchId); - mDecStepAsyncSndHdls = decoderSync(currRequests, decoderWaitEvent); - decoderWaitEvent.reset(); - - if (!mWorldConfig.isLastPipelineParallelRank()) - { - mAsyncSendWaitThread->notifyStart(); - } - } - else - { - for (auto const& requests : {currRequests.contextRequests, currRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - for (SizeType32 beam = 0; beam < llmReq->mSamplingConfig.beamWidth; ++beam) - { - llmReq->setNumPreDecodedTokens(0, beam); - } - if (llmReq->isGenerationToCompleteState()) - { - llmReq->setState(LlmRequestState::kGENERATION_COMPLETE); - terminateRequest(llmReq); - } - } - } - } - - (*mPauseRequests)(currRequests.generationRequests, mInflightReqIds, mReqIdsToPause, true, *mSeqSlotManager, - mKvCacheManager, mCrossKvCacheManager, mPeftCacheManager); - - if (!mReqIdsToTerminate.empty()) - { - for (auto const& requests : {currRequests.contextRequests, currRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - if (mReqIdsToTerminate.count(llmReq->mRequestId) != 0U) - { - if (!llmReq->isGenerationCompleteState()) - { - TLLM_LOG_DEBUG("Terminating request %lu with finish reason %d", llmReq->mRequestId, - static_cast(mReqIdsToTerminate[llmReq->mRequestId])); - terminateRequest(llmReq); - llmReq->finishByReason(mReqIdsToTerminate[llmReq->mRequestId]); - llmReq->clearGeneratedTokens(); - } - mReqIdsToTerminate.erase(llmReq->mRequestId); - } - } - } - } - - // Terminate draft requests whose logits have been sent by the background thread. - { - RequestVector doneSending; - { - std::lock_guard lk(mDraftRequestsMtx); - doneSending.swap(mDraftRequestsDoneSendingLogits); - } - for (auto const& llmReq : doneSending) - { - terminateRequest(llmReq); - } - } - - // Finished context requests have been moved to generationRequests by moveFinishedContextRequestsToGeneration - for (auto const& llmReq : currRequests.generationRequests) - { - // If a context-only request is finished, send its KV cache and mark it. - if (llmReq->isContextOnlyRequest() && llmReq->isContextFinished()) - { - // TODO: skip if sending layer-wise - { - TLLM_CHECK_WITH_INFO(mCacheTransceiver, - "Disaggregated serving is not enabled, please check the configuration of " - "cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendAsync(llmReq); - } - mSeqSlotManager->freeSequenceSlot(llmReq->mRequestId); - } - } - } - // report profile data - auto const bufferId = getFusedBufferId(); - auto const contextId = mBuffers[bufferId]->getContextIndex(); - if (mRuntime->hasLayerProfiler(contextId)) - { - mRuntime->reportToProfiler(contextId); - } - if (mCacheTransceiver) - { - mCacheTransceiver->checkContextTransferStatus(0, true); - } - ++mIterCounter; - - if (mKvCacheManager) - { - mKvCacheManager->flushIterationEvents(); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::storeContextBlocks(std::shared_ptr const& llmReq) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // TMJ - Note - // Make context blocks reusable immediately after context phase finishes. - // For chunked contexts, this occurs in step that processes last context chunk. - // isLastContextChunk() is always true for non-chunked contexts. - // This check is made in code that calls storeContextBlocks, so omitted here. - if (mKvCacheManager) - { - mKvCacheManager->storeContextBlocks(*llmReq); - } - if (mCrossKvCacheManager) - { - mCrossKvCacheManager->storeContextBlocks(*llmReq); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::storeNewBlock(std::shared_ptr const& llmReq) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // TMJ - Note - // Make context blocks reusable immediately after each generation step. - - if (mKvCacheManager) - { - mKvCacheManager->storeNewBlock(*llmReq); - } - if (mCrossKvCacheManager) - { - mCrossKvCacheManager->storeNewBlock(*llmReq); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::resetIterationStats() -{ - mLastIterationStatsIFB = IterationStatsIFB{mMicroBatchId}; -} - -void TrtGptModelInflightBatching::forwardAsync(RequestList const& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "TrtGptModelInflightBatching::forwardAsync"); - - TLLM_CUDA_CHECK(cudaSetDevice(mWorldConfig.getDevice())); - - try - { - verifyRequests(activeRequests); - if (mModelConfig.isTransformerBased() && getKVCacheManager() && mCacheTransceiver) - { - checkDisaggGenTransferStatus(activeRequests); - } - auto& currRequests = mMicroBatchScheduledRequests.at(mMicroBatchId); - - // Get a new set of requests for that context - // The scheduler will not include any requests that are (i) still in encoder state if encoder-decoder models OR - // (ii) already in flight for decoder models - TLLM_LOG_DEBUG("Running DECODER request scheduler"); - auto [fittingRequests, fittingDisaggGenInitRequests, requestsToPause] - = (*mCapacityScheduler)(activeRequests, mKvCacheManager, mPeftCacheManager, mCrossKvCacheManager); - // Remove from fitting requests the requests that cannot be scheduled due to disagg KV cache transfer - if (mModelConfig.isTransformerBased() && getKVCacheManager() && mCacheTransceiver) - { - prepareDisaggGenInitRequests(activeRequests, fittingDisaggGenInitRequests); - } - if (fittingRequests.empty() && fittingDisaggGenInitRequests.empty()) - { - TLLM_LOG_WARNING( - "CapacityScheduler didn't schedule any requests in iteration %lu, " - "probably because of insufficient resources such as KV cache, " - "will try wait for KV cache transfer to complete", - mIterCounter); - if (mCacheTransceiver) - { - mCacheTransceiver->checkContextTransferStatus(1, true); - // will free kvCache in next iteration. - } - } - std::tie(currRequests.contextRequests, currRequests.generationRequests) - = (*mMicroBatchScheduler)(fittingRequests, mInflightReqIds, mMaxBatchSizeRuntime, mMaxNumTokensRuntime); - TLLM_CHECK(currRequests.size() <= static_cast(getMaxBatchSize())); - - (*mPauseRequests)(requestsToPause, mInflightReqIds, mReqIdsToPause, false, *mSeqSlotManager, mKvCacheManager, - mCrossKvCacheManager, mPeftCacheManager); - - if (mUseSeamlessLookahead) - { - changeSpecDecMode(currRequests); - } - - if (!currRequests.empty()) - { - TLLM_LOG_DEBUG("Running DECODER model with batch size: %lu", currRequests.size()); - // For overlap don't store inflight requests, so they are not skipped in scheduler - if (!isTrtOverlap()) - { - NVTX3_SCOPED_RANGE(updateInflightReqIds); - // Add requests to in-flight set, so they can be skipped in other micro batches - for (auto const& llmReq : currRequests.contextRequests) - { - // Context requests that are chunking are not added to inflight set, so they are scheduled in the - // next micro batch. - if (llmReq->isLastContextChunk()) - { - TLLM_LOG_DEBUG( - "Context request with ID %lu added to DECODER model inflight set", llmReq->mRequestId); - mInflightReqIds.insert(llmReq->mRequestId); - } - } - for (auto const& llmReq : currRequests.generationRequests) - { - TLLM_LOG_DEBUG( - "Generation request with ID %lu added to DECODER model inflight set", llmReq->mRequestId); - mInflightReqIds.insert(llmReq->mRequestId); - } - } - - (*mAssignReqSeqSlots)(*mSeqSlotManager, currRequests.contextRequests, currRequests.generationRequests); - - if (mKvCacheManager) - { - (*mAllocateKvCache)(*mKvCacheManager, currRequests.contextRequests, currRequests.generationRequests, - mModelConfig, mCrossKvCacheManager); - } - - mPeftTables.at(mMicroBatchId) - = mPeftCacheManager->ensureBatch(currRequests.contextRequests, currRequests.generationRequests, true); - - // Do decoder setup before context phase if model needs to setup buffers for the context phase. - if (mModelConfig.getSpeculativeDecodingMode().needsDecoderPrologue()) - { - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - setupDecoderStep(currRequests.contextRequests, *mBuffers.at(contextBufferId), - mDecoderInputBuffers.at(getFusedBufferId())); - // WAR: Sync to ensure that the decoder setup is complete before the context phase starts. - // Without this, there may be a race condition between the decoder setup and the context phase - // which also leads to spurious test failure in trtGptModelRealDecoderTest. - mRuntime->getStream().synchronize(); - } - else - { - prepareDistGenBufferAndDecoder(currRequests.generationRequests); - } - sync_check_cuda_error(mRuntime->getStream().get()); - - executeBatch(currRequests); - if (mWorldConfig.isLastPipelineParallelRank() && mGuidedDecoder) - { - // XGrammar: build maskcache for context requests and perform maskgen for all requests - // These need to be overlapped with the kernel execution of forward step - mGuidedDecoder->build(currRequests); - } - - sync_check_cuda_error(mRuntime->getStream().get()); - - // Postpone decoder setup if model does not need to setup buffers for the context phase. - if (!mModelConfig.getSpeculativeDecodingMode().needsDecoderPrologue()) - { - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - setupDecoderStep(currRequests.contextRequests, *mBuffers.at(contextBufferId), - mDecoderInputBuffers.at(getFusedBufferId())); - } - - sync_check_cuda_error(mRuntime->getStream().get()); - - if (isTrtOverlap()) - { - // WAR: Because the decoder is not stateless (yet) a sync is needed between - // decoder execution and next decoder step preparation. - auto const prevMicroBatchId = getPrevMicroBatchId(mMicroBatchId); - auto& prevDecoderFinishedEvent = mDecoderFinishedEvents.at(prevMicroBatchId); - if (prevDecoderFinishedEvent) - { - prevDecoderFinishedEvent->synchronize(); - } - } - - auto& decoderFinishedEvent = mDecoderFinishedEvents.at(mMicroBatchId); - TLLM_CHECK_WITH_INFO(!decoderFinishedEvent.has_value(), "decoderFinishedEvent must be nullopt."); - decoderFinishedEvent = mWorldConfig.isLastPipelineParallelRank() - ? std::make_optional(decoderStepAsync(currRequests)) - : std::nullopt; - - sync_check_cuda_error(mRuntime->getStream().get()); - - mLastIterationStatsIFB = fillIterationStats(currRequests, requestsToPause); - for (auto const& requests : {currRequests.contextRequests, currRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - if (llmReq->isContextInitState()) - { - llmReq->moveToNextContextChunk(); - if (llmReq->getContextRemainingLength() == 0) - { - TLLM_LOG_DEBUG("[RANK %d] request with ID %lu finishes decoder ctx phase", - COMM_SESSION.getRank(), llmReq->mRequestId); - - llmReq->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - - // for encoder-decoder models, free encoder output buffers after decoder context phase is - // completed - if (llmReq->getEncoderTokens().has_value()) - { - llmReq->freeEncoderOutputBuffers(); - } - storeContextBlocks(llmReq); - - if (isTrtOverlap() && llmReq->willCompleteNextIteration()) - { - // This prohibits the request from being scheduled for another iteration if only one - // iteration is expected. - llmReq->setState(LlmRequestState::kGENERATION_TO_COMPLETE); - } - } - } - else if (llmReq->isGenerationInProgressState()) - { - storeNewBlock(llmReq); - TLLM_LOG_DEBUG("request with ID %lu forwards a step in decoder gen phase", llmReq->mRequestId); - } - } - } - - utils::moveFinishedContextRequestsToGeneration(currRequests); - } - else - { - mLastIterationStatsIFB = IterationStatsIFB{mMicroBatchId}; - } - - if (mWorldConfig.isPipelineParallel() && mWorldConfig.isLastPipelineParallelRank()) - { - mAsyncSendWaitThread->waitStop(); - if (!currRequests.empty()) - { - for (auto& hdl : mDecStepAsyncSndHdls) - { - TLLM_CHECK_WITH_INFO(hdl.get() == nullptr, "decoderSync handle must be nullptr."); - } - // Wait for decoding for requests in flight for the current micro batch - auto& decoderFinishedEvent = mDecoderFinishedEvents.at(mMicroBatchId); - mDecStepAsyncSndHdls = decoderSync(currRequests, decoderFinishedEvent); - decoderFinishedEvent.reset(); - - mAsyncSendWaitThread->notifyStart(); - } - } - - // Update the micro batch ID - mMicroBatchId = getNextMicroBatchId(mMicroBatchId); - } - // In case of error, we need to free the batch slot associated with those requests - catch (std::exception const&) - { - try - { - for (auto const& llmReq : activeRequests) - { - terminateRequest(llmReq); - } - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR("forwardAsync catch-all catch block that runs `terminateRequest` has failed with:"); - TLLM_LOG_EXCEPTION(e); - TLLM_LOG_ERROR("Rethrowing *outer* exception:"); - } - throw; - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::setRuntimeBatchSize(SizeType32 runtimeMaxBatchSize) -{ - mMaxBatchSizeTunerRecommended = runtimeMaxBatchSize; - mMaxBatchSizeRuntime = std::min(getMaxBatchSize(), runtimeMaxBatchSize); -} - -SizeType32 TrtGptModelInflightBatching::getRuntimeBatchSize() const -{ - return mMaxBatchSizeRuntime; -} - -void TrtGptModelInflightBatching::setRuntimeMaxNumTokens(SizeType32 runtimeMaxNumTokens) -{ - mMaxNumTokensTunerRecommended = runtimeMaxNumTokens; - mMaxNumTokensRuntime - = (mMaxNumTokensStatic) ? std::min(mMaxNumTokensStatic.value(), runtimeMaxNumTokens) : runtimeMaxNumTokens; -} - -void TrtGptModelInflightBatching::updatePeftCache(std::shared_ptr const& llmRequest) -{ - mPeftCacheManager->addRequestPeft(llmRequest, true); -} - -runtime::BufferManager const& TrtGptModelInflightBatching::getBufferManager() const -{ - return mRuntime->getBufferManager(); -} - -BufferManager::CudaStreamPtr TrtGptModelInflightBatching::getRuntimeStreamPtr() const -{ - return mRuntime->getStreamPtr(); -} - -void TrtGptModelInflightBatching::executeContext(SizeType32 runtimeContextId, SizeType32 bufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeContext); - - auto const& currBatchState = mBuffers[bufferId]->getBatchState(); - - bool hasCudaGraph = false; - // If batch state is context only, do not capture/launch graph and execute the engine as is. - if (isCudaGraphMode() && !currBatchState.isAnyContext()) - { - auto cudaGraphOpt = mCudaGraphExecutorCaches[bufferId].get(currBatchState); - // If graph exists for current batch state, launch it. - if (cudaGraphOpt.has_value()) - { - hasCudaGraph = true; - } - } - - // If there is no graph for current state, execute the engine. - if (!hasCudaGraph) - { - auto enqueueSuccessful = mRuntime->executeContext(runtimeContextId); - if (!enqueueSuccessful) - { - throw std::runtime_error("Executing TRT engine failed!"); - } - } - else - { - // Launch graph. - auto cudaGraphOpt = mCudaGraphExecutorCaches[bufferId].get(currBatchState); - cudaGraphOpt.value()->launch(mRuntime->getStream()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::setLayerProfiler() -{ - mRuntime->setLayerProfiler(); -} - -std::string TrtGptModelInflightBatching::getLayerProfileInfo() const -{ - return mRuntime->getLayerProfileInfo(); -} - -void TrtGptModelInflightBatching::verifyRequests(RequestList const& activeRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(verifyRequests); - - if (activeRequests.empty()) - { - return; - } - - auto const& firstRequest = activeRequests.front(); - auto const firstRequestId = firstRequest->mRequestId; - auto const firstBeamWidth = firstRequest->mSamplingConfig.beamWidth; - - for (auto const& llmReq : activeRequests) - { - auto const beamWidth = llmReq->mSamplingConfig.beamWidth; - auto const draftLength = llmReq->getNumDraftTokens(); - auto const maxDraftLength = mModelConfig.getMaxDecodingDraftTokens(); - - TLLM_CHECK_WITH_INFO(beamWidth == 1 || draftLength == 0, "Can't use speculative decoding with beam search."); - TLLM_CHECK_WITH_INFO(draftLength <= maxDraftLength, - "Number of draft tokens (%d) is larger than maximum number of draft tokens (%d)", draftLength, - maxDraftLength); - - // FIXME: Remove this check when varying beam width is supported - { - TLLM_CHECK_WITH_INFO(beamWidth == firstBeamWidth, - "All active requests must have same beam width, " - "but request %lu with beam width %d differs from first request %lu with beam width %d", - llmReq->mRequestId, beamWidth, firstRequestId, firstBeamWidth); - } - } - - if (firstBeamWidth != mOperatingBeamWidth) - { - changeBeamWidth(firstBeamWidth); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::executeBatch(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(executeBatch); - - if (!mCtxGenFusion) - { - if (!scheduledRequests.contextRequests.empty()) - { - auto const bufferId = getContextBufferId(); - executeStep(scheduledRequests.contextRequests, {}, bufferId); - } - if (!scheduledRequests.generationRequests.empty()) - { - auto const bufferId = getGenerationBufferId(); - executeStep({}, scheduledRequests.generationRequests, bufferId); - } - } - else - { - auto const bufferId = getFusedBufferId(); - executeStep(scheduledRequests.contextRequests, scheduledRequests.generationRequests, bufferId); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createRuntimeContexts() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRuntime->clearContexts(); - auto const numProfiles = mRuntime->getNbProfiles(); - for (auto i = 0; i < numProfiles; ++i) - { - mRuntime->addContext(i); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ -// TODO: move this somewhere else? -/** - * This function logic is also implemented in tensorrt_llm/_torch/pyexecutor/_util.py get_decoding_mode(). - */ -executor::DecodingMode getDecodingMode(SpeculativeDecodingMode specDecodingMode, - std::optional const& decodingModeOpt, runtime::SizeType32 const beamWidth) -{ - auto getDefaultDecodingMode = [beamWidth](std::optional const& decodingModeOpt) - { - if (decodingModeOpt.has_value() && !decodingModeOpt->isAuto()) - { - return decodingModeOpt.value(); - } - return (beamWidth == 1) ? executor::DecodingMode::TopKTopP() : executor::DecodingMode::BeamSearch(); - }; - - auto decodingMode = getDefaultDecodingMode(decodingModeOpt); - // Variable-Beam-Width-Search (special mode of Beam-Search) is enabled. - if (decodingMode.isBeamSearch() && decodingMode.isUseVariableBeamWidthSearch()) - { - TLLM_LOG_INFO("Variable-Beam-Width-Search is enabled"); - } - // Overwrite decoding mode when beam width is one. - if (beamWidth == 1 && decodingMode.isBeamSearch()) - { - TLLM_LOG_WARNING( - "Beam width is set to 1, but decoding mode is BeamSearch. Overwriting decoding mode to TopKTopP."); - decodingMode = executor::DecodingMode::TopKTopP(); - } - // Overwrite decoding mode when Medusa is used. - if (specDecodingMode.isMedusa() && !decodingMode.isMedusa()) - { - TLLM_LOG_WARNING("Model is Medusa, but decoding mode is not Medusa. Overwriting decoding mode to Medusa."); - decodingMode = executor::DecodingMode::Medusa(); - } - // Overwrite decoding mode when Medusa is not used. - if (!specDecodingMode.isMedusa() && decodingMode.isMedusa()) - { - TLLM_LOG_WARNING("Model is not Medusa, but decoding mode is Medusa. Overwriting decoding mode."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - // Overwrite decoding mode when lookahead decoding is used. - if (specDecodingMode.isLookaheadDecoding() && !decodingMode.isLookahead()) - { - TLLM_LOG_WARNING( - "Model is Lookahead, but decoding mode is not Lookahead. Overwriting decoding mode to Lookahead."); - decodingMode = executor::DecodingMode::Lookahead(); - } - // Overwrite decoding mode when lookahead decoding is not used. - if (!specDecodingMode.isLookaheadDecoding() && decodingMode.isLookahead()) - { - TLLM_LOG_WARNING( - "Model is not built with Lookahead decoding, but decoding mode is Lookahead. Overwriting decoding " - "mode."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - // Overwrite decoding mode when 'explicit draft tokens' is used. - if (specDecodingMode.isExplicitDraftTokens() && !decodingMode.isExplicitDraftTokens()) - { - TLLM_LOG_WARNING( - "Model is built with 'explicit draft tokens' decoding, but decoding mode is something else. Overwriting " - "decoding mode."); - decodingMode = executor::DecodingMode::ExplicitDraftTokens(); - } - // Overwrite decoding mode when 'explicit draft tokens' is not used. - if (!specDecodingMode.isExplicitDraftTokens() && decodingMode.isExplicitDraftTokens()) - { - TLLM_LOG_WARNING( - "Model is not built with 'explicit draft tokens' decoding, but decoding mode is set to it. Overwriting " - "decoding " - "mode to default."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - // Overwrite decoding mode when EAGLE is used. - if (specDecodingMode.isEagle() && !decodingMode.isEagle()) - { - TLLM_LOG_WARNING("Model is Eagle, but decoding mode is not Eagle. Overwriting decoding mode to Eagle."); - decodingMode = executor::DecodingMode::Eagle(); - } - // Overwrite decoding mode when Eagle is not used. - if (!specDecodingMode.isEagle() && decodingMode.isEagle()) - { - TLLM_LOG_WARNING("Model is not Eagle, but decoding mode is Eagle. Overwriting decoding mode."); - decodingMode = getDefaultDecodingMode(decodingModeOpt); - } - if (specDecodingMode.isDraftTokensExternal()) - { - TLLM_LOG_WARNING("Overwriting decoding mode to external draft token"); - decodingMode = executor::DecodingMode::ExternalDraftTokens(); - } - TLLM_LOG_DEBUG("DecodingMode: %s", decodingMode.getName()); - return decodingMode; -} -} // namespace - -void TrtGptModelInflightBatching::createDecoder(std::optional const& decodingModeOpt) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mDecoderState = std::make_unique(); - - if (mWorldConfig.isLastPipelineParallelRank()) - { - auto decoderType = mRuntime->getEngine().getTensorDataType("logits"); - - auto const decodingMode - = getDecodingMode(mModelConfig.getSpeculativeDecodingMode(), decodingModeOpt, mOperatingBeamWidth); - - if (decodingMode.isExplicitDraftTokens()) - { - // There are no logits in Explicit draft tokens model. - decoderType = mModelConfig.getDataType(); - // Decoder is not instantiated for bf16. We use half to get the same data size - // and explicitly pass dtype to redrafter that has bf16 kernels. - if (decoderType == nvinfer1::DataType::kBF16) - { - decoderType = nvinfer1::DataType::kHALF; - } - } - - mDecoder = std::make_unique(mRuntime->getStreamPtr()); - mDecoder->setup( - decodingMode, getMaxNumSequences(), mOperatingBeamWidth, decoderType, mModelConfig, mWorldConfig); - - mDecoderState->setup(getMaxNumSequences(), mOperatingBeamWidth, getMaxAttentionWindow(), getSinkTokenLen(), - getMaxSequenceLen(), decoderType, mModelConfig, mWorldConfig, mRuntime->getBufferManager()); - - if (!mModelConfig.getSpeculativeDecodingMode().isNone()) - { - mDecoderState->setupSpeculativeDecoding(mModelConfig.getSpeculativeDecodingMode(), - mModelConfig.getMaxDecodingTokens(), decoderType, mModelConfig, mWorldConfig, - mRuntime->getBufferManager()); - } - } - else - { - mDecoderState->setupCacheIndirection( - getMaxNumSequences(), mOperatingBeamWidth, getMaxAttentionWindow(), mRuntime->getBufferManager()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::createBuffers(executor::DecodingConfig const& decodingConfig, - std::optional> const& additionalModelOutputs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mBuffers.clear(); - for (SizeType32 i = 0; i < mNumBuffers; ++i) - { - mBuffers.emplace_back( - std::make_unique(getMaxBatchSize(), mOperatingBeamWidth, getMaxAttentionWindowVec(), - getMaxAttentionWindow(), getSinkTokenLen(), *mRuntime, mModelConfig, mWorldConfig, decodingConfig, - getGatherGenerationLogits(), getMaxNumTokens(), additionalModelOutputs, mPromptTableOffloading)); - } - - mDecoderInputBuffers.clear(); - mDecoderOutputBuffers.clear(); - for (SizeType32 i = 0; i < mNumMicroBatches; ++i) - { - mDecoderInputBuffers.emplace_back( - getMaxBatchSize(), mModelConfig.getMaxDecodingTokens(), mRuntime->getBufferManager()); - mDecoderInputBuffers.back().setupMedusaLogits(getMaxNumSequences(), mModelConfig); - mDecoderOutputBuffers.emplace_back(getMaxNumSequences(), mOperatingBeamWidth, getMaxSequenceLen(), - mModelConfig.getMaxDecodingTokens(), mRuntime->getBufferManager()); - mDecoderOutputBuffers.back().setupSpeculativeDecoding( - getMaxNumSequences(), mModelConfig.getMaxDecodingTokens(), mModelConfig); - } - - mSlotDecoderBuffers.clear(); - for (SizeType32 i = 0; i < getMaxNumSequences(); ++i) - { - mSlotDecoderBuffers.emplace_back(std::make_unique( - mOperatingBeamWidth, getMaxSequenceLen(), mRuntime->getBufferManager())); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::prepareDisaggGenInitRequests( - RequestList const& activeRequests, RequestVector& newGenReqs) -{ - NVTX3_SCOPED_RANGE(prepareDisaggGenInitRequests); - - // Allocate KV cache by treating them as context requests - (*mAllocateKvCache)(*mKvCacheManager, newGenReqs, {}, mModelConfig, mCrossKvCacheManager); - - // Initiate KV cache transfer - auto timeStart = std::chrono::steady_clock::now(); - - if (tc::getEnvDisaggBenchmarkGenOnly()) - { - TLLM_LOG_DEBUG("Disaggregated generation only benchmark mode is enabled"); - for (auto& req : newGenReqs) - { - req->setState(LlmRequestState::kDISAGG_GENERATION_TRANS_COMPLETE); - } - return; - } - - auto const genInitReqNum = std::count_if(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationInitState(); }); - - // Loop over the new disagg gen requests and trigger receive of KV cache - for (auto& newGenReq : newGenReqs) - { - TLLM_CHECK_WITH_INFO( - mCacheTransceiver, "Disaggregated serving is not enabled, please check the configuration."); - if (common::getEnvDisableKVCacheTransferOverlap()) - { - mCacheTransceiver->requestAndReceiveSync(newGenReq); - } - else - { - mCacheTransceiver->requestAndReceiveAsync(newGenReq); - } - } - if (!common::getEnvDisableKVCacheTransferOverlap()) - { - auto const blockTransfer = std::all_of(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "newGenReqs.size():%ld requests, activeRequests.size():%ld checkGenTransferStatus :%d original " - "gen_only_requests_num:%ld", - newGenReqs.size(), activeRequests.size(), blockTransfer, genInitReqNum); - mCacheTransceiver->checkGenTransferStatus(blockTransfer ? 1 : 0); - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "receiveDisaggGenCache time:%f ms, " - "blockTransfer:%d,genInitReqNum:%ld,newGenReqs.size():%ld,activeRequests.size():%ld", - duration, blockTransfer, genInitReqNum, newGenReqs.size(), activeRequests.size()); - } - - return; -} - -void TrtGptModelInflightBatching::checkDisaggGenTransferStatus(RequestList const& activeRequests) -{ - NVTX3_SCOPED_RANGE(checkDisaggGenTransferStatus); - - if (common::getEnvDisableKVCacheTransferOverlap()) - { - return; - } - - auto timeStart = std::chrono::steady_clock::now(); - - // TODO: - auto const needCheck = std::any_of(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - - if (needCheck) - { - auto const needCheckOne = std::all_of(activeRequests.begin(), activeRequests.end(), - [](auto const& req) { return req->isDisaggGenerationTransmissionInProgress(); }); - - int atLeastNum = needCheckOne ? 1 : 0; - TLLM_LOG_DEBUG( - mpi::MpiComm::world().getRank(), "noPreppared requests, checkGenTransferStatus atLeastNum:%d", atLeastNum); - - mCacheTransceiver->checkGenTransferStatus(atLeastNum); - - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "no Prepare checkDisaggGenTransferStatus time:%f ms, " - "needCheckOne:%d,needCheck:%ld,activeRequests.size():%ld", - duration, needCheckOne, needCheck, activeRequests.size()); - } -} - -void TrtGptModelInflightBatching::prepareDistGenBufferAndDecoder(RequestVector const& generationRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // set decoderStep for disagg_generation - RequestVector cacheTransCompleteRequests; - for (auto const& request : generationRequests) - { - if (request->isDisaggGenerationTransmissionComplete()) - { - cacheTransCompleteRequests.push_back((request)); - } - } - if (!cacheTransCompleteRequests.empty()) - { - auto timeStart = std::chrono::steady_clock::now(); - auto const bufferId = getFusedBufferId(); - auto& runtimeBuffers = *mBuffers[bufferId]; - runtimeBuffers.prepareStep(cacheTransCompleteRequests, {}, getMaxBeamWidth(), getMaxAttentionWindow(), - *mDecoderState, mKvCacheManager.get(), mCrossKvCacheManager.get(), mRnnStateManager.get(), - mPeftTables[mMicroBatchId], *mRuntime, mModelConfig, mWorldConfig, getGatherGenerationLogits(), - isTrtOverlap()); - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - setupDecoderStep( - cacheTransCompleteRequests, *mBuffers.at(contextBufferId), mDecoderInputBuffers.at(getFusedBufferId())); - sync_check_cuda_error(mRuntime->getStream().get()); - auto timeEnd = std::chrono::steady_clock::now(); - auto duration = std::chrono::duration(timeEnd - timeStart).count(); - TLLM_LOG_DEBUG(mpi::MpiComm::world().getRank(), - "prepareDistGenBufferAndDecoder time:%f ms , cacheTransCompleteRequests.size():%ld", duration, - cacheTransCompleteRequests.size()); - } - for (auto& request : cacheTransCompleteRequests) - { - request->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - request->setContextCurrentPosition(request->mPromptLen); - request->setDecodingIter(1); - auto const reqBeamWidth = request->mSamplingConfig.beamWidth; - auto firstGenTokens = request->getContextPhaseParams().value().getFirstGenTokens(); - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - request->addNewToken(firstGenTokens.at(beam), beam); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::debugIOTensors(RequestVector const& contextRequests, - RequestVector const& generationRequests, TensorMap const& inputMap, TensorMap const& outputMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CHECK(mDebugConfig); - - auto const& manager = mRuntime->getBufferManager(); - auto requestIds = utils::collectRequestIds(contextRequests, generationRequests); - - if (mDebugConfig->getDebugTensorsMaxIterations() > 0) - { - mLastIterationDebugTensors.clear(); - mLastIterationDebugTensors = utils::storeIOTensors(*mDebugConfig, requestIds, inputMap, outputMap, manager); - } - else - { - utils::dumpIOTensors(*mDebugConfig, mIterCounter, requestIds, inputMap, outputMap, mWorldConfig, manager); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::tuple const&, runtime::StringPtrMap&> -TrtGptModelInflightBatching::prepareBuffers( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(prepareBuffers); - - auto& runtimeBuffers = *mBuffers.at(bufferId); - - auto allNewTokens = mWorldConfig.isLastPipelineParallelRank() - ? RuntimeBuffers::OptionalRef(mDecoderState->getAllNewTokens()) - : std::nullopt; - - auto [optProfileId, inputMap, outputMap] = runtimeBuffers.prepareStep(contextRequests, generationRequests, - mOperatingBeamWidth, getMaxAttentionWindow(), *mDecoderState, mKvCacheManager.get(), mCrossKvCacheManager.get(), - mRnnStateManager.get(), mPeftTables[bufferId], *mRuntime, mModelConfig, mWorldConfig, - getGatherGenerationLogits(), isTrtOverlap(), allNewTokens); - - // For Variable-Beam-Width-Search - mRuntime->setCurrentBeamWidths( - tensorrt_llm::batch_manager::utils::getRequestBeamWidths(contextRequests, generationRequests)); - - mRuntime->setInputTensors(optProfileId, inputMap); - mRuntime->setOutputTensors(optProfileId, outputMap); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {optProfileId, inputMap, outputMap}; -} - -void TrtGptModelInflightBatching::prepareGraph(SizeType32 bufferId, SizeType32 optProfileId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(prepareGraph); - - auto const nextBatchState = mBuffers[bufferId]->getBatchState(); - auto cudaGraphOpt = mCudaGraphExecutorCaches[bufferId].get(nextBatchState); - // If graph is not found in the cache, capture it. - if (!cudaGraphOpt.has_value()) - { - // We need to prepare some tensors once again to properly set values for graph capture. - // Graph capture requires setting some tensors (e.g. past_kv_len) - // to the round_up(max_kv_cache_len, kKV_CACHE_LEN_CUDA_GRAPH_ROUND_SIZE) - // in order to capture the kernels with the large enough grid. - mBuffers[bufferId]->prepareBuffersForCudaGraph(getMaxSequenceLen()); - - auto cudaGraph = std::make_shared(); - cudaGraph->prepareNextGraph(mRuntime, optProfileId); - mCudaGraphExecutorCaches[bufferId].put(nextBatchState, cudaGraph); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::executeStep( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, - "executeStep: " + std::to_string(contextRequests.size()) + " ctx reqs, " - + std::to_string(generationRequests.size()) + " gen reqs"); - - if (mPromptTableOffloading) - { - prefetchNextPromptTableChunk(contextRequests, /* isFirstChunk */ true, bufferId); - } - - auto [optProfileId, inputMap, outputMap] = prepareBuffers(contextRequests, generationRequests, bufferId); - - if (mBuffers[bufferId]->transformerBuffers) - { - // Creation of context progress, or remains nullptr if not needed - std::shared_ptr progress = nullptr; - RequestVector layerWiseRequests; - if (common::getEnvDisaggLayerwise()) - { - for (auto const& request : contextRequests) - { - bool const enableLayerWise = request->isContextOnlyRequest() && request->isLastContextChunk(); - if (enableLayerWise) - { - layerWiseRequests.push_back(request); - } - } - } - // TODO: support layer-wise cross kv cache in encoder-decoder models - if (!layerWiseRequests.empty() && !mModelConfig.useCrossAttention()) - { - int const numLayers = mModelConfig.getNbAttentionLayers( - mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank()); - progress = std::make_shared(numLayers); - } - bufferCast(*mBuffers[bufferId]->transformerBuffers->contextProgressHost)[0] = progress.get(); - if (progress) - { - TLLM_CHECK_WITH_INFO(mCacheTransceiver, - "Disaggregated serving is not enabled, please check the configuration of cacheTransceiverConfig."); - mCacheTransceiver->respondAndSendLayerWise(layerWiseRequests, progress); - } - } - - if (mPromptTableOffloading) - { - prefetchNextPromptTableChunk(contextRequests, /* isFirstChunk */ false, bufferId); - } - - executeContext(optProfileId, bufferId); - - // If batch state has any context request, do not capture this graph. - if (isCudaGraphMode() && contextRequests.empty()) - { - // Capture graph of current batch state during engine execution. - // This is based on the assumptions that - // a) We can hide CPU graph capture behind the GPU engine execution. - // b) Batch size in the next iterations won't change and we can reuse the graph multiple times. - prepareGraph(bufferId, optProfileId); - } - - if (mDebugConfig) - { - debugIOTensors(contextRequests, generationRequests, inputMap, outputMap); - } - - if (mAdditionalModelOutputs.has_value() && !mAdditionalModelOutputs.value().empty()) - { - utils::copyAdditionalOutputs( - mAdditionalModelOutputs.value(), contextRequests, generationRequests, outputMap, getBufferManager()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::setupDecoderStep( - RequestVector const& contextRequests, RuntimeBuffers const& buffers, DecoderInputBuffers& inputBuffers) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(setupDecoderStep); - - if (mWorldConfig.isLastPipelineParallelRank() && !contextRequests.empty()) - { - auto const logitsType = mRuntime->getEngine().getTensorDataType("logits"); - - auto [batchSlots, samplingConfigs, lookaheadPrompt, lookaheadAlgoConfigs] - = (*mCreateNewDecoderRequests)(mModelConfig, mWorldConfig, mDecodingConfig, contextRequests, logitsType, - inputBuffers, *mDecoderState, mRuntime->getStream(), *mDecoder->getDecoderStream(), getMaxSequenceLen(), - mOperatingBeamWidth, buffers.mMedusaBuffers); - - auto const localBatchSize = batchSlots->getSize(); - if (localBatchSize > 0) - { - auto samplingConfig = SamplingConfig(samplingConfigs); - mDecoder->getUnderlyingDecoder().setup(samplingConfig, localBatchSize, batchSlots, - {mDecoderState->getJointDecodingOutput()}, mModelConfig.getDataType(), lookaheadPrompt, - lookaheadAlgoConfigs); - - auto const& stream = mDecoder->getDecoderStream(); - CudaEvent event{}; - stream->record(event); - mRuntime->getStreamPtr()->wait(event); - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::postProcessRequest( - LlmRequest& llmReq, std::vector const& numDroppedTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const seqSlot = llmReq.mSeqSlot.value(); - auto const reqBeamWidth = llmReq.getBeamWidthByIter(true); - auto const& bufferManager = getBufferManager(); - - if (llmReq.getReturnGenerationLogits() && !llmReq.getGenerationLogitsFragments().empty()) - { - TLLM_CHECK(!llmReq.isStreaming()); - auto const genBufferId = mCtxGenFusion ? getFusedBufferId() : getGenerationBufferId(); - auto& genRuntimeBuffers = *mBuffers.at(genBufferId); - - auto constexpr beforeDecoder = false; - utils::copyGenerationLogits( - genRuntimeBuffers.generationLogitsCache, bufferManager, llmReq, beforeDecoder, numDroppedTokens); - - bufferManager.getStream().synchronize(); - } - - if (reqBeamWidth == 1) - { - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return; - } - - // Update mDecoderBuffers->slotOutputIdsHost and synchronize - getDecoderSlotHostOutputs(seqSlot, llmReq.returnLogProbs(), llmReq.mSamplingConfig, llmReq.isStreaming()); - - auto const* outputIdsHostData = bufferCast(*mSlotDecoderBuffers[seqSlot]->outputIdsHost); - auto const* sequenceLengthsHostData = bufferCast(*mSlotDecoderBuffers[seqSlot]->sequenceLengthsHost); - auto const* cumLogProbsHostData = bufferCast(*mSlotDecoderBuffers[seqSlot]->cumLogProbsHost); - auto logProbsHost = mSlotDecoderBuffers[seqSlot]->logProbsHost; - auto const* logProbsHostData = bufferCast(*logProbsHost); - - auto const& outputIdsShape = mSlotDecoderBuffers[seqSlot]->outputIdsHost->getShape(); - auto const maxSeqLength = outputIdsShape.d[1]; - - std::vector> generatedTokens(reqBeamWidth); - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - auto const* const begin = outputIdsHostData + tc::flat_index2(beam, llmReq.mPromptLen, maxSeqLength); - auto const generatedLength = sequenceLengthsHostData[beam] - llmReq.mPromptLen; - auto const* const end = begin + generatedLength; - generatedTokens[beam].assign(begin, end); - - if (llmReq.returnLogProbs()) - { - llmReq.setCumLogProb(cumLogProbsHostData[beam], beam); - - auto const beginLogProbsOffset = reqBeamWidth == 1 ? llmReq.mPromptLen : 0; - auto const* const begin = logProbsHostData + beam * logProbsHost->getShape().d[1] + beginLogProbsOffset; - auto const* const end = begin + generatedLength; - LlmRequest::VecLogProbs logProbs(begin, end); - llmReq.setLogProbs(logProbs, beam); - } - } - - // store the generated tokens into the mTokensGathered buffer - llmReq.setGeneratedTokens(generatedTokens); - - if (llmReq.getReturnGenerationLogits() && llmReq.getGenerationLogitsHost() - && mWorldConfig.isLastPipelineParallelRank()) - { - reorderGenerationLogitsForBeamSearch( - llmReq, seqSlot, reqBeamWidth, maxSeqLength, outputIdsHostData, sequenceLengthsHostData); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::reorderGenerationLogitsForBeamSearch(LlmRequest& llmReq, SizeType32 seqSlot, - SizeType32 reqBeamWidth, SizeType32 maxSeqLength, TokenIdType const* outputIdsHostData, - SizeType32 const* sequenceLengthsHostData) -{ - // Reorder generation logits to match the gathered (finalized) beam ordering. - // During generation, logits are stored indexed by beam SLOT position. After beam search - // finalization (gatherTree), output_ids are reordered by tracing parentIds to reconstruct - // the correct beam paths. However, generation_logits are NOT reordered by gatherTree. - // We fix this here by tracing parentIds on the host to build the beam-slot mapping, - // then reindexing the logits accordingly. - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const promptLen = llmReq.mPromptLen; - - // Copy parentIds and ids (ungathered step IDs) from GPU to temporary host buffers. - // parentIds[slot][t] = the parent slot of beam slot `slot` at position t. - // ids[slot][t] = the token in beam slot `slot` at position t (before gather). - auto parentIdsDevice = ITensor::at(mDecoderState->getParentIds(), {seqSlot}); - auto idsDevice = mDecoderState->getIds(seqSlot); - - auto parentIdsHost = runtime::BufferManager::pinnedPool(parentIdsDevice->getShape(), nvinfer1::DataType::kINT32); - auto idsHost = runtime::BufferManager::pinnedPool(idsDevice->getShape(), nvinfer1::DataType::kINT32); - - mCopyBufferManager.copy(*parentIdsDevice, *parentIdsHost); - mCopyBufferManager.copy(*idsDevice, *idsHost); - mCopyBufferManager.getStream().synchronize(); - - auto const* parentIdsData = bufferCast(*parentIdsHost); - auto const* idsData = bufferCast(*idsHost); - - // For each final beam b, find the beam slot at the last generated step, then - // trace back through parentIds to build the slot trace for every generation step. - // slotTrace[beam][genStep] = the beam slot that produced the logits at that step. - auto const generationLogitsHost = llmReq.getGenerationLogitsHost(); - auto const& logitsShape = generationLogitsHost->getShape(); - // Non-streaming shape: [beamWidth, maxNewTokens, vocabSizePadded] - TLLM_CHECK_WITH_INFO(logitsShape.d[0] == reqBeamWidth, - "Generation logits beam dimension (%ld) does not match beam width (%d).", logitsShape.d[0], reqBeamWidth); - auto const maxNewTokens = logitsShape.d[1]; - auto const vocabSizePadded = logitsShape.d[2]; - - std::vector> slotTrace(reqBeamWidth, std::vector(maxNewTokens, 0)); - bool anyReorderNeeded = false; - - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - auto const seqLen = sequenceLengthsHostData[beam]; - auto const genLen = seqLen - promptLen; - if (genLen <= 0) - { - continue; - } - - // Find the starting beam slot at the last generated step by matching the - // backtracked token sequence against the gathered (finalized) output. - SizeType32 startSlot = -1; - for (SizeType32 s = 0; s < reqBeamWidth; ++s) - { - SizeType32 slot = s; - bool matches = true; - for (SizeType32 t = seqLen - 1; t >= promptLen; --t) - { - if (idsData[slot * maxSeqLength + t] != outputIdsHostData[beam * maxSeqLength + t]) - { - matches = false; - break; - } - if (t > promptLen) - { - slot = parentIdsData[slot * maxSeqLength + t]; - } - } - if (matches) - { - startSlot = s; - break; - } - } - - TLLM_CHECK_WITH_INFO(startSlot >= 0, - "Could not determine beam slot mapping for beam %d during generation logits reordering.", beam); - - // Build the slot trace: slotTrace[beam][g] = the pre-reassignment slot whose - // logits correspond to generation step g of this beam. - // - // The model runs BEFORE beam search reassigns beams to slots, so - // generationLogits[slot][g] was produced by the pre-reassignment slot — - // i.e. the slot the beam occupied in the *previous* step. - // parentIds[postSlot][promptLen+g] gives exactly that pre-reassignment slot, - // so taking the parentIds lookup before storing (rather than after) yields - // the correct source slot in a single pass. - SizeType32 slot = startSlot; - for (SizeType32 t = seqLen - 1; t >= promptLen; --t) - { - slot = parentIdsData[slot * maxSeqLength + t]; - slotTrace[beam][t - promptLen] = slot; - } - - // Check if any reordering is actually needed for this beam - auto& slotTraceIds = slotTrace[beam]; - anyReorderNeeded |= std::any_of( - slotTraceIds.begin(), slotTraceIds.begin() + genLen, [beam](SizeType32 s) { return s != beam; }); - } - - // Reorder the generation logits in-place using a per-step temporary buffer. - if (anyReorderNeeded) - { - auto const logitsDataType = generationLogitsHost->getDataType(); - auto const elemSize = runtime::BufferDataType(logitsDataType).getSize(); - auto const stepSize = static_cast(vocabSizePadded) * elemSize; - - // Temp buffer for one generation step across all beams: [beamWidth, vocabSizePadded] - auto tempLogits - = runtime::BufferManager::pinnedPool(ITensor::makeShape({reqBeamWidth, vocabSizePadded}), logitsDataType); - - auto* logitsPtr = static_cast(generationLogitsHost->data()); - auto* tempPtr = static_cast(tempLogits->data()); - - std::vector genLens(reqBeamWidth); - SizeType32 maxGenLen = 0; - for (SizeType32 b = 0; b < reqBeamWidth; ++b) - { - genLens[b] = std::max(SizeType32{0}, sequenceLengthsHostData[b] - promptLen); - maxGenLen = std::max(maxGenLen, genLens[b]); - } - - for (SizeType32 g = 0; g < maxGenLen; ++g) - { - // Check if any beam that generated this step needs reordering - bool stepNeedsReorder = false; - for (SizeType32 b = 0; b < reqBeamWidth; ++b) - { - if (g < genLens[b] && slotTrace[b][g] != b) - { - stepNeedsReorder = true; - break; - } - } - if (!stepNeedsReorder) - { - continue; - } - - // Copy all beams' logits at this step to the temp buffer - for (SizeType32 b = 0; b < reqBeamWidth; ++b) - { - // logits layout: [beamWidth, maxNewTokens, vocabSizePadded] - auto const offset = (static_cast(b) * maxNewTokens + g) * stepSize; - std::memcpy(tempPtr + static_cast(b) * stepSize, logitsPtr + offset, stepSize); - } - - // Reorder: logits[b][g] = temp[slotTrace[b][g]] - for (SizeType32 b = 0; b < reqBeamWidth; ++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); - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::getDecoderSlotHostOutputs( - SizeType32 seqSlot, bool returnLogProbs, SamplingConfig const& samplingConfig, bool streaming) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (mWorldConfig.isLastPipelineParallelRank()) - { - auto event = mDecoder->finalize(*mDecoderState, seqSlot, samplingConfig, streaming); - // Make sure that postprocessing is done before copying outputIds - mCopyBufferManager.getStream().wait(event.get()); - - auto sequenceLengths = mDecoderState->getSequenceLengths(seqSlot); - auto outputIds = mDecoderState->getGatheredIds(seqSlot); - auto cumLogProbs = mDecoderState->getCumLogProbs(seqSlot); - auto logProbs = mDecoderState->getLogProbs(seqSlot); - - mCopyBufferManager.copy(*sequenceLengths, *mSlotDecoderBuffers[seqSlot]->sequenceLengths); - mCopyBufferManager.copy(*outputIds, *mSlotDecoderBuffers[seqSlot]->outputIds); - if (returnLogProbs) - { - mCopyBufferManager.copy(*cumLogProbs, *mSlotDecoderBuffers[seqSlot]->cumLogProbs); - mCopyBufferManager.copy(*logProbs, *mSlotDecoderBuffers[seqSlot]->logProbs); - } - - if (mWorldConfig.isPipelineParallel()) - { - // Make sure that postprocessing is done before sending outputIds - event.synchronize(); - - auto const peerSend = 0; - mDecSlotAsyncSndHdls.emplace_back(std::make_unique( - outputIds, sequenceLengths, cumLogProbs, logProbs, returnLogProbs, *mMpiCommPipelinePara, peerSend)); - } - } - else - { - auto const peerRecv = mWorldConfig.getPipelineParallelRank() == 0 ? mWorldConfig.getPipelineParallelism() - 1 - : mWorldConfig.getPipelineParallelRank() - 1; - DecoderSlotAsyncSend::recv(*mSlotDecoderBuffers[seqSlot], returnLogProbs, *mMpiCommPipelinePara, peerRecv); - - auto const peerSend = mWorldConfig.getPipelineParallelRank() + 1; - if (peerSend != mWorldConfig.getPipelineParallelism() - 1) - { - mDecSlotAsyncSndHdls.emplace_back(std::make_unique( - *mSlotDecoderBuffers[seqSlot], returnLogProbs, *mMpiCommPipelinePara, peerSend)); - } - } - sync_check_cuda_error(mRuntime->getStream().get()); - - // Here copy stream is synchronized after receiving decoderSlotOutputIdsView either by copy or by receive - // before copying to host on copy stream - runtime::CudaEvent beforeEvent{}; - mRuntime->getStreamPtr()->record(beforeEvent); - mCopyBufferManager.getStream().wait(beforeEvent); - mCopyBufferManager.copy(*mSlotDecoderBuffers[seqSlot]->outputIds, *mSlotDecoderBuffers[seqSlot]->outputIdsHost); - mCopyBufferManager.copy( - *mSlotDecoderBuffers[seqSlot]->sequenceLengths, *mSlotDecoderBuffers[seqSlot]->sequenceLengthsHost); - - if (returnLogProbs) - { - mCopyBufferManager.copy( - *mSlotDecoderBuffers[seqSlot]->cumLogProbs, *mSlotDecoderBuffers[seqSlot]->cumLogProbsHost); - mCopyBufferManager.copy(*mSlotDecoderBuffers[seqSlot]->logProbs, *mSlotDecoderBuffers[seqSlot]->logProbsHost); - } - - // Make sure copy is done before continuing on host - mCopyBufferManager.getStream().synchronize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ -// Check if one of the request needs log probs, need to get from decoder and communicate -bool batchReturnLogProbs(ScheduledRequests const& scheduledRequests) -{ - auto pred = [](auto const& llmReq) { return llmReq->returnLogProbs(); }; - return std::any_of(scheduledRequests.contextRequests.begin(), scheduledRequests.contextRequests.end(), pred) - || std::any_of(scheduledRequests.generationRequests.begin(), scheduledRequests.generationRequests.end(), pred); -} -} // namespace - -runtime::CudaEvent TrtGptModelInflightBatching::decoderStepAsync(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(decoderStepAsync); - - auto& decoderInputBuffers = mDecoderInputBuffers.at(getFusedBufferId()); - - auto const contextBufferId = mCtxGenFusion ? getFusedBufferId() : getContextBufferId(); - auto& contextRuntimeBuffers = mBuffers.at(contextBufferId); - auto const logitsIndex = (*mHandleContextLogits)(decoderInputBuffers, scheduledRequests.contextRequests, - contextRuntimeBuffers->logits, contextRuntimeBuffers->numContextLogits, mModelConfig, - mRuntime->getBufferManager(), contextRuntimeBuffers->mMedusaBuffers); - - auto const genLogitsIndex = mCtxGenFusion ? logitsIndex : 0; - auto const genBufferId = mCtxGenFusion ? getFusedBufferId() : getGenerationBufferId(); - auto& genRuntimeBuffers = mBuffers.at(genBufferId); - (*mHandleGenerationLogits)(decoderInputBuffers, scheduledRequests.generationRequests, genRuntimeBuffers->logits, - genLogitsIndex, mModelConfig, mRuntime->getBufferManager(), *genRuntimeBuffers, - genRuntimeBuffers->mMedusaBuffers); - - if (mOperatingBeamWidth > 1) - { - copyCacheIndirectionFromOutputsToInputs(scheduledRequests, genBufferId); - } - - mLogitsPostProcessorIsApplied = (*mLogitsPostProcessor)(decoderInputBuffers, mReplicateLogitsPostProcessor, - mWorldConfig, mRuntime->getStreamPtr(), mLogitsPostProcessorBatched); - - if (mGuidedDecoder) - { - mGuidedDecoder->execute(decoderInputBuffers, mRuntime->getBufferManager()); - } - - auto const fusedBufferId = getFusedBufferId(); - auto& fusedRuntimeBuffers = mBuffers.at(fusedBufferId); - - (*mMakeDecodingBatchInputOutput)(decoderInputBuffers, *mDecoderState, mModelConfig, *fusedRuntimeBuffers); - - auto decoderFinishEvent = mDecoder->forwardAsync(*mDecoderState, decoderInputBuffers); - - auto const returnLogProbs = batchReturnLogProbs(scheduledRequests); - auto updateDecoderBuffersEvent = (*mUpdateDecoderBuffers)(mModelConfig, mDecoderOutputBuffers.at(fusedBufferId), - mRuntime->getBufferManager(), *mDecoderState, returnLogProbs, decoderFinishEvent); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return updateDecoderBuffersEvent; -} - -void TrtGptModelInflightBatching::copyCacheIndirectionFromOutputsToInputs( - ScheduledRequests const& scheduledRequests, SizeType32 genBufferId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(copyCacheIndirectionFromOutputsToInputs); - - auto& genRuntimeBuffers = *mBuffers.at(genBufferId); - auto* srcOffsetsPtr = bufferCast(*genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySrcOffsets); - auto* dstOffsetsPtr = bufferCast(*genRuntimeBuffers.cacheIndirDecoderIOBatchedCopyDstOffsets); - auto* copySizesPtr = bufferCast(*genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySizes); - - // Only `cacheIndirShape.d[2]` is used - auto const& cacheIndirShape = mDecoderState->getCacheIndirectionOutput()->getShape(); - auto const maxBeamWidth = cacheIndirShape.d[1]; - auto const maxAttentionWindow = cacheIndirShape.d[2]; - auto const slotOffset = maxBeamWidth * maxAttentionWindow; - - SizeType32 batchIdx{0}; - SizeType64 maxCopySize{0}; - auto& manager = mRuntime->getBufferManager(); - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - auto const seqSlot = llmReq->mSeqSlot.value(); - auto const copySize = reqBeamWidth * maxAttentionWindow; - srcOffsetsPtr[batchIdx] = seqSlot * slotOffset; - dstOffsetsPtr[batchIdx] = seqSlot * slotOffset; - copySizesPtr[batchIdx] = copySize; - maxCopySize = std::max(maxCopySize, copySize); - batchIdx++; - } - } - if (batchIdx != 0) - { - auto const srcOffsetsSlice - = ITensor::slice(genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySrcOffsets, 0, batchIdx); - auto const srcOffsetsSliceDeviceSlice - = ITensor::slice(genRuntimeBuffers.mCacheIndirDecoderIOBatchedCopySrcOffsetsSliceDevice, 0, batchIdx); - manager.copy(srcOffsetsSlice->data(), *srcOffsetsSliceDeviceSlice, - runtime::MemoryType::kGPU); // Explicitly move to device for faster access. - auto const dstOffsetsSlice - = ITensor::slice(genRuntimeBuffers.cacheIndirDecoderIOBatchedCopyDstOffsets, 0, batchIdx); - auto const dstOffsetsSliceDeviceSlice - = ITensor::slice(genRuntimeBuffers.mCacheIndirDecoderIOBatchedCopyDstOffsetsSliceDevice, 0, batchIdx); - manager.copy(dstOffsetsSlice->data(), *dstOffsetsSliceDeviceSlice, - runtime::MemoryType::kGPU); // Explicitly move to device for faster access. - auto const sizesSlice = ITensor::slice(genRuntimeBuffers.cacheIndirDecoderIOBatchedCopySizes, 0, batchIdx); - auto const copySizesDeviceSlice - = ITensor::slice(genRuntimeBuffers.mCacheIndirDecoderIOBatchedCopyCopySizesDevice, 0, batchIdx); - manager.copy(sizesSlice->data(), *copySizesDeviceSlice); // Explicitly move to device for faster access. - runtime::kernels::invokeCopyBatch(*mDecoderState->getCacheIndirectionOutput(), - *mDecoderState->getCacheIndirectionInput(), *srcOffsetsSliceDeviceSlice, *dstOffsetsSliceDeviceSlice, - *copySizesDeviceSlice, maxCopySize, manager.getStream()); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::vector> TrtGptModelInflightBatching::communicateDecoderBuffers( - bool returnLogProbs) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(communicateDecoderBuffers); - - auto& decoderOutputBuffers = mDecoderOutputBuffers.at(getFusedBufferId()); - - std::vector> asyncHandles; - if (mWorldConfig.isLastPipelineParallelRank()) - { - if (broadcastPostDecoder()) - { - DecoderStepAsyncSend::bcast(decoderOutputBuffers, *mDecoderState, returnLogProbs, mOperatingBeamWidth, - mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), *mMpiCommTensorPara, 0); - } - - if (mWorldConfig.isPipelineParallel()) - { - auto const peerSend = 0; - asyncHandles.emplace_back(std::make_unique(decoderOutputBuffers, *mDecoderState, - returnLogProbs, mOperatingBeamWidth, mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), - *mMpiCommPipelinePara, peerSend)); - } - } - else - { - auto const peerRecv = mWorldConfig.isFirstPipelineParallelRank() ? mWorldConfig.getPipelineParallelism() - 1 - : mWorldConfig.getPipelineParallelRank() - 1; - DecoderStepAsyncSend::recv(decoderOutputBuffers, *mDecoderState, returnLogProbs, mOperatingBeamWidth, - mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), *mMpiCommPipelinePara, peerRecv); - auto const peerSend = mWorldConfig.getPipelineParallelRank() + 1; - if (peerSend != mWorldConfig.getPipelineParallelism() - 1) - { - asyncHandles.emplace_back(std::make_unique(decoderOutputBuffers, *mDecoderState, - returnLogProbs, mOperatingBeamWidth, mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind(), - *mMpiCommPipelinePara, peerSend)); - } - } - TLLM_CHECK_WITH_INFO(asyncHandles.size() <= 2, "Up to two decoder step async handles expected"); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return asyncHandles; -} - -void TrtGptModelInflightBatching::updateRequests(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(updateRequests); - - auto const& decoderOutputBuffers = mDecoderOutputBuffers.at(getFusedBufferId()); - - auto const hostNewOutputTokensShape = decoderOutputBuffers.newOutputTokensHost->getShape(); - auto const* const hostNewOutputTokensData - = bufferCast(*decoderOutputBuffers.newOutputTokensHost); - auto const* const sequenceLengthsHostData = bufferCast(*decoderOutputBuffers.sequenceLengthsHost); - auto const* const decoderFinishedSumPtr = bufferCast(*decoderOutputBuffers.finishedSumHost); - auto const* const cumLogProbsPtr = bufferCast(*decoderOutputBuffers.cumLogProbsHost); - auto const* const logProbsPtr = bufferCast(*decoderOutputBuffers.logProbsHost); - auto const* const finishReasonsHostData - = bufferCast(*decoderOutputBuffers.finishReasonsHost); - - // Update only requests that ran through the decoder - for (auto const& llmReq : scheduledRequests.generationRequests) - { - if (llmReq->isGenerationCompleteState()) - { - continue; - } - auto const reqBeamWidth = llmReq->getBeamWidthByIter(true); - auto const seqSlot = llmReq->mSeqSlot.value(); - auto const currentNumOfTokens = llmReq->getMaxBeamNumTokens(); - - // Save the accepted token logits from target model - if (mModelConfig.getSpeculativeDecodingMode().isDraftTokensExternal() && llmReq->getReturnGenerationLogits() - && llmReq->hasDraftTokens()) - { - TLLM_CHECK_WITH_INFO(reqBeamWidth == 1, "Speculative decoding only works for beam width == 1"); - - SizeType32 numAcceptedTokens - = sequenceLengthsHostData[seqSlot * mOperatingBeamWidth + 0] - llmReq->getMaxBeamNumTokens(); - - auto const& generationLogitsHost = llmReq->getGenerationLogitsHost(); - auto shape = generationLogitsHost->getShape(); - shape.d[1] = numAcceptedTokens; - generationLogitsHost->reshape(shape); - } - - std::vector numNewTokens(reqBeamWidth); - std::vector numDroppedTokens(reqBeamWidth); - - // numGeneratedTokens is the number of tokens generated by the decoder. - // Some tokens might be dropped due to end token or rejected draft tokens. - auto const numGeneratedTokens = llmReq->getNumDraftTokens() + 1; - - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - // Sequence length is only advanced for accepted tokens. - auto const seqLen = sequenceLengthsHostData[seqSlot * mOperatingBeamWidth + beam]; - // Actual number of tokens that should be added to the request. - auto const numNewOutputTokens = seqLen - llmReq->getNumTokens(beam); - if (reqBeamWidth == 1) - { - TLLM_CHECK_WITH_INFO(numGeneratedTokens >= numNewOutputTokens, - "numNewOutputTokens must not be greater than numGeneratedTokens: " - "numGeneratedTokens %d < numNewOutputTokens %d", - numGeneratedTokens, numNewOutputTokens); - } - numNewTokens[beam] = std::min(numGeneratedTokens, numNewOutputTokens); - numDroppedTokens[beam] = numGeneratedTokens - numNewTokens[beam]; - for (SizeType32 step = 0; step < numNewTokens[beam]; ++step) - { - auto const newTokenIdx = tc::flat_index(hostNewOutputTokensShape.d, step, seqSlot, beam); - auto const newToken = hostNewOutputTokensData[newTokenIdx]; - llmReq->addNewToken(newToken, beam); - TLLM_LOG_DEBUG("request ID %ld beam %d newToken %d", llmReq->mRequestId, beam, newToken); - - if (llmReq->returnLogProbs()) - { - auto const cumLogProb = cumLogProbsPtr[seqSlot * mOperatingBeamWidth + beam]; - llmReq->setCumLogProb(cumLogProb, beam); - - auto const beginLogProbsOffset = reqBeamWidth == 1 ? llmReq->mPromptLen : 0; - SizeType32 offset - = (seqSlot * mOperatingBeamWidth + beam) * getMaxSequenceLen() + beginLogProbsOffset; - auto const generatedLength = seqLen - llmReq->mPromptLen; - std::vector logProbs(logProbsPtr + offset, logProbsPtr + offset + generatedLength); - llmReq->setLogProbs(logProbs, beam); - } - } - - auto const finishReason = finishReasonsHostData[seqSlot * mOperatingBeamWidth + beam]; - llmReq->setFinishedReason(finishReason.toFinishReason(), beam); - - TLLM_LOG_DEBUG("[RANK %d] decoderSync: request ID %lu beam %d tokens %s finished %d", - COMM_SESSION.getRank(), llmReq->mRequestId, beam, common::vec2str(llmReq->getTokens(beam)).c_str(), - static_cast(finishReason.toFinishReason())); - } - - // Set number of tokens predicted per runtime iteration. Will be > 1 for speculative decoding. - llmReq->updateNumTokensPerIteration(llmReq->getMaxBeamNumTokens() - currentNumOfTokens, mModelConfig); - - // Fill new draft tokens for the next step - if (decoderFinishedSumPtr[seqSlot] != reqBeamWidth - && (mModelConfig.getSpeculativeDecodingMode().predictsDraftTokens() - || mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind())) - { - auto const maxDraftTokensLen = mModelConfig.getMaxDecodingDraftTokens(); - auto prevDraftTokensLen = llmReq->getNumDraftTokens(); - - // We overallocate KV cache for EAGLE to the maxDecodingTokens + maxPathLen in order to fit both - // Base model verification (needs up to maxDecodingTokens) and - // Drafter (needs up to maxPathLen of accepted tokens and maxDecodingDraftTokens for new draft tokens). - if (mModelConfig.getSpeculativeDecodingMode().isEagle()) - { - prevDraftTokensLen = mModelConfig.getSpeculativeDecodingModule().getMaxDecodingTokens() - + mModelConfig.getSpeculativeDecodingModule().getMaxPathLen() - 1; - } - - auto nextDraftTokensLen = mModelConfig.getSpeculativeDecodingModule().getMaxDecodingDraftTokens(); - if (mModelConfig.getSpeculativeDecodingMode().variableDraftLength()) - { - auto const* const nextDraftTokensLengthsHostData - = bufferCast(*decoderOutputBuffers.nextDraftTokensLengthsHost); - nextDraftTokensLen = nextDraftTokensLengthsHostData[seqSlot]; - } - TLLM_CHECK(nextDraftTokensLen <= maxDraftTokensLen); - - auto const* const nextDraftTokensHostData - = bufferCast(*decoderOutputBuffers.nextDraftTokensHost); - auto draftTokensShared - = std::make_shared>(nextDraftTokensHostData + seqSlot * maxDraftTokensLen, - nextDraftTokensHostData + seqSlot * maxDraftTokensLen + nextDraftTokensLen); - - llmReq->setDraftTokens(draftTokensShared); - - // For all phases except context that does not have draft tokens - if (!llmReq->isGenerationCompleteState() && prevDraftTokensLen != 0 - && mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind()) - { - // -1 here is for current 'main' token - auto const acceptedTokensLen = llmReq->getMaxBeamNumTokens() - currentNumOfTokens - 1; - auto const rewindLength = prevDraftTokensLen - acceptedTokensLen; - - TLLM_LOG_DEBUG("request ID %lu (seqSlot %d): accepted %d of %d draft tokens, rewind %d tokens", - llmReq->mRequestId, seqSlot, acceptedTokensLen, prevDraftTokensLen, rewindLength); - TLLM_CHECK(0 <= acceptedTokensLen && acceptedTokensLen <= prevDraftTokensLen); - - // At this point, KV cache rows are already gathered and moved to the right location. - // We can safely rewind (draft - accepted) tokens - mKvCacheManager->rewindKVCache(llmReq->mRequestId, rewindLength); - } - } - - // Terminate if request has finished or if it is speculative decoding target model - if (decoderFinishedSumPtr[seqSlot] == reqBeamWidth - || (mModelConfig.getSpeculativeDecodingMode().isDraftTokensExternal() && llmReq->hasDraftTokens())) - { - postProcessRequest(*llmReq, numDroppedTokens); - - if (!mWorldConfig.isPipelineParallel() || !mWorldConfig.isLastPipelineParallelRank()) - { - if (llmReq->getReturnGenerationLogits() && mSpeculativeDecodingFastLogits && mIsLeaderInOrchMode) - { - std::lock_guard lk(mDraftRequestsMtx); - mDraftRequestsWaitingToSendLogits.push_back(llmReq); - } - else - { - terminateRequest(llmReq); - } - llmReq->setState(LlmRequestState::kGENERATION_COMPLETE); - } - else - { - llmReq->setState(LlmRequestState::kGENERATION_TO_COMPLETE); - } - } - else - { - // gather tokens in the case of streaming and beam search - if (llmReq->isStreaming() && llmReq->mSamplingConfig.beamWidth > 1) - { - postProcessRequest(*llmReq, numDroppedTokens); - } - if (llmReq->isContextInitState()) - { - llmReq->setState(LlmRequestState::kGENERATION_IN_PROGRESS); - } - - if (isTrtOverlap() && llmReq->willCompleteNextIteration()) - { - // This state prohibits the request from being scheduled for another iteration. It assumes that the next - // iteration has already been scheduled and the request can finish in the next call to updateRequests(). - llmReq->setState(LlmRequestState::kGENERATION_TO_COMPLETE); - } - } - - if (llmReq->getReturnPerfMetrics()) - { - llmReq->updatePerfMetrics(mIterCounter); - } - - llmReq->advanceDecodingIter(); - - if (mWorldConfig.isPipelineParallel() && mWorldConfig.isLastPipelineParallelRank()) - { - for (SizeType32 beam = 0; beam < reqBeamWidth; ++beam) - { - llmReq->setNumPreDecodedTokens(numNewTokens[beam], beam); - } - } - } - - if (mModelConfig.getSpeculativeDecodingMode().needsKVCacheRewind()) - { - SizeType32 numSequences{0}; - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - auto const reqBeamWidth = llmReq->mSamplingConfig.beamWidth; - numSequences += reqBeamWidth; - } - } - - TLLM_CHECK_WITH_INFO(mCtxGenFusion, "Current speculative decoding mode requires context-gen fusion IFB"); - rewindKVCacheBlocks(numSequences); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::vector> TrtGptModelInflightBatching::decoderSync( - ScheduledRequests const& scheduledRequests, std::optional const& decoderFinishEvent) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(decoderSync); - - if (mWorldConfig.isLastPipelineParallelRank()) - { - decoderFinishEvent->synchronize(); - } - - auto const returnLogProbs = batchReturnLogProbs(scheduledRequests); - auto asyncHandles = communicateDecoderBuffers(returnLogProbs); - - updateRequests(scheduledRequests); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return asyncHandles; -} - -void TrtGptModelInflightBatching::rewindKVCacheBlocks(SizeType32 numSequences) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const bufferId = getFusedBufferId(); - auto& runtimeBuffers = *mBuffers.at(bufferId); - auto& decoderOutputBuffers = mDecoderOutputBuffers.at(bufferId); - - auto localNbLayers = mModelConfig.getNbAttentionLayers( - mWorldConfig.getPipelineParallelism(), mWorldConfig.getPipelineParallelRank()); - if (mWorldConfig.isLastPipelineParallelRank() && mModelConfig.getSpeculativeDecodingMode().isEagle()) - { - // Do not correct the last kv caches, which are for EagleNet drafter. Those KV caches are managed separately. - auto eagleModulePtr - = std::dynamic_pointer_cast(mModelConfig.getSpeculativeDecodingModulePtr()); - localNbLayers -= eagleModulePtr->getNumTransformerLayers(); - } - - auto const tokensPerBlock = mModelConfig.getTokensPerBlock(); - auto const elemSize = BufferDataType(mModelConfig.getKvDataType()).getSize(); - auto const sizeInBytesPerKVHead = mModelConfig.getSizePerHead() * elemSize; - - auto const poolPointers = mKvCacheManager->getBlockPoolPointers(); - auto* const* pointerArrayPtr = bufferCast(*poolPointers); - auto const* offsetArrayPtr - = bufferCast(*runtimeBuffers.transformerBuffers->kvCacheBlockOffsetsDevice); - - auto commonRewindLen = mModelConfig.getSpeculativeDecodingModule().getMaxDecodingDraftTokens(); - SizeType32 const* rewindLens = nullptr; - if (mModelConfig.getSpeculativeDecodingMode().variableDraftLength()) - { - commonRewindLen = 0; - rewindLens = bufferCast(*decoderOutputBuffers.prevDraftTokensLengthsHost); - } - - tensorrt_llm::runtime::kernels::invokeUpdateKVBlockArrayDraftTokenLocation( - *mDecoderState->getAcceptedLengthsCumSum(), *mDecoderState->getAcceptedPackedPaths(), - *runtimeBuffers.sequenceLengthsDevice, pointerArrayPtr, offsetArrayPtr, localNbLayers, numSequences, - mRewindInputs.numKvHeads, sizeInBytesPerKVHead, commonRewindLen, rewindLens, *runtimeBuffers.seqSlots, - getMaxAttentionWindow(), mRewindInputs.maxBlocksPerSeq, tokensPerBlock, mRewindInputs.isUseOneMoreBlock, - mRuntime->getStreamPtr()->get()); - - sync_check_cuda_error(mRuntime->getStream().get()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -nvinfer1::DataType TrtGptModelInflightBatching::getLogitDataType() const -{ - return mModelConfig.getLogitsDtype(); -} - -TrtGptModelInflightBatching::SizeType32 TrtGptModelInflightBatching::numCachedCudaGraphs() const -{ - return std::accumulate(mCudaGraphExecutorCaches.begin(), mCudaGraphExecutorCaches.end(), SizeType32{0}, - [](SizeType32 sum, auto const& cache) { return sum + cache.size(); }); -} - -void TrtGptModelInflightBatching::changeBeamWidth(SizeType32 beamWidth) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK(mInflightReqIds.empty()); - - TLLM_CHECK_WITH_INFO(beamWidth <= getMaxBeamWidth(), - "Requested beam width %d is larger than configured max beam width %d", beamWidth, getMaxBeamWidth()); - TLLM_LOG_DEBUG("Changing operating beam width from %d to %d", mOperatingBeamWidth, beamWidth); - mOperatingBeamWidth = beamWidth; - - if (isCudaGraphMode()) - { - for (auto& cache : mCudaGraphExecutorCaches) - { - cache.clear(); - } - } - createBuffers(mDecodingConfig, mAdditionalModelOutputs); - createDecoder(mDecodingConfig.getDecodingMode()); - - if (static_cast(mKvCacheManager)) - { - auto const dims = mKvCacheManager->getOffsetTableDimensions(); - reshapeKvTensors(dims); - } - if (static_cast(mCrossKvCacheManager)) - { - auto const dims = mCrossKvCacheManager->getOffsetTableDimensions(); - reshapeKvTensors(dims); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::changeSpecDecMode(ScheduledRequests const& scheduledRequests) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if ((!mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() - && !mModelConfig.getSpeculativeDecodingMode().isNone()) - || scheduledRequests.empty() || mSeamlessLADMaxDraftLen == 0 || getGatherGenerationLogits() - || mModelConfig.isRnnBased()) - { - return; - } - - bool canUseLookahead = false; - auto maxNumRequestForLad = mDecodingConfig.getLookaheadDecodingMaxNumRequest(); - SizeType32 numRequests = scheduledRequests.contextRequests.size() + scheduledRequests.generationRequests.size(); - if (numRequests > maxNumRequestForLad) - { - if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding()) - { - canUseLookahead = false; - } - else - { - return; - } - } - { - bool useTopKTopP = false; - bool useBanWords = false; - bool useTempAccVocabPenalties = false; // use temperature and penalties that need to accumulate #vocab. - SizeType32 beamWidth = 1; - for (auto const& requests : {scheduledRequests.contextRequests, scheduledRequests.generationRequests}) - { - for (auto const& llmReq : requests) - { - useTopKTopP |= !(llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.topK, layers::DefaultDecodingParams::getTopK()) - || llmReq->mSamplingConfig.useDefaultValues(llmReq->mSamplingConfig.topK, 1)); - useTopKTopP |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.topP, layers::DefaultDecodingParams::getTopP()); - useBanWords |= llmReq->getBadWordsList().has_value(); - useBanWords |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.noRepeatNgramSize, layers::DefaultDecodingParams::getNoRepeatNgramSize()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.temperature, layers::DefaultDecodingParams::getTemperature()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.repetitionPenalty, layers::DefaultDecodingParams::getRepetitionPenalty()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.presencePenalty, layers::DefaultDecodingParams::getPresencePenalty()); - useTempAccVocabPenalties |= !llmReq->mSamplingConfig.useDefaultValues( - llmReq->mSamplingConfig.frequencyPenalty, layers::DefaultDecodingParams::getFrequencyPenalty()); - beamWidth = llmReq->mSamplingConfig.beamWidth; - if (useTopKTopP || useBanWords || useTempAccVocabPenalties || beamWidth > 1) - { - break; - } - } - canUseLookahead = !(useTopKTopP || useBanWords || useTempAccVocabPenalties || beamWidth > 1); - } - } - - // Change speculative decoding mode - auto const bufferId = mCtxGenFusion - ? getFusedBufferId() - : (!scheduledRequests.contextRequests.empty() ? getContextBufferId() : getGenerationBufferId()); - // TODO: enable lookahead for generation requests. - bool canChangeToLookahead = scheduledRequests.generationRequests.empty(); - if (mModelConfig.getSpeculativeDecodingMode().isNone() && canUseLookahead && canChangeToLookahead) - { - // None -> Lookahead - mModelConfig.enableSeamlessLookaheadDecoding(mSeamlessLADMaxDraftLen); - mDecodingConfig.enableSeamlessLookaheadDecoding(); - setupSpeculativeDecodingModule(mDecodingConfig); - mBuffers.at(bufferId)->mLookaheadBuffers->enableLookaheadDecoding( - getMaxBatchSize(), mModelConfig.getMaxDecodingTokens()); - mDecoderOutputBuffers.at(getFusedBufferId()) - .enableLookaheadDecoding(getMaxNumSequences(), mModelConfig.getMaxDecodingTokens()); - createDecoder(mDecodingConfig.getDecodingMode()); - } - else if (mModelConfig.getSpeculativeDecodingMode().isLookaheadDecoding() - && (!canUseLookahead || numRequests > maxNumRequestForLad)) - { - // Lookahead -> None - mModelConfig.disableSeamlessLookaheadDecoding(); - mDecodingConfig.setDecodingMode(executor::DecodingMode::Auto()); - mBuffers.at(bufferId)->mLookaheadBuffers->disableLookaheadDecoding(); - mDecoderOutputBuffers.at(getFusedBufferId()).disableLookaheadDecoding(getMaxNumSequences()); - mDecoder->disableLookahead( - scheduledRequests.generationRequests, mDecoderInputBuffers.at(getFusedBufferId()).setupBatchSlots); - mDecoderState->disableLookahead(scheduledRequests.generationRequests); - for (auto const& llmReq : scheduledRequests.generationRequests) - { - if (llmReq->getNumDraftTokens() > 0) - { - llmReq->discardDraftTokens(llmReq->getNumDraftTokens()); - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TrtGptModelInflightBatching::getCurrentIterationStats(executor::IterationStats& stats) const -{ - stats.iter = mIterCounter; - - // Max batch size and max num tokens can be tuned at runtime - stats.maxBatchSizeStatic = getMaxBatchSize(); - stats.maxBatchSizeTunerRecommended = mMaxBatchSizeTunerRecommended; - stats.maxBatchSizeRuntime = mMaxBatchSizeRuntime; - stats.maxNumTokensStatic = mMaxNumTokensStatic.value_or(0); - stats.maxNumTokensTunerRecommended = mMaxNumTokensTunerRecommended; - stats.maxNumTokensRuntime = mMaxNumTokensRuntime.value_or(0); - - // KVCacheManager statistics - auto const& kvCacheManager = getKVCacheManager(); - if (kvCacheManager) - { - executor::KvCacheStats kvStats{}; - auto kvCacheStats = kvCacheManager->getKvCacheStats(); - kvStats.maxNumBlocks = kvCacheStats.maxNumBlocks; - kvStats.freeNumBlocks = kvCacheStats.freeNumBlocks; - kvStats.usedNumBlocks = kvCacheStats.usedNumBlocks; - kvStats.tokensPerBlock = kvCacheStats.toksPerBlock; - kvStats.allocTotalBlocks = kvCacheStats.allocTotalBlocks; - kvStats.allocNewBlocks = kvCacheStats.allocNewBlocks; - kvStats.reusedBlocks = kvCacheStats.reusedBlocks; - kvStats.missedBlocks = kvCacheStats.missedBlocks; - kvStats.cacheHitRate = kvCacheStats.cacheHitRate; - stats.kvCacheStats = kvStats; - } - auto const& crossKvCacheManager = getCrossKVCacheManager(); - if (crossKvCacheManager) - { - executor::KvCacheStats kvStats{}; - auto kvCacheStats = crossKvCacheManager->getKvCacheStats(); - kvStats.maxNumBlocks = kvCacheStats.maxNumBlocks; - kvStats.freeNumBlocks = kvCacheStats.freeNumBlocks; - kvStats.usedNumBlocks = kvCacheStats.usedNumBlocks; - kvStats.tokensPerBlock = kvCacheStats.toksPerBlock; - kvStats.allocTotalBlocks = kvCacheStats.allocTotalBlocks; - kvStats.allocNewBlocks = kvCacheStats.allocNewBlocks; - kvStats.reusedBlocks = kvCacheStats.reusedBlocks; - kvStats.missedBlocks = kvCacheStats.missedBlocks; - kvStats.cacheHitRate = kvCacheStats.cacheHitRate; - stats.crossKvCacheStats = kvStats; - } - executor::InflightBatchingStats modelStats{}; - modelStats.numScheduledRequests = mLastIterationStatsIFB.scheduledRequests.size(); - modelStats.numContextRequests = mLastIterationStatsIFB.numCtxRequests; - modelStats.numGenRequests = mLastIterationStatsIFB.numGenRequests; - modelStats.numPausedRequests = mLastIterationStatsIFB.pausedRequests.size(); - modelStats.avgNumDecodedTokensPerIter = mLastIterationStatsIFB.avgNumDecodedTokensPerIter; - modelStats.numCtxTokens = mLastIterationStatsIFB.numCtxTokens; - modelStats.microBatchId = mLastIterationStatsIFB.microBatchId; - stats.inflightBatchingStats = modelStats; -} - -void TrtGptModelInflightBatching::getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const -{ - stats.iter = mIterCounter; - for (auto& requestStat : stats.requestStats) - { - requestStat.scheduled - = mLastIterationStatsIFB.scheduledRequests.count(static_cast(requestStat.id)); - requestStat.paused = mLastIterationStatsIFB.pausedRequests.count(static_cast(requestStat.id)); - } -} - -executor::DebugTensorsPerIteration TrtGptModelInflightBatching::getCurrentDebugTensors() const -{ - executor::DebugTensorsPerIteration debugTensors; - debugTensors.iter = mIterCounter; - - for (auto const& [name, tensor] : mLastIterationDebugTensors) - { - debugTensors.debugTensors.emplace(name, executor::detail::ofITensor(tensor)); - } - - return debugTensors; -} - -nvinfer1::DataType TrtGptModelInflightBatching::getTensorDataType(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorDataType(name.c_str()); -} - -nvinfer1::Dims TrtGptModelInflightBatching::getTensorShape(std::string const& name) const -{ - auto const& engine = mRuntime->getEngine(); - return engine.getTensorShape(name.c_str()); -} - -SizeType32 TrtGptModelInflightBatching::getMaxCapacityBatchSize(SizeType32 inputLength, SizeType32 outputLength) const -{ - return mKvCacheManager->getMaxCapacityBatchSize(inputLength, outputLength); -} - -/* - * Manages prefetching of prompt table chunks using a double-buffer strategy - * - * Function Flow: - * 1. First Chunk Processing (isFirstChunk == true): - * - Uses blocking prefetch on main runtime stream - * - Ensures initial data is ready before computation starts - * - * 2. Subsequent Chunks (isFirstChunk == false): - * - Uses non-blocking prefetch on separate copy stream - * - Overlaps data transfer with computation - * - * Synchronization: - * - First prefetch: No wait needed (fresh start) - * - Later prefetches: Wait for previous copy to complete - * - Uses mPtableCopyDoneEvent to track completion - * - * Key Functions: - * 1. prefetchNextPromptTableChunk: - * - Calls the correct function based on position in code (before or after prepareBuffers()) - * - Waits for previous copy to complete if not the first chunk - * - * 2. remapInputTokensForPromptTable: - * - Identifies tokens that need prompt table embeddings (tokens that are greater than vocabSize) - * - Remaps IDs to match chunked prompt table layout - * - * 3. copyPromptTableToGpuInChunk: - * - Handles actual transfer from CPU pinned memory to GPU - * - Uses appropriate buffer manager based on isFirstChunk - */ -void TrtGptModelInflightBatching::prefetchNextPromptTableChunk( - RequestVector const& contextRequests, bool isFirstChunk, SizeType32 bufferId) -{ - auto& promptTuningBuffers = mBuffers[bufferId]->promptTuningBuffers; - - if (!isFirstChunk) - { - // Only switch buffer after prepareBuffer() - promptTuningBuffers->switchChunkPtableBuffer(); - } - - SizeType32 contextId = 0; - for (auto const& llmReq : contextRequests) - { - if (llmReq->isFirstContextChunk() && isFirstChunk) - { - // For first chunk: Blocking prefetch on runtime stream to ensure data is ready - remapInputTokensForPromptTable(llmReq, true, bufferId, contextId); - } - else if (!isFirstChunk) // prefetching for subsequent chunks - { - // For the first prefetch chunk, don't need to wait for previous prefetch to complete - // For subsequent chunks: Need to wait for previous prefetch to complete - if (!llmReq->isFirstContextChunk()) - { - mRuntime->getBufferManager().getStream().wait(mPtableCopyDoneEvent); - } - - // Non-blocking prefetch on copy stream to prepare next chunk in pong buffer - if (llmReq->getContextRemainingLength() > 0) - { - remapInputTokensForPromptTable(llmReq, false, bufferId, contextId); - } - } - - ++contextId; - } -} - -void TrtGptModelInflightBatching::remapInputTokensForPromptTable( - std::shared_ptr const& llmReq, bool isFirstChunk, SizeType32 bufferId, SizeType32 contextId) -{ - NVTX3_SCOPED_RANGE_WITH_NAME(range, "remapInputTokensForPromptTable"); - auto& promptTuningBuffers = mBuffers[bufferId]->promptTuningBuffers; - auto const chunkSize = llmReq->getContextChunkSize(); - auto& inputTokensMutable = llmReq->getTokensMutable(0); - auto vocabSize = mModelConfig.getVocabSize(); - - if (isFirstChunk) - { - promptTuningBuffers->initializeChunkPtableBuffers( - mRuntime->getBufferManager(), mModelConfig, chunkSize, llmReq); - } - - size_t processChunkSize; - size_t beginPos; - - if (!isFirstChunk) - { - processChunkSize = std::min(chunkSize, llmReq->getContextRemainingLength() - chunkSize); - } - else - { - processChunkSize = std::min(chunkSize, llmReq->getContextRemainingLength()); - } - - if (!isFirstChunk) - { - // For prefetching next chunk - if (llmReq->getContextRemainingLength() - chunkSize <= 0) - { - promptTuningBuffers->updateBufferStartPosition(promptTuningBuffers->getChunkPtableCurrentIndex(), 0); - return; // No more chunks to prefetch - } - beginPos = llmReq->getContextCurrentPosition() + chunkSize; - } - else - { - // For current chunk - beginPos = llmReq->getContextCurrentPosition(); - } - - TLLM_CHECK_WITH_INFO(beginPos + processChunkSize <= inputTokensMutable.size(), - "Invalid chunk access: beginPos(%zu) + processChunkSize(%zu) > totalSize(%zu)", beginPos, processChunkSize, - inputTokensMutable.size()); - - auto inputTokensChunk = inputTokensMutable.begin() + beginPos; - std::vector outOfVocabTokens; - SizeType32 ptableTokenId = vocabSize; - for (size_t i = 0; i < processChunkSize; i++) - { - if (inputTokensChunk[i] >= vocabSize) - { - outOfVocabTokens.push_back(inputTokensChunk[i]); - inputTokensChunk[i] = ptableTokenId++; - } - } - - copyPromptTableToGpuInChunk(llmReq, outOfVocabTokens, isFirstChunk, bufferId, contextId); -} - -void TrtGptModelInflightBatching::copyPromptTableToGpuInChunk(std::shared_ptr const& llmReq, - std::vector const& outOfVocabTokens, bool isFirstChunk, SizeType32 bufferId, SizeType32 contextId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE_WITH_NAME(range, "copyPromptTableToGpuInChunk"); - auto& promptTuningBuffers = mBuffers[bufferId]->promptTuningBuffers; - - if (outOfVocabTokens.empty()) - { - return; - } - - auto const& promptTable = llmReq->getPromptEmbeddingTable(); - TLLM_CHECK_WITH_INFO(promptTable.has_value(), "promptTable is empty but there's fake_prompt"); - TLLM_CHECK_WITH_INFO(promptTable.value() != nullptr, "promptTable value is null but there's fake_prompt"); - - auto currentBufferManager = isFirstChunk ? mRuntime->getBufferManager() : mCopyBufferManager; - auto const hiddenSize = mModelConfig.getHiddenSize(); - auto numRows = outOfVocabTokens.size(); - std::size_t sliceSize = static_cast(numRows * hiddenSize); - auto currentIndex = promptTuningBuffers->getChunkPtableCurrentIndex(); - - // Calculate the offset based on current position - size_t srcOffset = llmReq->mPtableCurrentPosition * hiddenSize; - size_t dstOffset = promptTuningBuffers->getChunkPtableBufferStartPosition(currentIndex, contextId); - - auto gpuBuffer = promptTuningBuffers->getChunkPtableBuffer(currentIndex); - - // First view as 1D tensor of elements - auto totalElements = promptTable.value()->getSize(); - auto table1D = runtime::ITensor::view( - promptTable.value(), runtime::ITensor::makeShape({static_cast(totalElements)})); - - TLLM_CHECK_WITH_INFO(srcOffset + sliceSize <= totalElements, - "Buffer bounds violation: Trying to access up to %zu elements but buffer only has %zu elements (offset: %zu, " - "slice size: %zu)", - srcOffset + sliceSize, totalElements, srcOffset, sliceSize); - - auto table1DShared = runtime::ITensor::SharedPtr(table1D.release()); - auto pTableView = runtime::ITensor::slice(table1DShared, srcOffset, sliceSize); - - auto gpuBufferSlice = runtime::ITensor::slice(gpuBuffer, dstOffset, numRows); - - currentBufferManager.copy(*pTableView, *gpuBufferSlice); - - promptTuningBuffers->updateBufferStartPosition(currentIndex, outOfVocabTokens.size()); - - llmReq->mPtableCurrentPosition += outOfVocabTokens.size(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h b/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h deleted file mode 100644 index ca9f7c7f4a7f..000000000000 --- a/cpp/tensorrt_llm/batch_manager/trtGptModelInflightBatching.h +++ /dev/null @@ -1,637 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026 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/batch_manager/kvCacheType.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/worldConfig.h" -#include "trtGptModel.h" - -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime; -class GptDecoderBatched; -class AllReduceBuffers; -class NcclCommunicator; -class SpeculativeDecodingMode; - -namespace decoder -{ -class DecoderState; -} // namespace decoder - -namespace decoder_batch -{ -class Input; -class Output; -} // namespace decoder_batch - -} // namespace tensorrt_llm::runtime - -namespace tensorrt_llm::mpi -{ -class MpiWaitThread; -} // namespace tensorrt_llm::mpi - -namespace tensorrt_llm::batch_manager -{ -class BaseCacheTransceiver; -} - -namespace tensorrt_llm::batch_manager -{ - -namespace kv_cache_manager -{ -class KVCacheManager; -struct OffsetTableDimensions; -} // namespace kv_cache_manager - -namespace rnn_state_manager -{ -class RnnStateManager; -} // namespace rnn_state_manager - -class SequenceSlotManager; -class DecoderStepAsyncSend; -class DecoderSlotAsyncSend; -class DecoderInputBuffers; -class DecoderOutputBuffers; -class SlotDecoderBuffers; -class LlmRequest; -class RuntimeBuffers; -class BasePeftCacheManager; -class GuidedDecoder; -class TrtGptModelTest; - -// Algorithms -class CapacityScheduler; -class MicroBatchScheduler; -class PauseRequests; -class AssignReqSeqSlots; -class AllocateKvCache; -class HandleContextLogits; -class HandleGenerationLogits; -class GenerateRequestOptions; -class LogitsPostProcessor; -class MakeDecodingBatchInputOutput; -class CreateNewDecoderRequests; -class UpdateDecoderBuffers; - -namespace utils -{ -class CudaGraphExecutorCache; -} // namespace utils - -struct RewindInputs -{ - SizeType32 maxBlocksPerSeq; - bool isUseOneMoreBlock; - SizeType32 numKvHeads; -}; - -class TrtGptModelInflightBatching : public TrtGptModel -{ - using BaseKVCacheManager = kv_cache_manager::BaseKVCacheManager; - using OffsetTableDimensions = kv_cache_manager::OffsetTableDimensions; - using KVCacheManager = kv_cache_manager::KVCacheManager; - using KvCacheType = kv_cache_manager::CacheType; - using KvCacheConfig = executor::KvCacheConfig; - using RnnStateManager = rnn_state_manager::RnnStateManager; - using LlmRequestPtr = std::shared_ptr; - -public: - class IterationStatsIFB - { - public: - explicit IterationStatsIFB(SizeType32 microBatchId) - : microBatchId{microBatchId} - { - } - - SizeType32 microBatchId; - SizeType32 numCtxRequests{}; - SizeType32 numGenRequests{}; - SizeType32 numCtxTokens{}; - float avgNumDecodedTokensPerIter{}; - ReqIdsSet scheduledRequests; - ReqIdsSet pausedRequests; - }; - - using SizeType32 = tensorrt_llm::runtime::SizeType32; - using TokenIdType = tensorrt_llm::runtime::TokenIdType; - using BufferManager = tensorrt_llm::runtime::BufferManager; - using PeftTable = PeftCacheManager::PeftTable; - using TensorMap = runtime::StringPtrMap; - using TensorPtr = runtime::ITensor::SharedPtr; - - TrtGptModelInflightBatching(std::shared_ptr logger, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, runtime::RawEngine const& rawEngine, bool ctxGenFusion, - executor::ExecutorConfig const& executorConfig, bool isLeaderInOrchMode); - - ~TrtGptModelInflightBatching() override; - - /// @brief Calculate the cache size per token for the disaggregated serving. - /// @param modelConfig Model configuration. - /// @param worldConfig World configuration. - /// @param maxAttentionWindowVec Maximum attention window vector. (may have fewer elements than numLayers, in which - /// case it cycles) - /// @param isCrossAttention Whether the attention is cross attention. - /// @param kvFactor KV factor. - /// @return Cache size per token for the disaggregated layers. Note that window size is not included in the result - /// here. - [[nodiscard]] static std::map calculateCacheSizePerTokenForDisagg( - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - std::vector const& maxAttentionWindowVec, bool isCrossAttention, SizeType32 kvFactor); - - void terminateRequest(LlmRequestPtr const& llmRequest, bool pause = false) override; - - /// @brief Terminate request in the next forwardSync call that includes the request. - /// @details This function does not terminate requests immediately. It will add the requests to the - /// mReqIdsToTerminate set. The requests will be terminated in the next forwardSync call that - /// includes the request in the batch. - void terminateRequestSync(LlmRequestPtr const& llmRequest, executor::FinishReason finishReason) override; - - /// @brief Function that waits for the decoding of requests in flight. - /// When the requests have finished or using speculative decoding, the state of requests - /// will become LlmRequestState::kGENERATION_COMPLETE. Else, it will be set to - /// LlmRequestState::kGENERATION_IN_PROGRESS. - void forwardSync() override; - - /// @brief Function that tries to advance the active requests. - /// Depending on resources available, it's possible that not all requests will get advanced. - /// Requests that may be in state LlmRequestState::kCONTEXT_INIT become - /// LlmRequestState::kGENERATION_IN_PROGRESS or LlmRequestState::kGENERATION_TO_COMPLETE. - /// @param activeRequests The list of request to try to advance. - void forwardAsync(RequestList const& activeRequests) override; - - /// @brief Override the runtime batch size for the model - void setRuntimeBatchSize(SizeType32 runtimeMaxBatchSize) override; - - /// @brief Get the runtime batch size for the model - [[nodiscard]] SizeType32 getRuntimeBatchSize() const override; - - /// @brief Override the runtime max num tokens for the model - void setRuntimeMaxNumTokens(SizeType32 runtimeMaxNumTokens) override; - - void updatePeftCache(std::shared_ptr const& llmRequest) override; - - [[nodiscard]] IterationStatsIFB getLastIterationStats() const - { - return mLastIterationStatsIFB; - } - - [[nodiscard]] TrtGptModelType getModelType() const override - { - return mCtxGenFusion ? TrtGptModelType::InflightFusedBatching : TrtGptModelType::InflightBatching; - }; - - [[nodiscard]] runtime::BufferManager const& getBufferManager() const override; - [[nodiscard]] runtime::BufferManager::CudaStreamPtr getRuntimeStreamPtr() const override; - - void getCurrentIterationStats(executor::IterationStats& stats) const override; - void getCurrentRequestStats(executor::RequestStatsPerIteration& stats) const override; - [[nodiscard]] executor::DebugTensorsPerIteration getCurrentDebugTensors() const override; - - [[nodiscard]] executor::IterationType getIterCounter() const noexcept override - { - return mIterCounter; - } - - [[nodiscard]] static bool executorConfigIsValid( - runtime::ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig); - [[nodiscard]] static executor::ExecutorConfig fixExecutorConfig( - runtime::ModelConfig const& modelConfig, executor::ExecutorConfig const& executorConfig); - - void prepareDisaggGenInitRequests(RequestList const& activeRequests, RequestVector& newGenReques); - void checkDisaggGenTransferStatus(RequestList const& activeRequests); - void prepareDistGenBufferAndDecoder(RequestVector const& generationRequests); - - void resetIterationStats() override; - - runtime::SpeculativeDecodingMode getSpeculativeDecodingMode() const noexcept - { - return mModelConfig.getSpeculativeDecodingMode(); - } - - [[nodiscard]] SizeType32 numCachedCudaGraphs() const; - -private: - friend class TrtGptModelTest; - - [[nodiscard]] SizeType32 getContextBufferId() const - { - return mMicroBatchId; - } - - [[nodiscard]] SizeType32 getGenerationBufferId() const - { - return mNumMicroBatches + mMicroBatchId; - } - - [[nodiscard]] SizeType32 getFusedBufferId() const - { - return mMicroBatchId; - } - - [[nodiscard]] SizeType32 getNextMicroBatchId(SizeType32 bufferId) const - { - return (bufferId + 1) % mNumMicroBatches; - } - - [[nodiscard]] SizeType32 getPrevMicroBatchId(SizeType32 bufferId) const - { - return (bufferId + mNumMicroBatches - 1) % mNumMicroBatches; - } - - //! @brief Store full kv cache blocks contributed by req. - //! These blocks become reusable from next step. - void storeContextBlocks(std::shared_ptr const& req); - - //! @brief Store newest kv cache block for reuse. - //! The block become reusable from next step. - void storeNewBlock(std::shared_ptr const& req); - - //! @brief Set LayerProfiler to collect performance per layer. - void setLayerProfiler() override; - - //! @brief Print profile information per layer. - std::string getLayerProfileInfo() const override; - - std::tuple prepareBuffers( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId); - - //! @brief Capture graph of current batch state during engine execution. - //! This is based on the assumptions that - //! a) We can hide CPU graph capture behind the GPU engine execution. - //! b) Batch size in the next iterations won't change and we can reuse the graph multiple times. - void prepareGraph(SizeType32 bufferId, SizeType32 optProfileId); - - void executeContext(SizeType32 runtimeContextId, SizeType32 bufferId); - void executeBatch(ScheduledRequests const& scheduledRequests); - void executeStep( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId); - - void debugIOTensors(RequestVector const& contextRequests, RequestVector const& generationRequests, - TensorMap const& inputMap, TensorMap const& outputMap); - - void createRuntimeContexts(); - void createDecoder(std::optional const& decodingModeOpt); - void createBuffers(executor::DecodingConfig const& decodingConfig, - std::optional> const& additionalModelOutputs); - std::unique_ptr createKvCacheManager(KvCacheConfig const& kvCacheConfig, KvCacheType kvCacheType, - uint64_t freePrimaryMemBytes, uint64_t freeSecondaryMemBytes, size_t extraCostMemory, - bool const failFastOnAttentionWindowTooLarge = false); - void createRnnStateManager(); - void createCustomAllReduceWorkspace(); - void createRuntimePerfKnobsTensor(executor::ExtendedRuntimePerfKnobConfig const& extendedRuntimePerfKnobConfig); - - /// @brief Verify draft token length and beam width of all active requests. - /// May change operating beam width if all requests agree on same beam width. - void verifyRequests(RequestList const& activeRequests); - - /// @brief Change the operating beam width. - /// Only possible if no requests are currently in-flight. - /// @param beamWidth New operating beam width. Must be smaller than initial maxBeamWidth. - void changeBeamWidth(SizeType32 beamWidth); - - SizeType32 getOperatingBeamWidth() const override - { - return mOperatingBeamWidth; - } - - /// @details Should be called after setting up the current batch in executeBatch to get the correct number of - /// context tokens. - IterationStatsIFB fillIterationStats( - ScheduledRequests const& scheduledRequests, RequestVector const& requestsToPause); - - /// @brief Function that sets up the TensorRT execution context that is going to be used for execution. If multiple - /// TensorRT optimization profiles are built in the engine, it selects the corresponding context that is going to be - /// used, and prepares the input and output tensors so that both buffers and the context is ready for the execution. - /// @return The TensorRT execution context index that has been setup. - void setupContext( - RequestVector const& contextRequests, RequestVector const& generationRequests, SizeType32 bufferId); - - void setupDecoderStep( - RequestVector const& contextRequests, RuntimeBuffers const& buffers, DecoderInputBuffers& inputBuffers); - runtime::CudaEvent decoderStepAsync(ScheduledRequests const& scheduledRequests); - std::vector> decoderSync( - ScheduledRequests const& scheduledRequests, std::optional const& decoderFinishEvent); - - std::vector> communicateDecoderBuffers(bool returnLogProbs); - void updateRequests(ScheduledRequests const& scheduledRequests); - - /// @brief It gathers the logits if they need to be returned, calls getDecoderSlotHostOutputs, - /// and overwrites the llmRequest tokens buffer. - /// Called either on request finishing, or at every step when doing beam search and streaming. - void postProcessRequest(LlmRequest& llmReq, std::vector const& numDroppedTokens); - /// @brief Reorders generation logits to match finalized beam paths after gatherTree. - /// During beam search, logits are stored by beam slot. After finalization, output_ids are - /// reordered by parentIds, but logits are not. This method traces parentIds on the host - /// to build the slot mapping and reindexes the logits accordingly. - void reorderGenerationLogitsForBeamSearch(LlmRequest& llmReq, SizeType32 seqSlot, SizeType32 reqBeamWidth, - SizeType32 maxSeqLength, TokenIdType const* outputIdsHostData, SizeType32 const* sequenceLengthsHostData); - /// @brief Calls gatherTree (via finalize) and transmits the received data across ranks if PP>1 - void getDecoderSlotHostOutputs( - SizeType32 seqSlot, bool returnLogProbs, runtime::SamplingConfig const& samplingConfig, bool streaming); - void rewindKVCacheBlocks(SizeType32 numSequences); - void setupSpeculativeDecodingModule(executor::DecodingConfig const& decodingConfig); - - /// @brief Copies the content of the cache indirection outputs to the cache indirection inputs. - /// @param[in] scheduledRequests The requests to copy the cache indirections for. - /// @param[in] genBufferId The id of the generation buffers for those requests. - void copyCacheIndirectionFromOutputsToInputs(ScheduledRequests const& scheduledRequests, SizeType32 genBufferId); - - [[nodiscard]] bool getGatherGenerationLogits() const override - { - return getModelConfig().computeGenerationLogits() || mGatherGenerationLogits; - } - - [[nodiscard]] runtime::ModelConfig const& getModelConfig() const override - { - return mModelConfig; - } - - [[nodiscard]] runtime::WorldConfig const& getWorldConfig() const override - { - return mWorldConfig; - } - - [[nodiscard]] SizeType32 getNumMicroBatches() const override - { - return mNumMicroBatches; - } - - [[nodiscard]] nvinfer1::DataType getLogitDataType() const override; - - [[nodiscard]] nvinfer1::DataType getTensorDataType(std::string const& name) const override; - - [[nodiscard]] nvinfer1::Dims getTensorShape(std::string const& name) const override; - - void reshapeKvTensors(OffsetTableDimensions const& dims); - - [[nodiscard]] bool hasSpeculativeDecodingFastLogits() const noexcept override - { - return mSpeculativeDecodingFastLogits; - } - - [[nodiscard]] bool hasGuidedDecoder() const noexcept override - { - return static_cast(mGuidedDecoder); - } - - using BlocksPerWindow = std::map>; - /// @brief Based on the KV-cache manager's capacity and configuration, we adjust the maximum supported attention - /// window. - /// - /// @param blocksPerWindow map of window size to number of blocks. - /// @param failFastOnAttentionWindowTooLarge if true, the function will report a runtime error if the attention - /// window is too large to fit even a single sequence in the KV cache. - /// @return pair of new blocks per window and new maxAttentionWindowVec - [[nodiscard]] std::pair> clampWindowSizesToFitAtLeastOneSequence( - BlocksPerWindow const& blocksPerWindow, bool const failFastOnAttentionWindowTooLarge = false); - - /// @brief Change the speculative decoding mode. - void changeSpecDecMode(ScheduledRequests const& scheduledRequests); - - void prefetchNextPromptTableChunk(RequestVector const& contextRequests, bool isFirstChunk, SizeType32 bufferId); - - void remapInputTokensForPromptTable( - std::shared_ptr const& llmReq, bool isCurrentChunk, SizeType32 bufferId, SizeType32 contextId); - - void copyPromptTableToGpuInChunk(std::shared_ptr const& llmReq, - std::vector const& outOfVocabTokens, bool useCurrentBuffer, SizeType32 bufferId, SizeType32 contextId); - -protected: - std::shared_ptr getKVCacheManager() override - { - return mKvCacheManager; - } - - [[nodiscard]] std::shared_ptr getKVCacheManager() const override - { - return mKvCacheManager; - } - - std::shared_ptr getCrossKVCacheManager() - { - return mCrossKvCacheManager; - } - - [[nodiscard]] std::shared_ptr getCrossKVCacheManager() const - { - return mCrossKvCacheManager; - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() override - { - return mPeftCacheManager; - } - - [[nodiscard]] std::shared_ptr getPeftCacheManager() const override - { - return mPeftCacheManager; - } - - void setLogitsPostProcessorBatched(std::optional logitsPostProcessorBatched) override - { - mLogitsPostProcessorBatched = logitsPostProcessorBatched; - } - - void setReplicateLogitsPostProcessor(bool replicateLogitsPostProcessor) override - { - mReplicateLogitsPostProcessor = replicateLogitsPostProcessor; - } - - [[nodiscard]] bool getReplicateLogitsPostProcessor() const override - { - return mReplicateLogitsPostProcessor; - } - - SizeType32 getMaxCapacityBatchSize(SizeType32 inputLength, SizeType32 outputLength) const override; - -private: - /******************** Configs ********************/ - // Parameters of the model (TRT engine) - runtime::ModelConfig mModelConfig; - // Parameters of the execution environment - runtime::WorldConfig mWorldConfig; - // Device ID of this instance - int mDevice{-1}; - // Config for (speculative) decoding - executor::DecodingConfig mDecodingConfig; - // Performance knobs for the engine. - executor::ExtendedRuntimePerfKnobConfig mExtendedRuntimePerfKnobConfig; - TensorPtr mExtendedRuntimePerfKnobsHost; - // Config for debugging output - std::optional mDebugConfig; - // List of additional outputs for each request - std::optional> mAdditionalModelOutputs; - - /******************** Components ********************/ - std::shared_ptr mLogger; - // Runner for the TRT engine. The engine produces logits. - std::unique_ptr mRuntime; - // Decoder that generates new tokens from the logits. - std::unique_ptr mDecoder; - // Decoder state for all requests - std::unique_ptr mDecoderState; - // Synchronization handles for decoder - std::vector> mDecoderFinishedEvents; - - // Manager that maps requests to slots - std::shared_ptr mSeqSlotManager; - // KV cache manager for attention layers (optional) - std::shared_ptr mKvCacheManager; - // KV cache manager for cross attention in enc-dec models (optional) - std::shared_ptr mCrossKvCacheManager = nullptr; - // RNN state manager for recurrent layers (optional) - std::unique_ptr mRnnStateManager; - // PEFT cache manager for LoRA tasks (optional) - std::shared_ptr mPeftCacheManager; - // BufferManager using a separate stream for async copy operations. - runtime::BufferManager mCopyBufferManager; - // Event for async data transfers - runtime::CudaEvent mPtableCopyDoneEvent; - - /******************** Logits Post-Processor ********************/ - std::optional mLogitsPostProcessorBatched; - bool mReplicateLogitsPostProcessor{true}; - // Set if any request invoked a logits processor in current step - bool mLogitsPostProcessorIsApplied{false}; - - constexpr bool broadcastPostDecoder() - { - return mWorldConfig.isTensorParallel() && !mReplicateLogitsPostProcessor && mLogitsPostProcessorIsApplied; - } - - std::unique_ptr mGuidedDecoder; - - /******************** Pipeline parallelism ********************/ - std::unique_ptr mMpiCommPipelinePara; - std::vector> mDecStepAsyncSndHdls; - std::vector> mDecSlotAsyncSndHdls; - std::unique_ptr mAsyncSendWaitThread; - - /******************** Tensor parallelism ********************/ - std::unique_ptr mMpiCommTensorPara; - std::unique_ptr mAllReduceBuffers; - - /******************** Runtime parameters ********************/ - // Flag to select fused or unfused context+generation execution - bool mCtxGenFusion; - // ID of current micro batch, changes after each iteration - SizeType32 mMicroBatchId{0}; - // Number of micro batches. Multiple batches are used for overlapping setup and execution, - // and in pipeline parallelism. - SizeType32 mNumMicroBatches; - // Number of buffers to be added to mBuffers. - SizeType32 mNumBuffers; - // Current operating beam width. Can be changed with changeBeamWidth function. - SizeType32 mOperatingBeamWidth; - // Runtime batch size optimized during execution for microBatchScheduler: - /// The max batch size recommended by the dynamic tuner - SizeType32 mMaxBatchSizeTunerRecommended; - /// The min of mMaxBatchSize and mMaxBatchSizeTunerRecommended - SizeType32 mMaxBatchSizeRuntime; - // Runtime max num tokens optimized during execution for microBatchScheduler: - /// Build time max num tokens - std::optional mMaxNumTokensStatic; - /// The max num tokens recommended by the dynamic tuner - SizeType32 mMaxNumTokensTunerRecommended; - /// The min of mMaxNumTokens and mMaxNumTokensTunerRecommended - std::optional mMaxNumTokensRuntime; - // Controls if generation logits should be gathered, so that returnGenerationLogits can be requested. - bool mGatherGenerationLogits{false}; - // offloading and prefetching the prompt tuning table (only effective in chunked prefill mode) - bool mPromptTableOffloading; - - /******************** Buffers ********************/ - // Buffers for each micro batch. Unfused path (mCtxGenFusion==false) uses two times the buffers. - std::vector> mBuffers; - // Decoder input buffers for each micro batch. - std::vector mDecoderInputBuffers; - // Decoder output buffers for each micro batch. - std::vector mDecoderOutputBuffers; - // Buffers for each slot in the decoder - std::vector> mSlotDecoderBuffers; - // PEFT table for each micro batch - std::vector mPeftTables; - - /******************** Book keeping ********************/ - // List of requests in each micro batch - std::vector mMicroBatchScheduledRequests; - // Set of in-flight requests of *all* micro batches - ReqIdsSet mInflightReqIds; - // Requests that should be terminated (requested from outside the model) - std::unordered_map mReqIdsToTerminate; - // Requests that the scheduler selected to be paused - ReqIdsSet mReqIdsToPause; - // Stats collected in last iteration - IterationStatsIFB mLastIterationStatsIFB{-1}; - // Iteration counter used to distinguish debug output - executor::IterationType mIterCounter{0}; - // Debug tensors of last itreation - TensorMap mLastIterationDebugTensors; - // Cuda graph instances for each microbatch. - std::vector mCudaGraphExecutorCaches; - - /******************** Cache transceiver ********************/ - std::unique_ptr mCacheTransceiver; - - /******************** Spec dec ***********************/ - std::unique_ptr mDraftModelSendLogitsThread; - bool mSpeculativeDecodingFastLogits; - std::atomic mDraftModelThreadShouldExit{false}; - bool mIsLeaderInOrchMode{false}; - // List of completed draft requests which logits will need to be sent to the target model. - // Guarded by mDraftRequestsMtx (shared with the background logits sender thread). - RequestVector mDraftRequestsWaitingToSendLogits; - // Draft requests whose logits have been sent — pending termination by main thread. - // Guarded by mDraftRequestsMtx. - RequestVector mDraftRequestsDoneSendingLogits; - std::mutex mDraftRequestsMtx; - SizeType32 mSeamlessLADMaxDraftLen{0}; - bool mUseSeamlessLookahead{false}; - RewindInputs mRewindInputs; - - /******************** Algorithms ********************/ - // Algorithms are reentrant, they are assigned a state at - // construction time and it is not modified through execution, hence they are const. - // Schedulers that select which requests to run in each iteration - std::unique_ptr mCapacityScheduler; - std::unique_ptr mMicroBatchScheduler; - std::unique_ptr mPauseRequests; - std::unique_ptr mAssignReqSeqSlots; - std::unique_ptr mAllocateKvCache; - std::unique_ptr mHandleContextLogits; - std::unique_ptr mHandleGenerationLogits; - std::unique_ptr mLogitsPostProcessor; - std::unique_ptr mMakeDecodingBatchInputOutput; - std::unique_ptr mCreateNewDecoderRequests; - std::unique_ptr mUpdateDecoderBuffers; -}; - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/updateDecoderBuffers.cpp b/cpp/tensorrt_llm/batch_manager/updateDecoderBuffers.cpp deleted file mode 100644 index ead120135f3a..000000000000 --- a/cpp/tensorrt_llm/batch_manager/updateDecoderBuffers.cpp +++ /dev/null @@ -1,78 +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/updateDecoderBuffers.h" -#include "tensorrt_llm/batch_manager/decoderBuffers.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/runtime/decoderState.h" -#include "tensorrt_llm/runtime/iTensor.h" - -namespace tensorrt_llm::batch_manager -{ - -using BufferManager = tensorrt_llm::runtime::BufferManager; -using TensorPtr = runtime::ITensor::SharedPtr; -using ITensor = runtime::ITensor; -using SizeType32 = tensorrt_llm::runtime::SizeType32; - -runtime::CudaEvent UpdateDecoderBuffers::operator()(runtime::ModelConfig const& modelConfig, - DecoderOutputBuffers& decoderOutputBuffers, runtime::BufferManager const& copyBufferManager, - runtime::decoder::DecoderState const& decoderState, bool returnLogProbs, - runtime::CudaEvent const& decoderFinishEvent) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(updateDecoderBuffers); - - // Chain copy after decoder event, using a different stream - copyBufferManager.getStream().wait(decoderFinishEvent); - - copyBufferManager.copy(*decoderState.getAllNewTokens(), *decoderOutputBuffers.newOutputTokensHost); - copyBufferManager.copy(*decoderState.getSequenceLengths(), *decoderOutputBuffers.sequenceLengthsHost); - - auto const finishedSumDevice = decoderState.getFinishedSum(); - copyBufferManager.copy(*finishedSumDevice, *decoderOutputBuffers.finishedSumHost); - auto const finishReasonsDevice = decoderState.getFinishReasons(); - copyBufferManager.copy(*finishReasonsDevice, *decoderOutputBuffers.finishReasonsHost); - - if (returnLogProbs) - { - copyBufferManager.copy(*decoderState.getCumLogProbs(), *decoderOutputBuffers.cumLogProbsHost); - copyBufferManager.copy(*decoderState.getLogProbs(), *decoderOutputBuffers.logProbsHost); - } - - if (modelConfig.getSpeculativeDecodingMode().predictsDraftTokens()) - { - // TODO: keep data on device for next iteration - copyBufferManager.copy(*decoderState.getNextDraftTokens(), *decoderOutputBuffers.nextDraftTokensHost); - - if (modelConfig.getSpeculativeDecodingMode().variableDraftLength()) - { - copyBufferManager.copy( - *decoderState.getNextDraftTokensLengths(), *decoderOutputBuffers.nextDraftTokensLengthsHost); - copyBufferManager.copy( - *decoderState.getPrevDraftTokensLengths(), *decoderOutputBuffers.prevDraftTokensLengthsHost); - } - } - - runtime::CudaEvent copyEvent{}; - copyBufferManager.getStream().record(copyEvent); - // Store the event for later sync. Sync stream before calling next decoder. Sync host before updating requests. - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return copyEvent; -} - -} // namespace tensorrt_llm::batch_manager diff --git a/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h b/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h index e4732a75f649..2a0e14c29989 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h +++ b/cpp/tensorrt_llm/batch_manager/utils/debugUtils.h @@ -25,7 +25,6 @@ namespace tensorrt_llm::runtime { -class TllmRuntime; } // namespace tensorrt_llm::runtime namespace tensorrt_llm::batch_manager::utils diff --git a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp index bdb12886337c..a3e54a6b0f9b 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp +++ b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.cpp @@ -16,7 +16,6 @@ */ #include "inflightBatchingUtils.h" -#include "tensorrt_llm/runtime/runtimeKernels.h" namespace tensorrt_llm::batch_manager::utils { @@ -88,169 +87,6 @@ void moveFinishedContextRequestsToGeneration(ScheduledRequests& scheduledRequest TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -void copyGenerationLogits(RuntimeBuffers::GenerationLogitsCache& generationLogitsCache, - runtime::BufferManager const& bufferManager, LlmRequest& llmReq, bool beforeDecoder, - std::vector const& numDroppedTokens) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO( - !beforeDecoder || numDroppedTokens.empty(), "numDroppedTokens are only possible after decoder."); - - auto const reqBeamWidth = llmReq.getBeamWidthByIter(); - TLLM_CHECK_WITH_INFO(numDroppedTokens.empty() || numDroppedTokens.size() == static_cast(reqBeamWidth), - "Dropped tokens have to be defined for all beams."); - - auto const fragmentSize = llmReq.getGenerationLogitsFragmentsSize(); - - // Merge logits fragments on device - auto const& transposeBufferPtr = generationLogitsCache.transposedLogits; - auto const& cachePointerDevice = generationLogitsCache.fragmentPointerDevice; - auto const& cachePointerHost = generationLogitsCache.getFragmentPointerHost(); - tensorrt_llm::runtime::kernels::mergeLogitsFragments(bufferManager, *transposeBufferPtr, - llmReq.getGenerationLogitsFragments(), *cachePointerDevice, *cachePointerHost, 0, 1, reqBeamWidth, - bufferManager.getStream(), 0); - llmReq.clearGenerationLogitsFragments(); - - // Copy logits to host - for (SizeType32 beam = 0; beam < reqBeamWidth; beam++) - { - auto const droppedSize = !numDroppedTokens.empty() ? numDroppedTokens.at(beam) : 0; - // Ignore logits of dropped tokens - auto const beamFragmentSize = fragmentSize - droppedSize; - // If this function is called before the decoder, the request does not contain the generated token of the - // current iteration, so we add 1 to the number of tokens. - auto const numGenerationToken - = static_cast(beforeDecoder) + llmReq.getNumTokens(beam) - llmReq.mPromptLen; - auto const hostOffset = numGenerationToken - beamFragmentSize; - - // [beamWidth, GENERATION_LOGITS_BUFFER_LENGTH, vocabSizePadded] -> [beamFragmentSize, vocabSizePadded] - auto beamDeviceTensorPtr = ITensor::slice(transposeBufferPtr, {beam, 0}, beamFragmentSize); - // [beamWidth, mMaxNewTokens, vocabSizePadded] -> [beamFragmentSize, vocabSizePadded] - auto beamHostTensorPtr = ITensor::slice(llmReq.getGenerationLogitsHost(), {beam, hostOffset}, beamFragmentSize); - bufferManager.copy(*beamDeviceTensorPtr, *beamHostTensorPtr); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -namespace -{ - -std::pair findOutputTensor(std::string const& outputTensorName, - std::vector const& additionalModelOutputs, - RuntimeBuffers::TensorMap const& outputMap, bool isContext) -{ - auto const aoIter = std::find_if(additionalModelOutputs.cbegin(), additionalModelOutputs.cend(), - [&outputTensorName](auto const& ao) { return ao.name == outputTensorName; }); - TLLM_CHECK_WITH_INFO(aoIter != additionalModelOutputs.cend(), "Additional %s output tensor not found: %s", - isContext ? "context" : "generation", outputTensorName.c_str()); - - auto const gatherContext = aoIter->gatherContext; - if (isContext) - { - TLLM_CHECK_WITH_INFO( - gatherContext, "Additional context output tensor not gathered: %s", outputTensorName.c_str()); - } - - auto const tensorIt = outputMap.find(outputTensorName); - TLLM_CHECK_WITH_INFO(tensorIt != outputMap.end(), "Additional %s output tensor not found: %s", - isContext ? "context" : "generation", outputTensorName.c_str()); - - return {tensorIt->second, gatherContext}; -} - -} // namespace - -void copyAdditionalOutputs(std::vector const& additionalModelOutputs, - RequestVector const& contextRequests, RequestVector const& generationRequests, - RuntimeBuffers::TensorMap const& outputMap, runtime::BufferManager const& manager) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // One index shared across all output tensors that have gatherContext - SizeType32 srcTensorIndexWithContext{0}; - // One index shared across all output tensors that do not have gatherContext - SizeType32 srcTensorIndexWithoutContext{0}; - - for (auto const& llmReq : contextRequests) - { - auto numContextTokens = llmReq->getContextChunkSize(); - for (auto const& outputTensor : llmReq->getAdditionalContextOutputs()) - { - auto const& [tensor, gatherContext] - = findOutputTensor(outputTensor.first, additionalModelOutputs, outputMap, true); - - auto const srcTensorIndex = srcTensorIndexWithContext; - auto srcView = ITensor::slice(tensor, srcTensorIndex, numContextTokens); - auto dstView = ITensor::slice(outputTensor.second, llmReq->getContextCurrentPosition(), numContextTokens); - manager.copy(*srcView, *dstView); - } - srcTensorIndexWithContext += numContextTokens; - srcTensorIndexWithoutContext += 1; - - // Copy output of last token to generation outputs - if (llmReq->isLastContextChunk()) - { - for (auto const& outputTensor : llmReq->getAdditionalGenerationOutputs()) - { - auto const& [tensor, gatherContext] - = findOutputTensor(outputTensor.first, additionalModelOutputs, outputMap, false); - - auto const srcTensorIndex = gatherContext ? srcTensorIndexWithContext : srcTensorIndexWithoutContext; - auto srcView = ITensor::slice(tensor, srcTensorIndex - 1, 1); - for (SizeType32 beam = 0; beam < llmReq->getBeamWidthByIter(); beam++) - { - auto dstView = ITensor::slice(outputTensor.second, {beam, 0}, 1); - manager.copy(*srcView, *dstView); - } - } - } - } - - for (auto const& llmReq : generationRequests) - { - auto const reqBeamWidth = llmReq->getBeamWidthByIter(); - for (auto const& outputTensor : llmReq->getAdditionalGenerationOutputs()) - { - auto const& [tensor, gatherContext] - = findOutputTensor(outputTensor.first, additionalModelOutputs, outputMap, false); - - auto const srcTensorIndex = gatherContext ? srcTensorIndexWithContext : srcTensorIndexWithoutContext; - for (SizeType32 beam = 0; beam < reqBeamWidth; beam++) - { - auto const generatedLength = llmReq->getNumTokens(beam) - llmReq->getPromptLen(); - TLLM_CHECK(generatedLength >= 1); - auto srcView = ITensor::slice(tensor, srcTensorIndex + beam, 1); - auto dstView = ITensor::slice(outputTensor.second, {beam, generatedLength}, 1); - manager.copy(*srcView, *dstView); - } - } - srcTensorIndexWithContext += reqBeamWidth; - srcTensorIndexWithoutContext += reqBeamWidth; - } - - // Check final indices - for (auto const& outputTensor : additionalModelOutputs) - { - auto const& outputTensorName = outputTensor.name; - auto const gatherContext = outputTensor.gatherContext; - - auto const tensorIt = outputMap.find(outputTensorName); - TLLM_CHECK_WITH_INFO( - tensorIt != outputMap.end(), "Additional output tensor not found: %s", outputTensorName.c_str()); - - auto const& outputShape = tensorIt->second->getShape(); - auto const outputSize = outputShape.d[0]; - auto const finalIndex = gatherContext ? srcTensorIndexWithContext : srcTensorIndexWithoutContext; - - TLLM_CHECK_WITH_INFO(finalIndex == outputSize, "Additional %s output tensor final index mismatch %d != %ld: %s", - gatherContext ? "context" : "generation", finalIndex, outputSize, outputTensorName.c_str()); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - void terminateRequest(SequenceSlotManager& seqSlotManager, LlmRequest& llmReq, SizeType32 maxInputLen, OptionalRef kvCacheManager, OptionalRef crossKvCacheManager, @@ -299,109 +135,4 @@ std::vector getRequestBeamWidths( return beamWidths; } -void CudaGraphExecutor::create(cudaGraph_t const& graph) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - assert(mInstance == nullptr); - TLLM_CUDA_CHECK(cudaGraphInstantiate(&mInstance, graph, nullptr, nullptr, 0)); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void CudaGraphExecutor::uploadToStream(runtime::CudaStream const& stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - assert(hasInstance()); - TLLM_CUDA_CHECK(cudaGraphUpload(mInstance, stream.get())); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void CudaGraphExecutor::launch(runtime::CudaStream const& stream) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_CUDA_CHECK(cudaGraphLaunch(mInstance, stream.get())); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -bool CudaGraphExecutor::update(cudaGraph_t const& graph) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - return cudaGraphExecUpdate(mInstance, graph, nullptr) != cudaSuccess; -} - -void CudaGraphExecutor::clear() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - if (mInstance != nullptr) - { - TLLM_CUDA_CHECK(cudaGraphExecDestroy(mInstance)); - mInstance = nullptr; - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void CudaGraphExecutor::prepareNextGraph(std::unique_ptr& runtime, SizeType32 nextContextId) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto& stream = runtime->getStream(); - - cudaGraph_t nextGraph; - TLLM_CUDA_CHECK(cudaStreamBeginCapture(stream.get(), cudaStreamCaptureModeThreadLocal)); - runtime->executeContext(nextContextId); - TLLM_CUDA_CHECK(cudaStreamEndCapture(stream.get(), &nextGraph)); - - if (hasInstance()) - { - if (update(nextGraph)) - { - clear(); - create(nextGraph); - } - } - else - { - create(nextGraph); - } - - TLLM_CUDA_CHECK(cudaGraphDestroy(nextGraph)); - uploadToStream(stream); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::optional> CudaGraphExecutorCache::get(BatchState const& state) -{ - auto it = mMap.find(state); - if (it == mMap.end()) - { - return std::nullopt; - } - mCache.splice(mCache.begin(), mCache, it->second); - return it->second->second; -} - -void CudaGraphExecutorCache::put(BatchState const& state, std::shared_ptr const& value) -{ - auto it = mMap.find(state); - if (it != mMap.end()) - { - mCache.erase(it->second); - } - mCache.emplace_front(state, value); - mMap[state] = mCache.begin(); - - if (static_cast(mMap.size()) > mCapacity) - { - auto lastState = mCache.back().first; - mCache.pop_back(); - mMap.erase(lastState); - } -} - -void CudaGraphExecutorCache::clear() -{ - // Releasing the shared_ptrs runs ~CudaGraphExecutor, which calls - // cudaGraphExecDestroy on each cached instance. - mMap.clear(); - mCache.clear(); -} - } // namespace tensorrt_llm::batch_manager::utils diff --git a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h index fe0c4e505218..374ae398f781 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h +++ b/cpp/tensorrt_llm/batch_manager/utils/inflightBatchingUtils.h @@ -20,7 +20,6 @@ #include "tensorrt_llm/batch_manager/common.h" #include "tensorrt_llm/batch_manager/kvCacheManager.h" #include "tensorrt_llm/batch_manager/peftCacheManager.h" -#include "tensorrt_llm/batch_manager/runtimeBuffers.h" #include "tensorrt_llm/batch_manager/sequenceSlotManager.h" #include "tensorrt_llm/common/optionalRef.h" #include "tensorrt_llm/runtime/iTensor.h" @@ -50,17 +49,6 @@ void sortRequests(RequestVector& contextRequests, RequestVector& generationReque //! @param scheduledRequests The scheduled context and generation requests. void moveFinishedContextRequestsToGeneration(ScheduledRequests& scheduledRequests); -//! @param beforeDecoder Whether the function is called before the decoder. If it is true, correct the output offset. -//! @param numDroppedTokens The number of dropped tokens for each beam (e.g. when the requests finished early). -//! Generation logits for dropped tokens are ignored. -void copyGenerationLogits(RuntimeBuffers::GenerationLogitsCache& generationLogitsCache, - runtime::BufferManager const& bufferManager, LlmRequest& llmReq, bool beforeDecoder, - std::vector const& numDroppedTokens = {}); - -void copyAdditionalOutputs(std::vector const& additionalModelOutputs, - RequestVector const& contextRequests, RequestVector const& generationRequests, - RuntimeBuffers::TensorMap const& outputMap, runtime::BufferManager const& manager); - void terminateRequest(SequenceSlotManager& seqSlotManager, LlmRequest& llmRequest, SizeType32 maxInputLen, OptionalRef kvCacheManager = std::nullopt, OptionalRef crossKvCacheManager = std::nullopt, @@ -68,66 +56,4 @@ void terminateRequest(SequenceSlotManager& seqSlotManager, LlmRequest& llmReques std::vector getRequestBeamWidths( RequestVector const& contextRequests, RequestVector const& generationRequests); - -class CudaGraphExecutor -{ -public: - CudaGraphExecutor() = default; - - ~CudaGraphExecutor() - { - try - { - clear(); - } - catch (std::exception& e) - { - TLLM_LOG_EXCEPTION(e); - } - } - - bool hasInstance() const - { - return mInstance != nullptr; - } - - void clear(); - void prepareNextGraph(std::unique_ptr& runtime, SizeType32 nextContextId); - void launch(runtime::CudaStream const& stream); - -private: - void create(cudaGraph_t const& graph); - bool update(cudaGraph_t const& graph); - void uploadToStream(runtime::CudaStream const& stream); - - cudaGraphExec_t mInstance = nullptr; -}; - -class CudaGraphExecutorCache -{ - /// @brief LRU cache to store cuda graph instances. -public: - explicit CudaGraphExecutorCache(runtime::SizeType32 capacity) - : mCapacity(capacity) - { - } - - std::optional> get(BatchState const& state); - - void put(BatchState const& state, std::shared_ptr const& value); - - void clear(); - - [[nodiscard]] runtime::SizeType32 size() const noexcept - { - return static_cast(mCache.size()); - } - -private: - using BatchStateGraphExecutorPair = std::pair>; - using GraphExecutorLruCache = std::list; - SizeType32 mCapacity; - GraphExecutorLruCache mCache; - std::unordered_map mMap; -}; } // namespace tensorrt_llm::batch_manager::utils diff --git a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp index 941c1b655073..30799aff2967 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp +++ b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.cpp @@ -124,7 +124,7 @@ void draftModelSendLogitsThread(int device, std::atomic* draftModelThreadS } void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost, - executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype) + executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, tensorrt_llm::DataType logitsDtype) { #if ENABLE_MULTI_DEVICE auto const& worldComm = tensorrt_llm::mpi::MpiComm::world(); diff --git a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h index 637f8a850610..73a00bc04876 100644 --- a/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h +++ b/cpp/tensorrt_llm/batch_manager/utils/logitsThread.h @@ -46,6 +46,6 @@ void draftModelSendLogitsThread(int device, std::atomic* draftModelThreadS std::mutex* draftRequestsMtx); void targetModelReceiveLogits(runtime::ITensor::SharedPtr& draftLogitsHost, - executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, nvinfer1::DataType logitsDtype); + executor::SpeculativeDecodingFastLogitsInfo const& fastLogitsInfo, tensorrt_llm::DataType logitsDtype); } // namespace tensorrt_llm::batch_manager::utils diff --git a/cpp/tensorrt_llm/common/attentionOp.cpp b/cpp/tensorrt_llm/common/attentionOp.cpp index 0ea70da0d73b..a34e6988adff 100644 --- a/cpp/tensorrt_llm/common/attentionOp.cpp +++ b/cpp/tensorrt_llm/common/attentionOp.cpp @@ -747,7 +747,7 @@ size_t AttentionOp::getFmhaMultiCtasKvScratchSize() const noexcept return partialStatsSize + partialOSize; } -size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t max_num_seq, int32_t input_seq_length, +size_t AttentionOp::getWorkspaceSizeForContext(tensorrt_llm::DataType type, int32_t max_num_seq, int32_t input_seq_length, int32_t cross_kv_length, int32_t max_num_tokens, int32_t total_kv_len) const noexcept { if (max_num_tokens == 0) @@ -900,7 +900,7 @@ size_t AttentionOp::getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t return context_workspace_size; } -size_t AttentionOp::getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32_t max_num_seq, +size_t AttentionOp::getWorkspaceSizeForGeneration(tensorrt_llm::DataType type, int32_t max_num_seq, int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept { if (max_num_tokens == 0) @@ -2763,7 +2763,7 @@ int AttentionOp::initialize() noexcept if (mEnableContextFMHA) { mEnableContextFMHA = false; - if (!(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16)) + if (!(mType == tensorrt_llm::DataType::kHALF || mType == tensorrt_llm::DataType::kBF16)) { TLLM_LOG_WARNING("Fall back to unfused MHA because of unsupported data type."); } @@ -2808,7 +2808,7 @@ int AttentionOp::initialize() noexcept "mFP8ContextFMHA must enable if FP4 KV cache is enabled"); TLLM_CHECK(isRoPE() == (mRotaryEmbeddingDim != 0)); - TLLM_CHECK_WITH_INFO((mSM >= 80) || (mType != nvinfer1::DataType::kBF16), + TLLM_CHECK_WITH_INFO((mSM >= 80) || (mType != tensorrt_llm::DataType::kBF16), "Unsupported data type, pre SM 80 GPUs do not support bfloat16"); // Pre-check whether the head size is supported by MMHA. @@ -2860,11 +2860,11 @@ int AttentionOp::initialize() noexcept // Pre-checked during constructing. Data_type data_type, data_type_kv; - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { data_type = DATA_TYPE_FP16; } - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { data_type = DATA_TYPE_BF16; } @@ -3019,13 +3019,13 @@ int AttentionOp::initialize() noexcept Data_type kvDataType = DATA_TYPE_FP32; Data_type outputDataType = DATA_TYPE_FP32; - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { qDataType = DATA_TYPE_FP16; kvDataType = DATA_TYPE_FP16; outputDataType = DATA_TYPE_FP16; } - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { qDataType = DATA_TYPE_BF16; kvDataType = DATA_TYPE_BF16; @@ -3110,7 +3110,7 @@ int AttentionOp::initialize() noexcept } mEnableXQA = (mEnableXQA || mIsSpecDecodingEnabled) - && (mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16) && mUseKVCache; + && (mType == tensorrt_llm::DataType::kHALF || mType == tensorrt_llm::DataType::kBF16) && mUseKVCache; if (mEnableXQA) { @@ -3120,12 +3120,12 @@ int AttentionOp::initialize() noexcept fixedParams.isMLA = mIsGenerationMLA; // TODO: support more combinations. // Update Q and O dtype. - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { fixedParams.inputDataType = DATA_TYPE_FP16; fixedParams.outputDataType = DATA_TYPE_FP16; } - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { fixedParams.inputDataType = DATA_TYPE_BF16; fixedParams.outputDataType = DATA_TYPE_BF16; diff --git a/cpp/tensorrt_llm/common/attentionOp.h b/cpp/tensorrt_llm/common/attentionOp.h index 5fcf901f8a79..6e8441894012 100644 --- a/cpp/tensorrt_llm/common/attentionOp.h +++ b/cpp/tensorrt_llm/common/attentionOp.h @@ -56,10 +56,10 @@ class AttentionOp [[nodiscard]] size_t getFmhaMultiCtasKvScratchSize() const noexcept; [[nodiscard]] int getHeadSize(bool checkInit = true) const; [[nodiscard]] int getMaxNumSeqLenTile(int batch_beam_size = 1) const; - [[nodiscard]] size_t getWorkspaceSizeForContext(nvinfer1::DataType type, int32_t nbReq, int32_t max_input_length, + [[nodiscard]] size_t getWorkspaceSizeForContext(tensorrt_llm::DataType type, int32_t nbReq, int32_t max_input_length, int32_t cross_kv_length = 0, int32_t max_num_tokens = 0, int32_t total_kv_len = 0) const noexcept; // total_num_seq is the sum of beam_width for multiple requests - [[nodiscard]] size_t getWorkspaceSizeForGeneration(nvinfer1::DataType type, int32_t total_num_seq, + [[nodiscard]] size_t getWorkspaceSizeForGeneration(tensorrt_llm::DataType type, int32_t total_num_seq, int32_t max_attention_window_size, int32_t max_num_tokens, int32_t max_blocks_per_sequence) const noexcept; template @@ -181,14 +181,14 @@ class AttentionOp if (this->context_lengths && batch_size > 0) { ss << "context_lengths: " - << *(runtime::ITensor::wrap((void*) this->context_lengths, nvinfer1::DataType::kINT32, + << *(runtime::ITensor::wrap((void*) this->context_lengths, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))) << std::endl; } if (this->sequence_lengths && batch_size > 0) { ss << "sequence_lengths: " - << *(runtime::ITensor::wrap((void*) this->sequence_lengths, nvinfer1::DataType::kINT32, + << *(runtime::ITensor::wrap((void*) this->sequence_lengths, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))) << std::endl; } @@ -478,7 +478,7 @@ class AttentionOp int mTpSize = 1; int mTpRank = 0; bool mUnfuseQkvGemm = false; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; int32_t mMaxContextLength = 0; int32_t mMaxSeqLen = 0; int32_t mMaxNumRequests = 0; diff --git a/cpp/tensorrt_llm/common/opUtils.cpp b/cpp/tensorrt_llm/common/opUtils.cpp index a738e28377a1..925cfb55dc55 100644 --- a/cpp/tensorrt_llm/common/opUtils.cpp +++ b/cpp/tensorrt_llm/common/opUtils.cpp @@ -32,18 +32,18 @@ TRTLLM_NAMESPACE_BEGIN #if ENABLE_MULTI_DEVICE -std::unordered_map* getDtypeMap() +std::unordered_map* getDtypeMap() { - static std::unordered_map dtypeMap = { - {nvinfer1::DataType::kFLOAT, ncclFloat32}, - {nvinfer1::DataType::kHALF, ncclFloat16}, - {nvinfer1::DataType::kBF16, ncclBfloat16}, - {nvinfer1::DataType::kFP8, ncclInt8}, - {nvinfer1::DataType::kBOOL, ncclInt8}, - {nvinfer1::DataType::kINT32, ncclInt32}, - {nvinfer1::DataType::kINT64, ncclInt64}, - {nvinfer1::DataType::kUINT8, ncclUint8}, - {nvinfer1::DataType::kINT8, ncclInt8}, + static std::unordered_map dtypeMap = { + {tensorrt_llm::DataType::kFLOAT, ncclFloat32}, + {tensorrt_llm::DataType::kHALF, ncclFloat16}, + {tensorrt_llm::DataType::kBF16, ncclBfloat16}, + {tensorrt_llm::DataType::kFP8, ncclInt8}, + {tensorrt_llm::DataType::kBOOL, ncclInt8}, + {tensorrt_llm::DataType::kINT32, ncclInt32}, + {tensorrt_llm::DataType::kINT64, ncclInt64}, + {tensorrt_llm::DataType::kUINT8, ncclUint8}, + {tensorrt_llm::DataType::kINT8, ncclInt8}, }; return &dtypeMap; } diff --git a/cpp/tensorrt_llm/common/opUtils.h b/cpp/tensorrt_llm/common/opUtils.h index 72e5a5ea3e09..79f467dc91b9 100644 --- a/cpp/tensorrt_llm/common/opUtils.h +++ b/cpp/tensorrt_llm/common/opUtils.h @@ -21,7 +21,7 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/workspace.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -62,14 +62,14 @@ void read(char const*& buffer, T& val) buffer += sizeof(T); } -inline cudaDataType_t trtToCublasDtype(nvinfer1::DataType type) +inline cudaDataType_t trtToCublasDtype(tensorrt_llm::DataType type) { switch (type) { - case nvinfer1::DataType::kFLOAT: return CUDA_R_32F; - case nvinfer1::DataType::kHALF: return CUDA_R_16F; + case tensorrt_llm::DataType::kFLOAT: return CUDA_R_32F; + case tensorrt_llm::DataType::kHALF: return CUDA_R_16F; #if defined(NV_TENSORRT_MAJOR) && NV_TENSORRT_MAJOR >= 9 - case nvinfer1::DataType::kBF16: return CUDA_R_16BF; + case tensorrt_llm::DataType::kBF16: return CUDA_R_16BF; #endif default: TLLM_THROW("Not supported data type for cuBLAS"); } @@ -214,7 +214,7 @@ inline bool isBuilding() } \ } while (0) -std::unordered_map* getDtypeMap(); +std::unordered_map* getDtypeMap(); std::shared_ptr getComm(std::set const& group); diff --git a/cpp/tensorrt_llm/common/safetensors.cpp b/cpp/tensorrt_llm/common/safetensors.cpp index 9171f79e44e5..8bd91ccfbd51 100644 --- a/cpp/tensorrt_llm/common/safetensors.cpp +++ b/cpp/tensorrt_llm/common/safetensors.cpp @@ -18,7 +18,7 @@ #include "nlohmann/json.hpp" #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -30,7 +30,7 @@ TRTLLM_NAMESPACE_BEGIN namespace common::safetensors { -using nvinfer1::DataType; +using tensorrt_llm::DataType; static DataType convertDataTypeStrToEnum(std::string const& str) { diff --git a/cpp/tensorrt_llm/common/safetensors.h b/cpp/tensorrt_llm/common/safetensors.h index e31225f1be24..bdecf95e5909 100644 --- a/cpp/tensorrt_llm/common/safetensors.h +++ b/cpp/tensorrt_llm/common/safetensors.h @@ -18,7 +18,7 @@ #include "tensorrt_llm/common/assert.h" #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/logger.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -34,13 +34,13 @@ class INdArray [[nodiscard]] virtual void const* data() const = 0; [[nodiscard]] virtual int ndim() const = 0; [[nodiscard]] virtual std::vector const& dims() const = 0; - [[nodiscard]] virtual nvinfer1::DataType dtype() const = 0; + [[nodiscard]] virtual tensorrt_llm::DataType dtype() const = 0; - [[nodiscard]] nvinfer1::Dims trtDims() const + [[nodiscard]] tensorrt_llm::Dims trtDims() const { - nvinfer1::Dims dims; + tensorrt_llm::Dims dims; dims.nbDims = ndim(); - TLLM_CHECK(dims.nbDims <= nvinfer1::Dims::MAX_DIMS); + TLLM_CHECK(dims.nbDims <= tensorrt_llm::Dims::MAX_DIMS); memset(dims.d, 0, sizeof(dims.d)); for (int i = 0; i < dims.nbDims; ++i) { diff --git a/cpp/tensorrt_llm/executor/CMakeLists.txt b/cpp/tensorrt_llm/executor/CMakeLists.txt index ca1ab298d7f4..bef39826847c 100644 --- a/cpp/tensorrt_llm/executor/CMakeLists.txt +++ b/cpp/tensorrt_llm/executor/CMakeLists.txt @@ -19,6 +19,7 @@ set(TARGET_DIR ${CMAKE_CURRENT_SOURCE_DIR}) # keep this list sorted alphabetically set(SRCS + kvCacheEvent.cpp cache_transmission/mpi_utils/connection.cpp cache_transmission/agent_utils/connection.cpp cache_transmission/transferAgent.cpp @@ -27,9 +28,7 @@ set(SRCS contextPhaseParams.cpp debugConfig.cpp decodingConfig.cpp - executor.cpp executorConfig.cpp - executorImpl.cpp executorKVCacheEventManager.cpp extendedRuntimePerfKnobConfig.cpp guidedDecodingConfig.cpp @@ -60,7 +59,6 @@ set(SRCS types.cpp requestUtils.cpp contextPhaseParams.cpp - disaggServerUtil.cpp cacheTransceiverConfig.cpp) if(NOT WIN32) diff --git a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu index 6e8c68d7efa7..2722d3dbf78a 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu +++ b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.cu @@ -27,7 +27,7 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -467,7 +467,7 @@ void concatKVCache(runtime::ITensor::SharedPtr* inputBlocks, int inputBlockNum, blockInfos[outputBlockNum * inputAllRankNum + oi] = fillBlockInfo(oCacheState, outputBlocks[oi], oRank); } runtime::BufferManager::IBufferPtr blockInfosDeviceBuffer - = bufferManager.gpu(sizeof(BlockInfo) * (blockInfos.size()), nvinfer1::DataType::kUINT8); + = bufferManager.gpu(sizeof(BlockInfo) * (blockInfos.size()), tensorrt_llm::DataType::kUINT8); bufferManager.copy((blockInfos.data()), *blockInfosDeviceBuffer, runtime::MemoryType::kCPU); BlockInfo* iBlockInfoDevice = static_cast*>(blockInfosDeviceBuffer->data()); @@ -594,7 +594,7 @@ void concatKVCacheDispatch(runtime::ITensor::SharedPtr* inputBlocks, int inputBl } } -nvinfer1::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState) +tensorrt_llm::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState) { int64_t blockSize = static_cast(cacheState.getModelConfig().mNbKvHeadsPerLayer[0] @@ -1131,7 +1131,7 @@ void splitKVCache(std::map> size_t cacheBlockSizeSum = 0; size_t inputBlockLayerNumSum = 0; auto cacheDataType - = isIndexerKCache ? nvinfer1::DataType::kUINT8 : kVCacheBlocksPerWindow.begin()->second.front()->getDataType(); + = isIndexerKCache ? tensorrt_llm::DataType::kUINT8 : kVCacheBlocksPerWindow.begin()->second.front()->getDataType(); for (auto const& [window, blocks] : kVCacheBlocksPerWindow) { @@ -1170,7 +1170,7 @@ void splitKVCache(std::map> bool const isWindow = windowSizes.size() > 1; runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(T*)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -1182,7 +1182,7 @@ void splitKVCache(std::map> windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), blockNumInwindow.begin(), blockNumInwindow.end()); windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), layersInWindow.begin(), layersInWindow.end()); - windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), nvinfer1::DataType::kINT32); + windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), tensorrt_llm::DataType::kINT32); bufferManager.copy(windowInfoHostBuffer.data(), *windowInfoDeviceBuffer, runtime::MemoryType::kCPU); for (auto layerNum : layersInWindow) @@ -1405,7 +1405,7 @@ void splitKVCacheDispatch(std::mapsecond.front()->getDataType(); + = isIndexerKCache ? tensorrt_llm::DataType::kUINT8 : kVCacheBlocksPerWindow.begin()->second.front()->getDataType(); auto dataSize = tensorrt_llm::common::getDTypeSize(dataType); switch (dataSize) @@ -1513,7 +1513,7 @@ void concatKVCache(std::vector const& inputSplitBlo } cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); bool const isWindow = windowSizes.size() > 1; @@ -1525,7 +1525,7 @@ void concatKVCache(std::vector const& inputSplitBlo windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), blockNumInwindow.begin(), blockNumInwindow.end()); windowInfoHostBuffer.insert(windowInfoHostBuffer.end(), layersInWindow.begin(), layersInWindow.end()); - windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), nvinfer1::DataType::kINT32); + windowInfoDeviceBuffer = bufferManager.gpu(windowInfoHostBuffer.size(), tensorrt_llm::DataType::kINT32); bufferManager.copy(windowInfoHostBuffer.data(), *windowInfoDeviceBuffer, runtime::MemoryType::kCPU); } constexpr int subWarpSize = 8; diff --git a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h index c7036b219612..80816bc5c3a0 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h +++ b/cpp/tensorrt_llm/executor/cache_transmission/cacheSplitConcat.h @@ -27,7 +27,7 @@ #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::executor::kv_cache { @@ -78,7 +78,7 @@ void concatKVCacheDispatch(runtime::ITensor::SharedPtr* inputBlocks, int inputBl runtime::ITensor::SharedPtr* outputBlocks, int outputBlockNum, int selfRank, kv_cache::CacheState const& selfCacheState, runtime::BufferManager const& bufferManager); -nvinfer1::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState); +tensorrt_llm::Dims makeShapeFromCacheState(kv_cache::CacheState const& cacheState); void splitKVCacheDispatch(std::map> const& kVCacheBlocksPerWindow, std::vector& ouputSplitBlocks, kv_cache::CacheState const& peerCacheState, @@ -147,7 +147,7 @@ void concatRnnSsmStateDispatch(std::vector const& i void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType ssmDataType, + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, runtime::BufferManager const& bufferManager); /** @@ -156,7 +156,7 @@ void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType convDataType, + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager); /** @@ -165,7 +165,7 @@ void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType ssmDataType, runtime::BufferManager const& bufferManager); + size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, runtime::BufferManager const& bufferManager); /** * @brief Concat conv state from per-source buffers into unified pool blocks (section-aware). @@ -173,6 +173,6 @@ void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType convDataType, runtime::BufferManager const& bufferManager); + size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager); } // namespace tensorrt_llm::executor::rnn_cache diff --git a/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu b/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu index b697f2eeb3d6..2beaa659b167 100644 --- a/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu +++ b/cpp/tensorrt_llm/executor/cache_transmission/rnnCacheSplitConcat.cu @@ -31,7 +31,7 @@ #include "tensorrt_llm/runtime/iBuffer.h" #include "tensorrt_llm/runtime/iTensor.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -366,7 +366,7 @@ void splitRnnConvState(std::vector const& inputConv // Allocate and copy pointer array to device runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -497,7 +497,7 @@ void splitRnnSsmState(std::vector const& inputSsmBl cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -625,7 +625,7 @@ void concatRnnConvState(std::vector const& inputSpl cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -747,7 +747,7 @@ void concatRnnSsmState(std::vector const& inputSpli cachePtrs.insert(cachePtrs.end(), prefixLayerNum.begin(), prefixLayerNum.end()); runtime::BufferManager::IBufferPtr PtrsDeviceBuffer - = bufferManager.gpu(cachePtrs.size(), nvinfer1::DataType::kINT64); + = bufferManager.gpu(cachePtrs.size(), tensorrt_llm::DataType::kINT64); TLLM_CHECK(PtrsDeviceBuffer->getSizeInBytes() == cachePtrs.size() * sizeof(uint64_t)); bufferManager.copy(cachePtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); @@ -1349,7 +1349,7 @@ void splitUnifiedPoolSsm(runtime::ITensor::SharedPtr const& pool, std::vector(PtrsDeviceBuffer->data()); @@ -1482,10 +1482,10 @@ void splitUnifiedPoolConv(runtime::ITensor::SharedPtr const& pool, std::vector(PtrsDeviceBuffer->data()); @@ -1607,7 +1607,7 @@ void concatUnifiedPoolSsm(runtime::ITensor::SharedPtr const& pool, std::vector(PtrsDeviceBuffer->data()); @@ -1732,10 +1732,10 @@ void concatUnifiedPoolConv(runtime::ITensor::SharedPtr const& pool, std::vector< sectionInfo.insert(sectionInfo.end(), sectionDimsDomainTP.begin(), sectionDimsDomainTP.end()); sectionInfo.insert(sectionInfo.end(), sectionOffsetsLocal.begin(), sectionOffsetsLocal.end()); - auto PtrsDeviceBuffer = bufferManager.gpu(allPtrs.size(), nvinfer1::DataType::kINT64); + auto PtrsDeviceBuffer = bufferManager.gpu(allPtrs.size(), tensorrt_llm::DataType::kINT64); bufferManager.copy(allPtrs.data(), *PtrsDeviceBuffer, runtime::MemoryType::kCPU); - auto sectionInfoBuffer = bufferManager.gpu(sectionInfo.size(), nvinfer1::DataType::kINT32); + auto sectionInfoBuffer = bufferManager.gpu(sectionInfo.size(), tensorrt_llm::DataType::kINT32); bufferManager.copy(sectionInfo.data(), *sectionInfoBuffer, runtime::MemoryType::kCPU); T** outputPtrsDev = static_cast(PtrsDeviceBuffer->data()); @@ -1811,7 +1811,7 @@ void concatUnifiedPoolConv(runtime::ITensor::SharedPtr const& pool, std::vector< void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType ssmDataType, runtime::BufferManager const& bufferManager) + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(ssmDataType); switch (dataSize) @@ -1835,7 +1835,7 @@ void splitUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector& outputSplitBlocks, kv_cache::CacheState const& destCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, - size_t ssmBytes, size_t blockSizeBytes, nvinfer1::DataType convDataType, + size_t ssmBytes, size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(convDataType); @@ -1860,7 +1860,7 @@ void splitUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType ssmDataType, runtime::BufferManager const& bufferManager) + size_t blockSizeBytes, tensorrt_llm::DataType ssmDataType, runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(ssmDataType); switch (dataSize) @@ -1884,7 +1884,7 @@ void concatUnifiedPoolSsmDispatch(runtime::ITensor::SharedPtr const& pool, void concatUnifiedPoolConvDispatch(runtime::ITensor::SharedPtr const& pool, std::vector const& realBlockIndices, std::vector const& inputSplitBlocks, kv_cache::CacheState const& srcCacheState, kv_cache::CacheState const& selfCacheState, int selfIdx, size_t ssmBytes, - size_t blockSizeBytes, nvinfer1::DataType convDataType, runtime::BufferManager const& bufferManager) + size_t blockSizeBytes, tensorrt_llm::DataType convDataType, runtime::BufferManager const& bufferManager) { auto dataSize = tensorrt_llm::common::getDTypeSize(convDataType); switch (dataSize) diff --git a/cpp/tensorrt_llm/executor/disaggServerUtil.cpp b/cpp/tensorrt_llm/executor/disaggServerUtil.cpp deleted file mode 100644 index 6be2e4fb8ae8..000000000000 --- a/cpp/tensorrt_llm/executor/disaggServerUtil.cpp +++ /dev/null @@ -1,555 +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/executor/disaggServerUtil.h" -#include "tensorrt_llm/common/utils.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -namespace tensorrt_llm::executor::disagg_executor -{ - -class DisaggExecutorOrchestrator::Impl -{ -public: - Impl(std::vector const& ctxEnginePaths, - std::vector const& genEnginePaths, - std::vector const& ctxExecutorConfigs, - std::vector const& genExecutorConfigs, bool hasContextAwaitThreads, - bool hasGenAwaitThreads) - : mhasContextAwaitThreads(hasContextAwaitThreads) - , mhasGenAwaitThreads(hasGenAwaitThreads) - { - TLLM_CHECK(ctxEnginePaths.size() == ctxExecutorConfigs.size()); - TLLM_CHECK(genEnginePaths.size() == genExecutorConfigs.size()); - TLLM_CHECK(!(ctxEnginePaths.empty() || genEnginePaths.empty())); - int worldRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - mIsOrchestrator = (worldRank == 0); - auto contextNum = ctxEnginePaths.size(); - mContextReqIdToGlobalId = std::vector>(contextNum); - mContextMapMutexs = std::vector(contextNum); - auto genNum = genEnginePaths.size(); - mGenerationReqIdToGlobalId = std::vector>(genNum); - mGenerationMapMutexs = std::vector(genNum); - - for (size_t cN = 0; cN < contextNum; cN++) - { - mContextExecutors.push_back(std::make_unique( - ctxEnginePaths[cN], texec::ModelType::kDECODER_ONLY, ctxExecutorConfigs[cN])); - } - - for (size_t gN = 0; gN < genNum; gN++) - { - mGenerationExecutors.push_back(std::make_unique( - genEnginePaths[gN], texec::ModelType::kDECODER_ONLY, genExecutorConfigs[gN])); - } - - if (mIsOrchestrator) - { - if (mhasContextAwaitThreads) - { - for (size_t contextIdx = 0; contextIdx < contextNum; contextIdx++) - { - mContextThreads.emplace_back( - [this, contextIdx]() { this->waitResponseAndAppendThreadFun(true, contextIdx); }); - } - } - if (mhasGenAwaitThreads) - { - - for (size_t genIdx = 0; genIdx < genNum; genIdx++) - { - mGenerationThreads.emplace_back( - [this, genIdx]() { this->waitResponseAndAppendThreadFun(false, genIdx); }); - } - } - } - tensorrt_llm::mpi::MpiComm::world().barrier(); - } - - std::vector enqueueContext(std::vector const& requests, - std::optional selectContextId = std::nullopt, bool batch = false) - { - - std::vector globalReqIds; - for (auto const& request : requests) - { - globalReqIds.push_back(generatedGlobalId()); - TLLM_CHECK(request.getRequestType() == tensorrt_llm::executor::RequestType::REQUEST_TYPE_CONTEXT_ONLY); - } - - if (batch) - { - size_t contextId = selectContextId.has_value() ? selectContextId.value() : selectContextExecutor(); - auto contextReqIds = mContextExecutors[contextId]->enqueueRequests(requests); - { - std::scoped_lock lock{mContextMapMutexs[contextId]}; - for (size_t i = 0; i < requests.size(); ++i) - { - mContextReqIdToGlobalId[contextId][contextReqIds[i]] = globalReqIds[i]; - } - } - } - else - { - for (size_t i = 0; i < requests.size(); ++i) - { - size_t contextId = selectContextId.has_value() ? selectContextId.value() : selectContextExecutor(); - - auto contextReqId = mContextExecutors[contextId]->enqueueRequest(requests[i]); - { - std::scoped_lock lock{mContextMapMutexs[contextId]}; - mContextReqIdToGlobalId[contextId][contextReqId] = globalReqIds[i]; - } - } - } - return globalReqIds; - } - - void enqueueGeneration(std::vector const& requests, std::vector const& globalRequestIds, - std::optional selectGenIdx = std::nullopt, bool batch = false) - { - - TLLM_CHECK(globalRequestIds.size() == requests.size()); - - for (auto const& request : requests) - { - - TLLM_CHECK(request.getRequestType() == tensorrt_llm::executor::RequestType::REQUEST_TYPE_GENERATION_ONLY); - } - if (batch) - { - size_t genIdx = selectGenIdx.has_value() ? selectGenIdx.value() : selectGenerationExecutor(); - auto genReqIds = mGenerationExecutors[genIdx]->enqueueRequests(requests); - { - std::scoped_lock lock{mGenerationMapMutexs[genIdx]}; - for (size_t i = 0; i < requests.size(); ++i) - { - mGenerationReqIdToGlobalId[genIdx][genReqIds[i]] = globalRequestIds[i]; - } - } - } - else - { - for (size_t i = 0; i < requests.size(); ++i) - { - size_t genIdx = selectGenIdx.has_value() ? selectGenIdx.value() : selectGenerationExecutor(); - - auto genReqId = mGenerationExecutors[genIdx]->enqueueRequest(requests[i]); - { - std::scoped_lock lock{mGenerationMapMutexs[genIdx]}; - mGenerationReqIdToGlobalId[genIdx][genReqId] = globalRequestIds[i]; - } - } - } - } - - std::vector awaitContextResponses( - std::optional contextIdx, std::optional const& timeout) - { - - std::vector responses; - - if (mhasContextAwaitThreads) - { - - std::unique_lock lock(mResponsesContextMtx); - auto pred = [&mShutdown = mShutdown, &resp = this->mContextResponses]() -> bool - { return !resp.empty() || mShutdown; }; - auto storeResponses = [&resp = this->mContextResponses, &responses]() - { - responses = std::move(resp); - resp.clear(); - }; - if (timeout) - { - if (mContextResponsesCV.wait_for(lock, timeout.value(), pred)) - { - storeResponses(); - } - } - else - { - mContextResponsesCV.wait(lock, pred); - storeResponses(); - } - TLLM_CHECK_WITH_INFO( - !contextIdx.has_value(), "contextIdx should not be provided when mhasContextAwaitThreads is true"); - - return responses; - } - - if (contextIdx.has_value()) - { - TLLM_CHECK(!mhasContextAwaitThreads); - auto responseFromExecutor = mContextExecutors[contextIdx.value()]->awaitResponses(timeout); - for (auto&& resp : responseFromExecutor) - { - - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mContextMapMutexs.at(contextIdx.value())}; - globalId = mContextReqIdToGlobalId.at(contextIdx.value()).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - return responses; - } - TLLM_CHECK(timeout.has_value()); - auto timeouP = timeout.value() / mContextExecutors.size(); - for (size_t ci = 0; ci < mContextExecutors.size(); ci++) - { - auto responseFromExecutor = mContextExecutors.at(ci)->awaitResponses(timeouP); - for (auto&& resp : responseFromExecutor) - { - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mContextMapMutexs.at(ci)}; - globalId = mContextReqIdToGlobalId.at(ci).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - } - - return responses; - }; - - std::vector awaitGenerationResponses( - std::optional genIdx, std::optional const& timeout) - { - - std::vector responses; - - if (mhasGenAwaitThreads) - { - - std::unique_lock lock(mResponseGenerationMtx); - auto pred = [&mShutdown = mShutdown, &resp = this->mGenerationResponses]() -> bool - { return !resp.empty() || mShutdown; }; - auto storeResponses = [&resp = this->mGenerationResponses, &responses]() - { - responses = std::move(resp); - resp.clear(); - }; - if (timeout) - { - if (mGenerationResponsesCv.wait_for(lock, timeout.value(), pred)) - { - storeResponses(); - } - } - else - { - mGenerationResponsesCv.wait(lock, pred); - storeResponses(); - } - TLLM_CHECK_WITH_INFO(!genIdx.has_value(), "genIdx should not be provided when mhasGenAwaitThreads is true"); - return responses; - } - - if (genIdx.has_value()) - { - TLLM_CHECK(!mhasGenAwaitThreads); - auto responseFromExecutor = mGenerationExecutors[genIdx.value()]->awaitResponses(timeout); - for (auto&& resp : responseFromExecutor) - { - - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mGenerationMapMutexs.at(genIdx.value())}; - globalId = mGenerationReqIdToGlobalId.at(genIdx.value()).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - return responses; - } - TLLM_CHECK(timeout.has_value()); - auto timeouP = timeout.value() / mGenerationExecutors.size(); - - for (size_t gi = 0; gi < mGenerationExecutors.size(); gi++) - { - auto responseFromExecutor = mGenerationExecutors.at(gi)->awaitResponses(timeouP); - for (auto&& resp : responseFromExecutor) - { - - auto reqId = resp.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mGenerationMapMutexs.at(gi)}; - globalId = mGenerationReqIdToGlobalId.at(gi).at(reqId); - } - TLLM_CHECK(globalId != 0); - responses.emplace_back(std::move(resp), globalId); - } - } - - return responses; - }; - - [[nodiscard]] bool canEnqueue() const - { - return mIsOrchestrator; - } - - [[nodiscard]] std::vector> const& getContextExecutors() const - { - return mContextExecutors; - } - - [[nodiscard]] std::vector> const& getGenExecutors() const - { - return mGenerationExecutors; - } - - ~Impl() - { - - mShutdown = true; - - mContextResponsesCV.notify_all(); - mGenerationResponsesCv.notify_all(); - for (auto&& executor : mContextExecutors) - { - executor->shutdown(); - } - for (auto&& executor : mGenerationExecutors) - { - executor->shutdown(); - } - - if (mIsOrchestrator) - { - if (mhasContextAwaitThreads) - { - for (auto&& contextThread : mContextThreads) - { - if (contextThread.joinable()) - { - contextThread.join(); - } - } - } - if (mhasGenAwaitThreads) - { - for (auto&& genThread : mGenerationThreads) - { - if (genThread.joinable()) - { - genThread.join(); - } - } - } - } - } - -private: - IdType generatedGlobalId() - { - return (++mLastId % UINT64_MAX); - }; - - size_t selectContextExecutor() - { - static size_t selectContextId = 0; - auto contextId = (selectContextId++) % mContextExecutors.size(); - if (selectContextId >= mContextExecutors.size()) - { - selectContextId = 0; - } - return contextId; - } - - size_t selectGenerationExecutor() - { - static size_t selectGenerationId = 0; - auto generationIdx = (selectGenerationId++) % mGenerationExecutors.size(); - if (selectGenerationId >= mGenerationExecutors.size()) - { - selectGenerationId = 0; - } - return generationIdx; - } - - void appendNewContextResponse(std::vector&& newResponses) - { - { - std::scoped_lock lock(mResponsesContextMtx); - for (auto&& response : newResponses) - { - mContextResponses.emplace_back(std::move(response)); - } - } - mContextResponsesCV.notify_all(); - } - - void appendNewGenerationResponse(std::vector&& newResponses) - { - { - std::scoped_lock lock(mResponseGenerationMtx); - for (auto&& response : newResponses) - { - mGenerationResponses.emplace_back(std::move(response)); - } - } - mGenerationResponsesCv.notify_all(); - } - - void waitResponseAndAppendThreadFun(bool isContext, int executorIdx) - { - - tensorrt_llm::common::setThreadName("waitResponseAndAppendThreadFun"); - - auto& executor = isContext ? mContextExecutors[executorIdx] : mGenerationExecutors[executorIdx]; - - while (!mShutdown) - { - auto responses = executor->awaitResponses(); - - if (responses.empty()) - { - continue; - } - std::vector responseWithIds; - if (isContext) - { - for (auto&& response : responses) - { - auto reqId = response.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mContextMapMutexs.at(executorIdx)}; - globalId = mContextReqIdToGlobalId.at(executorIdx).at(reqId); - } - TLLM_CHECK(globalId != 0); - responseWithIds.emplace_back(std::move(response), globalId); - } - if (responseWithIds.size() > 0) - { - appendNewContextResponse(std::move(responseWithIds)); - } - } - else - { - - for (auto&& response : responses) - { - auto reqId = response.getRequestId(); - IdType globalId{0}; - - { - std::scoped_lock lock{mGenerationMapMutexs.at(executorIdx)}; - globalId = mGenerationReqIdToGlobalId.at(executorIdx).at(reqId); - } - TLLM_CHECK(globalId != 0); - responseWithIds.emplace_back(std::move(response), globalId); - } - if (responseWithIds.size() > 0) - { - appendNewGenerationResponse(std::move(responseWithIds)); - } - } - } - }; - - std::vector> mContextExecutors; - std::vector> mGenerationExecutors; - std::vector mContextThreads; - std::vector mGenerationThreads; - - std::atomic mLastId{0}; - std::vector> mContextReqIdToGlobalId; - std::vector> mGenerationReqIdToGlobalId; - std::vector mContextMapMutexs; - std::vector mGenerationMapMutexs; - std::vector mContextResponses; - std::condition_variable mContextResponsesCV; - std::mutex mResponsesContextMtx; - - std::vector mGenerationResponses; - std::condition_variable mGenerationResponsesCv; - std::mutex mResponseGenerationMtx; - std::atomic mShutdown{false}; - std::atomic mhasContextAwaitThreads{false}; - std::atomic mhasGenAwaitThreads{false}; - bool mIsOrchestrator{false}; -}; - -DisaggExecutorOrchestrator::DisaggExecutorOrchestrator(std::vector const& ctxEnginePaths, - std::vector const& genEnginePaths, - std::vector const& ctxExecutorConfigs, - std::vector const& genExecutorConfigs, bool hasContextAwaitThreads, - bool hasGenAwaitThreads) - : mImpl(std::make_unique(ctxEnginePaths, genEnginePaths, ctxExecutorConfigs, - genExecutorConfigs, hasContextAwaitThreads, hasGenAwaitThreads)) -{ -} - -std::vector DisaggExecutorOrchestrator::enqueueContext( - std::vector const& requests, std::optional selectContextId, bool batch) -{ - return mImpl->enqueueContext(requests, selectContextId, batch); -} - -void DisaggExecutorOrchestrator::enqueueGeneration(std::vector const& requests, - std::vector const& globalRequestIds, std::optional selectGenIdx, bool batch) -{ - mImpl->enqueueGeneration(requests, globalRequestIds, selectGenIdx, batch); -} - -std::vector DisaggExecutorOrchestrator::awaitContextResponses( - std::optional const& timeout, std::optional contextIdx) -{ - return mImpl->awaitContextResponses(contextIdx, timeout); -} - -std::vector DisaggExecutorOrchestrator::awaitGenerationResponses( - std::optional const& timeout, std::optional genIdx) -{ - return mImpl->awaitGenerationResponses(genIdx, timeout); -} - -bool DisaggExecutorOrchestrator::canEnqueue() const -{ - return mImpl->canEnqueue(); -}; - -std::vector> const& DisaggExecutorOrchestrator::getContextExecutors() const -{ - return mImpl->getContextExecutors(); -} - -std::vector> const& DisaggExecutorOrchestrator::getGenExecutors() const -{ - return mImpl->getGenExecutors(); -} - -DisaggExecutorOrchestrator::~DisaggExecutorOrchestrator() = default; - -} // namespace tensorrt_llm::executor::disagg_executor diff --git a/cpp/tensorrt_llm/executor/executor.cpp b/cpp/tensorrt_llm/executor/executor.cpp deleted file mode 100644 index 091bb5128230..000000000000 --- a/cpp/tensorrt_llm/executor/executor.cpp +++ /dev/null @@ -1,144 +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 -#include - -namespace tensorrt_llm::executor -{ - -Executor::Executor(std::filesystem::path const& modelPath, ModelType modelType, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(modelPath, std::nullopt, modelType, executorConfig)) -{ -} - -Executor::Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, - ModelType modelType, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(decoderModelPath, encoderModelPath, modelType, executorConfig)) -{ -} - -Executor::Executor(BufferView const& engineBuffer, std::string const& jsonConfigStr, ModelType modelType, - ExecutorConfig const& executorConfig, std::optional> const& managedWeights) - : mImpl(std::make_unique( - engineBuffer, jsonConfigStr, std::nullopt, std::nullopt, modelType, executorConfig, managedWeights)) -{ -} - -Executor::Executor(BufferView const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, - BufferView const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, ModelType modelType, - ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(decoderEngineBuffer, decoderJsonConfigStr, encoderEngineBuffer, - encoderJsonConfigStr, modelType, executorConfig, std::nullopt)) -{ -} - -Executor::Executor(std::shared_ptr model, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(std::move(model), std::nullopt, executorConfig)) -{ -} - -Executor::Executor( - std::shared_ptr encoderModel, std::shared_ptr decoderModel, ExecutorConfig const& executorConfig) - : mImpl(std::make_unique(std::move(decoderModel), std::move(encoderModel), executorConfig)) -{ -} - -Executor::~Executor() = default; - -IdType Executor::enqueueRequest(Request const& llmRequest) -{ - return mImpl->enqueueRequest(llmRequest); -} - -std::vector Executor::enqueueRequests(std::vector const& llmRequests) -{ - return mImpl->enqueueRequests(llmRequests); -} - -std::vector Executor::awaitResponses(std::optional const& timeout) -{ - return mImpl->awaitResponses(timeout); -} - -std::vector Executor::awaitResponses( - IdType const& requestId, std::optional const& timeout) -{ - return mImpl->awaitResponses(requestId, timeout); -} - -std::vector> Executor::awaitResponses( - std::vector const& requestIds, std::optional const& timeout) -{ - return mImpl->awaitResponses(requestIds, timeout); -} - -SizeType32 Executor::getNumResponsesReady(std::optional const& requestId) const -{ - return mImpl->getNumResponsesReady(requestId); -} - -void Executor::cancelRequest(IdType requestId) -{ - return mImpl->cancelRequest(requestId); -} - -void Executor::shutdown() -{ - return mImpl->shutdown(); -} - -std::deque Executor::getLatestIterationStats() -{ - return mImpl->getLatestIterationStats(); -} - -std::deque Executor::getLatestRequestStats() -{ - return mImpl->getLatestRequestStats(); -} - -std::deque Executor::getLatestDebugTensors() -{ - return mImpl->getLatestDebugTensors(); -} - -bool Executor::canEnqueueRequests() const -{ - return mImpl->canEnqueueRequests(); -} - -bool Executor::isParticipant() const -{ - return mImpl->isParticipant(); -} - -std::optional> Executor::getKVCacheEventManager() const -{ - return mImpl->getKVCacheEventManager(); -} - -KVCacheEvent::KVCacheEvent( - size_t eventId, KVCacheEventData data, SizeType32 windowSize, std::optional attentionDpRank) - : eventId{eventId} - , data{std::move(data)} - , windowSize{windowSize} - , attentionDpRank{attentionDpRank} -{ -} - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/executorImpl.cpp b/cpp/tensorrt_llm/executor/executorImpl.cpp deleted file mode 100644 index 2fb20b7ea572..000000000000 --- a/cpp/tensorrt_llm/executor/executorImpl.cpp +++ /dev/null @@ -1,2773 +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/executor/executorImpl.h" -#include "tensorrt_llm/batch_manager/trtEncoderModel.h" -#include "tensorrt_llm/batch_manager/trtGptModelFactory.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaProfilerUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/timestampUtils.h" -#include "tensorrt_llm/common/utils.h" -#include "tensorrt_llm/executor/dataTransceiverState.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/orchestratorUtils.h" -#include "tensorrt_llm/executor/requestUtils.h" -#include "tensorrt_llm/executor/serialization.h" -#include "tensorrt_llm/executor/serializeUtils.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/executor/version.h" -#include "tensorrt_llm/runtime/loraCache.h" -#include "tensorrt_llm/runtime/memoryCounters.h" -#include "tensorrt_llm/runtime/utils/mpiTags.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::executor -{ - -namespace -{ - -[[nodiscard]] bool executorConfigIsValid( - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, runtime::ModelConfig const& modelConfig) -{ - // Make sure logic in this function matches fixExecutorConfig - if (executorConfig.getEnableChunkedContext()) - { - if (modelConfig.isRnnBased() || !modelConfig.isKVCacheEnabled() || !modelConfig.getPagedContextFMHA()) - { - return false; - } - } - return true; -} - -[[nodiscard]] ::tensorrt_llm::executor::ExecutorConfig fixExecutorConfig( - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, runtime::ModelConfig const& modelConfig) -{ - // Make sure logic in this function matches executorConfigIsValid - auto fixedExecutorConfig = executorConfig; - // Disable chunked context when not supported - if (executorConfig.getEnableChunkedContext()) - { - if (modelConfig.isRnnBased() || !modelConfig.isKVCacheEnabled() || !modelConfig.getPagedContextFMHA()) - { - fixedExecutorConfig.setEnableChunkedContext(false); - TLLM_LOG_WARNING( - "Chunked context is not supported for this configuration and will be disabled. " - "Related configs: RNNBased: %d, KVCacheEnabled: %d, PagedContextFMHA: %d", - modelConfig.isRnnBased(), modelConfig.isKVCacheEnabled(), modelConfig.getPagedContextFMHA()); - } - } - return fixedExecutorConfig; -} - -SizeType32 getNumChildRequests(Request const& request) -{ - auto samplingConfig = request.getSamplingConfig(); - return samplingConfig.getBeamWidth() > 1 ? 0 : samplingConfig.getNumReturnSequences().value_or(1) - 1; -} - -} // namespace - -/// @brief Version of TRT-LLM as defined in tensorrt_llm/version.py -char const* version() noexcept -{ - return kTensorRtLlmVersion; -} - -class CancelledRequestsAsyncSend -{ -public: - CancelledRequestsAsyncSend(std::shared_ptr const& commSession, - std::unordered_set const& cancelledReqIds, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mNumReq = static_cast(cancelledReqIds.size()); - TLLM_LOG_DEBUG("start send %ld cancelled requests to rank %d", mNumReq, peer); - mRequest1 - = commSession->sendAsync(&mNumReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kCancelledRequestsNumReq); - if (mNumReq > 0) - { - mIds.assign(cancelledReqIds.begin(), cancelledReqIds.end()); - mRequest2 = commSession->sendAsync( - mIds.data(), mIds.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kCancelledRequestsIds); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - ~CancelledRequestsAsyncSend() - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRequest1->wait(); - if (mRequest2) - { - mRequest2->wait(); - } - TLLM_LOG_DEBUG("end send cancelled requests"); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - CancelledRequestsAsyncSend(CancelledRequestsAsyncSend const& executor) = delete; - CancelledRequestsAsyncSend& operator=(CancelledRequestsAsyncSend const& executor) = delete; - CancelledRequestsAsyncSend(CancelledRequestsAsyncSend&&) = delete; - CancelledRequestsAsyncSend& operator=(CancelledRequestsAsyncSend&&) = delete; - - static std::unordered_set cancelledRequestsRecv( - std::shared_ptr const& commSession, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start recv cancelled requests from rank %d", peer); - std::unordered_set cancelledReqIds; - int64_t numReq{0}; - commSession->recv(&numReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kCancelledRequestsNumReq); - TLLM_LOG_DEBUG("recv %ld cancelled requests", numReq); - if (numReq > 0) - { - std::vector buffer(numReq); - commSession->recv( - buffer.data(), buffer.size(), mpi::MpiType::kUINT64, peer, mpi::MpiTag::kCancelledRequestsIds); - cancelledReqIds = std::unordered_set(buffer.begin(), buffer.end()); - } - TLLM_LOG_DEBUG("end recv cancelled requests from rank %d", peer); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return cancelledReqIds; - } - -private: - int64_t mNumReq; - std::vector mIds; - std::shared_ptr mRequest1; - std::shared_ptr mRequest2; -}; - -class RequestWithIdAsyncSend -{ -public: - RequestWithIdAsyncSend(std::shared_ptr const& commSession, - std::vector const& reqWithIds, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start send requests to rank %d", peer); - mNumReq = static_cast(reqWithIds.size()); - mRequest1 = commSession->sendAsync(&mNumReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdNumReq); - if (mNumReq > 0) - { - mPacked = RequestWithId::serializeReqWithIds(reqWithIds); - mVecSize = static_cast(mPacked.size()); - mRequest2 - = commSession->sendAsync(&mVecSize, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdVecSize); - mRequest3 = commSession->sendAsync( - mPacked.data(), mPacked.size(), mpi::MpiType::kCHAR, peer, mpi::MpiTag::kRequestWithIdPacked); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - ~RequestWithIdAsyncSend() - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - mRequest1->wait(); - if (mRequest2) - { - mRequest2->wait(); - } - if (mRequest3) - { - mRequest3->wait(); - } - TLLM_LOG_DEBUG("end send requests"); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - } - - RequestWithIdAsyncSend(RequestWithIdAsyncSend const& executor) = delete; - RequestWithIdAsyncSend& operator=(RequestWithIdAsyncSend const& executor) = delete; - RequestWithIdAsyncSend(RequestWithIdAsyncSend&&) = delete; - RequestWithIdAsyncSend& operator=(RequestWithIdAsyncSend&&) = delete; - - static std::vector requestWithIdRecv( - std::shared_ptr const& commSession, int peer) - { - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - TLLM_LOG_DEBUG("start recv requests from rank %d", peer); - std::vector reqWithIds; - int64_t numReq{0}; - commSession->recv(&numReq, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdNumReq); - if (numReq > 0) - { - std::vector buffer; - int64_t vecSize = 0; - commSession->recv(&vecSize, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kRequestWithIdVecSize); - buffer.resize(vecSize); - commSession->recv( - buffer.data(), buffer.size(), mpi::MpiType::kCHAR, peer, mpi::MpiTag::kRequestWithIdPacked); - reqWithIds = RequestWithId::deserializeReqWithIds(buffer); - } - TLLM_LOG_DEBUG("end recv requests from rank %d", peer); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return reqWithIds; - } - -private: - int64_t mNumReq; - int64_t mVecSize; - std::vector mPacked; - std::shared_ptr mRequest1; - std::shared_ptr mRequest2; - std::shared_ptr mRequest3; -}; - -void Executor::Impl::loadModel(std::optional const& modelPathOpt, - std::optional const& engineBufferOpt, runtime::GptJsonConfig const& jsonConfig, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, bool isEncoder, - std::optional> const& managedWeightsOpt) -{ - auto const gpusPerNode = jsonConfig.getGpusPerNode(); - auto const tp = jsonConfig.getTensorParallelism(); - auto const pp = jsonConfig.getPipelineParallelism(); - auto const cp = jsonConfig.getContextParallelism(); - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - auto worldConfig = runtime::WorldConfig::mpi(gpusPerNode, tp, pp, cp, parallelConfig.getDeviceIds()); - - TLLM_CHECK_WITH_INFO(modelPathOpt.has_value() || engineBufferOpt.has_value(), - "Either engine path or deserialized engine buffer should be given to load the model properly."); - auto rawEngine = engineBufferOpt.has_value() - ? runtime::RawEngine(engineBufferOpt.value().data(), engineBufferOpt.value().size()) - : runtime::RawEngine(modelPathOpt.value() / jsonConfig.engineFilename(worldConfig)); - - if (rawEngine.getType() != tensorrt_llm::runtime::RawEngine::FilePath) - { - if (modelPathOpt.has_value()) - { - rawEngine.setPath(modelPathOpt.value() / jsonConfig.engineFilename(worldConfig)); - if (managedWeightsOpt.has_value()) - { - TLLM_LOG_WARNING( - "Executor::Impl::loadModel: managedWeightsOpt argument is ignored when loading engine from file."); - } - } - else if (managedWeightsOpt.has_value()) - { - rawEngine.setManagedWeightsMap(managedWeightsOpt.value()); - } - } - - auto const& modelConfig = jsonConfig.getModelConfig(); - - if (isEncoder) - { - mEncoderModel = createEncoderModel(rawEngine, modelConfig, worldConfig, executorConfig); - } - else - { - mModel = createModel(rawEngine, modelConfig, worldConfig, executorConfig); - } -}; - -Executor::Impl::Impl(std::filesystem::path const& modelPath, - std::optional const& encoderModelPath, ModelType const modelType, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto decoderJsonConfig = runtime::GptJsonConfig::parse(modelPath / "config.json"); - - // for now, assume encoder & decoder models share the same MPI config - auto const tp = decoderJsonConfig.getTensorParallelism(); - auto const pp = decoderJsonConfig.getPipelineParallelism(); - auto const cp = decoderJsonConfig.getContextParallelism(); - initializeCommAndWorkers(tp, pp, cp, executorConfig, modelType, modelPath, std::nullopt, decoderJsonConfig); - - if (mIsWorker) - { - if (modelType == ModelType::kENCODER_DECODER) - { - if (encoderModelPath.has_value()) - { - auto const encoderJsonConfig = runtime::GptJsonConfig::parse(encoderModelPath.value() / "config.json"); - - auto const encoderMaxInputLen = encoderJsonConfig.getModelConfig().getMaxInputLen(); - auto const encoderHiddenSize = encoderJsonConfig.getModelConfig().getHiddenSize() - * encoderJsonConfig.getTensorParallelism(); // recover full hidden size - // add encoder info to decoder for encoder-decoder models - // note: GptJsonConfig can no longer have modelConfig as const member since it must be mutable here - decoderJsonConfig.getModelConfigMutable().setMaxEncoderLen(encoderMaxInputLen); - decoderJsonConfig.getModelConfigMutable().setEncoderHiddenSize(encoderHiddenSize); - - loadModel( - encoderModelPath.value(), std::nullopt, encoderJsonConfig, executorConfig, true, std::nullopt); - } - else - { - TLLM_LOG_WARNING("Encoder model path not provided. Skipping Encoder Run."); - } - } - loadModel(modelPath, std::nullopt, decoderJsonConfig, executorConfig, false, std::nullopt); - } - initialize(executorConfig); -} - -Executor::Impl::Impl(BufferView const& engineBufferView, std::string const& jsonConfigStr, - std::optional const& encoderEngineBufferView, std::optional const& encoderJsonConfigStr, - ModelType const modelType, ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, - std::optional> const& managedWeightsOpt) -{ - auto decoderJsonConfig = runtime::GptJsonConfig::parse(jsonConfigStr); - - // for now, assume encoder & decoder models share the same MPI config - auto const tp = decoderJsonConfig.getTensorParallelism(); - auto const pp = decoderJsonConfig.getPipelineParallelism(); - auto const cp = decoderJsonConfig.getContextParallelism(); - initializeCommAndWorkers(tp, pp, cp, executorConfig, modelType, std::nullopt, std::nullopt, decoderJsonConfig); - - if (mIsWorker) - { - if (modelType == ModelType::kENCODER_DECODER) - { - TLLM_CHECK(encoderEngineBufferView.has_value() && encoderJsonConfigStr.has_value()); - TLLM_CHECK_WITH_INFO( - !managedWeightsOpt.has_value(), "Managed weights are not supported for enc-dec models"); - - auto const encoderJsonConfig = runtime::GptJsonConfig::parse(encoderJsonConfigStr.value()); - - auto const encoderMaxInputLen = encoderJsonConfig.getModelConfig().getMaxInputLen(); - auto const encoderHiddenSize = encoderJsonConfig.getModelConfig().getHiddenSize() - * encoderJsonConfig.getTensorParallelism(); // recover full hidden size - // add encoder info to decoder for encoder-decoder models - // note: GptJsonConfig can no longer have modelConfig as const member since it must be mutable here - decoderJsonConfig.getModelConfigMutable().setMaxEncoderLen(encoderMaxInputLen); - decoderJsonConfig.getModelConfigMutable().setEncoderHiddenSize(encoderHiddenSize); - - loadModel( - std::nullopt, encoderEngineBufferView.value(), encoderJsonConfig, executorConfig, true, std::nullopt); - } - loadModel(std::nullopt, engineBufferView, decoderJsonConfig, executorConfig, false, managedWeightsOpt); - } - initialize(executorConfig); -} - -Executor::Impl::Impl(std::shared_ptr model, std::optional> encoderModel, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto const& worldConfig = model->getWorldConfig(); - auto const tp = worldConfig.getTensorParallelism(); - auto const pp = worldConfig.getPipelineParallelism(); - auto const cp = worldConfig.getContextParallelism(); - auto const modelType = encoderModel.has_value() ? ModelType::kENCODER_DECODER : ModelType::kDECODER_ONLY; - initializeCommAndWorkers(tp, pp, cp, executorConfig, modelType, std::nullopt, worldConfig); - if (modelType == ModelType::kENCODER_DECODER) - { - mEncoderModel = encoderModel.value(); - } - mModel = std::move(model); - initialize(executorConfig); -} - -Executor::Impl::~Impl() -{ - shutdown(); -} - -void Executor::Impl::initialize(::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - mShutdown = false; - mShutdownCalled = false; - mIterStatsMaxIterations = executorConfig.getIterStatsMaxIterations(); - mRequestStatsMaxIterations = executorConfig.getRequestStatsMaxIterations(); - mDebugTensorsMaxIterations - = executorConfig.getDebugConfig() ? executorConfig.getDebugConfig()->getDebugTensorsMaxIterations() : 0; - TLLM_CHECK_WITH_INFO(mDebugTensorsMaxIterations == 0 || mCommMode == CommunicationMode::kLEADER, - "debugTensorsMaxIterations > 0 is only allowed in leader mode."); - mBatchingType = executorConfig.getBatchingType(); - mIsSchedulerMaxUtilization = (executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy() - == CapacitySchedulerPolicy::kMAX_UTILIZATION); - mIsSchedulerGuaranteedNoEvict = (executorConfig.getSchedulerConfig().getCapacitySchedulerPolicy() - == CapacitySchedulerPolicy::kGUARANTEED_NO_EVICT); - mIsChunkedContext = executorConfig.getEnableChunkedContext(); - mPromptTableOffloading = executorConfig.getPromptTableOffloading(); - mMaxQueueSize = executorConfig.getMaxQueueSize(); - - mLastReqId = 1; - - auto const& logitsProcConfig = executorConfig.getLogitsPostProcessorConfig(); - if (logitsProcConfig.has_value()) - { - mLogitsPostProcessorMap = logitsProcConfig.value().getProcessorMap().value_or(LogitsPostProcessorMap{}); - initializeLogitsPostProcessorBatched(logitsProcConfig.value()); - if (!logitsProcConfig.value().getReplicate()) - { - mModel->setReplicateLogitsPostProcessor(false); - } - } - - auto const& commComm = COMM_SESSION; - int32_t const commSize = commComm.getSize(); - if (mIsWorker) - { - if (commSize > 1) - { - auto const& worldConfig = mModel->getWorldConfig(); - auto const& commSession = COMM_SESSION; - auto const& rank = commSession.getRank(); - auto const& tp = worldConfig.getTensorParallelism(); - auto const& cp = worldConfig.getContextParallelism(); - - mCommTensorParallel = std::make_shared( - commSession.split(rank / tp, worldConfig.getTensorParallelRank())); - mCommContextParallel = std::make_shared( - commSession.split(rank / (tp * cp) * tp + rank % tp, worldConfig.getContextParallelRank())); - mCommPipelineParallel = std::make_shared( - commSession.split(rank % (tp * cp), worldConfig.getPipelineParallelRank())); - - if (worldConfig.isPipelineParallel()) - { - mRequestWithIdWaitThread = std::make_unique( - "requestWithIdWaitThread", [this]() { mRequestWithIdAsyncSndHdl.reset(nullptr); }); - mCancelledRequestsWaitThread = std::make_unique( - "cancelledRequestsWaitThread", [this]() { mCancelledRequestsAsyncSndHdl.reset(nullptr); }); - if (mIsLeader) - { - mRequestWithIdLeaderThread - = std::make_unique(&Executor::Impl::requestWithIdLeaderThread, this); - mCancelledRequestsLeaderThread - = std::make_unique(&Executor::Impl::cancelledRequestsLeaderThread, this); - } - } - } - // Launch the execution thread - mMaxNumActiveRequests = mModel->getMaxNumSequences(); - mExecutionThread = std::thread(&Impl::executionLoop, this); - } - - mEnableBlockReuse = executorConfig.getKvCacheConfig().getEnableBlockReuse(); - - auto const& dynamicBatchConfig = executorConfig.getSchedulerConfig().getDynamicBatchConfig(); - if (dynamicBatchConfig) - { - if (mIsWorker) - { - if (mModel->getModelConfig().isTransformerBased() && mModel->getModelConfig().isKVCacheEnabled()) - { - mDynamicBatchTuner = std::make_shared(dynamicBatchConfig.value()); - } - else - { - TLLM_LOG_WARNING("Dynamic batch tuner can only support transformer models that use KV cache."); - } - } - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -std::shared_ptr Executor::Impl::createModel(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto const gptModelType = [&executorConfig, &modelConfig]() - { - switch (executorConfig.getBatchingType()) - { - case BatchingType::kSTATIC: - TLLM_THROW( - "Static batching type is deprecated. Please use in-flight batching with " - "CapacitySchedulerPolicy::kSTATIC_BATCH instead."); - case BatchingType::kINFLIGHT: - return modelConfig.isRnnBased() ? batch_manager::TrtGptModelType::InflightBatching - : batch_manager::TrtGptModelType::InflightFusedBatching; - default: TLLM_THROW("Invalid batching strategy"); - } - }(); - - bool const isLeaderInOrchMode = (mCommMode == CommunicationMode::kORCHESTRATOR) && mIsLeader; - auto const& fixedExecutorConfig = executorConfigIsValid(executorConfig, modelConfig) - ? executorConfig - : fixExecutorConfig(executorConfig, modelConfig); - - return batch_manager::TrtGptModelFactory::create( - rawEngine, modelConfig, worldConfig, gptModelType, fixedExecutorConfig, isLeaderInOrchMode); -} - -std::shared_ptr Executor::Impl::createEncoderModel(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig) -{ - auto fixedExecutorConfig = ExecutorConfig{}; - fixedExecutorConfig.setSchedulerConfig(executorConfig.getSchedulerConfig()); - return std::make_shared( - modelConfig, worldConfig, rawEngine, std::make_shared(), fixedExecutorConfig); -} - -void Executor::Impl::setOrchLeaderComm( - SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig const& parallelConfig) -{ -#if ENABLE_MULTI_DEVICE - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - if (optOrchestratorConfig.value().getIsOrchestrator()) - { - TLLM_CHECK_WITH_INFO(mWorldRank == 0, "Rank 0 must be orchestrator"); - } - - TLLM_CHECK_WITH_INFO(parallelConfig.getParticipantIds(), - "When not spawning processes in orchestrator mode, participant IDs must be provided"); - auto participantIds = parallelConfig.getParticipantIds().value(); - - TLLM_CHECK_WITH_INFO(static_cast(participantIds.size()) == tp * pp * cp, - "When specifying participantIds, participantIds size must be equal to tp*pp*cp"); - - bool isLeader = (mWorldRank == participantIds.front()); - bool isOrchestrator = (mWorldRank == 0); - - // OrchLeaderComm rank 0 is orchestrator, rank 1 is leader - mOrchRank = 0; - mLeaderRank = 1; - - // Create a leaderOrch comm - std::vector leaderOrchRanks{0, participantIds.front()}; - - MPI_Group worldGroup = nullptr; - MPICHECK(MPI_Comm_group(MPI_COMM_WORLD, &worldGroup)); // NOLINT - int worldGroupRank = 0; - MPI_Group_rank(worldGroup, &worldGroupRank); - - int worldSize = 0; - MPICHECK(MPI_Group_size(worldGroup, &worldSize)); // NOLINT - TLLM_CHECK_WITH_INFO(participantIds.front() < worldSize, "Not enough ranks in world"); - - MPI_Group leaderOrchCommGroup = nullptr; - MPICHECK( - MPI_Group_incl(worldGroup, leaderOrchRanks.size(), leaderOrchRanks.data(), &leaderOrchCommGroup)); // NOLINT - int leaderOrchGroupRank = 0; - int leaderOrchGroupSize = 0; - MPI_Group_rank(leaderOrchCommGroup, &leaderOrchGroupRank); - MPI_Group_size(leaderOrchCommGroup, &leaderOrchGroupSize); - - if (isOrchestrator || isLeader) - { - MPI_Comm leaderOrchComm = nullptr; - MPICHECK(MPI_Comm_create_group( - MPI_COMM_WORLD, leaderOrchCommGroup, participantIds.front(), &leaderOrchComm)); // NOLINT - mOrchLeaderComm = std::make_shared(leaderOrchComm, false); - } - else - { - mOrchLeaderComm = nullptr; - } -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::initializeCommAndWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, std::optional modelType, - std::optional const& modelPath, std::optional const& worldConfig, - std::optional const& decoderGptJsonConfig) -{ - if (modelType.has_value() && modelType.value() == ModelType::kENCODER_DECODER) - { - TLLM_CHECK_WITH_INFO(pp == 1, - "Encoder-Decoder C++ runtime doesn't support Pipeline Parallelism currently. Please switch to Python " - "runtime for PP mode, if necessary."); - } - - tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE); - mWorldRank = tensorrt_llm::mpi::MpiComm::world().getRank(); - mUsePipelineParallel = pp > 1; - - auto parallelConfig = executorConfig.getParallelConfig().value_or(ParallelConfig()); - validateParallelConfig(parallelConfig, modelType, modelPath); - - mCommMode = parallelConfig.getCommunicationMode(); - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - - mRecvPollPeriodMs = executorConfig.getRecvPollPeriodMs(); - - // Need to create communicator between orchestrator and leader if not spawning processes in orchestrator mode - if (mCommMode == CommunicationMode::kORCHESTRATOR && !optOrchestratorConfig.value().getSpawnProcesses()) - { - setOrchLeaderComm(tp, pp, cp, parallelConfig); - } - - if (mCommMode == CommunicationMode::kORCHESTRATOR && optOrchestratorConfig.value().getIsOrchestrator()) - { - initializeOrchestrator(tp, pp, cp, executorConfig, parallelConfig, modelType.value(), modelPath.value()); - } - else - { - initializeWorkers(tp, pp, cp, parallelConfig, worldConfig, decoderGptJsonConfig); - } -} - -void Executor::Impl::validateParallelConfig(ParallelConfig const& parallelConfig, std::optional modelType, - std::optional const& modelPath) -{ - TLLM_CHECK_WITH_INFO(parallelConfig.getCommunicationType() == CommunicationType::kMPI, - "Only CommunicationType kMPI is supported for now."); - - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - - if (parallelConfig.getCommunicationMode() == CommunicationMode::kORCHESTRATOR) - { - TLLM_CHECK_WITH_INFO( - optOrchestratorConfig, "OrchestratorConfig must be set when using ORCHESTRATOR communication mode."); - - TLLM_CHECK_WITH_INFO(modelPath, "OrchestratorMode only supports reading model weight from disk currently."); - - TLLM_CHECK_WITH_INFO(modelType, "OrchestratorMode requires modelType to be specified."); - } -} - -void Executor::Impl::initializeOrchestrator(SizeType32 tp, SizeType32 pp, SizeType32 cp, - ::tensorrt_llm::executor::ExecutorConfig const& executorConfig, ParallelConfig parallelConfig, ModelType modelType, - std::filesystem::path const& modelPath) -{ -#if ENABLE_MULTI_DEVICE - namespace su = tensorrt_llm::executor::serialize_utils; - - auto const& worldComm = tensorrt_llm::mpi::MpiComm::world(); - int32_t const worldSize = worldComm.getSize(); - - auto orchestratorConfig = parallelConfig.getOrchestratorConfig().value(); - - mIsWorker = false; - mIsLeader = false; - mIsPipelineLeader = false; - mIsOrchestrator = true; - - // Verify that worldSize is 1 - if (orchestratorConfig.getSpawnProcesses()) - { - TLLM_CHECK_WITH_INFO(worldSize == 1, - "When using the orchestrator mode and isOrchestrator is true, expect MPI worldSize to be 1."); - - // Spawn the worker threads - auto workerExecPath = orchestratorConfig.getWorkerExecutablePath(); - MPI_Comm intercomm = nullptr; - MPI_Info mpiInfo = nullptr; - MPICHECK(MPI_Info_create(&mpiInfo)); - MPICHECK(MPI_Info_set(mpiInfo, "env", "FORCE_NCCL_ALL_REDUCE_STRATEGY")); - - // Binding policy is not inherited for dynamically spawned jobs, resulting in the worker being bound - // to a single core. Override the setting to avoid perf issue - see https://nvbugs/4574329 - MPICHECK(MPI_Info_set(mpiInfo, "bind_to", "none")); - - MPICHECK(MPI_Comm_spawn(workerExecPath.c_str(), MPI_ARGV_NULL, tp * pp * cp, mpiInfo, 0, MPI_COMM_SELF, - &intercomm, MPI_ERRCODES_IGNORE)); - - mOrchLeaderComm = std::make_shared(intercomm, true); - // With intercomm, leader is rank 0 in the local group - mLeaderRank = 0; - mOrchRank = 0; - - // Copy the executor config, but set the orchestrator flag to false - auto newOrchConfig = OrchestratorConfig(false, orchestratorConfig.getWorkerExecutablePath()); - parallelConfig.setOrchestratorConfig(newOrchConfig); - auto execConfig = executorConfig; - execConfig.setParallelConfig(parallelConfig); - - // Serialize and send the executorConfig, the modelType and the modelPath - std::ostringstream oStream; - su::serialize(modelPath.string(), oStream); - su::serialize(modelType, oStream); - su::serialize(execConfig, oStream); - - auto str = oStream.str(); - std::vector buffer(str.begin(), str.end()); - auto bufferSize = static_cast(buffer.size()); - mOrchLeaderComm->bcast(&bufferSize, 1, mpi::MpiType::kINT64, MPI_ROOT); - mOrchLeaderComm->bcast(buffer.data(), buffer.size(), mpi::MpiType::kCHAR, MPI_ROOT); - - // Wait for workers to have created their executor instance - MPICHECK(MPI_Barrier(intercomm)); - } - - // Spawn the thread responsible for sending new requests to the leader of the model - mOrchSendReqThread = std::thread(&Impl::orchSendReqThread, this); - - // Spawn the thread responsible for receiving new responses from the leader of the model - mOrchRecvThread - = std::thread([&]() { this->orchRecvThread(mpi::MpiTag::kOrchestratorId, mpi::MpiTag::kOrchestratorData); }); - -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::initializeWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig& parallelConfig, - std::optional const& worldConfig, - std::optional const& decoderGptJsonConfig) -{ - auto const& worldComm = tensorrt_llm::mpi::MpiComm::world(); - int32_t const worldSize = worldComm.getSize(); - - auto const& orchestratorConfig = parallelConfig.getOrchestratorConfig(); - mIsOrchestrator = mCommMode == CommunicationMode::kORCHESTRATOR && orchestratorConfig.value().getIsOrchestrator(); - - TLLM_CHECK_WITH_INFO(mCommMode != CommunicationMode::kORCHESTRATOR || orchestratorConfig.has_value(), - "When using ORCHESTRATOR mode, orchestrator config must be set"); - - if (mCommMode == CommunicationMode::kORCHESTRATOR && !orchestratorConfig.value().getSpawnProcesses()) - { - TLLM_CHECK_WITH_INFO(parallelConfig.getParticipantIds(), - "When not spawning processes in orchestrator mode, participant IDs must be provided"); - - // Check that rank 0 is reserved for the orchestrator - auto const participantIds = parallelConfig.getParticipantIds().value(); - for (auto const& participantId : participantIds) - { - TLLM_CHECK_WITH_INFO(participantId != 0, "Rank 0 is reserved for the orchestrator"); - } - } - - // Participant ids - std::vector participantIds; - if (!parallelConfig.getParticipantIds()) - { - TLLM_CHECK_WITH_INFO(worldSize == tp * pp * cp, - "With communicationMode kLEADER, MPI worldSize is expected to be equal to tp*pp*cp when " - "participantIds are not specified"); - - participantIds.resize(tp * pp * cp); - std::iota(participantIds.begin(), participantIds.end(), 0); - } - else - { - if (mCommMode == CommunicationMode::kORCHESTRATOR && orchestratorConfig.value().getSpawnProcesses()) - { - TLLM_THROW( - "Participant ids should not be set when using CommunicationMode::kORCHESTRATOR with " - "spawnProcesses=true"); - } - participantIds = parallelConfig.getParticipantIds().value(); - TLLM_CHECK_WITH_INFO(static_cast(participantIds.size()) == tp * pp * cp, - tensorrt_llm::common::fmtstr("When specifying participantIds, participantIds size (%lu) must be equal to " - "tp*pp*cp (tp is %u, pp is %u, cp is %u)", - participantIds.size(), tp, pp, cp)); - } - - // If deviceIds are specified, check that they match tp*pp*cp - if (parallelConfig.getDeviceIds()) - { - auto deviceIds = parallelConfig.getDeviceIds().value(); - auto const hasNumNodes = parallelConfig.getNumNodes().has_value(); - if (hasNumNodes || static_cast(deviceIds.size()) != tp * pp * cp) - { - auto const numNodes = hasNumNodes ? parallelConfig.getNumNodes().value() : tensorrt_llm::mpi::getNumNodes(); - TLLM_CHECK_WITH_INFO(static_cast(deviceIds.size() * numNodes) == tp * pp * cp, - tensorrt_llm::common::fmtstr("When specifying deviceIds, deviceIds (%lu) * numNodes (%u) must be equal " - "to tp*pp*cp (tp is %u, pp is %u, cp is %u)", - deviceIds.size(), numNodes, tp, pp, cp)); - } - } - - // Bool that indicates if current process is worker for this model or not - auto participantIt = std::find(participantIds.begin(), participantIds.end(), mWorldRank); - mIsWorker = participantIt != participantIds.end(); - // Bool that indicates if current ranks is leader for this model - mIsLeader = (mWorldRank == participantIds.front()); - mIsPipelineLeader = (mWorldRank == participantIds[tp * (pp - 1)]); - -#if ENABLE_MULTI_DEVICE - if (mIsWorker) - { - // Create a session, but only assign to COMM_SESSION for ranks participating in this model - MPI_Group worldGroup = MPI_GROUP_NULL; - MPICHECK(MPI_Comm_group(MPI_COMM_WORLD, &worldGroup)); // NOLINT - MPI_Group sessionGroup = MPI_GROUP_NULL; - if (pp > 1) - { - // reverse participantIds to move leader to last pp rank. retain order in each tp group - std::reverse(participantIds.begin(), participantIds.end()); - if (tp > 1) - { - for (SizeType32 ppRank = 0; ppRank < pp; ppRank++) - { - std::reverse(participantIds.begin() + ppRank * tp, participantIds.begin() + (ppRank + 1) * tp); - } - } - } - MPICHECK(MPI_Group_incl(worldGroup, participantIds.size(), participantIds.data(), &sessionGroup)); // NOLINT - MPI_Comm sessionComm = MPI_COMM_NULL; - MPICHECK( - MPI_Comm_create_group(MPI_COMM_WORLD, sessionGroup, 1000 + participantIds.front(), &sessionComm)); // NOLINT - - tensorrt_llm::mpi::MpiComm::setSession(tensorrt_llm::mpi::MpiComm(sessionComm, false)); - } - - if (mIsLeader && mCommMode == CommunicationMode::kORCHESTRATOR) - { - auto optOrchestratorConfig = parallelConfig.getOrchestratorConfig(); - if (orchestratorConfig.has_value() && orchestratorConfig.value().getSpawnProcesses()) - { - mOrchLeaderComm = optOrchestratorConfig.value().getOrchLeaderComm(); - } - else - { - // mOrchLeaderComm has already been created - } - TLLM_CHECK(mOrchLeaderComm.get() != nullptr); - - TLLM_CHECK(worldConfig.has_value() || decoderGptJsonConfig.has_value()); - if (worldConfig.has_value()) - { - mDeviceId = worldConfig->getDevice(); - } - else - { - auto gpusPerNode = decoderGptJsonConfig->getGpusPerNode(); - auto worldConfig = runtime::WorldConfig::mpi(gpusPerNode, tp, pp, cp, parallelConfig.getDeviceIds()); - mDeviceId = worldConfig.getDevice(); - } - // Spawn the thread responsible for receiving new requests from the orchestrator - mLeaderRecvReqThread = std::thread(&Impl::leaderRecvReqThread, this); - - // Spawn the thread responsible for sending new responses to the orchestrator - mLeaderSendThread = std::thread([&]() - { this->leaderSendThread(mSendQueue, mpi::MpiTag::kOrchestratorId, mpi::MpiTag::kOrchestratorData); }); - } -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::initializeLogitsPostProcessorBatched(LogitsPostProcessorConfig const& logitsProcConfig) -{ - if (logitsProcConfig.getProcessorBatched().has_value()) - { - mLogitsPostProcessorBatched - = [cb = logitsProcConfig.getProcessorBatched().value()]( - std::vector const& reqIdsVec, - std::vector& logitsVec, - std::vector> const& beamTokensVec, - CudaStreamPtr const& cudaStreamPtr, - std::vector> const& clientIdsVec) - { - std::vector cbLogitsVec; - cbLogitsVec.reserve(logitsVec.size()); - for (auto& logits : logitsVec) - { - cbLogitsVec.emplace_back(executor::detail::ofITensor(logits)); - } - - cb(reqIdsVec, cbLogitsVec, beamTokensVec, cudaStreamPtr, clientIdsVec); - }; - - mModel->setLogitsPostProcessorBatched(mLogitsPostProcessorBatched); - } -} - -IdType Executor::Impl::enqueueRequest(Request const& request) -{ - return enqueueRequests({&request, 1}).at(0); -} - -std::vector Executor::Impl::enqueueRequests(std::vector const& requests) -{ - return enqueueRequests({requests.data(), requests.size()}); -} - -std::vector Executor::Impl::enqueueRequests(common::ArrayView const& requests) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called, cannot enqueue requests"); - checkParallelApiUsage(__func__); - - TLLM_LOG_DEBUG("Enqueuing %lu requests", requests.size()); - std::vector requestWithIds; - requestWithIds.reserve(requests.size()); - - // First check valid of request in enqueue thread, so Exceptions can be thrown to user. - for (auto const& req : requests) - { - auto logitsPostProcessorName = req.getLogitsPostProcessorName(); - if (logitsPostProcessorName && logitsPostProcessorName.value() != Request::kBatchedPostProcessorName) - { - getLogitsPostProcessor(*logitsPostProcessorName); - } - } - - std::vector ids; - { - auto now = std::chrono::steady_clock::now(); - for (auto const& req : requests) - { - ids.emplace_back(generateReqId(req)); - TLLM_LOG_DEBUG("Enqueue new request with id %d", ids.back()); - - std::vector childReqIds; - auto numChildRequests = getNumChildRequests(req); - if (numChildRequests > 0) - { - childReqIds.reserve(numChildRequests); - for (int childId = 0; childId < numChildRequests; childId++) - { - childReqIds.emplace_back(generateLocalReqId()); - TLLM_LOG_DEBUG("Add new child request with id %d", childReqIds.back()); - } - } - requestWithIds.emplace_back(RequestWithId{req, ids.back(), std::move(childReqIds), now}); - } - } - - if (mCommMode == CommunicationMode::kLEADER) - { - { - std::scoped_lock const lck(mQueuedReqMtx); - if (mMaxQueueSize) - { - auto const maxQueueSize = mMaxQueueSize.value(); - - auto totalRequestSize = 0; - for (auto&& reqWithId : requestWithIds) - { - totalRequestSize += (getNumChildRequests(reqWithId.req) + 1); - } - - if (maxQueueSize > 0 && mQueuedRequests.size() + totalRequestSize > static_cast(maxQueueSize)) - { - TLLM_THROW("Maximum queue size of %d has been reached, please try again later", maxQueueSize); - } - } - - for (auto&& req : requestWithIds) - { - insertRequestInOrder(mQueuedRequests, std::move(req)); - } - } - mQueuedReqCv.notify_one(); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - MpiMessage message(MpiId::PENDING_REQUEST); - message.data = PendingRequestData{std::move(requestWithIds)}; - mSendQueue.push(std::move(message)); - } - return ids; -} - -std::vector Executor::Impl::awaitResponses(std::optional const& timeout) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::unique_lock lck(mResponsesMtx); - auto pred = [this]() -> bool { return !mResponses.empty() || mShutdown; }; - auto storeResponses = [this]() - { - std::vector responses; - for (auto it = mResponses.begin(); it != mResponses.end();) - { - responses.insert(responses.end(), it->second.begin(), it->second.end()); - addTerminatedReqId(it->second, it->first); - it = mResponses.erase(it); - } - return responses; - }; - - std::vector responses; - if (timeout) - { - if (mResponsesCv.wait_for(lck, timeout.value(), pred)) - { - responses = storeResponses(); - } - } - else - { - mResponsesCv.wait(lck, pred); - responses = storeResponses(); - } - return responses; -} - -std::vector Executor::Impl::awaitResponses( - IdType const& reqId, std::optional const& timeout) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::unique_lock lck(mResponsesMtx); - auto pred = [this, reqId]() -> bool - { return (mResponses.find(reqId) != mResponses.end() && !mResponses.at(reqId).empty()) || mShutdown; }; - auto storeIdResponse = [this, reqId]() - { - std::vector responses; - responses.swap(mResponses.at(reqId)); - mResponses.erase(reqId); - addTerminatedReqId(responses, reqId); - return responses; - }; - - // We don't process a terminated request again. Terminated request is defined as a response - // with isFinal = true for a given requestId. - if (mTerminatedReqIds.contains(reqId)) - { - if (mResponses.find(reqId) != mResponses.end()) - { - TLLM_THROW("ReqId should already be removed from responses!"); - } - std::string const err = "ReqId " + std::to_string(reqId) + " has already been processed and was terminated."; - TLLM_LOG_ERROR("%s", err.c_str()); - - return {Response(reqId, err)}; - } - - std::vector responses; - if (timeout) - { - if (mResponsesCv.wait_for(lck, timeout.value(), pred)) - { - responses = storeIdResponse(); - } - } - else - { - mResponsesCv.wait(lck, pred); - responses = storeIdResponse(); - } - return responses; -} - -std::vector> Executor::Impl::awaitResponses( - std::vector const& requestIds, std::optional const& timeout) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::vector> responses; - responses.reserve(requestIds.size()); - if (timeout) - { - auto const start_time = std::chrono::high_resolution_clock::now(); - for (auto const requestId : requestIds) - { - auto const elapsed_ms = std::chrono::duration_cast( - std::chrono::high_resolution_clock::now() - start_time); - responses.emplace_back(awaitResponses( - requestId, timeout.value() > elapsed_ms ? timeout.value() - elapsed_ms : std::chrono::milliseconds{0})); - } - } - else - { - for (auto const requestId : requestIds) - { - responses.emplace_back(awaitResponses(requestId)); - } - } - return responses; -} - -SizeType32 Executor::Impl::getNumResponsesReady(std::optional const& optId) const -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::scoped_lock lck(mResponsesMtx); - SizeType32 numResponsesReady = 0; - if (optId) - { - auto const reqId = optId.value(); - auto const respIt = mResponses.find(reqId); - if (respIt != mResponses.end()) - { - numResponsesReady = static_cast(respIt->second.size()); - } - } - else - { - for (auto const& [id, responses] : mResponses) - { - numResponsesReady += static_cast(responses.size()); - } - } - return numResponsesReady; -} - -void Executor::Impl::shutdown() -{ - // Cannot call shutdown multiple times - if (mShutdownCalled) - { - return; - } - mShutdownCalled = true; - - if (!mShutdown) - { - if (mCommMode == CommunicationMode::kLEADER && mIsLeader) - { - // Enqueue a request to indicate to other ranks to terminate - enqueueTerminateRequest(); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - if (mIsOrchestrator) - { - // Send to the leader the termination signal - mShutdown = true; - mResponsesCv.notify_all(); - - mSendQueue.push(MpiMessage(MpiId::TERMINATION)); - - // Wait for sender thread to exit - if (mOrchSendReqThread.joinable()) - { - mOrchSendReqThread.join(); - } - // Wait for recv response thread to exit - if (mOrchRecvThread.joinable()) - { - mOrchRecvThread.join(); - } - } - else if (mIsLeader) - { - // Wait for sender thread to exit - if (mLeaderRecvReqThread.joinable()) - { - mLeaderRecvReqThread.join(); - } - // Wait for send response thread to exit - if (mLeaderSendThread.joinable()) - { - mLeaderSendThread.join(); - } - } - } - } - - // Wait for execution thread to terminate - if (mExecutionThread.joinable()) - { - mExecutionThread.join(); - } - - // If we overwrote COMM_SESSION with split, free it now. Otherwise, since - // COMM_SESSION is a global static object, it will be destroyed in an - // undefined order and can cause crashes on program exit. - if (mIsWorker) - { - tensorrt_llm::mpi::MpiComm::setSession(tensorrt_llm::mpi::MpiComm(MPI_COMM_WORLD, false)); - } -} - -void Executor::Impl::cancelRequest(IdType requestId) -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - - // Check if the request is terminated already. If so, return - { - std::scoped_lock lckResp(mResponsesMtx); - if (mTerminatedReqIds.contains(requestId)) - { - TLLM_LOG_INFO("Ignoring already terminated request %lu", requestId); - return; - } - } - - if (mCommMode == CommunicationMode::kLEADER) - { - std::scoped_lock lck(mCancelReqMtx); - auto& selCancelledReqIds = mUsePipelineParallel ? mPipelineCancelledReqIds : mCancelledReqIds; - selCancelledReqIds.insert(requestId); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - MpiMessage message(MpiId::CANCEL_REQUEST); - std::vector cancelledReqIds{requestId}; - message.data = RequestIdsData{std::move(cancelledReqIds)}; - mSendQueue.push(std::move(message)); - } -} - -std::deque Executor::Impl::getLatestIterationStats() -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - std::scoped_lock lck(mIterStatsMtx); - return std::exchange(mIterationStats, {}); -} - -std::deque Executor::Impl::getLatestRequestStats() -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - checkParallelApiUsage(__func__); - - std::scoped_lock lck(mRequestStatsMtx); - return std::exchange(mRequestStats, {}); -} - -std::deque Executor::Impl::getLatestDebugTensors() -{ - TLLM_CHECK_WITH_INFO(!mShutdownCalled, "Shutdown called"); - if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - TLLM_LOG_WARNING("getLatestDebugTensors is not supported in ORCHESTRATOR mode yet"); - return {}; - } - if (mEncoderModel) - { - TLLM_LOG_WARNING("getLatestDebugTensors is not supported for encoder model yet"); - } - std::scoped_lock lck(mDebugTensorsMtx); - return std::exchange(mDebugTensors, {}); -} - -bool Executor::Impl::canEnqueueRequests() const -{ - return !mShutdownCalled - && ((mCommMode == CommunicationMode::kLEADER && mIsLeader) - || (mCommMode == CommunicationMode::kORCHESTRATOR && mIsOrchestrator)); -} - -bool Executor::Impl::isParticipant() const -{ - return mIsWorker; -} - -std::optional> Executor::Impl::getKVCacheEventManager() const -{ - if (!mModel) - { - return std::nullopt; - } - auto cacheEventManager = mModel->getKVCacheManager(); - return cacheEventManager ? std::optional(std::make_shared(cacheEventManager)) : std::nullopt; -} - -void Executor::Impl::requestWithIdLeaderThread() -{ - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - auto constexpr peer = 0; - while (true) - { - int64_t numActiveRequests; - mCommPipelineParallel->recv( - &numActiveRequests, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kExecutorNumActiveRequests); - if (numActiveRequests < 0) - { - break; - } - - bool lowestPriorityActiveHasValue; - std::optional lowestPriorityActive; - mCommPipelineParallel->recv(&lowestPriorityActiveHasValue, 1, mpi::MpiType::kBOOL, peer, - mpi::MpiTag::kExecutorLowestPriorityActiveHasValue); - if (lowestPriorityActiveHasValue) - { - PriorityType lowestPriorityActiveValue; - mCommPipelineParallel->recv( - &lowestPriorityActiveValue, 1, mpi::MpiType::kFLOAT, peer, mpi::MpiTag::kExecutorLowestPriorityActive); - lowestPriorityActive = lowestPriorityActiveValue; - } - - auto reqWithIds = getLeaderNewReqWithIds(numActiveRequests, lowestPriorityActive); - setupDynamicLogitsPostProcessors(reqWithIds); - auto requestWithIdAsyncSndHdl - = std::make_unique(mCommPipelineParallel, reqWithIds, peer); - requestWithIdAsyncSndHdl.reset(nullptr); - } -} - -void Executor::Impl::cancelledRequestsLeaderThread() -{ - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - auto constexpr peer = 0; - while (true) - { - bool shouldExit; - mCommPipelineParallel->recv(&shouldExit, 1, mpi::MpiType::kBOOL, peer, mpi::MpiTag::kExecutorShouldExit); - if (shouldExit) - { - break; - } - - std::unique_ptr cancelledRequestsAsyncSndHdl; - { - std::scoped_lock lck(mCancelReqMtx); - cancelledRequestsAsyncSndHdl - = std::make_unique(mCommPipelineParallel, mPipelineCancelledReqIds, peer); - mPipelineCancelledReqIds.clear(); - } - cancelledRequestsAsyncSndHdl.reset(nullptr); - } -} - -std::vector Executor::Impl::getLeaderNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive) -{ - std::unique_lock lck(mQueuedReqMtx); - mQueuedReqCv.wait(lck, [&]() { return (!mQueuedRequests.empty() || numActiveRequests > 0 || mShutdown); }); - - std::vector reqWithIds; - - if (mQueuedRequests.empty() || mShutdown) - { - return reqWithIds; - } - - if (mQueuedRequests.front().id == kTerminateReqId) - { - reqWithIds.emplace_back(std::move(mQueuedRequests.front())); - mQueuedRequests.pop_front(); - return reqWithIds; - } - - auto const& firstRequest = mQueuedRequests.front(); - auto const firstBeamWidth = firstRequest.req.getSamplingConfig().getBeamWidth(); - auto const operatingBeamWidth = numActiveRequests > 0 ? mModel->getOperatingBeamWidth() : firstBeamWidth; - - auto const tryInsertQueuedRequestIntoReqWithIds = [this, &reqWithIds, operatingBeamWidth]() -> bool - { - auto& nextRequest = mQueuedRequests.front(); - auto const beamWidth = nextRequest.req.getSamplingConfig().getBeamWidth(); - if (beamWidth != operatingBeamWidth) - { - TLLM_LOG_INFO( - "Can't dequeue request with ID %ld because beam width %d differs from operating beam width %d.", - nextRequest.id, beamWidth, operatingBeamWidth); - return false; - } - - TLLM_LOG_DEBUG("Dequeue request with ID %ld", nextRequest.id); - reqWithIds.emplace_back(std::move(nextRequest)); - mQueuedRequests.pop_front(); - return true; - }; - - auto const maxNewRequests = static_cast(std::max(mMaxNumActiveRequests - numActiveRequests, 0)); - for (size_t req = 0; !mQueuedRequests.empty() && req < maxNewRequests;) - { - req += (getNumChildRequests(mQueuedRequests.front().req) + 1); - if (req > maxNewRequests) - { - break; - } - if (!tryInsertQueuedRequestIntoReqWithIds()) - { - break; - } - } - - if (lowestPriorityActive) - { - while (!mQueuedRequests.empty() && mQueuedRequests.front().req.getPriority() > (*lowestPriorityActive)) - { - if (!tryInsertQueuedRequestIntoReqWithIds()) - { - break; - } - } - } - return reqWithIds; -} - -std::vector Executor::Impl::getNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - auto const& worldConfig = mModel->getWorldConfig(); - - if (worldConfig.isPipelineParallel()) - { - mRequestWithIdWaitThread->waitStop(); - } - - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - std::vector reqWithIds; - if (mIsPipelineLeader) - { - if (!worldConfig.isPipelineParallel()) - { - reqWithIds = getLeaderNewReqWithIds(numActiveRequests, lowestPriorityActive); - setupDynamicLogitsPostProcessors(reqWithIds); - } - else - { - auto const peer = worldConfig.getPipelineParallelism() - 1; - auto numActiveRequestsValue = static_cast(numActiveRequests); - auto request1 = mCommPipelineParallel->sendAsync( - &numActiveRequestsValue, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kExecutorNumActiveRequests); - bool lowestPriorityActiveHasValue = lowestPriorityActive.has_value(); - auto request2 = mCommPipelineParallel->sendAsync(&lowestPriorityActiveHasValue, 1, mpi::MpiType::kBOOL, - peer, mpi::MpiTag::kExecutorLowestPriorityActiveHasValue); - auto request3 = lowestPriorityActiveHasValue - ? mCommPipelineParallel->sendAsync(&lowestPriorityActive.value(), 1, mpi::MpiType::kFLOAT, peer, - mpi::MpiTag::kExecutorLowestPriorityActive) - : nullptr; - request1->wait(); - request2->wait(); - if (request3) - { - request3->wait(); - } - reqWithIds = RequestWithIdAsyncSend::requestWithIdRecv(mCommPipelineParallel, peer); - } - if (worldConfig.isTensorParallel() || worldConfig.isContextParallel()) - { - auto packed = RequestWithId::serializeReqWithIds(reqWithIds); - if (worldConfig.isTensorParallel()) - { - mCommTensorParallel->bcast(packed, 0); - } - if (worldConfig.isContextParallel()) - { - mCommContextParallel->bcast(packed, 0); - } - } - } - else - { - if (worldConfig.isFirstPipelineParallelRank()) - { - std::vector buffer; - mCommTensorParallel->bcast(buffer, 0); - mCommContextParallel->bcast(buffer, 0); - reqWithIds = RequestWithId::deserializeReqWithIds(buffer); - } - else - { - auto const peer = worldConfig.getPipelineParallelRank() - 1; - reqWithIds = RequestWithIdAsyncSend::requestWithIdRecv(mCommPipelineParallel, peer); - } - } - if (!worldConfig.isLastPipelineParallelRank()) - { - auto const peer = worldConfig.getPipelineParallelRank() + 1; - mRequestWithIdAsyncSndHdl = std::make_unique(mCommPipelineParallel, reqWithIds, peer); - mRequestWithIdWaitThread->notifyStart(); - } - TLLM_CUDA_CHECK(cudaSetDevice(mModel->getWorldConfig().getDevice())); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return reqWithIds; -} - -std::tuple Executor::Impl::fetchNewRequests( - SizeType32 numActiveRequests, std::optional lowestPriorityActive) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_SCOPED_RANGE(fetchNewRequests); - - // If grab requests from queue, do exchange between ranks - auto reqWithIds = getNewReqWithIds(numActiveRequests, lowestPriorityActive); - RequestList newRequests; - double newActiveRequestsQueueLatencyMS{0.}; - for (auto& reqWithId : reqWithIds) - { - if (reqWithId.id == kTerminateReqId) - { - mShutdown = true; - mResponsesCv.notify_all(); - return {}; - } - - try - { - std::optional llmRequestLogitsPostProcessor; - bool applyLogitsPostProcessorBatched{false}; - if (mModel->getWorldConfig().isLastPipelineParallelRank()) - { - auto logitsPostProcessorName = reqWithId.req.getLogitsPostProcessorName(); - if (logitsPostProcessorName) - { - if (logitsPostProcessorName.value() == Request::kBatchedPostProcessorName) - { - TLLM_CHECK_WITH_INFO( - mLogitsPostProcessorBatched, "Batched logits post processor is not defined."); - applyLogitsPostProcessorBatched = true; - } - else - { - if (logitsPostProcessorName->compare(0, - std::char_traits::length(Request::kDynamicPostProcessorNamePrefix), - Request::kDynamicPostProcessorNamePrefix) - == 0) - { - TLLM_CHECK_WITH_INFO(!mModel->getReplicateLogitsPostProcessor() - || mModel->getWorldConfig().getTensorParallelism() == 1, - "Dynamic logits postprocessor must be used with replicate=false or no tensor " - "parallelism."); - } - if (mModel->getWorldConfig().isFirstTensorParallelRank() - || mModel->getReplicateLogitsPostProcessor()) - { - llmRequestLogitsPostProcessor = getLogitsPostProcessor(logitsPostProcessorName.value()); - } - else - { - llmRequestLogitsPostProcessor - = [](IdType reqId, RtTensorPtr& logits, BeamTokens const& beamTokens, - CudaStreamPtr const& cudaStreamPtr, std::optional clientId) {}; - } - } - } - } - auto newLlmReq = std::make_shared( - reqWithId.id, reqWithId.req, llmRequestLogitsPostProcessor, applyLogitsPostProcessorBatched); - - auto numReturnSequences = newLlmReq->getNumSubRequests(); - if (numReturnSequences > 1) - { - TLLM_CHECK(reqWithId.childReqIds.size() == static_cast(numReturnSequences - 1)); - mChildReqIdsMap[reqWithId.id] = reqWithId.childReqIds; - } - - for (auto seqIdx = 0; seqIdx < numReturnSequences; seqIdx++) - { - auto newReq - = seqIdx == 0 ? newLlmReq : newLlmReq->createChildRequest(reqWithId.childReqIds.at(seqIdx - 1)); - - // If static batching and streaming, disable streaming and exclude input - if (mBatchingType == BatchingType::kSTATIC && newReq->isStreaming()) - { - newReq->setStreaming(false); - newReq->setExcludeInputFromOutput(true); - } - - // Validate the request parameters - newReq->validate(mModel->getMaxInputLen(), mModel->getMaxSequenceLen(), mModel->getMaxDraftLen(), - mModel->getVocabSizePadded(), - mEncoderModel ? std::optional(mEncoderModel->getMaxInputLen()) : std::nullopt, - mEnableBlockReuse); - - TLLM_CHECK_WITH_INFO(!mEncoderModel || !mIsSchedulerMaxUtilization, - "Encoder or Encoder-Decoder model don't support max utilization scheduler yet. Only max requests " - "or guaranteed no evict."); - - // When streaming is enabled and scheduling policy permits evict/restart, need to guard against the case - // where the sequence is truncated on eviction (to respect maxInputLen limits), resulting in loss of - // some tokens that have been streamed out. In this case, resuming generation may result in different - // completion for locations whose tokens have already been returned. There is no way to protect against - // this, so disallowing. - if (newReq->isStreaming() && !mIsSchedulerGuaranteedNoEvict && !mIsChunkedContext) - { - auto const maxReqSeqLen = newReq->mPromptLen + newReq->mMaxNewTokens; - auto const maxRestartLen = maxReqSeqLen - 1; - TLLM_CHECK_WITH_INFO(maxRestartLen <= mModel->getMaxInputLen(), - "Request sequence length is potentially greater than max input length. This cannot be run " - "unless streaming is disabled, context chunking is enabled or the GUARANTEED_NO_EVICT " - "scheduling policy is used"); - } - - // Create the encoder output tensor - if (mEncoderModel) - { - TLLM_CHECK_WITH_INFO(mModel || (!mModel && newReq->getReturnEncoderOutput()), - "Encoder-Decoder models allow optionally returning encoder output. But if it is Encoder-only " - "models, please make sure returnEncoderOutput is always true."); - - // gpu buffers for passing to the next phase - newReq->allocEncoderOutput(mEncoderModel->getBufferManager(), mEncoderModel->getLogitDataType()); - newReq->allocEncoderHiddenStates( - mEncoderModel->getBufferManager(), mEncoderModel->getLogitDataType()); - // pinned buffers for returning results to host - if (newReq->getReturnEncoderOutput()) - { - newReq->allocEncoderOutputHost( - mEncoderModel->getHiddenSize() * mEncoderModel->getWorldConfig().getTensorParallelism(), - mEncoderModel->getLogitDataType()); - } - } - - if (!mEncoderModel && newReq->getEncoderInputFeatures()) - { - TLLM_LOG_INFO("Allocating buffers for encoder output"); - // gpu buffers for passing to the next phase - newReq->allocEncoderOutput(mModel->getBufferManager(), mModel->getLogitDataType()); - newReq->allocEncoderHiddenStates(mModel->getBufferManager(), mModel->getLogitDataType()); - } - - // Create the context logits tensor - if (newReq->getReturnContextLogits()) - { - TLLM_CHECK_WITH_INFO(mModel->getModelConfig().computeContextLogits(), - "Return context logit need to build engine with gather_context_logits"); - newReq->allocContextLogitsHost(mModel->getVocabSizePadded(), mModel->getLogitDataType()); - } - - // Create the generation logits tensor - if (newReq->getReturnGenerationLogits()) - { - TLLM_CHECK_WITH_INFO(mModel->getGatherGenerationLogits(), - "To return generation logits, gather_generation_logits must be enabled in ExecutorConfig"); - - if (mModel->getModelConfig().getSpeculativeDecodingMode().isDraftTokensExternal() - && newReq->hasDraftTokens()) - { - newReq->allocTargetModelAcceptedTokenLogitsHost( - mModel->getVocabSizePadded(), mModel->getLogitDataType()); - } - else - { - newReq->allocGenerationLogitsHost(mModel->getVocabSizePadded(), mModel->getLogitDataType()); - } - } - - if (mModel->getWorldConfig().isLastPipelineParallelRank() && newReq->getGuidedDecodingParams()) - { - TLLM_CHECK_WITH_INFO(mModel->hasGuidedDecoder(), - "Request is specified with GuidedDecodingParams, but GuidedDecoder is not setup. Please " - "provide a valid GuidedDecodingConfig to setup GuidedDecoder."); - } - - if (mModel->getWorldConfig().isLastPipelineParallelRank() && newReq->hasAdditionalOutputs()) - { - newReq->allocAdditionalOutputs([this](std::string const& name) - { return mModel->getTensorDataType(name); }, - [this](std::string const& name) { return mModel->getTensorShape(name); }); - } - - mModel->updatePeftCache(newReq); - - newRequests.emplace_back(std::move(newReq)); - } - - auto queuedEnd = std::chrono::steady_clock::now(); - auto reqQueueLatencyMS - = std::chrono::duration(queuedEnd - reqWithId.queuedStart).count(); - newActiveRequestsQueueLatencyMS += reqQueueLatencyMS; - } - catch (runtime::LoraExpectedException const& e) - { - if (mIsLeader) - { - // In case of an expected LoRA exception (e.g. cache full, cache miss), log a warning and enqueue - // response - TLLM_LOG_WARNING("%s", e.what()); - enqueueNewResponses({{reqWithId.id, e.what(), reqWithId.req.getClientId()}}); - } - } - catch (std::exception const& e) - { - if (mIsLeader) - { - // In case of error, create a response with error for this request - auto err = std::string("Encountered an error when fetching new request: ") + e.what(); - TLLM_LOG_ERROR("%s", err.c_str()); - enqueueNewResponses({{reqWithId.id, err, reqWithId.req.getClientId()}}); - } - } - } - TLLM_LOG_DEBUG("[RANK %d] num new requests fetched from queue: %d", COMM_SESSION.getRank(), newRequests.size()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return {newRequests, newActiveRequestsQueueLatencyMS}; -} - -void Executor::Impl::terminateActiveRequests(RequestList& activeRequests, std::string const& err) -{ - TLLM_LOG_ERROR("%s", err.c_str()); - - // Create a response for all requests and add to queue - for (auto it = activeRequests.cbegin(); it != activeRequests.cend();) - { - auto llmReq = (*it); - - llmReq->setState(batch_manager::LlmRequestState::kGENERATION_COMPLETE); - mModel->terminateRequest(llmReq); - - if (mIsLeader) - { - enqueueNewResponses({{llmReq->mRequestId, err, llmReq->mClientId}}); - } - - // Remove from the requestList - it = activeRequests.erase(it); - } -} - -void Executor::Impl::forwardSync(RequestList& activeRequests) -{ - TLLM_LOG_TRACE("[RANK %d] %s start", COMM_SESSION.getRank(), __PRETTY_FUNCTION__); - try - { - if (mEncoderModel) - { - mEncoderModel->forwardSync(); - } - mModel->forwardSync(); - } - catch (std::exception const& e) - { - std::string const err = std::string("Encountered an error in forwardSync function: ") + e.what(); - terminateActiveRequests(activeRequests, err); - } - TLLM_LOG_TRACE("[RANK %d] %s stop", COMM_SESSION.getRank(), __PRETTY_FUNCTION__); -} - -// The function is used to change the state of a request to context_init from encoder_init for enc-dec model whose -// encoder is skipped. The encoder output is populated accordingly with input features given through model executor of -// decoder. -void Executor::Impl::prepRequestsForEncoderSkip(RequestList& activeRequests) -{ - - for (auto& req : activeRequests) - { - - if (req->isEncoderInitState() && req->getEncoderInputFeatures()) - { - TLLM_LOG_INFO("Changing state of request and setting encoder output to skip encoder run"); - req->setState(batch_manager::LlmRequestState::kCONTEXT_INIT); - req->setEncoderOutput(req->getEncoderInputFeatures()); - } - } -} - -void Executor::Impl::finishTimedOutRequests(RequestList const& activeRequests) -{ - if (mIsLeader) - { - for (auto const& request : activeRequests) - { - if (request->isTimedOut() && !request->isFinished()) - { - // workaround to cancelRequest since it throws an error if - // mCommMode == CommunicationMode::kORCHESTRATOR && !mIsOrchestrator - { - std::scoped_lock lck(mCancelReqMtx); - auto& selCancelledReqIds = mUsePipelineParallel ? mPipelineCancelledReqIds : mCancelledReqIds; - selCancelledReqIds.insert(request->mRequestId); - } - } - } - } -} - -void Executor::Impl::forwardAsync(RequestList& activeRequests) -{ - try - { - TLLM_LOG_DEBUG("num active requests in scope: %d", activeRequests.size()); - - if (mDynamicBatchTuner) - { - auto const averageInputLength = static_cast(mDynamicBatchTuner->getAverageInputLength()); - auto const averageOutputLength = static_cast(mDynamicBatchTuner->getAverageOutputLength()); - auto const maxCapacityBatchSize = mModel->getMaxCapacityBatchSize(averageInputLength, averageOutputLength); - - if (mDynamicBatchTuner->isBatchSizeTuningEnabled()) - { - auto runtimeBatchSize = mDynamicBatchTuner->getRuntimeBatchSize(maxCapacityBatchSize); - mModel->setRuntimeBatchSize(runtimeBatchSize); - } - - if (mDynamicBatchTuner->isMaxNumTokensTuningEnabled()) - { - auto runtimeBatchSize = mModel->getRuntimeBatchSize(); - auto runtimeMaxNumTokens = mDynamicBatchTuner->getRuntimeMaxNumTokens(runtimeBatchSize); - mModel->setRuntimeMaxNumTokens(runtimeMaxNumTokens); - } - } - - if (mEncoderModel) - { - mEncoderModel->forwardAsync(activeRequests); - auto const& encoderStream = *(mEncoderModel->getRuntimeStreamPtr()); - auto const& decoderStream = *(mModel->getRuntimeStreamPtr()); - runtime::CudaEvent encoderFinished; - encoderStream.record(encoderFinished); - decoderStream.wait(encoderFinished); - } - - if (!mEncoderModel) - { - prepRequestsForEncoderSkip(activeRequests); - } - - mModel->forwardAsync(activeRequests); - } - catch (std::exception const& e) - { - std::string err = std::string("Encountered an error in forwardAsync function: ") + e.what(); - terminateActiveRequests(activeRequests, err); - } -} - -IterationStats Executor::Impl::getCurrentIterationStats(RequestList const& activeRequests, double iterLatencyMS, - SizeType32 numNewActiveRequests, double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests) -{ - IterationStats stats; - // Timestamp - stats.timestamp = tensorrt_llm::common::getCurrentTimestamp(); - stats.numNewActiveRequests = numNewActiveRequests; - stats.iterLatencyMS = iterLatencyMS; - stats.newActiveRequestsQueueLatencyMS = newActiveRequestsQueueLatencyMS; - // Active request count - stats.numActiveRequests = static_cast(activeRequests.size()); - // Queued request count - { - std::scoped_lock lck(mQueuedReqMtx); - stats.numQueuedRequests = static_cast(mQueuedRequests.size()); - } - stats.numCompletedRequests = numCompletedRequests; - // Max number of requests - stats.maxNumActiveRequests = mMaxNumActiveRequests; - // Runtime memory allocation statistics - auto const& memoryCounters = runtime::MemoryCounters::getInstance(); - stats.gpuMemUsage = memoryCounters.getGpu(); - stats.cpuMemUsage = memoryCounters.getCpu(); - stats.pinnedMemUsage = memoryCounters.getPinned(); - - // Model specific stats - mModel->getCurrentIterationStats(stats); - return stats; -} - -RequestStatsPerIteration Executor::Impl::getCurrentRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests) -{ - std::vector requestStatsVec; - - auto includeDisServingStats = [](LlmRequestPtr const& request, tensorrt_llm::executor::RequestStats& requestStats) - { - auto requestType = request->getLlmRequestType(); - if (requestType == batch_manager::LlmRequestType::LLMREQUEST_TYPE_CONTEXT_ONLY - || requestType == batch_manager::LlmRequestType::LLMREQUEST_TYPE_GENERATION_ONLY) - { - requestStats.disServingStats - = executor::DisServingRequestStats{request->getKvCacheTransferTimeMS(), request->getKvCacheSize()}; - } - }; - - for (auto const& request : activeRequests) - { - RequestStats requestStats; - requestStats.id = request->mRequestId; - requestStats.stage = request->getRequestStage(); - requestStats.contextPrefillPosition = request->getContextCurrentPosition(); - requestStats.numGeneratedTokens = request->getMaxBeamNumTokens() - request->getOrigPromptLen(); - requestStats.avgNumDecodedTokensPerIter = request->getAvgDecodedTokensPerIter(); - includeDisServingStats(request, requestStats); - requestStats.allocTotalBlocksPerRequest = request->getAllocTotalBlocksPerRequest(); - requestStats.allocNewBlocksPerRequest = request->getAllocNewBlocksPerRequest(); - requestStats.reusedBlocksPerRequest = request->getReusedBlocksPerRequest(); - requestStats.missedBlocksPerRequest = request->getMissedBlocksPerRequest(); - requestStats.kvCacheHitRatePerRequest = request->getKVCacheHitRatePerRequest(); - requestStatsVec.emplace_back(requestStats); - } - - { - std::unique_lock lck(mQueuedReqMtx); - for (auto const& request : mQueuedRequests) - { - // Still waiting for the first scheduling - RequestStats requestStats; - requestStats.id = static_cast(request.id); - requestStats.stage = executor::RequestStage::kQUEUED; - requestStats.contextPrefillPosition = 0; - requestStats.numGeneratedTokens = 0; - requestStats.avgNumDecodedTokensPerIter = 0; - requestStats.allocTotalBlocksPerRequest = 0; - requestStats.allocNewBlocksPerRequest = 0; - requestStats.reusedBlocksPerRequest = 0; - requestStats.missedBlocksPerRequest = 0; - requestStats.kvCacheHitRatePerRequest = 0; - requestStatsVec.emplace_back(requestStats); - } - } - - for (auto const& request : finishedRequests) - { - // Still waiting for the first scheduling - RequestStats requestStats; - requestStats.id = static_cast(request->mRequestId); - requestStats.stage = executor::RequestStage::kGENERATION_COMPLETE; - requestStats.contextPrefillPosition = request->getContextCurrentPosition(); - requestStats.numGeneratedTokens = request->getMaxBeamNumTokens() - request->getOrigPromptLen(); - requestStats.avgNumDecodedTokensPerIter = request->getAvgDecodedTokensPerIter(); - includeDisServingStats(request, requestStats); - requestStats.allocTotalBlocksPerRequest = request->getAllocTotalBlocksPerRequest(); - requestStats.allocNewBlocksPerRequest = request->getAllocNewBlocksPerRequest(); - requestStats.reusedBlocksPerRequest = request->getReusedBlocksPerRequest(); - requestStats.missedBlocksPerRequest = request->getMissedBlocksPerRequest(); - requestStats.kvCacheHitRatePerRequest = request->getKVCacheHitRatePerRequest(); - requestStatsVec.emplace_back(requestStats); - } - - RequestStatsPerIteration stats{0, std::move(requestStatsVec)}; - - // Model specific stats - mModel->getCurrentRequestStats(stats); - return stats; -} - -void Executor::Impl::appendCurrentIterStats(IterationStats&& currentIterStats) -{ - std::scoped_lock lck(mIterStatsMtx); - if (mIterationStats.size() >= mIterStatsMaxIterations) - { - mIterationStats.pop_front(); - } - mIterationStats.emplace_back(std::move(currentIterStats)); -} - -void Executor::Impl::appendMultipleIterStats(std::vector&& currentIterStatsVec) -{ - std::scoped_lock lck(mIterStatsMtx); - if (mIterationStats.size() + currentIterStatsVec.size() > mIterStatsMaxIterations) - { - size_t removeCount = mIterationStats.size() + currentIterStatsVec.size() - mIterStatsMaxIterations; - for (size_t i = 0; i < removeCount; i++) - { - mIterationStats.pop_front(); - } - } - mIterationStats.insert(mIterationStats.end(), std::make_move_iterator(currentIterStatsVec.begin()), - std::make_move_iterator(currentIterStatsVec.end())); -} - -void Executor::Impl::updateIterationStats(RequestList const& activeRequests, double iterLatencyMS, - SizeType32 numNewActiveRequests, double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests, - bool flushToOrchestrator) -{ - NVTX3_SCOPED_RANGE(updateIterationStats); - if (mIterStatsMaxIterations > 0 && mIsLeader) - { - auto currentIterStats = getCurrentIterationStats( - activeRequests, iterLatencyMS, numNewActiveRequests, newActiveRequestsQueueLatencyMS, numCompletedRequests); - // Send the stats to the orchestrator - if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - bool hasSchedThisIter = (currentIterStats.inflightBatchingStats - && currentIterStats.inflightBatchingStats->numScheduledRequests > 0) - || (currentIterStats.staticBatchingStats - && currentIterStats.staticBatchingStats->numScheduledRequests > 0); - appendCurrentIterStats(std::move(currentIterStats)); - if (hasSchedThisIter || flushToOrchestrator) - { - std::deque iterStatsQueue; - { - std::scoped_lock lck(mIterStatsMtx); - iterStatsQueue = std::exchange(mIterationStats, {}); - } - MpiMessage message(MpiId::ITER_STATS); - std::vector iterStates( - std::make_move_iterator(iterStatsQueue.begin()), std::make_move_iterator(iterStatsQueue.end())); - message.data = IterStatsData{std::move(iterStates)}; - mSendQueue.push(std::move(message)); - } - } - else - { - // Add current iteration stats - appendCurrentIterStats(std::move(currentIterStats)); - } - } -} - -void Executor::Impl::appendCurrentRequestStats(RequestStatsPerIteration&& currentRequestStats) -{ - std::scoped_lock lck(mRequestStatsMtx); - if (mRequestStats.size() >= mRequestStatsMaxIterations) - { - mRequestStats.pop_front(); - } - mRequestStats.emplace_back(std::move(currentRequestStats)); -} - -void Executor::Impl::appendMultipleRequestStats(std::vector&& currentRequestStatsVec) -{ - std::scoped_lock lck(mRequestStatsMtx); - if (mRequestStats.size() + currentRequestStatsVec.size() > mRequestStatsMaxIterations) - { - size_t removeCount = mRequestStats.size() + currentRequestStatsVec.size() - mRequestStatsMaxIterations; - for (size_t i = 0; i < removeCount; i++) - { - mRequestStats.pop_front(); - } - } - mRequestStats.insert(mRequestStats.end(), std::make_move_iterator(currentRequestStatsVec.begin()), - std::make_move_iterator(currentRequestStatsVec.end())); -} - -void Executor::Impl::updateRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests, bool flushToOrchestrator) -{ - NVTX3_SCOPED_RANGE(updateRequestStats); - if (mRequestStatsMaxIterations > 0 && mIsLeader) - { - // Add current iteration request stats - auto currentRequestStats = getCurrentRequestStats(activeRequests, finishedRequests); - // Send the stats to the orchestrator - if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - bool hasScheduledReqs = false; - if (!flushToOrchestrator) - { - size_t activeSize = activeRequests.size(); - TLLM_CHECK_WITH_INFO(currentRequestStats.requestStats.size() >= activeSize, - "currentRequestStats num is %ld should >= activeRequest num:%zu", - currentRequestStats.requestStats.size(), activeSize); - hasScheduledReqs = std::any_of(currentRequestStats.requestStats.begin(), - currentRequestStats.requestStats.begin() + static_cast(activeSize), - [](RequestStats const& requestStat) { return requestStat.scheduled; }); - } - appendCurrentRequestStats(std::move(currentRequestStats)); - if (hasScheduledReqs || flushToOrchestrator) - { - std::deque requestStatsQueue; - { - std::scoped_lock lck(mRequestStatsMtx); - requestStatsQueue = std::exchange(mRequestStats, {}); - } - std::vector requestIterStates( - std::make_move_iterator(requestStatsQueue.begin()), - std::make_move_iterator(requestStatsQueue.end())); - MpiMessage message(MpiId::REQUEST_ITER_STATS); - message.data = RequestStatsPerIterationData{std::move(requestIterStates)}; - mSendQueue.push(std::move(message)); - } - } - else - { - // Add current iteration stats - appendCurrentRequestStats(std::move(currentRequestStats)); - } - } -} - -void Executor::Impl::appendCurrentDebugTensors() -{ - if (mDebugTensorsMaxIterations > 0) - { - std::scoped_lock lck(mDebugTensorsMtx); - if (mDebugTensors.size() >= mDebugTensorsMaxIterations) - { - mDebugTensors.pop_front(); - } - mDebugTensors.emplace_back(mModel->getCurrentDebugTensors()); - } -} - -void Executor::Impl::terminateCancelledRequests(RequestList& activeRequests) -{ - NVTX3_SCOPED_RANGE(terminateCancelledRequests); - auto const& worldConfig = mModel->getWorldConfig(); - auto const broadcastCancelledRequests = [this, &activeRequests, &worldConfig] - { - auto const& commSession = COMM_SESSION; - - if (worldConfig.isPipelineParallel()) - { - mCancelledRequestsWaitThread->waitStop(); - } - - if (commSession.getSize() > 1 && !activeRequests.empty()) - { - if (mIsPipelineLeader) - { - if (worldConfig.isPipelineParallel()) - { - auto const peer = worldConfig.getPipelineParallelism() - 1; - bool shouldExit = false; - mCommPipelineParallel->send( - &shouldExit, 1, mpi::MpiType::kBOOL, peer, mpi::MpiTag::kExecutorShouldExit); - auto pipelineCancelledReqIds - = CancelledRequestsAsyncSend::cancelledRequestsRecv(mCommPipelineParallel, peer); - mCancelledReqIds.insert(pipelineCancelledReqIds.begin(), pipelineCancelledReqIds.end()); - } - - auto numCancelledRequests = static_cast(mCancelledReqIds.size()); - if (worldConfig.isTensorParallel()) - { - mCommTensorParallel->bcastValue(numCancelledRequests, 0); - if (numCancelledRequests > 0) - { - std::vector cancelledReqIdsVec(mCancelledReqIds.begin(), mCancelledReqIds.end()); - mCommTensorParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - } - } - if (worldConfig.isContextParallel()) - { - mCommContextParallel->bcastValue(numCancelledRequests, 0); - if (numCancelledRequests > 0) - { - std::vector cancelledReqIdsVec(mCancelledReqIds.begin(), mCancelledReqIds.end()); - mCommContextParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - } - } - } - // If not leader - else - { - if (worldConfig.isFirstPipelineParallelRank()) - { - int64_t numCancelledRequests = 0; - mCommTensorParallel->bcastValue(numCancelledRequests, 0); - mCommContextParallel->bcastValue(numCancelledRequests, 0); - if (numCancelledRequests > 0) - { - std::vector cancelledReqIdsVec(numCancelledRequests); - mCommTensorParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - mCommContextParallel->bcast( - cancelledReqIdsVec.data(), cancelledReqIdsVec.size(), mpi::MpiType::kUINT64, 0); - mCancelledReqIds - = std::unordered_set(cancelledReqIdsVec.begin(), cancelledReqIdsVec.end()); - } - } - else - { - auto const peer = worldConfig.getPipelineParallelRank() - 1; - mCancelledReqIds = CancelledRequestsAsyncSend::cancelledRequestsRecv(mCommPipelineParallel, peer); - } - } - if (!worldConfig.isLastPipelineParallelRank()) - { - auto const peer = worldConfig.getPipelineParallelRank() + 1; - mCancelledRequestsAsyncSndHdl - = std::make_unique(mCommPipelineParallel, mCancelledReqIds, peer); - mCancelledRequestsWaitThread->notifyStart(); - } - } - }; - - std::unique_lock lck{mCancelReqMtx, std::defer_lock}; - if (!worldConfig.isPipelineParallel()) - { - lck.lock(); - } - - broadcastCancelledRequests(); - - if (!mCancelledReqIds.empty()) - { - // Loop over active requests and terminate those that have been cancelled - std::unordered_set terminatedReqIds; - for (auto& req : activeRequests) - { - auto reqId = req->isChild() ? req->getParentRequestId() : req->mRequestId; - if (mCancelledReqIds.find(reqId) != mCancelledReqIds.end()) - { - auto finishReason = req->isTimedOut() ? FinishReason::kTIMED_OUT : FinishReason::kCANCELLED; - mModel->terminateRequestSync(req, finishReason); - // Parent and child requests share the same request id. - // Mark it terminated first and remove from the set later. - terminatedReqIds.insert(reqId); - } - } - - for (auto const& reqId : terminatedReqIds) - { - mCancelledReqIds.erase(reqId); - } - } -} - -void Executor::Impl::terminateContextFinishedRequests(InTransList& inTransmissionRequests) -{ - NVTX3_SCOPED_RANGE(terminateContextFinishedRequests); - for (auto it = inTransmissionRequests.begin(); it != inTransmissionRequests.end();) - { - auto& item = *it; - auto req = item.request; - if (req->isDisaggContextCompleteState()) - { - // If pinnedBlockIds were tracked, unpin them. Otherwise, just terminate. - auto kvMgr = mModel->getKVCacheManager(); - if (kvMgr && !item.pinnedBlockIds.empty()) - { - kvMgr->unpinBlocksById(item.pinnedBlockIds); - } - else - { - mModel->terminateRequest(req); - } - it = inTransmissionRequests.erase(it); - } - else - { - ++it; - } - } -} - -void Executor::Impl::appendNewResponses(std::vector&& newResponses) -{ - { - std::scoped_lock lck(mResponsesMtx); - for (auto& response : newResponses) - { - mResponses[response.getRequestId()].emplace_back(std::move(response)); - } - } - mResponsesCv.notify_all(); -} - -Executor::Impl::RequestList Executor::Impl::populateNewResponses( - RequestList& activeRequests, InTransList& inTransmissionRequests, std::vector& newResponses) -{ - NVTX3_SCOPED_RANGE(populateNewResponses); - RequestList finishedRequests; - for (auto it = activeRequests.begin(); it != activeRequests.end();) - { - auto const& llmReq = (*it); - bool const requestDone = llmReq->isFinished(); - // Only leader should store responses - if (mIsLeader) - { - auto response = llmReq->createResponse(mModel->hasSpeculativeDecodingFastLogits(), mWorldRank); - if (response) - { - newResponses.emplace_back(std::move(response.value())); - } - } - // Remove from active requests if last response has been generated - if (requestDone) - { - // move the in transmission requests to another tracker - if (llmReq->isDisaggContextTransmissionState()) - { - std::vector pinnedBlockIds{}; - auto kvMgr = mModel->getKVCacheManager(); - if (kvMgr && kvMgr->isEnableBlockReuse() && !kvMgr->getBlockManager().isVariableWindow()) - { - pinnedBlockIds = kvMgr->storeBlocksForReuse(llmReq->mRequestId, llmReq, /*pinBlocks=*/true); - mModel->terminateRequest(llmReq); - } - inTransmissionRequests.push_back(InTransmissionItem{*it, pinnedBlockIds}); - } - finishedRequests.push_back(*it); - it = activeRequests.erase(it); - } - else - { - ++it; - } - } - return finishedRequests; -} - -void Executor::Impl::executionLoop() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - tensorrt_llm::common::setThreadName("executionLoop"); - - auto const& worldConfig = mModel->getWorldConfig(); - TLLM_CUDA_CHECK(cudaSetDevice(worldConfig.getDevice())); - - auto const [profileIterIdxs, stopIterIdxs] = tensorrt_llm::common::populateIterationIndexes( - kPROFILE_START_STOP_ENV_VAR_NAME, kLEGACY_PROFILE_START_STOP_ENV_VAR_NAME); - - SizeType32 numNewActiveRequests{0}; - std::chrono::time_point iterStart; - std::chrono::time_point iterEnd; - bool firstIteration{true}; - RequestList activeRequests; - InTransList inTransmissionRequests; - std::vector newResponses; - while (!mShutdown || !activeRequests.empty()) - { - double iterLatencyMS{0.0}; - double newActiveRequestsQueueLatencyMS{0.0}; - bool reportFinishedRequests = true; - RequestList finishedRequests; - if (!activeRequests.empty()) - { - finishTimedOutRequests(activeRequests); - terminateCancelledRequests(activeRequests); - forwardSync(activeRequests); - finishedRequests = populateNewResponses(activeRequests, inTransmissionRequests, newResponses); - cleanupDynamicLogitsPostProcessors(finishedRequests); - auto const iterCounter = mModel->getIterCounter(); - auto const stopIter = !stopIterIdxs.empty() && (stopIterIdxs.count(iterCounter - 1) > 0); - if (stopIter) - { - cudaProfilerStop(); - } - - // When there are no active or inflight requests, we need to update the stats before calling - // fetchNewRequests to make sure that the stats are reported accurately. - if (activeRequests.empty() && (!firstIteration)) - { - mModel->resetIterationStats(); - updateIterationStats(activeRequests, iterLatencyMS, numNewActiveRequests, - newActiveRequestsQueueLatencyMS, static_cast(finishedRequests.size()), true); - updateRequestStats(activeRequests, finishedRequests, true); - reportFinishedRequests = false; - } - if (!newResponses.empty()) - { - enqueueNewResponses(std::move(newResponses)); - newResponses.clear(); - } - iterEnd = std::chrono::steady_clock::now(); - iterLatencyMS = std::chrono::duration(iterEnd - iterStart).count(); - } - - if (!inTransmissionRequests.empty()) - { - terminateContextFinishedRequests(inTransmissionRequests); - } - - if (!mShutdown) - { - auto const iterCounter = mModel->getIterCounter(); - auto const profileIter = !profileIterIdxs.empty() && (profileIterIdxs.count(iterCounter) > 0); - if (profileIter) - { - cudaProfilerStart(); - } - iterStart = std::chrono::steady_clock::now(); - std::optional lowestPriority = std::nullopt; - if (!activeRequests.empty()) - { - lowestPriority = activeRequests.back()->priority(); - } - - auto [newRequests, newActiveRequestsQueueLatency] - = fetchNewRequests(static_cast(activeRequests.size()), lowestPriority); - newActiveRequestsQueueLatencyMS = newActiveRequestsQueueLatency; - numNewActiveRequests = newRequests.size(); - - if (firstIteration) - { - firstIteration = false; - } - - for (auto const& newRequest : newRequests) - { - insertRequestInOrder(activeRequests, newRequest); - } - - // Update dynamic tuning stats - if (mDynamicBatchTuner) - { - for (auto const& req : activeRequests) - { - auto const inputLength = req->mPromptLen; - auto const outputLength = req->mMaxNewTokens; - mDynamicBatchTuner->updateStats(inputLength, outputLength); - } - } - } - if (!activeRequests.empty()) - { - forwardAsync(activeRequests); - updateIterationStats(activeRequests, iterLatencyMS, numNewActiveRequests, newActiveRequestsQueueLatencyMS, - static_cast(finishedRequests.size()), false); - // Finished requests were reported once. Avoid reporting it twice. - if (reportFinishedRequests) - { - updateRequestStats(activeRequests, finishedRequests, false); - } - else - { - updateRequestStats(activeRequests, {}, false); - } - appendCurrentDebugTensors(); - } - } - - if (mCancelledRequestsWaitThread) - { - mCancelledRequestsWaitThread.reset(nullptr); - } - if (mRequestWithIdWaitThread) - { - mRequestWithIdWaitThread.reset(nullptr); - } - if (worldConfig.isPipelineParallel() && mIsPipelineLeader) - { - auto const peer = worldConfig.getPipelineParallelism() - 1; - int64_t numActiveRequests = -1; - mCommPipelineParallel->send( - &numActiveRequests, 1, mpi::MpiType::kINT64, peer, mpi::MpiTag::kExecutorNumActiveRequests); - bool shouldExit = true; - mCommPipelineParallel->send(&shouldExit, 1, mpi::MpiType::kBOOL, peer, mpi::MpiTag::kExecutorShouldExit); - } - if (mRequestWithIdLeaderThread) - { - mRequestWithIdLeaderThread->join(); - mRequestWithIdLeaderThread.reset(nullptr); - } - if (mCancelledRequestsLeaderThread) - { - mCancelledRequestsLeaderThread->join(); - mCancelledRequestsLeaderThread.reset(nullptr); - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void Executor::Impl::enqueueTerminateRequest() -{ - { - std::scoped_lock lck(mQueuedReqMtx); - Request dummyReq({1}, 1); - RequestWithId reqWithId{std::move(dummyReq), kTerminateReqId}; - mQueuedRequests.emplace_back(reqWithId); - } - mQueuedReqCv.notify_one(); -} - -void Executor::Impl::enqueueNewResponses(std::vector&& newResponses) -{ - TLLM_CHECK_WITH_INFO(mIsLeader, "Only leader should store responses"); - - if (mCommMode == CommunicationMode::kLEADER) - { - appendNewResponses(std::move(newResponses)); - } - else if (mCommMode == CommunicationMode::kORCHESTRATOR) - { - MpiMessage message(MpiId::RESPONSE); - message.data = ResponseData{std::move(newResponses)}; - mSendQueue.push(std::move(message)); - } -} - -// Orchestrator thread sending new requests to leader of the model -void Executor::Impl::orchSendReqThread() -{ - tensorrt_llm::common::setThreadName("orchSendReq"); - - while (true) - { - auto message = mSendQueue.pop(); - - if (message.id == MpiId::TERMINATION) - { - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorId); - TLLM_LOG_INFO("Orchestrator sendReq thread exiting"); - break; - } - if (message.id == MpiId::PENDING_REQUEST) - { - auto& reqWithIds = std::get(message.data); - auto packed = RequestWithId::serializeReqWithIds(reqWithIds.requests); - - TLLM_LOG_DEBUG("Orchestrator sendReq thread sending %d pending requests", reqWithIds.requests.size()); - // Temporary WAR to indicate to client that we cannot send the serialized request - // because it exceeds int32_t size limit. - // TODO: Should fix as part of https://jirasw.nvidia.com/browse/TRTLLM-708 - if (packed.size() > std::numeric_limits::max()) - { - for (auto const& reqWithId : reqWithIds.requests) - { - { - std::scoped_lock lck(mResponsesMtx); - mResponses[reqWithId.id].emplace_back(reqWithId.id, - "Request is too large, or you are enqueuing too many requests at once " - "to be sent via MPI_Send, please try to enqueue the request(s) again. " - "This issue will be resolved in a future version of TRT-LLM."); - } - mResponsesCv.notify_all(); - } - } - else - { - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorId); - mOrchLeaderComm->send( - packed.data(), packed.size(), mpi::MpiType::kCHAR, mLeaderRank, mpi::MpiTag::kOrchestratorData); - } - } - else if (message.id == MpiId::CANCEL_REQUEST) - { - auto& data = std::get(message.data); - - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorId); - mOrchLeaderComm->send( - data.ids.data(), data.ids.size(), mpi::MpiType::kUINT64, mLeaderRank, mpi::MpiTag::kOrchestratorData); - } - else - { - TLLM_THROW("Invalid message id"); - } - } -} - -// Leader thread receiving new requests from orchestrator -void Executor::Impl::leaderRecvReqThread() -{ - tensorrt_llm::common::setThreadName("leaderRecvReq"); - TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); -#if ENABLE_MULTI_DEVICE - auto& selCancelledReqIds = mUsePipelineParallel ? mPipelineCancelledReqIds : mCancelledReqIds; - while (true) - { - if (mRecvPollPeriodMs > 0) - { - mOrchLeaderComm->recvPoll(mOrchRank, mpi::MpiTag::kOrchestratorId, mRecvPollPeriodMs); - } - - // Blocking is okay: terminate message is expected to arrive here - MPI_Message msg = nullptr; - MPI_Status status; - mOrchLeaderComm->mprobe(mOrchRank, mpi::MpiTag::kOrchestratorId, &msg, &status); - - int32_t count = 0; - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); // NOLINT - TLLM_CHECK(count == 1); - - MpiId mpiId{}; - MPICHECK(MPI_Mrecv(&mpiId, count, MPI_UINT64_T, &msg, &status)); // NOLINT - - // EXIT condition from receiving TERMINATE msg - if (mpiId == MpiId::TERMINATION) - { - // Enqueue a request to indicate to other ranks to terminate - enqueueTerminateRequest(); - - // Send message to orchestrator to indicate to terminate orch recv thread - mSendQueue.push(MpiMessage(mpiId)); - TLLM_LOG_INFO("Leader recvReq thread exiting"); - break; - } - if (mpiId == MpiId::PENDING_REQUEST) - { - mOrchLeaderComm->mprobe(mOrchRank, mpi::MpiTag::kOrchestratorData, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_CHAR, &count)); // NOLINT - std::vector buffer(count); - MPICHECK(MPI_Mrecv(buffer.data(), count, MPI_CHAR, &msg, &status)); // NOLINT - - auto requestWithIds = RequestWithId::deserializeReqWithIds(buffer); - TLLM_LOG_DEBUG("Leader recvReq thread receiving %d pending requests", requestWithIds.size()); - { - std::scoped_lock lck(mQueuedReqMtx); - if (mMaxQueueSize) - { - auto const maxQueueSize = mMaxQueueSize.value(); - if (maxQueueSize > 0 && mQueuedRequests.size() >= static_cast(maxQueueSize)) - { - auto err = tensorrt_llm::common::fmtstr( - "Maximum queue size of %d has been reached, please try again later", maxQueueSize); - TLLM_LOG_ERROR("%s", err.c_str()); - std::vector responses; - responses.reserve(requestWithIds.size()); - for (auto const& reqWithId : requestWithIds) - { - responses.emplace_back(reqWithId.id, err); - } - enqueueNewResponses(std::move(responses)); - continue; - } - } - for (auto&& req : requestWithIds) - { - req.queuedStart = std::chrono::steady_clock::now(); - insertRequestInOrder(mQueuedRequests, std::move(req)); - } - } - mQueuedReqCv.notify_one(); - } - else if (mpiId == MpiId::CANCEL_REQUEST) - { - // Prepare receiving data - mOrchLeaderComm->mprobe(mOrchRank, mpi::MpiTag::kOrchestratorData, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); // NOLINT - std::vector cancelledReqIds(count); - MPICHECK(MPI_Mrecv(cancelledReqIds.data(), count, MPI_UINT64_T, &msg, &status)); // NOLINT - - std::scoped_lock lck(mCancelReqMtx); - selCancelledReqIds.insert(cancelledReqIds.begin(), cancelledReqIds.end()); - } - else - { - TLLM_THROW("Invalid message id"); - } - } -#endif // ENABLE_MULTI_DEVICE -} - -// Leader thread sending responses to orchestrator -void Executor::Impl::leaderSendThread(MpiMessageQueue& sendQueue, mpi::MpiTag idTag, mpi::MpiTag dataTag) -{ - tensorrt_llm::common::setThreadName("leaderSend"); - TLLM_CUDA_CHECK(cudaSetDevice(mDeviceId)); - -#if ENABLE_MULTI_DEVICE - while (true) - { - auto message = sendQueue.pop(); - - if (message.id == MpiId::TERMINATION) - { - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mOrchRank, idTag); - TLLM_LOG_INFO("Leader sendThread exiting"); - break; - } - if (message.id == MpiId::RESPONSE || message.id == MpiId::ITER_STATS - || message.id == MpiId ::REQUEST_ITER_STATS) - { - std::vector buffer; - if (message.id == MpiId::RESPONSE) - { - auto& responseData = std::get(message.data); - TLLM_LOG_DEBUG("Leader sendResp thread sending %d responses", responseData.responses.size()); - buffer = Serialization::serialize(responseData.responses); - } - else if (message.id == MpiId::ITER_STATS) - { - auto& iterStatsData = std::get(message.data); - TLLM_LOG_DEBUG("Leader sendResp thread sending iter stats"); - buffer = Serialization::serialize(iterStatsData.iterStatsVec); - } - else if (message.id == MpiId::REQUEST_ITER_STATS) - { - auto& requestIterStatsData = std::get(message.data); - TLLM_LOG_DEBUG("Leader sendResp thread sending iter request stats"); - buffer = Serialization::serialize(requestIterStatsData.requestStatsPerIterationVec); - } - mOrchLeaderComm->send(&message.id, 1, mpi::MpiType::kUINT64, mOrchRank, idTag); - mOrchLeaderComm->send(buffer.data(), buffer.size(), mpi::MpiType::kCHAR, mOrchRank, dataTag); - } - else - { - TLLM_THROW("Invalid message id"); - } - } -#endif // ENABLE_MULTI_DEVICE -} - -void Executor::Impl::orchRecvThread(mpi::MpiTag idTag, mpi::MpiTag dataTag) -{ - tensorrt_llm::common::setThreadName("orchRecv"); - -#if ENABLE_MULTI_DEVICE - while (true) - { - if (mRecvPollPeriodMs > 0) - { - mOrchLeaderComm->recvPoll(mOrchRank, mpi::MpiTag::kOrchestratorId, mRecvPollPeriodMs); - } - - MPI_Message msg = nullptr; - MPI_Status status; - mOrchLeaderComm->mprobe(mLeaderRank, idTag, &msg, &status); - - int32_t count = 0; - MPICHECK(MPI_Get_count(&status, MPI_UINT64_T, &count)); // NOLINT - TLLM_CHECK(count == 1); - - MpiId mpiId{}; - MPICHECK(MPI_Mrecv(&mpiId, count, MPI_UINT64_T, &msg, &status)); // NOLINT - - if (mpiId == MpiId::TERMINATION) - { - TLLM_LOG_INFO("Orchestrator recv thread exiting"); - break; - } - if (mpiId == MpiId::RESPONSE || mpiId == MpiId::ITER_STATS || mpiId == MpiId::REQUEST_ITER_STATS) - { - mOrchLeaderComm->mprobe(mLeaderRank, dataTag, &msg, &status); - MPICHECK(MPI_Get_count(&status, MPI_CHAR, &count)); // NOLINT - - std::vector buffer(count); - MPICHECK(MPI_Mrecv(buffer.data(), count, MPI_CHAR, &msg, &status)); // NOLINT - - if (mpiId == MpiId::RESPONSE) - { - auto newResponses = Serialization::deserializeResponses(buffer); - TLLM_LOG_DEBUG("Orchestrator recv thread receiving %d responses", newResponses.size()); - appendNewResponses(std::move(newResponses)); - } - else if (mpiId == MpiId::ITER_STATS) - { - appendMultipleIterStats(Serialization::deserializeIterationStatsVec(buffer)); - } - else if (mpiId == MpiId::REQUEST_ITER_STATS) - { - appendMultipleRequestStats(Serialization::deserializeRequestStatsPerIterationVec(buffer)); - } - } - else - { - TLLM_THROW("Invalid message id"); - } - } -#endif // ENABLE_MULTI_DEVICE -} - -Executor::Impl::LlmRequestLogitsPostProcessor Executor::Impl::getLogitsPostProcessor(std::string const& name) -{ - auto const postProcIt = mLogitsPostProcessorMap.find(name); - TLLM_CHECK_WITH_INFO( - postProcIt != mLogitsPostProcessorMap.end(), "LogitsPostProcessor %s not found.", name.c_str()); - auto executorLogitsPostProcessor = postProcIt->second; - return [executorLogitsPostProcessor](IdType reqId, RtTensorPtr& logits, BeamTokens const& beamTokens, - CudaStreamPtr const& cudaStreamPtr, std::optional clientId) - { - auto logitsTensor = executor::detail::ofITensor(logits); - executorLogitsPostProcessor(reqId, logitsTensor, beamTokens, cudaStreamPtr, clientId); - }; -} - -void Executor::Impl::setupDynamicLogitsPostProcessors(std::vector& newReqWithIds) -{ - for (auto& reqWithId : newReqWithIds) - { - auto logitsPostProcessor = reqWithId.req.getLogitsPostProcessor(); - if (logitsPostProcessor) - { - std::string const name = Request::kDynamicPostProcessorNamePrefix + std::to_string(reqWithId.id); - mLogitsPostProcessorMap[name] = logitsPostProcessor.value(); - reqWithId.req.setLogitsPostProcessor(std::nullopt); - reqWithId.req.setLogitsPostProcessorName(name); - } - } -} - -void Executor::Impl::cleanupDynamicLogitsPostProcessors(RequestList const& finishedRequests) -{ - for (auto& req : finishedRequests) - { - std::string const name = Request::kDynamicPostProcessorNamePrefix + std::to_string(req->mRequestId); - auto const postProcIt = mLogitsPostProcessorMap.find(name); - if (postProcIt != mLogitsPostProcessorMap.end()) - { - mLogitsPostProcessorMap.erase(name); - } - } -} - -void Executor::Impl::addTerminatedReqId(std::vector const& responses, IdType const& reqId) -{ - for (auto const& response : responses) - { - if (response.hasError() || (!response.hasError() && response.getResult().isFinal)) - { - mTerminatedReqIds.insert(reqId); - if (mChildReqIdsMap.find(reqId) != mChildReqIdsMap.end()) - { - for (auto childReqId : mChildReqIdsMap.at(reqId)) - { - mTerminatedReqIds.insert(childReqId); - } - mChildReqIdsMap.erase(reqId); - } - } - } -} - -void Executor::Impl::checkParallelApiUsage(std::string const& methodName) const -{ - // If leader mode, and not leader, throw error - if (mCommMode == CommunicationMode::kLEADER && !mIsLeader) - { - // Non-leader are not expected to call cancelRequest - TLLM_THROW("With LEADER communication mode, only leader rank is expected to call %s", methodName.c_str()); - } - if (mCommMode == CommunicationMode::kORCHESTRATOR && !mIsOrchestrator) - { - TLLM_THROW( - "With ORCHESTRATOR communication mode, only orchestrator rank is expected to call %s", methodName.c_str()); - } -} - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/executorImpl.h b/cpp/tensorrt_llm/executor/executorImpl.h deleted file mode 100644 index 6e545dbf6def..000000000000 --- a/cpp/tensorrt_llm/executor/executorImpl.h +++ /dev/null @@ -1,385 +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/common/arrayView.h" -#include "tensorrt_llm/executor/dynamicBatchTuner.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/intervalSet.h" -#include "tensorrt_llm/executor/model.h" -#include "tensorrt_llm/executor/orchestratorUtils.h" -#include "tensorrt_llm/executor/requestWithId.h" -#include "tensorrt_llm/executor/types.h" -#include "tensorrt_llm/runtime/gptJsonConfig.h" -#include "tensorrt_llm/runtime/modelConfig.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tensorrt_llm/runtime/worldConfig.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::executor -{ - -class RequestWithIdAsyncSend; -class CancelledRequestsAsyncSend; - -class MpiMessageQueue -{ -public: - void push(MpiMessage&& message) - { - std::lock_guard const lock(mMutex); - mQueue.push(std::move(message)); - mCv.notify_one(); - } - - MpiMessage pop() - { - std::unique_lock lock(mMutex); - mCv.wait(lock, [this] { return !mQueue.empty(); }); - MpiMessage message = std::move(mQueue.front()); - mQueue.pop(); - return message; - } - -private: - std::queue mQueue; - std::mutex mMutex; - std::condition_variable mCv; -}; - -class Executor::Impl - -{ - using LlmRequestPtr = std::shared_ptr; - using RequestList = std::list; - - // When block reuse is enabled for context worker for disaggregated serving, - // we need to store the pinned block ids so that we can unpin them when - // the request is finished. - struct InTransmissionItem - { - LlmRequestPtr request; - std::vector pinnedBlockIds; - }; - - using InTransList = std::list; - -public: - Impl(std::filesystem::path const& modelPath, std::optional const& encoderModelPath, - [[maybe_unused]] ModelType modelType, ExecutorConfig const& executorConfig); - - Impl(BufferView const& engineBufferView, std::string const& jsonConfigStr, - std::optional const& encoderEngineBufferView, - std::optional const& encoderJsonConfigStr, [[maybe_unused]] ModelType modelType, - ExecutorConfig const& executorConfig, std::optional> const& managedWeightsOpt); - - Impl(std::shared_ptr model, std::optional> encoderModel, - ExecutorConfig const& executorConfig); - - ~Impl(); - - Impl(Impl const& executor) = delete; - Impl& operator=(Impl const& executor) = delete; - Impl(Impl&&) = delete; - Impl& operator=(Impl&&) = delete; - - IdType enqueueRequest(Request const& request); - - std::vector enqueueRequests(std::vector const& requests); - - std::vector enqueueRequests(common::ArrayView const& requests); - - std::vector awaitResponses(std::optional const& timeout = std::nullopt); - - std::vector awaitResponses( - IdType const& reqId, std::optional const& optTimeout = std::nullopt); - - std::vector> awaitResponses( - std::vector const& requestIds, std::optional const& timeout); - - SizeType32 getNumResponsesReady(std::optional const& optId = std::nullopt) const; - - void cancelRequest(IdType requestId); - - void shutdown(); - - std::deque getLatestIterationStats(); - std::deque getLatestRequestStats(); - std::deque getLatestDebugTensors(); - - bool canEnqueueRequests() const; - - bool isParticipant() const; - - std::optional> getKVCacheEventManager() const; - -private: - using RtTensorPtr = runtime::ITensor::SharedPtr; - using CudaStreamPtr = runtime::BufferManager::CudaStreamPtr; - using LlmRequestLogitsPostProcessor - = std::function)>; - - void initialize(ExecutorConfig const& executorConfig); - - void loadModel(std::optional const& modelPath, std::optional const& engineBuffer, - runtime::GptJsonConfig const& jsonConfig, ExecutorConfig const& executorConfig, bool isEncoder, - std::optional> const& managedWeightsOpt); - - std::shared_ptr createModel(runtime::RawEngine const& rawEngine, runtime::ModelConfig const& modelConfig, - runtime::WorldConfig const& worldConfig, ExecutorConfig const& executorConfig); - - std::shared_ptr createEncoderModel(runtime::RawEngine const& rawEngine, - runtime::ModelConfig const& modelConfig, runtime::WorldConfig const& worldConfig, - ExecutorConfig const& executorConfig); - - void setOrchLeaderComm(SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig const& parallelConfig); - - void initializeCommAndWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, ExecutorConfig const& executorConfig, - std::optional modelType = std::nullopt, - std::optional const& modelPath = std::nullopt, - std::optional const& worldConfig = std::nullopt, - std::optional const& decoderGptJsonConfig = std::nullopt); - - static void validateParallelConfig(ParallelConfig const& parallelConfig, std::optional modelType, - std::optional const& modelPath); - - void initializeOrchestrator(SizeType32 tp, SizeType32 pp, SizeType32 cp, ExecutorConfig const& executorConfig, - ParallelConfig parallelConfig, ModelType modelType, std::filesystem::path const& modelPath); - - void initializeWorkers(SizeType32 tp, SizeType32 pp, SizeType32 cp, ParallelConfig& parallelConfig, - std::optional const& worldConfig = std::nullopt, - std::optional const& decoderGptJsonConfig = std::nullopt); - - void initializeLogitsPostProcessorBatched(LogitsPostProcessorConfig const& logitsProcConfig); - - IdType generateReqId(Request const& request) - { - // If the request has a disaggregated request id, prefer it. - if (request.getDisaggRequestId().has_value() && request.getDisaggRequestId().value() > kMaxLocalReqId) - { - return request.getDisaggRequestId().value(); - } - // Otherwise, generate a local request id in range [1, kMaxLocalReqId). - return generateLocalReqId(); - } - - IdType generateLocalReqId() - { - return (mLastReqId++ % kMaxLocalReqId); - } - - std::vector getLeaderNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive); - std::vector getNewReqWithIds( - SizeType32 numActiveRequests, std::optional lowestPriorityActive); - - std::tuple fetchNewRequests( - SizeType32 numActiveRequests, std::optional lowestPriorityActive); - - void forwardSync(RequestList& activeRequests); - - void forwardAsync(RequestList& activeRequests); - - void prepRequestsForEncoderSkip(RequestList& activeRequests); - - void terminateActiveRequests(RequestList& activeRequests, std::string const& err); - - IterationStats getCurrentIterationStats(RequestList const& activeRequests, double iterLatencyMS, - SizeType32 numNewActiveRequests, double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests); - - void appendCurrentIterStats(IterationStats&& currentIterStats); - void appendMultipleIterStats(std::vector&& currentIterStatsVec); - void updateIterationStats(RequestList const& activeRequests, double iterLatencyMS, SizeType32 numNewActiveRequests, - double newActiveRequestsQueueLatencyMS, SizeType32 numCompletedRequests, bool flushToOrchestrator); - void appendCurrentRequestStats(RequestStatsPerIteration&& currentRequestStats); - void appendMultipleRequestStats(std::vector&& currentRequestStatsVec); - RequestStatsPerIteration getCurrentRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests); - void updateRequestStats( - RequestList const& activeRequests, RequestList const& finishedRequests, bool flushToOrchestrator); - - void appendCurrentDebugTensors(); - - void terminateCancelledRequests(RequestList& activeRequests); - - void terminateContextFinishedRequests(InTransList& inTransmissionRequests); - - void appendNewResponses(std::vector&& newResponses); - - /// @brief Populates new responses from active requests. - /// Active requests that have completed are erased from activeRequests - /// and returned for bookkeeping. - /// @return A list of requests that have completed. - RequestList populateNewResponses( - RequestList& activeRequests, InTransList& inTransmissionRequests, std::vector& newResponses); - - void executionLoop(); - - void enqueueTerminateRequest(); - void enqueueNewResponses(std::vector&& newResponses); - - LlmRequestLogitsPostProcessor getLogitsPostProcessor(std::string const& name); - void setupDynamicLogitsPostProcessors(std::vector& newReqWithIds); - void cleanupDynamicLogitsPostProcessors(RequestList const& finishedRequests); - - void orchSendReqThread(); - void orchRecvThread(mpi::MpiTag idTag, mpi::MpiTag dataTag); - void leaderRecvReqThread(); - void leaderSendThread(MpiMessageQueue& sendQueue, mpi::MpiTag idTag, mpi::MpiTag dataTag); - - void addTerminatedReqId(std::vector const& responses, IdType const& reqId); - - // Check that the current process is the leader or orchestrator - void checkParallelApiUsage(std::string const& methodName) const; - - // These functions wait for MPI async sends on separate threads - void requestWithIdWaitThread(); - void cancelledRequestsWaitThread(); - // These functions send data from leader to pipeline leader on separate threads - void requestWithIdLeaderThread(); - void cancelledRequestsLeaderThread(); - - /// @brief mark requests that have timed out before ever being executed as finished. - /// uses cancellation based on communication mode. - /// - /// @param activeRequests [in] List of active requests to check for timeouts - void finishTimedOutRequests(RequestList const& activeRequests); - - // The model to execute - std::shared_ptr mModel = nullptr; - std::shared_ptr mEncoderModel = nullptr; - - // The maximum number of activeRequests - SizeType32 mMaxNumActiveRequests; - - // Thread the executes the main loop - std::thread mExecutionThread; - - // Atomic that indicates threads should shutdown - std::atomic mShutdown; - - // Atomic that indicates if shutdown method has been called - std::atomic mShutdownCalled = false; - - // Queued requests - std::mutex mQueuedReqMtx; - std::condition_variable mQueuedReqCv; - std::deque mQueuedRequests; - std::optional mMaxQueueSize; - - // Cancelled requests - std::mutex mCancelReqMtx; - std::unordered_set mCancelledReqIds; - std::unordered_set mPipelineCancelledReqIds; - - // Ready responses - std::unordered_map> mResponses; - mutable std::mutex mResponsesMtx; - std::condition_variable mResponsesCv; - - // Since the request IDs are generated sequentially, IntervalSet is preferred over unordered_set for its efficient - // memory usage to stores request ID intervals rather than individual request ID numbers. - IntervalSet mTerminatedReqIds; - - std::unordered_map> mChildReqIdsMap; - - // Iteration stats - IterationType mIterStatsMaxIterations; - std::mutex mIterStatsMtx; - std::deque mIterationStats; - - // Request stats - IterationType mRequestStatsMaxIterations; - std::mutex mRequestStatsMtx; - std::deque mRequestStats; - - // Debug - IterationType mDebugTensorsMaxIterations; - std::mutex mDebugTensorsMtx; - std::deque mDebugTensors; - - IdType mLastReqId = 1; - - static constexpr IdType kTerminateReqId = 0; - // Request id > kMaxLocalReqId is reserved for disaggregated requests. - // This max ID is also in Python side. - static constexpr IdType kMaxLocalReqId = 1ULL << 42U; - - BatchingType mBatchingType; - bool mIsSchedulerMaxUtilization; - bool mIsSchedulerGuaranteedNoEvict; - bool mIsChunkedContext; - bool mPromptTableOffloading; - - CommunicationMode mCommMode; - bool mIsWorker = false; - bool mIsLeader = false; - bool mIsPipelineLeader = false; - bool mUsePipelineParallel = false; - - std::unordered_map mLogitsPostProcessorMap; - std::optional mLogitsPostProcessorBatched; - - bool mIsOrchestrator = false; - std::shared_ptr mOrchLeaderComm; - - std::thread mOrchSendReqThread; - std::thread mOrchRecvThread; - std::thread mLeaderRecvReqThread; - std::thread mLeaderSendThread; - - int32_t mRecvPollPeriodMs = 0; - - int32_t mLeaderRank = -1; - int32_t mOrchRank = 0; - int32_t mWorldRank = -1; - int32_t mDeviceId = 0; - - MpiMessageQueue mSendQueue; - - std::shared_ptr mCommTensorParallel; - std::shared_ptr mCommPipelineParallel; - std::shared_ptr mCommContextParallel; - std::unique_ptr mRequestWithIdAsyncSndHdl; - std::unique_ptr mCancelledRequestsAsyncSndHdl; - std::unique_ptr mRequestWithIdLeaderThread; - std::unique_ptr mCancelledRequestsLeaderThread; - std::unique_ptr mRequestWithIdWaitThread; - std::unique_ptr mCancelledRequestsWaitThread; - - // for validating requests - bool mEnableBlockReuse; - - inline static std::string const kPROFILE_START_STOP_ENV_VAR_NAME = "TLLM_PROFILE_START_STOP"; - inline static std::string const kLEGACY_PROFILE_START_STOP_ENV_VAR_NAME = "TLLM_GPTM_PROFILE_START_STOP"; - - std::shared_ptr mDynamicBatchTuner; -}; - -} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/kvCacheEvent.cpp b/cpp/tensorrt_llm/executor/kvCacheEvent.cpp new file mode 100644 index 000000000000..5158eb93c983 --- /dev/null +++ b/cpp/tensorrt_llm/executor/kvCacheEvent.cpp @@ -0,0 +1,44 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2025-2026 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. + */ + +// Definition of the (backend-agnostic) KVCacheEvent constructor. It previously +// lived in executor.cpp alongside the legacy TensorRT-engine Executor class +// implementation; that file was removed with the TensorRT backend, so the +// constructor is relocated here so the retained KV-cache event path +// (kvCacheEventManager, serialization, nanobind) continues to link. + +#include +#include + +namespace tensorrt_llm::executor +{ + +char const* version() noexcept +{ + return kTensorRtLlmVersion; +} + +KVCacheEvent::KVCacheEvent( + size_t eventId, KVCacheEventData data, SizeType32 windowSize, std::optional attentionDpRank) + : eventId{eventId} + , data{std::move(data)} + , windowSize{windowSize} + , attentionDpRank{attentionDpRank} +{ +} + +} // namespace tensorrt_llm::executor diff --git a/cpp/tensorrt_llm/executor/model.h b/cpp/tensorrt_llm/executor/model.h index 52fedf1d1113..03fca2007374 100644 --- a/cpp/tensorrt_llm/executor/model.h +++ b/cpp/tensorrt_llm/executor/model.h @@ -87,7 +87,7 @@ class Model [[nodiscard]] virtual SizeType32 getMaxDraftLen() const = 0; [[nodiscard]] virtual SizeType32 getNumMicroBatches() const = 0; [[nodiscard]] virtual SizeType32 getOperatingBeamWidth() const = 0; - [[nodiscard]] virtual nvinfer1::DataType getLogitDataType() const = 0; + [[nodiscard]] virtual tensorrt_llm::DataType getLogitDataType() const = 0; [[nodiscard]] virtual runtime::WorldConfig const& getWorldConfig() const = 0; [[nodiscard]] virtual runtime::ModelConfig const& getModelConfig() const = 0; [[nodiscard]] virtual runtime::BufferManager const& getBufferManager() const = 0; @@ -95,8 +95,8 @@ class Model [[nodiscard]] virtual IterationType getIterCounter() const noexcept = 0; [[nodiscard]] virtual bool hasSpeculativeDecodingFastLogits() const noexcept = 0; [[nodiscard]] virtual bool getGatherGenerationLogits() const = 0; - [[nodiscard]] virtual nvinfer1::DataType getTensorDataType(std::string const& name) const = 0; - [[nodiscard]] virtual nvinfer1::Dims getTensorShape(std::string const& name) const = 0; + [[nodiscard]] virtual tensorrt_llm::DataType getTensorDataType(std::string const& name) const = 0; + [[nodiscard]] virtual tensorrt_llm::Dims getTensorShape(std::string const& name) const = 0; /// @brief Function that provides per iteration stats specific to a certain model /// @param stats The json object to write stats to diff --git a/cpp/tensorrt_llm/executor/serialization.cpp b/cpp/tensorrt_llm/executor/serialization.cpp index ce081e10c603..abe7bd2a5d5f 100644 --- a/cpp/tensorrt_llm/executor/serialization.cpp +++ b/cpp/tensorrt_llm/executor/serialization.cpp @@ -626,8 +626,8 @@ kv_cache::CacheState Serialization::deserializeCacheState(std::istream& is) auto hasRnnConfig = su::deserialize(is); std::optional rnnModelConfig; std::vector rnnLayerNumPerPP; - nvinfer1::DataType convStateDataType{nvinfer1::DataType::kFLOAT}; - nvinfer1::DataType ssmStateDataType{nvinfer1::DataType::kFLOAT}; + tensorrt_llm::DataType convStateDataType{tensorrt_llm::DataType::kFLOAT}; + tensorrt_llm::DataType ssmStateDataType{tensorrt_llm::DataType::kFLOAT}; if (hasRnnConfig) { CacheState::RnnModelConfig rnnCfg; @@ -641,8 +641,8 @@ kv_cache::CacheState Serialization::deserializeCacheState(std::istream& is) rnnCfg.mNumHeads = su::deserialize(is); rnnCfg.mConvSectionLayout = static_cast(su::deserialize(is)); - convStateDataType = su::deserialize(is); - ssmStateDataType = su::deserialize(is); + convStateDataType = su::deserialize(is); + ssmStateDataType = su::deserialize(is); rnnLayerNumPerPP = su::deserialize>(is); rnnModelConfig = std::move(rnnCfg); } diff --git a/cpp/tensorrt_llm/executor/tensor.cpp b/cpp/tensorrt_llm/executor/tensor.cpp index c38feb0e34b8..4c495b19f56f 100644 --- a/cpp/tensorrt_llm/executor/tensor.cpp +++ b/cpp/tensorrt_llm/executor/tensor.cpp @@ -53,17 +53,17 @@ DataType Tensor::getDataType() const } switch (mTensor->getDataType()) { - case nvinfer1::DataType::kBOOL: return DataType::kBOOL; - case nvinfer1::DataType::kINT8: return DataType::kINT8; - case nvinfer1::DataType::kINT32: return DataType::kINT32; - case nvinfer1::DataType::kUINT8: return DataType::kUINT8; - case nvinfer1::DataType::kFP8: return DataType::kFP8; - case nvinfer1::DataType::kHALF: return DataType::kFP16; - case nvinfer1::DataType::kFLOAT: return DataType::kFP32; - case nvinfer1::DataType::kBF16: return DataType::kBF16; - case nvinfer1::DataType::kINT64: return DataType::kINT64; - case nvinfer1::DataType::kINT4: [[fallthrough]] /* do nothing */; - case nvinfer1::DataType::kFP4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kBOOL: return DataType::kBOOL; + case tensorrt_llm::DataType::kINT8: return DataType::kINT8; + case tensorrt_llm::DataType::kINT32: return DataType::kINT32; + case tensorrt_llm::DataType::kUINT8: return DataType::kUINT8; + case tensorrt_llm::DataType::kFP8: return DataType::kFP8; + case tensorrt_llm::DataType::kHALF: return DataType::kFP16; + case tensorrt_llm::DataType::kFLOAT: return DataType::kFP32; + case tensorrt_llm::DataType::kBF16: return DataType::kBF16; + case tensorrt_llm::DataType::kINT64: return DataType::kINT64; + case tensorrt_llm::DataType::kINT4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kFP4: [[fallthrough]] /* do nothing */; default: TLLM_THROW("Unsupported data type"); } } @@ -135,19 +135,19 @@ tr::ITensor::Shape toDims(Shape const& shape) return dims; } -nvinfer1::DataType toDataType(DataType dataType) +tensorrt_llm::DataType toDataType(DataType dataType) { switch (dataType) { - case DataType::kBOOL: return nvinfer1::DataType::kBOOL; - case DataType::kUINT8: return nvinfer1::DataType::kUINT8; - case DataType::kINT8: return nvinfer1::DataType::kINT8; - case DataType::kINT32: return nvinfer1::DataType::kINT32; - case DataType::kINT64: return nvinfer1::DataType::kINT64; - case DataType::kBF16: return nvinfer1::DataType::kBF16; - case DataType::kFP8: return nvinfer1::DataType::kFP8; - case DataType::kFP16: return nvinfer1::DataType::kHALF; - case DataType::kFP32: return nvinfer1::DataType::kFLOAT; + case DataType::kBOOL: return tensorrt_llm::DataType::kBOOL; + case DataType::kUINT8: return tensorrt_llm::DataType::kUINT8; + case DataType::kINT8: return tensorrt_llm::DataType::kINT8; + case DataType::kINT32: return tensorrt_llm::DataType::kINT32; + case DataType::kINT64: return tensorrt_llm::DataType::kINT64; + case DataType::kBF16: return tensorrt_llm::DataType::kBF16; + case DataType::kFP8: return tensorrt_llm::DataType::kFP8; + case DataType::kFP16: return tensorrt_llm::DataType::kHALF; + case DataType::kFP32: return tensorrt_llm::DataType::kFLOAT; case DataType::kUNKNOWN: TLLM_THROW("Unsupported data type"); } diff --git a/cpp/tensorrt_llm/executor_worker/CMakeLists.txt b/cpp/tensorrt_llm/executor_worker/CMakeLists.txt deleted file mode 100644 index 2feb6dfe5790..000000000000 --- a/cpp/tensorrt_llm/executor_worker/CMakeLists.txt +++ /dev/null @@ -1,26 +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. -set(SRCS executorWorker.cpp) - -include_directories(${PROJECT_SOURCE_DIR}/include) - -set(EXECUTOR_WORKER_TARGET executorWorker) - -add_executable(${EXECUTOR_WORKER_TARGET} ${SRCS}) - -target_link_libraries(${EXECUTOR_WORKER_TARGET} - PUBLIC ${SHARED_TARGET} nvinfer_plugin_tensorrt_llm) - -target_compile_features(${EXECUTOR_WORKER_TARGET} PRIVATE cxx_std_17) diff --git a/cpp/tensorrt_llm/executor_worker/executorWorker.cpp b/cpp/tensorrt_llm/executor_worker/executorWorker.cpp deleted file mode 100644 index aa1b06c2cb74..000000000000 --- a/cpp/tensorrt_llm/executor_worker/executorWorker.cpp +++ /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. - */ - -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/executor.h" -#include "tensorrt_llm/executor/serialization.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include - -namespace tle = tensorrt_llm::executor; - -int main(int argc, char* argv[]) -{ -#if ENABLE_MULTI_DEVICE - - if (std::getenv("FORCE_NCCL_ALL_REDUCE_STRATEGY") != nullptr) - { - TLLM_LOG_INFO("FORCE_NCCL_ALL_REDUCE_STRATEGY env variable detected in worker"); - } - - // Register the TRT-LLM plugins - initTrtLlmPlugins(); - - tensorrt_llm::mpi::initialize(tensorrt_llm::mpi::MpiThreadSupport::THREAD_MULTIPLE, true); - - MPI_Comm parentComm; - MPI_Comm_get_parent(&parentComm); - if (parentComm == MPI_COMM_NULL) - { - TLLM_LOG_ERROR("TRT-LLM worker has no parent!"); - return -1; - } - - int size; - MPI_Comm_remote_size(parentComm, &size); - if (size != 1) - { - TLLM_LOG_ERROR("Parent size is %d, must be 1", size); - return -1; - } - - // Since parentComm is an intercommunicator, input root - // is the rank of the parent process in his group - // (always 0 as the parent size is checked before) - - // Receive from the parent the executor configuration - int64_t bufferSize; - MPICHECK(MPI_Bcast(&bufferSize, 1, MPI_INT64_T, 0, parentComm)); - std::vector buffer(bufferSize); - MPICHECK(MPI_Bcast(buffer.data(), bufferSize, MPI_CHAR, 0, parentComm)); - std::istringstream is(std::string(buffer.begin(), buffer.end())); - auto modelPath = tle::Serialization::deserializeString(is); - auto modelType = tle::Serialization::deserializeModelType(is); - auto executorConfig = tle::Serialization::deserializeExecutorConfig(is); - - // Create the orchestrator config for workers - auto orchLeaderComm = std::make_shared(parentComm, true); - auto parallelConfig = executorConfig.getParallelConfig(); - TLLM_CHECK_WITH_INFO(parallelConfig.has_value(), "Parallel config should have a value."); - TLLM_CHECK_WITH_INFO( - parallelConfig.value().getOrchestratorConfig().has_value(), "Orchestrator config should have a value."); - auto orchConfig = parallelConfig.value().getOrchestratorConfig().value(); - TLLM_CHECK_WITH_INFO(parallelConfig.has_value(), "Parallel config should have a value."); - auto newOrchConfig = tle::OrchestratorConfig(false, orchConfig.getWorkerExecutablePath(), orchLeaderComm); - parallelConfig.value().setOrchestratorConfig(newOrchConfig); - executorConfig.setParallelConfig(parallelConfig.value()); - // In orchestrator mode, the spawned threads will wait for termination signal from orchestrator - auto executor = tle::Executor(modelPath, modelType, executorConfig); - - // Wait for all workers to have created their instances - MPI_Barrier(parentComm); - TLLM_LOG_INFO("Executor instance created by worker"); - -#endif // ENABLE_MULTI_DEVICE - - return 0; -} diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu index 5be8b1c2ff78..e9b6850fba6b 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.cu @@ -816,7 +816,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) bool use_float4 = (params.allreduce_in_k != nullptr) && (params.hidden_dim * params.nranks == 6144) && (params.hidden_dim_k * params.nranks == 1024); - if (params.dtype == nvinfer1::DataType::kHALF) + if (params.dtype == tensorrt_llm::DataType::kHALF) { if (use_float4) { @@ -827,7 +827,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) minimax_reduce_rms_kernel_launcher(params); } } - else if (params.dtype == nvinfer1::DataType::kBF16) + else if (params.dtype == tensorrt_llm::DataType::kBF16) { if (use_float4) { @@ -838,7 +838,7 @@ void dispatch_dtype(MiniMaxReduceRMSParams const& params) minimax_reduce_rms_kernel_launcher<__nv_bfloat16, NRanks>(params); } } - else if (params.dtype == nvinfer1::DataType::kFLOAT) + else if (params.dtype == tensorrt_llm::DataType::kFLOAT) { if (use_float4) { diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h index b0cfd0ca074c..bf5775f96de9 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/MiniMaxReduceRMSKernel.h @@ -15,7 +15,7 @@ */ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -59,7 +59,7 @@ struct MiniMaxReduceRMSParams { int nranks{}; int rank{}; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; int size_q{}; // numel of Q (num_token * head_dim_q) int hidden_dim{}; // head_dim_q int size_k{}; // numel of K (num_token * head_dim_k) diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu index 5a3edda04a70..e01878e6c13a 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.cu @@ -795,15 +795,15 @@ void allreduce_fusion_op(AllReduceFusionParams const& params) } #define DISPATCH_DTYPE(NRanks) \ - if (params.dtype == nvinfer1::DataType::kHALF) \ + if (params.dtype == tensorrt_llm::DataType::kHALF) \ { \ DISPATCH_PATTERN(half, NRanks); \ } \ - else if (params.dtype == nvinfer1::DataType::kBF16) \ + else if (params.dtype == tensorrt_llm::DataType::kBF16) \ { \ DISPATCH_PATTERN(__nv_bfloat16, NRanks); \ } \ - else if (params.dtype == nvinfer1::DataType::kFLOAT) \ + else if (params.dtype == tensorrt_llm::DataType::kFLOAT) \ { \ DISPATCH_PATTERN(float, NRanks); \ } \ diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h index 6d2074a6589e..769776273ef6 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h @@ -16,7 +16,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -124,7 +124,7 @@ struct AllReduceFusionParams { int nranks; int rank; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; int size; int hidden_dim; void** workspace; diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu index f1d5c08bda6b..22cdd2de7ab5 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.cu @@ -1357,7 +1357,7 @@ std::vector splitNumber(size_t number) } LowPrecisionAllReduceParams LowPrecisionAllReduceParams::deserialize( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size) + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size) { // Get appropriate static buffer @@ -1401,7 +1401,7 @@ LowPrecisionAllReduceParams LowPrecisionAllReduceParams::deserialize( } LowPrecisionAllReduceParams LowPrecisionAllReduceParams::deserialize_hier( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size) + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size) { // Get appropriate static buffer @@ -1616,7 +1616,7 @@ int32_t max_workspace_size_lowprecision(int32_t tp_size) } void customLowPrecisionAllReduce( - kernels::LowPrecisionAllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream) + kernels::LowPrecisionAllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(lowPrecisionConfigurationSupported(params.ranks_per_node, params.elts_total), "Low Precision Custom all-reduce configuration unsupported"); @@ -1625,10 +1625,10 @@ void customLowPrecisionAllReduce( switch (dataType) { - case nvinfer1::DataType::kFLOAT: lowPrecisionAllReduceDispatchType(params, stream); break; - case nvinfer1::DataType::kHALF: lowPrecisionAllReduceDispatchType(params, stream); break; + case tensorrt_llm::DataType::kFLOAT: lowPrecisionAllReduceDispatchType(params, stream); break; + case tensorrt_llm::DataType::kHALF: lowPrecisionAllReduceDispatchType(params, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: lowPrecisionAllReduceDispatchType<__nv_bfloat16>(params, stream); break; + case tensorrt_llm::DataType::kBF16: lowPrecisionAllReduceDispatchType<__nv_bfloat16>(params, stream); break; #endif default: TLLM_THROW("Unsupported dataType for customAllReduce"); } diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h index 5fc87ef1a523..fb7d4cf679c4 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/customLowPrecisionAllReduceKernels.h @@ -20,7 +20,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/customAllReduceKernels.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -111,15 +111,15 @@ struct LowPrecisionAllReduceParams uint64_t* ag_notify_peer_inside_numa_flags[LP_ALLREDUCE_MAX_BLOCKS * 4]; // 3*flags , 3 is other rank inside numa static LowPrecisionAllReduceParams deserialize( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size); + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size); static LowPrecisionAllReduceParams deserialize_hier( - size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, int token_num, int hidden_size); + size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size); }; bool lowPrecisionConfigurationSupported(size_t msg_size, size_t n_ranks); void customLowPrecisionAllReduce( - kernels::LowPrecisionAllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream); + kernels::LowPrecisionAllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream); int32_t max_workspace_size_lowprecision(int32_t tp_size); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu index eb44f1638a19..0dbc86764d58 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.cu @@ -856,9 +856,9 @@ void oneshotAllreduceFusionOp(AllReduceFusionParams const& params) }; #undef LAUNCH_ALLREDUCE_KERNEL #undef DISPATCH_ALLREDUCE_PATTERN - bool launched = (params.dType == nvinfer1::DataType::kBF16 && dispatchImpl((__nv_bfloat16*) nullptr)) - || (params.dType == nvinfer1::DataType::kFLOAT && dispatchImpl((float*) nullptr)) - || (params.dType == nvinfer1::DataType::kHALF && dispatchImpl((__nv_half*) nullptr)); + bool launched = (params.dType == tensorrt_llm::DataType::kBF16 && dispatchImpl((__nv_bfloat16*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kFLOAT && dispatchImpl((float*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kHALF && dispatchImpl((__nv_half*) nullptr)); if (!launched) { TLLM_CHECK_WITH_INFO(false, "Failed to dispatch MNNVL AllReduceOneShot kernel."); @@ -1246,9 +1246,9 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) #undef LAUNCH_ALLREDUCE_KERNEL - bool launched = (params.dType == nvinfer1::DataType::kFLOAT && dispatchAR((float*) nullptr)) - || (params.dType == nvinfer1::DataType::kBF16 && dispatchAR((__nv_bfloat16*) nullptr)) - || (params.dType == nvinfer1::DataType::kHALF && dispatchAR((__nv_half*) nullptr)); + bool launched = (params.dType == tensorrt_llm::DataType::kFLOAT && dispatchAR((float*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kBF16 && dispatchAR((__nv_bfloat16*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kHALF && dispatchAR((__nv_half*) nullptr)); if (!launched) { TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] Failed to dispatch twoshotAllreduce kernel."); @@ -1388,9 +1388,9 @@ void twoshotAllreduceFusionOp(AllReduceFusionParams const& params) return true; }; - launched = (params.dType == nvinfer1::DataType::kFLOAT && dispatchRN((float*) nullptr)) - || (params.dType == nvinfer1::DataType::kBF16 && dispatchRN((__nv_bfloat16*) nullptr)) - || (params.dType == nvinfer1::DataType::kHALF && dispatchRN((__nv_half*) nullptr)); + launched = (params.dType == tensorrt_llm::DataType::kFLOAT && dispatchRN((float*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kBF16 && dispatchRN((__nv_bfloat16*) nullptr)) + || (params.dType == tensorrt_llm::DataType::kHALF && dispatchRN((__nv_half*) nullptr)); if (!launched) { TLLM_CHECK_WITH_INFO(false, "[MNNVL AllReduceTwoShot] Failed to dispatch rmsnorm lamport kernel."); diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h index 2a228e815b8d..f7789f64ac50 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/mnnvlAllreduceKernels.h @@ -18,7 +18,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/kernels/communicationKernels/allReduceFusionKernels.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -41,7 +41,7 @@ struct AllReduceFusionParams int nRanks; //!< Total number of participating ranks in the AllReduce operation int rank; //!< Current rank ID - nvinfer1::DataType dType; //!< Data type of the tensors (e.g., FP16, BF16, FP32) + tensorrt_llm::DataType dType; //!< Data type of the tensors (e.g., FP16, BF16, FP32) int numTokens; //!< Number of tokens in the input tensor int tokenDim; //!< Hidden Dimension void** bufferPtrsDev; //!< Unicast Device pointers to communication buffers for each rank diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu index 306d42677e2f..cd480869a53e 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.cu @@ -442,11 +442,11 @@ void moereduction_allreduce_fusion_op(MoeReductionAllReduceFusionParams const& p #define MOE_DISPATCH1(DTYPE, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ return moereduction_allreduce_fusion_kernel_launcher(params); #define MOE_DISPATCH0(NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ - if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kHALF) \ + if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kHALF) \ { \ MOE_DISPATCH1(half, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } \ - else if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kBF16) \ + else if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kBF16) \ { \ MOE_DISPATCH1(__nv_bfloat16, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } @@ -727,13 +727,13 @@ void moefinalize_allreduce_fusion_op(MoeFinalizeAllReduceFusionParams const& par #define MOE_FINALIZE_DISPATCH1(DTYPE, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ return moefinalize_allreduce_fusion_kernel_launcher(params); #define MOE_FINALIZE_DISPATCH0(NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT) \ - if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kHALF \ - && params.scale_dtype == nvinfer1::DataType::kHALF) \ + if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kHALF \ + && params.scale_dtype == tensorrt_llm::DataType::kHALF) \ { \ MOE_FINALIZE_DISPATCH1(half, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } \ - else if (params.nranks == NRANKS && params.dtype == nvinfer1::DataType::kBF16 \ - && params.scale_dtype == nvinfer1::DataType::kBF16) \ + else if (params.nranks == NRANKS && params.dtype == tensorrt_llm::DataType::kBF16 \ + && params.scale_dtype == tensorrt_llm::DataType::kBF16) \ { \ MOE_FINALIZE_DISPATCH1(__nv_bfloat16, NRANKS, RESIDUAL_OUT, NORM_OUT, QUANT_OUT); \ } diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h index 556dd4e5cd24..e526a70268b3 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAllReduceFusionKernels.h @@ -16,7 +16,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -44,7 +44,7 @@ struct AllReduceFusionParams { int nranks; int rank; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; // size = token_num * hidden_dim int size; int hidden_dim; @@ -94,7 +94,7 @@ struct MoeFinalizeAllReduceFusionParams : public AllReduceFusionParams // Refer to kernel implementation on layout of those params // number of active experts on current device int top_k; - nvinfer1::DataType scale_dtype; + tensorrt_llm::DataType scale_dtype; // [num_tokens, top_k] void* expert_scale_factor = nullptr; void* shared_expert_output = nullptr; diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu index 91cb5725fede..b5e0df7b1fd6 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.cu @@ -129,25 +129,25 @@ using tensorrt_llm::common::launchWithPdlWhenEnabled; #define SWITCH_DTYPE(dtype, TYPE, ...) \ switch (dtype) \ { \ - case nvinfer1::DataType::kHALF: \ + case tensorrt_llm::DataType::kHALF: \ { \ using TYPE = half; \ __VA_ARGS__; \ break; \ } \ - case nvinfer1::DataType::kBF16: \ + case tensorrt_llm::DataType::kBF16: \ { \ using TYPE = __nv_bfloat16; \ __VA_ARGS__; \ break; \ } \ - case nvinfer1::DataType::kFLOAT: \ + case tensorrt_llm::DataType::kFLOAT: \ { \ using TYPE = float; \ __VA_ARGS__; \ break; \ } \ - case nvinfer1::DataType::kFP8: \ + case tensorrt_llm::DataType::kFP8: \ { \ using TYPE = __nv_fp8_e4m3; \ __VA_ARGS__; \ @@ -1318,7 +1318,7 @@ void moe_a2a_combine_launch(MoeA2ACombineParams const& params) // When use_low_precision is set the recv buffers contain FP8 data regardless of params.dtype, // so dispatch the FP8 accumulation kernel in that case. - auto const effective_dtype = params.use_low_precision ? nvinfer1::DataType::kFP8 : params.dtype; + auto const effective_dtype = params.use_low_precision ? tensorrt_llm::DataType::kFP8 : params.dtype; // Launch appropriate kernel with compact macros SWITCH_DTYPE(effective_dtype, TKernelType, { diff --git a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h index 317ff4d2240c..b88e3170c87f 100644 --- a/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h +++ b/cpp/tensorrt_llm/kernels/communicationKernels/moeAlltoAllKernels.h @@ -16,7 +16,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -153,7 +153,7 @@ struct MoeA2ACombineParams void* output_data; // Output buffer [local_num_tokens, elements_per_token] // Payload information int elements_per_token; // Number of elements per token - nvinfer1::DataType dtype; // Data type of the payload (used for combine kernel dispatch) + tensorrt_llm::DataType dtype; // Data type of the payload (used for combine kernel dispatch) bool use_low_precision; // If true, prepare kernel quantizes payload→FP8; combine kernel accumulates FP8→output dtype diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu index 81e947977797..5ac96f1f3c82 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.cu @@ -117,16 +117,16 @@ void cudaGraphGroupedGemmTemplate(cutlass::gemm::GemmCoord* problemSizesPtr, int template void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - nvinfer1::DataType dataType, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) + tensorrt_llm::DataType dataType, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { cudaGraphGroupedGemmTemplate( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, hostMaxProblemSizesPtr, stream); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { cudaGraphGroupedGemmTemplate( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, @@ -141,7 +141,7 @@ void cudaGraphGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int pro void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, bool isLoraIn, - nvinfer1::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) + tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream) { if (isLoraIn) { @@ -283,17 +283,17 @@ void cudaGraphSplitKGroupedGemmTemplate(cutlass::gemm::GemmCoord* problemSizesPt template void cudaGraphSplitKGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - nvinfer1::DataType dataType, int splitKSlices, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, + tensorrt_llm::DataType dataType, int splitKSlices, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { cudaGraphSplitKGroupedGemmTemplate( problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, splitKSlices, hostMaxProblemSizesPtr, splitKOffsetsGpu, stream); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { cudaGraphSplitKGroupedGemmTemplate(problemSizesPtr, problemCount, ptrAGpu, ptrBGpu, ptrCGpu, ptrDGpu, ldaGpu, ldbGpu, ldcGpu, lddGpu, @@ -308,7 +308,7 @@ void cudaGraphSplitKGroupedGemmType(cutlass::gemm::GemmCoord* problemSizesPtr, i void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream) { if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h index 0eecccb78852..b447bba3a785 100644 --- a/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h +++ b/cpp/tensorrt_llm/kernels/cuda_graph_grouped_gemm.h @@ -18,7 +18,7 @@ #include "cutlass/gemm_coord.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -45,7 +45,7 @@ namespace kernels */ void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, bool isLoraIn, - nvinfer1::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream); + tensorrt_llm::DataType dataType, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, cudaStream_t stream); /** * @brief CUDA Graph compatible wrapper for split-K grouped GEMM operations. @@ -55,7 +55,7 @@ void cudaGraphGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problem */ void cudaGraphSplitKGroupedGemm(cutlass::gemm::GemmCoord* problemSizesPtr, int problemCount, void** ptrAGpu, void** ptrBGpu, void** ptrCGpu, void** ptrDGpu, int64_t* ldaGpu, int64_t* ldbGpu, int64_t* ldcGpu, int64_t* lddGpu, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cutlass::gemm::GemmCoord* hostMaxProblemSizesPtr, int64_t* splitKOffsetsGpu, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu b/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu index 9cf2b51eb583..e4e476507403 100644 --- a/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu +++ b/cpp/tensorrt_llm/kernels/customAllReduceKernels.cu @@ -1187,14 +1187,14 @@ bool is_lamport_supported(int token_num, int hidden_size) return true; } -bool is_lamport_supported(nvinfer1::DataType dataType, int token_num, int hidden_size) +bool is_lamport_supported(tensorrt_llm::DataType dataType, int token_num, int hidden_size) { switch (dataType) { - case nvinfer1::DataType::kFLOAT: return is_lamport_supported(token_num, hidden_size); - case nvinfer1::DataType::kHALF: return is_lamport_supported(token_num, hidden_size); + case tensorrt_llm::DataType::kFLOAT: return is_lamport_supported(token_num, hidden_size); + case tensorrt_llm::DataType::kHALF: return is_lamport_supported(token_num, hidden_size); #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: return is_lamport_supported<__nv_bfloat16>(token_num, hidden_size); + case tensorrt_llm::DataType::kBF16: return is_lamport_supported<__nv_bfloat16>(token_num, hidden_size); #endif default: return false; } @@ -1658,7 +1658,7 @@ static __global__ void __launch_bounds__(512, 1) twoShotAllReduceKernel(AllReduc update_barrier_flag(params.barrier_flag_ptr, params.barrier_flag_counter_ptr); } -bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, nvinfer1::DataType type) +bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, tensorrt_llm::DataType type) { size_t elts_per_thread = 16 / common::getDTypeSize(type); int const msg_align = (algo == AllReduceStrategyType::TWOSHOT) ? n_ranks * elts_per_thread : elts_per_thread; @@ -1894,7 +1894,7 @@ void AllReduceDispatchType(AllReduceParams& params, AllReduceStrategyType strat, } } -AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, +AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size, AllReduceFusionOp op) { void* const* buffer_ptrs = reinterpret_cast(buffer); @@ -1933,7 +1933,7 @@ AllReduceParams AllReduceParams::deserialize(int64_t* buffer, size_t tpSize, siz return params; } -void customAllReduce(kernels::AllReduceParams& params, nvinfer1::DataType dataType, AllReduceStrategyType strat, +void customAllReduce(kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, AllReduceStrategyType strat, AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(configurationSupported(strat, params.elts_total, params.ranks_per_node, dataType), @@ -1943,10 +1943,10 @@ void customAllReduce(kernels::AllReduceParams& params, nvinfer1::DataType dataTy switch (dataType) { - case nvinfer1::DataType::kFLOAT: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; - case nvinfer1::DataType::kHALF: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; + case tensorrt_llm::DataType::kFLOAT: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; + case tensorrt_llm::DataType::kHALF: AllReduceDispatchType(params, strat, config, fusionOp, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: AllReduceDispatchType<__nv_bfloat16>(params, strat, config, fusionOp, stream); break; #endif @@ -1991,22 +1991,22 @@ void launchResidualRmsNormKernel(kernels::AllReduceParams& params, cudaStream_t } void residualRmsNorm( - kernels::AllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp) + kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp) { sync_check_cuda_error(stream); switch (dataType) { - case nvinfer1::DataType::kFLOAT: launchResidualRmsNormKernel(params, stream, fusionOp); break; - case nvinfer1::DataType::kHALF: launchResidualRmsNormKernel(params, stream, fusionOp); break; + case tensorrt_llm::DataType::kFLOAT: launchResidualRmsNormKernel(params, stream, fusionOp); break; + case tensorrt_llm::DataType::kHALF: launchResidualRmsNormKernel(params, stream, fusionOp); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: launchResidualRmsNormKernel<__nv_bfloat16>(params, stream, fusionOp); break; + case tensorrt_llm::DataType::kBF16: launchResidualRmsNormKernel<__nv_bfloat16>(params, stream, fusionOp); break; #endif default: TLLM_THROW("Unsupported dataType for customAllReduce"); } sync_check_cuda_error(stream); } -void lamportInitialize(void* buffer, size_t size, nvinfer1::DataType dataType, cudaStream_t stream) +void lamportInitialize(void* buffer, size_t size, tensorrt_llm::DataType dataType, cudaStream_t stream) { sync_check_cuda_error(stream); if (size == 0) @@ -2015,14 +2015,14 @@ void lamportInitialize(void* buffer, size_t size, nvinfer1::DataType dataType, c } switch (dataType) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: reduce_fusion::lamport_initialize_kernel_launcher(buffer, size, stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: reduce_fusion::lamport_initialize_kernel_launcher(buffer, size, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: reduce_fusion::lamport_initialize_kernel_launcher<__nv_bfloat16>(buffer, size, stream); break; #endif diff --git a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h index f7151f1cd0ab..93f67ffdd911 100644 --- a/cpp/tensorrt_llm/kernels/customAllReduceKernels.h +++ b/cpp/tensorrt_llm/kernels/customAllReduceKernels.h @@ -17,7 +17,7 @@ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -178,23 +178,23 @@ struct AllReduceParams AllReduceFusionParams fusion_params; - static AllReduceParams deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, nvinfer1::DataType dataType, + static AllReduceParams deserialize(int64_t* buffer, size_t tpSize, size_t tpRank, tensorrt_llm::DataType dataType, int token_num, int hidden_size, AllReduceFusionOp op); }; -bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, nvinfer1::DataType type); +bool configurationSupported(AllReduceStrategyType algo, size_t msg_size, size_t n_ranks, tensorrt_llm::DataType type); -void customAllReduce(kernels::AllReduceParams& params, nvinfer1::DataType dataType, AllReduceStrategyType strat, +void customAllReduce(kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, AllReduceStrategyType strat, AllReduceStrategyConfig config, AllReduceFusionOp fusionOp, cudaStream_t stream); void residualRmsNorm( - kernels::AllReduceParams& params, nvinfer1::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp); + kernels::AllReduceParams& params, tensorrt_llm::DataType dataType, cudaStream_t stream, AllReduceFusionOp fusionOp); -void lamportInitialize(void* buffer, size_t size, nvinfer1::DataType dataType, cudaStream_t stream); +void lamportInitialize(void* buffer, size_t size, tensorrt_llm::DataType dataType, cudaStream_t stream); namespace reduce_fusion { -bool is_lamport_supported(nvinfer1::DataType dataType, int token_num, int hidden_size); +bool is_lamport_supported(tensorrt_llm::DataType dataType, int token_num, int hidden_size); } } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h index dbbed4e08c97..6632f273cc35 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h @@ -17,7 +17,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include "cutlass/half.h" #include @@ -38,34 +38,34 @@ namespace kernels namespace cutlass_kernels { /////////////////////////////////////////////////////////////////////////////////////////////////// -// nvinfer1::DataType to Cutlass +// tensorrt_llm::DataType to Cutlass /////////////////////////////////////////////////////////////////////////////////////////////////// -template +template struct CutlassType { using type = void; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::half_t; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::bfloat16_t; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::float_e4m3_t; }; template <> -struct CutlassType +struct CutlassType { using type = cutlass::float_e2m1_t; }; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h index 216877a4ffc7..159ea447c6e7 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h @@ -27,7 +27,7 @@ #include #endif #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -958,8 +958,8 @@ struct GemmProfilerBackend using Config = cutlass_extensions::CutlassGemmConfig; using GemmToProfile = MoeGemmId; - void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, nvinfer1::DataType dtype, - nvinfer1::DataType wtype, nvinfer1::DataType otype, int num_experts, int k, int64_t hidden_size, + void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, tensorrt_llm::DataType dtype, + tensorrt_llm::DataType wtype, tensorrt_llm::DataType otype, int num_experts, int k, int64_t hidden_size, int64_t unpadded_hidden_size, int64_t inter_size, int64_t group_size, ActivationType activation_type, bool bias, bool use_lora, bool min_latency_mode, bool need_weights, MOEParallelismConfig parallelism_config, bool const enable_alltoall) @@ -986,13 +986,13 @@ struct GemmProfilerBackend mSM = common::getSMVersion(); mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - if (dtype == nvinfer1::DataType::kFP8 - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + if (dtype == tensorrt_llm::DataType::kFP8 + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } - else if ((dtype == nvinfer1::DataType::kFP4 || dtype == nvinfer1::DataType::kINT64) - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + else if ((dtype == tensorrt_llm::DataType::kFP4 || dtype == tensorrt_llm::DataType::kINT64) + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; } @@ -1023,9 +1023,9 @@ struct GemmProfilerBackend int mSampleIndex = 0; - nvinfer1::DataType mDType{}; - nvinfer1::DataType mWType{}; - nvinfer1::DataType mOType{}; + tensorrt_llm::DataType mDType{}; + tensorrt_llm::DataType mWType{}; + tensorrt_llm::DataType mOType{}; // This will be a unique value for every iteration of warmup and actual bench constexpr static int64_t NUM_ROUTING_SAMPLES = 16; diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h index 8214e84e3b7f..0f2a128eb7d9 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_lora_device_path.h @@ -18,7 +18,7 @@ #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -56,7 +56,7 @@ struct MoeLoraDevicePathModule; // stream: CUDA stream to launch onto. using MoeLoraDeviceRunFn = void (*)(MoeLoraDevicePathModule const& mod, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, int64_t splitk_slices, void const* input_base, - void* output_base, nvinfer1::DataType data_type, cudaStream_t stream); + void* output_base, tensorrt_llm::DataType data_type, cudaStream_t stream); // Per-module device-resident scratch for the MoE LoRA capture-safe path. // Pointers refer to device memory unless noted. diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h index e902e2c9d6d3..25abc9b0f911 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/include/moe_util_kernels.h @@ -25,7 +25,7 @@ #ifdef ENABLE_FP4 #include #endif -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include diff --git a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu index a99f42003e47..6ffe6fc26636 100644 --- a/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu +++ b/cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu @@ -3670,7 +3670,7 @@ void CutlassMoeFCRunner -constexpr nvinfer1::DataType moeLoraNvInferType() +constexpr tensorrt_llm::DataType moeLoraNvInferType() { if constexpr (std::is_same_v) { - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; } #if defined(ENABLE_BF16) else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; } #endif else if constexpr (std::is_same_v) { - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; } else { @@ -3886,7 +3886,7 @@ auto CutlassMoeFCRunner(); + tensorrt_llm::DataType const data_type = moeLoraNvInferType(); // The device-path GEMM skips rank-0 rows, but the bias/reorder paths // read lora_fc1_result_ for every valid row. Zero the buffer first so @@ -3979,7 +3979,7 @@ void CutlassMoeFCRunner(); + tensorrt_llm::DataType const data_type = moeLoraNvInferType(); // As in loraFC1, zero the output so rank-0 rows the GEMM skips do not // feed stale data into the downstream add. @@ -4701,18 +4701,18 @@ std::map> GemmProfilerBackend::getProfile size_t k = mK; size_t num_expanded_tokens = mMinLatencyMode ? maxM * mNumExpertsPerNode : maxM * k; - TLLM_CHECK(mDType != nvinfer1::DataType::kINT4); + TLLM_CHECK(mDType != tensorrt_llm::DataType::kINT4); // nvllm still uses int64 because torch doesn't have fp4 yet. - bool is_4bit_act = mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64; - bool is_4bit_weight = mWType == nvinfer1::DataType::kINT4 || mWType == nvinfer1::DataType::kFP4 - || mWType == nvinfer1::DataType::kINT64; + bool is_4bit_act = mDType == tensorrt_llm::DataType::kFP4 || mDType == tensorrt_llm::DataType::kINT64; + bool is_4bit_weight = mWType == tensorrt_llm::DataType::kINT4 || mWType == tensorrt_llm::DataType::kFP4 + || mWType == tensorrt_llm::DataType::kINT64; TLLM_CHECK_WITH_INFO(!is_4bit_act || is_4bit_weight, "Cannot have 4-bit activation with non-4-bit weight"); float dtype_bytes = is_4bit_act ? 0.5f - : static_cast(mWType == nvinfer1::DataType::kINT4 ? getDTypeSize(mOType) : getDTypeSize(mDType)); + : static_cast(mWType == tensorrt_llm::DataType::kINT4 ? getDTypeSize(mOType) : getDTypeSize(mDType)); float weight_bytes = is_4bit_weight ? 0.5f : static_cast(getDTypeSize(mWType)); size_t output_bytes = getDTypeSize(mOType); - size_t gemm_output_bytes = (mOType == nvinfer1::DataType::kFP8) + size_t gemm_output_bytes = (mOType == tensorrt_llm::DataType::kFP8) ? sizeof(TmaWarpSpecializedGroupedGemmInput::OutputTypeAdaptor_t<__nv_fp8_e4m3>) : output_bytes; @@ -4754,18 +4754,18 @@ std::map> GemmProfilerBackend::getProfile // TODO Make quant 2 & 4 bigger for FP8 if we ever change to scaling per expert bool is_int_w_quant - = (mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4) && mGroupSize <= 0; + = (mWType == tensorrt_llm::DataType::kINT8 || mWType == tensorrt_llm::DataType::kINT4) && mGroupSize <= 0; bool is_int_groupwise_w_quant - = (mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4) && mGroupSize > 0; - bool is_fp8_act_quant = mDType == nvinfer1::DataType::kFP8; - bool is_fp8_w_quant = mWType == nvinfer1::DataType::kFP8; + = (mWType == tensorrt_llm::DataType::kINT8 || mWType == tensorrt_llm::DataType::kINT4) && mGroupSize > 0; + bool is_fp8_act_quant = mDType == tensorrt_llm::DataType::kFP8; + bool is_fp8_w_quant = mWType == tensorrt_llm::DataType::kFP8; // nvllm still uses int64 because torch doesn't have fp4 yet. - // bool is_fp4_act_quant = mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64; - bool is_fp4_w_quant = mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64; + // bool is_fp4_act_quant = mDType == tensorrt_llm::DataType::kFP4 || mDType == tensorrt_llm::DataType::kINT64; + bool is_fp4_w_quant = mWType == tensorrt_llm::DataType::kFP4 || mWType == tensorrt_llm::DataType::kINT64; bool is_w4afp8_quant = is_int_groupwise_w_quant && is_fp8_act_quant; // bool is_wfp4afp8_quant = is_fp4_w_quant && is_fp8_act_quant; - bool is_wfp4a16_quant = (mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) - && mWType == nvinfer1::DataType::kUINT8; + bool is_wfp4a16_quant = (mDType == tensorrt_llm::DataType::kHALF || mDType == tensorrt_llm::DataType::kBF16) + && mWType == tensorrt_llm::DataType::kUINT8; // Int sizes size_t quant_1_size = is_int_w_quant ? fc1_out_size * num_experts_per_node * dtype_bytes : 0; @@ -4985,19 +4985,19 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr GET_WS_PTR(float const*, w4a8_alpha); #undef GET_WS_PTR - if ((mWType == nvinfer1::DataType::kINT8 || mWType == nvinfer1::DataType::kINT4 - || mWType == nvinfer1::DataType::kUINT8) + if ((mWType == tensorrt_llm::DataType::kINT8 || mWType == tensorrt_llm::DataType::kINT4 + || mWType == tensorrt_llm::DataType::kUINT8) && mGroupSize < 0) { TLLM_CHECK(quant_1 && quant_2); mQuantParams = QuantParams::Int(quant_1, quant_2); } - else if (mWType == nvinfer1::DataType::kINT4 || mWType == nvinfer1::DataType::kUINT8) + else if (mWType == tensorrt_llm::DataType::kINT4 || mWType == tensorrt_llm::DataType::kUINT8) { TLLM_CHECK(quant_1 && quant_2); - if (mDType == nvinfer1::DataType::kFP8 - || (mWType == nvinfer1::DataType::kUINT8 - && (mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16))) + if (mDType == tensorrt_llm::DataType::kFP8 + || (mWType == tensorrt_llm::DataType::kUINT8 + && (mDType == tensorrt_llm::DataType::kHALF || mDType == tensorrt_llm::DataType::kBF16))) { TLLM_CHECK(w4a8_alpha); mQuantParams = QuantParams::GroupWise( @@ -5008,14 +5008,14 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr mQuantParams = QuantParams::GroupWise(mGroupSize, quant_1, quant_2, nullptr, nullptr, quant_3, quant_4); } } - else if (mWType == nvinfer1::DataType::kFP8) + else if (mWType == tensorrt_llm::DataType::kFP8) { TLLM_CHECK(quant_1 && quant_2 && quant_3); mQuantParams = QuantParams::FP8(static_cast(quant_1), static_cast(quant_2), static_cast(quant_3), static_cast(quant_4)); } - else if (mDType == nvinfer1::DataType::kFP8 - && (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64)) + else if (mDType == tensorrt_llm::DataType::kFP8 + && (mWType == tensorrt_llm::DataType::kFP4 || mWType == tensorrt_llm::DataType::kINT64)) { TLLM_CHECK(quant_1 && quant_2 && quant_3 && quant_4 && quant_5 && quant_6); mQuantParams = QuantParams::FP8MXFP4(static_cast(quant_1), @@ -5024,8 +5024,8 @@ void GemmProfilerBackend::prepareQuantParams(int num_tokens, char* workspace_ptr static_cast(quant_5), static_cast(quant_6)); } - else if ((mDType == nvinfer1::DataType::kFP4 || mDType == nvinfer1::DataType::kINT64) - && (mWType == nvinfer1::DataType::kFP4 || mWType == nvinfer1::DataType::kINT64)) + else if ((mDType == tensorrt_llm::DataType::kFP4 || mDType == tensorrt_llm::DataType::kINT64) + && (mWType == tensorrt_llm::DataType::kFP4 || mWType == tensorrt_llm::DataType::kINT64)) { // nvllm still uses int64 because torch doesn't have fp4 yet. TLLM_CHECK(quant_1 && quant_2 && quant_3 && quant_4 && quant_5 && quant_6); @@ -5045,9 +5045,9 @@ void GemmProfilerBackend::prepareTmaWsInputs(int num_tokens, char* workspace_ptr return; } - bool use_w4afp8 = (mDType == nvinfer1::DataType::kFP8 && mWType == nvinfer1::DataType::kINT4); - bool use_wfp4a16 = ((mDType == nvinfer1::DataType::kHALF || mDType == nvinfer1::DataType::kBF16) - && mWType == nvinfer1::DataType::kUINT8); + bool use_w4afp8 = (mDType == tensorrt_llm::DataType::kFP8 && mWType == tensorrt_llm::DataType::kINT4); + bool use_wfp4a16 = ((mDType == tensorrt_llm::DataType::kHALF || mDType == tensorrt_llm::DataType::kBF16) + && mWType == tensorrt_llm::DataType::kUINT8); bool const use_finalize_fusion = fusion == TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE; bool const finalize_fusion_not_supported = !mInterface->use_fused_finalize_ || mMinLatencyMode || use_wfp4a16 || mGemmToProfile != GemmToProfile::GEMM_2; diff --git a/cpp/tensorrt_llm/kernels/gptKernels.h b/cpp/tensorrt_llm/kernels/gptKernels.h index e13e9bca4d6a..6e8dc7a17830 100644 --- a/cpp/tensorrt_llm/kernels/gptKernels.h +++ b/cpp/tensorrt_llm/kernels/gptKernels.h @@ -227,11 +227,11 @@ struct BuildDecoderInfoParams std::string toString() const { std::stringstream ss; - auto printTensor = [&ss](char const* name, void* ptr, nvinfer1::Dims shape) + auto printTensor = [&ss](char const* name, void* ptr, tensorrt_llm::Dims shape) { ss << name << ": "; if (ptr) - ss << *(runtime::ITensor::wrap((void*) ptr, nvinfer1::DataType::kINT32, shape)); + ss << *(runtime::ITensor::wrap((void*) ptr, tensorrt_llm::DataType::kINT32, shape)); else ss << "nullptr"; ss << std::endl; diff --git a/cpp/tensorrt_llm/kernels/groupGemm.cu b/cpp/tensorrt_llm/kernels/groupGemm.cu index 5b8c0d929150..f87824e280e4 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.cu +++ b/cpp/tensorrt_llm/kernels/groupGemm.cu @@ -63,7 +63,7 @@ template problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - nvinfer1::DataType dataType, cudaStream_t stream) + tensorrt_llm::DataType dataType, cudaStream_t stream) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); using ElementA = cutlassType; @@ -178,20 +178,20 @@ template problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - nvinfer1::DataType dataType, cudaStream_t stream) + tensorrt_llm::DataType dataType, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { groupedGemm_(problem_sizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, dataType, stream); } - else if (dataType == nvinfer1::DataType::kFLOAT) + else if (dataType == tensorrt_llm::DataType::kFLOAT) { TLLM_CHECK_WITH_INFO(false, "not support float input/output"); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { groupedGemm_(problem_sizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkspaceSize, @@ -203,7 +203,7 @@ void groupedGemmType_(std::vector problem_sizes, std:: void groupedGemm(std::vector problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int minKN, cudaStream_t stream) + bool isLoraIn, tensorrt_llm::DataType dataType, int minKN, cudaStream_t stream) { TLLM_LOG_TRACE("%s start, isLoraIn: %d, minKN = %d", __PRETTY_FUNCTION__, static_cast(isLoraIn), minKN); if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/groupGemm.h b/cpp/tensorrt_llm/kernels/groupGemm.h index dbc1e498b7b2..c526e08e986c 100644 --- a/cpp/tensorrt_llm/kernels/groupGemm.h +++ b/cpp/tensorrt_llm/kernels/groupGemm.h @@ -17,7 +17,7 @@ #include "cutlass/gemm_coord.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" TRTLLM_NAMESPACE_BEGIN @@ -29,7 +29,7 @@ int64_t getGroupedGemmParamsWorkSpaceSize(int64_t problem_count); void groupedGemm(std::vector problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkspace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int minKN, cudaStream_t stream); + bool isLoraIn, tensorrt_llm::DataType dataType, int minKN, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu index 409968bb510d..90ec7ae960eb 100644 --- a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu +++ b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.cu @@ -655,9 +655,9 @@ void GroupRMSNormBaseKernelLauncher(GroupRMSParams& params) switch (params.dtype) { - case nvinfer1::DataType::kHALF: GROUP_RMS_NORM_DISPATCH(half); break; - case nvinfer1::DataType::kBF16: GROUP_RMS_NORM_DISPATCH(__nv_bfloat16); break; - case nvinfer1::DataType::kFLOAT: GROUP_RMS_NORM_DISPATCH(float); break; + case tensorrt_llm::DataType::kHALF: GROUP_RMS_NORM_DISPATCH(half); break; + case tensorrt_llm::DataType::kBF16: GROUP_RMS_NORM_DISPATCH(__nv_bfloat16); break; + case tensorrt_llm::DataType::kFLOAT: GROUP_RMS_NORM_DISPATCH(float); break; default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNorm"); } @@ -750,9 +750,9 @@ void GroupRMSNormKernelLargeBatchLauncher(GroupRMSParams& params) switch (params.dtype) { - case nvinfer1::DataType::kHALF: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(half); break; - case nvinfer1::DataType::kBF16: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(__nv_bfloat16); break; - case nvinfer1::DataType::kFLOAT: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(float); break; + case tensorrt_llm::DataType::kHALF: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(half); break; + case tensorrt_llm::DataType::kBF16: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(__nv_bfloat16); break; + case tensorrt_llm::DataType::kFLOAT: GROUP_RMS_NORM_LARGE_BATCH_DISPATCH(float); break; default: TLLM_CHECK_WITH_INFO(false, "Unsupported data type for GroupRMSNormV2"); } @@ -813,15 +813,15 @@ void GroupRMSNormKernelLauncherWithHeuristic(GroupRMSParams& params) // Choose the appropriate DType switch (params.dtype) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: base_warps = calculateNumWarpsBase(params); large_batch_warps = calculateNumWarpsLargeBatch(params).num_warps_to_launch; break; - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: base_warps = calculateNumWarpsBase<__nv_bfloat16, n>(params); large_batch_warps = calculateNumWarpsLargeBatch<__nv_bfloat16, n>(params).num_warps_to_launch; break; - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: base_warps = calculateNumWarpsBase(params); large_batch_warps = calculateNumWarpsLargeBatch(params).num_warps_to_launch; break; diff --git a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h index 335adf44ed67..70425f924217 100644 --- a/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h +++ b/cpp/tensorrt_llm/kernels/groupRmsNormKernels/groupRmsNormKernels.h @@ -15,7 +15,7 @@ */ #pragma once #include "tensorrt_llm/common/assert.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -44,7 +44,7 @@ struct GroupRMSParams float eps; float weight_bias; bool enable_weights; - nvinfer1::DataType dtype; + tensorrt_llm::DataType dtype; cudaStream_t stream; }; diff --git a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h index 132990603db9..0b02686f5e53 100644 --- a/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h +++ b/cpp/tensorrt_llm/kernels/internal_cutlass_kernels/include/moe_kernels.h @@ -26,7 +26,7 @@ #ifdef ENABLE_FP4 #include #endif -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include @@ -869,8 +869,8 @@ struct GemmProfilerBackend GEMM_2 }; - void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, nvinfer1::DataType dtype, - nvinfer1::DataType wtype, nvinfer1::DataType otype, int num_experts, int k, int64_t hidden_size, + void init(CutlassMoeFCRunnerInterface& runner, GemmToProfile gemm_to_profile, tensorrt_llm::DataType dtype, + tensorrt_llm::DataType wtype, tensorrt_llm::DataType otype, int num_experts, int k, int64_t hidden_size, int64_t inter_size, int64_t group_size, ActivationType activation_type, bool bias, bool use_lora, bool min_latency_mode, bool need_weights, MOEParallelismConfig parallelism_config) { @@ -895,13 +895,13 @@ struct GemmProfilerBackend mSorter.updateNumExperts(mNumExpertsPerNode); mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NONE; - if (dtype == nvinfer1::DataType::kFP8 - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + if (dtype == tensorrt_llm::DataType::kFP8 + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::MXFPX; } - else if ((dtype == nvinfer1::DataType::kFP4 || dtype == nvinfer1::DataType::kINT64) - && (wtype == nvinfer1::DataType::kFP4 || wtype == nvinfer1::DataType::kINT64)) + else if ((dtype == tensorrt_llm::DataType::kFP4 || dtype == tensorrt_llm::DataType::kINT64) + && (wtype == tensorrt_llm::DataType::kFP4 || wtype == tensorrt_llm::DataType::kINT64)) { mScalingType = TmaWarpSpecializedGroupedGemmInput::FpXBlockScalingType::NVFP4; } @@ -932,9 +932,9 @@ struct GemmProfilerBackend int mSampleIndex = 0; - nvinfer1::DataType mDType{}; - nvinfer1::DataType mWType{}; - nvinfer1::DataType mOType{}; + tensorrt_llm::DataType mDType{}; + tensorrt_llm::DataType mWType{}; + tensorrt_llm::DataType mOType{}; // This will be a unique value for every iteration of warmup and actual bench constexpr static int64_t NUM_ROUTING_SAMPLES = 16; diff --git a/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu b/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu index 3b91cf3f1776..24d353a6b04b 100644 --- a/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu +++ b/cpp/tensorrt_llm/kernels/kvCachePartialCopy.cu @@ -89,42 +89,42 @@ void kvCacheBlockPartialCopy(IBuffer& dst, IBuffer const& src, unsigned int numL TLLM_CHECK_WITH_INFO(dataType == dst.getDataType(), "src and dst dataType does not match"); switch (dataType) { - case nvinfer1::DataType::kINT64: + case tensorrt_llm::DataType::kINT64: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kINT32: + case tensorrt_llm::DataType::kINT32: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: hostKVCacheBlockPartialCopy<__nv_bfloat16>( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; #endif - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kBOOL: + case tensorrt_llm::DataType::kBOOL: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kUINT8: + case tensorrt_llm::DataType::kUINT8: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; - case nvinfer1::DataType::kINT8: + case tensorrt_llm::DataType::kINT8: hostKVCacheBlockPartialCopy( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: + case tensorrt_llm::DataType::kFP8: hostKVCacheBlockPartialCopy<__nv_fp8_e4m3>( dst, src, numLayers, numHeads, tokensPerBlock, numHidden, numTokensToCopy, kvFactor, stream); break; diff --git a/cpp/tensorrt_llm/kernels/lora/dora.cpp b/cpp/tensorrt_llm/kernels/lora/dora.cpp index 43dbf4fdccb0..883d02df9291 100644 --- a/cpp/tensorrt_llm/kernels/lora/dora.cpp +++ b/cpp/tensorrt_llm/kernels/lora/dora.cpp @@ -20,13 +20,13 @@ #include "tensorrt_llm/common/memoryUtils.h" #include "tensorrt_llm/kernels/doraScaling.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include using tensorrt_llm::kernels::DoraImpl; -DoraImpl::DoraImpl(std::vector const& outHiddenSizes, nvinfer1::DataType type) +DoraImpl::DoraImpl(std::vector const& outHiddenSizes, tensorrt_llm::DataType type) : mType(type) { mCumModuleSizes.resize(outHiddenSizes.size()); @@ -73,14 +73,14 @@ int DoraImpl::run(int64_t numTokens, void const* input, void const* const* loraW auto const* deviceCumModuleSizes = reinterpret_cast(workspace); auto const* deviceScalePtrs = reinterpret_cast((&deviceCumModuleSizes[numModules])); - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { tokenPerChannelScale(numel, numModules, numTokens, deviceCumModuleSizes, reinterpret_cast(input), reinterpret_cast(deviceScalePtrs), reinterpret_cast(outputs[0]), stream); } #ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { tokenPerChannelScale(numel, numModules, numTokens, deviceCumModuleSizes, reinterpret_cast(input), reinterpret_cast(deviceScalePtrs), diff --git a/cpp/tensorrt_llm/kernels/lora/dora.h b/cpp/tensorrt_llm/kernels/lora/dora.h index fc21fe669366..02cd68e7c1f0 100644 --- a/cpp/tensorrt_llm/kernels/lora/dora.h +++ b/cpp/tensorrt_llm/kernels/lora/dora.h @@ -17,7 +17,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cudaUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" TRTLLM_NAMESPACE_BEGIN @@ -28,7 +28,7 @@ class DoraImpl public: DoraImpl() = delete; - DoraImpl(std::vector const& outHiddenSizes, nvinfer1::DataType type); + DoraImpl(std::vector const& outHiddenSizes, tensorrt_llm::DataType type); ~DoraImpl() = default; @@ -41,7 +41,7 @@ class DoraImpl private: std::vector mCumModuleSizes; std::vector mHostBuf; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; }; } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/lora/lora.cpp b/cpp/tensorrt_llm/kernels/lora/lora.cpp index 61f6af00fedc..212111d5a65c 100644 --- a/cpp/tensorrt_llm/kernels/lora/lora.cpp +++ b/cpp/tensorrt_llm/kernels/lora/lora.cpp @@ -48,7 +48,7 @@ void _getProblemParams(cublasOperation_t& transa, cublasOperation_t& transb, int // TODO should reuse the function in gemmPlugin void _runGemm(int const M, int const N, int const K, bool const transA, bool const transB, - nvinfer1::DataType const type, CublasGemmWrapperPtr const& cublasWrapperPtr, void const* act, void const* weight, + tensorrt_llm::DataType const type, CublasGemmWrapperPtr const& cublasWrapperPtr, void const* act, void const* weight, void* output, std::optional const& heuristic, void* workspace, cudaStream_t stream) { cublasWrapperPtr->setStream(stream); @@ -65,7 +65,7 @@ void _runGemm(int const M, int const N, int const K, bool const transA, bool con } LoraImpl::LoraImpl(int in_hidden_size, std::vector out_hidden_sizes, bool transA, bool transB, - int num_lora_modules, nvinfer1::DataType type, int max_low_rank, std::shared_ptr cublasWrapper) + int num_lora_modules, tensorrt_llm::DataType type, int max_low_rank, std::shared_ptr cublasWrapper) : mInHiddenSize(in_hidden_size) , mTransA(transA) , mTransB(transB) @@ -82,16 +82,16 @@ LoraImpl::LoraImpl(int in_hidden_size, std::vector out_hidden_sizes, bool t void LoraImpl::setGemmConfig() { TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - if (mType == nvinfer1::DataType::kHALF) + if (mType == tensorrt_llm::DataType::kHALF) { mCublasWrapper->setFP16GemmConfig(); } - else if (mType == nvinfer1::DataType::kFLOAT) + else if (mType == tensorrt_llm::DataType::kFLOAT) { mCublasWrapper->setFP32GemmConfig(); } #ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) + else if (mType == tensorrt_llm::DataType::kBF16) { mCublasWrapper->setBF16GemmConfig(); } @@ -121,7 +121,7 @@ int64_t getGemmWorkSpaceSize(int64_t numTokens, int64_t maxLoraModuleNum, int64_ } size_t LoraImpl::getWorkspaceSize( - int64_t const numTokens, int64_t const numReqs, nvinfer1::DataType const type) const noexcept + int64_t const numTokens, int64_t const numReqs, tensorrt_llm::DataType const type) const noexcept { TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); auto const typeSize = tensorrt_llm::common::getDTypeSize(type); diff --git a/cpp/tensorrt_llm/kernels/lora/lora.h b/cpp/tensorrt_llm/kernels/lora/lora.h index 7215a7af74d4..9924f6d68f08 100644 --- a/cpp/tensorrt_llm/kernels/lora/lora.h +++ b/cpp/tensorrt_llm/kernels/lora/lora.h @@ -19,7 +19,7 @@ #include "tensorrt_llm/common/config.h" #include "tensorrt_llm/common/cublasMMWrapper.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -37,9 +37,9 @@ class LoraImpl { public: LoraImpl(int in_hidden_size, std::vector out_hidden_sizes, bool transA, bool transB, int num_lora_modules, - nvinfer1::DataType type, int max_low_rank, std::shared_ptr cublasWrapper); + tensorrt_llm::DataType type, int max_low_rank, std::shared_ptr cublasWrapper); - [[nodiscard]] size_t getWorkspaceSize(int64_t numTokens, int64_t numReqs, nvinfer1::DataType type) const noexcept; + [[nodiscard]] size_t getWorkspaceSize(int64_t numTokens, int64_t numReqs, tensorrt_llm::DataType type) const noexcept; void setBestTactic(std::optional config); int run(int64_t numTokens, int64_t numReqs, void const* input, int32_t const* loraRanks, void const* const* loraWeightsPtr, int weightIndex, void* const* outputs, void* workspace, cudaStream_t stream); @@ -54,7 +54,7 @@ class LoraImpl private: bool mTransA; bool mTransB; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; int mNumLoraModules; // @fixme: seems this is shared across multiple clones. diff --git a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu index c3276ea487bc..5c6a0e74f184 100644 --- a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu +++ b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.cu @@ -358,7 +358,7 @@ void launchLoraGroupGEMMParamFillRowReorderFusion(int32_t* in_sizes, int32_t* ou int64_t a_base, int64_t d_base, int64_t d_prime_base, int32_t const* slot_counts, int32_t const* slot_ranks, int64_t const* slot_offsets, int32_t const* module_out_sizes, int64_t const* module_out_prefix, int64_t const* b_ptrs, int64_t const* b_prime_ptrs, void const* input, int64_t const* sorted_ids, - int32_t module_count, nvinfer1::DataType dtype, cudaStream_t stream) + int32_t module_count, tensorrt_llm::DataType dtype, cudaStream_t stream) { // Determine block dimensions (1D) // Requirements: 1) >= max_lora_count * module_count 2) >= 256 3) divisible by 32 diff --git a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h index 3043054ca4b5..835b9f8bed96 100644 --- a/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h +++ b/cpp/tensorrt_llm/kernels/lora/loraGroupGEMMParamFillRowReorderFusion.h @@ -17,7 +17,7 @@ #pragma once #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -70,7 +70,7 @@ void launchLoraGroupGEMMParamFillRowReorderFusion(int32_t* in_sizes, int32_t* ou int64_t a_base, int64_t d_base, int64_t d_prime_base, int32_t const* slot_counts, int32_t const* slot_ranks, int64_t const* slot_offsets, int32_t const* module_out_sizes, int64_t const* module_out_prefix, int64_t const* b_ptrs, int64_t const* b_prime_ptrs, void const* input, int64_t const* sorted_ids, - int32_t module_count, nvinfer1::DataType dtype, cudaStream_t stream); + int32_t module_count, tensorrt_llm::DataType dtype, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu index 6397396ea6f6..755d9a74ea78 100644 --- a/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu +++ b/cpp/tensorrt_llm/kernels/splitkGroupGemm.cu @@ -203,20 +203,20 @@ template const& problemSizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkSpaceSize, - nvinfer1::DataType dataType, int splitKSlices, cudaStream_t stream) + tensorrt_llm::DataType dataType, int splitKSlices, cudaStream_t stream) { - if (dataType == nvinfer1::DataType::kHALF) + if (dataType == tensorrt_llm::DataType::kHALF) { splitkGroupedGemm_(problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, gemmWorkSpaceSize, splitKSlices, stream); } - else if (dataType == nvinfer1::DataType::kFLOAT) + else if (dataType == tensorrt_llm::DataType::kFLOAT) { TLLM_CHECK_WITH_INFO(false, "not support float input/output"); } #ifdef ENABLE_BF16 - else if (dataType == nvinfer1::DataType::kBF16) + else if (dataType == tensorrt_llm::DataType::kBF16) { splitkGroupedGemm_( problemSizes, ptrA, ptrB, ptrC, ptrD, gemmParamsWorkSpace, gemmParamsWorkSpaceSize, gemmWorkSpace, @@ -228,7 +228,7 @@ void splitkGroupedGemmType_(std::vector const& problem void splitkGroupedGemm(std::vector const& problemSizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkSpace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkSpaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream) + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream) { TLLM_LOG_TRACE("%s start, isLoraIn: %d, minKN = %d", __PRETTY_FUNCTION__, static_cast(isLoraIn), minKN); if (isLoraIn) diff --git a/cpp/tensorrt_llm/kernels/splitkGroupGemm.h b/cpp/tensorrt_llm/kernels/splitkGroupGemm.h index 6ada8255292e..bcde457db7d2 100644 --- a/cpp/tensorrt_llm/kernels/splitkGroupGemm.h +++ b/cpp/tensorrt_llm/kernels/splitkGroupGemm.h @@ -17,7 +17,7 @@ #include "cutlass/gemm_coord.h" #include "tensorrt_llm/common/config.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -30,7 +30,7 @@ int64_t getSplitkGroupedGemmParamsWorkSpaceSize(int64_t problem_count); void splitkGroupedGemm(std::vector const& problem_sizes, std::vector const& ptrA, std::vector const& ptrB, std::vector const& ptrC, std::vector const& ptrD, void* gemmParamsWorkspace, int64_t gemmParamsWorkSpaceSize, void* gemmWorkSpace, int64_t gemmWorkspaceSize, - bool isLoraIn, nvinfer1::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream); + bool isLoraIn, tensorrt_llm::DataType dataType, int splitKSlices, int minKN, cudaStream_t stream); } // namespace kernels diff --git a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h index 302c278f7a2a..1d8cdb038145 100644 --- a/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h +++ b/cpp/tensorrt_llm/kernels/unfusedAttentionKernels.h @@ -242,42 +242,42 @@ struct QKVPreprocessingParams { ss << "seq_lens: " << *(runtime::ITensor::wrap( - (void*) seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + (void*) seq_lens, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); } if (cache_seq_lens && batch_size > 0) { ss << "cache_seq_lens: " << *(runtime::ITensor::wrap( - (void*) cache_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + (void*) cache_seq_lens, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); } if (encoder_seq_lens && batch_size > 0) { ss << "encoder_seq_lens: " << *(runtime::ITensor::wrap( - (void*) encoder_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + (void*) encoder_seq_lens, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); } if (cu_seq_lens && batch_size > 0) { ss << "cu_seq_lens: " << *(runtime::ITensor::wrap( - (void*) cu_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + (void*) cu_seq_lens, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); } if (cu_kv_seq_lens && batch_size > 0) { ss << "cu_kv_seq_lens: " << *(runtime::ITensor::wrap( - (void*) cu_kv_seq_lens, nvinfer1::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); + (void*) cu_kv_seq_lens, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size}))); } if (sparse_kv_offsets) { ss << "sparse_kv_offsets: " - << *(runtime::ITensor::wrap((void*) sparse_kv_offsets, nvinfer1::DataType::kINT32, + << *(runtime::ITensor::wrap((void*) sparse_kv_offsets, tensorrt_llm::DataType::kINT32, runtime::ITensor::makeShape({batch_size + 1}))); } if (rotary_embedding_inv_freq && batch_size > 0 && rotary_embedding_dim > 0) { ss << "rotary_embedding_inv_freq: " - << *(runtime::ITensor::wrap((void*) rotary_embedding_inv_freq, nvinfer1::DataType::kFLOAT, + << *(runtime::ITensor::wrap((void*) rotary_embedding_inv_freq, tensorrt_llm::DataType::kFLOAT, runtime::ITensor::makeShape({batch_size, rotary_embedding_dim / 2}))); } ss << "rotary_coef_cache_buffer: " << rotary_coef_cache_buffer << std::endl; diff --git a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp index 3e19f9ebe72a..4131abc5eacd 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp +++ b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.cpp @@ -80,13 +80,13 @@ namespace kernels::ub { void allreduce2_userbuff_inplace_launcher(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { allreduce2_userbuff_inplace_impl(handler, offset, elements, dataType, comm, stream); } int allgather2_userbuff_residual_launcher(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable) { return allgather2_userbuff_residual_impl( @@ -95,7 +95,7 @@ int allgather2_userbuff_residual_launcher(int const handler, size_t const offset int allreduce2_userbuff_rmsnorm_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return allreduce2_userbuff_rmsnorm_impl(handler, offset, out_handler, out_offset, elements, hidden_size, beta, gamma, eps, residual_in, residual_out, dataType, comm, stream); @@ -103,7 +103,7 @@ int allreduce2_userbuff_rmsnorm_launcher(int const handler, size_t const offset, int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return allreduce2_userbuff_inplace_rmsnorm_quant_impl(handler, offset, out_handler, out_offset, elements, @@ -113,7 +113,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(handler, offset, out_handler, out_offset, scale_handler, scale_offset, elements, hidden_size, beta, gamma, eps, scalefactor, residual_in, residual_out, dataType, comm, @@ -165,12 +165,12 @@ TRTLLM_NAMESPACE_BEGIN namespace kernels::ub { void allreduce2_userbuff_inplace_launcher(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { } int allgather2_userbuff_residual_launcher(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable) { return 0; @@ -178,7 +178,7 @@ int allgather2_userbuff_residual_launcher(int const handler, size_t const offset int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return 0; @@ -187,7 +187,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { return 0; } diff --git a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h index e8a48e2c680b..63ec755be52f 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h +++ b/cpp/tensorrt_llm/kernels/userbuffers/ub_interface.h @@ -40,24 +40,24 @@ namespace kernels::ub using ::tensorrt_llm::runtime::ub::communicator; void allreduce2_userbuff_inplace_launcher(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream = 0); + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream = 0); int allgather2_userbuff_residual_launcher(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable = false); int allreduce2_userbuff_rmsnorm_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); } // namespace kernels::ub TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu index 8cb5814e0398..08c8e9cc2b8f 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu +++ b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.cu @@ -1774,11 +1774,11 @@ int allgather2_userbuff_residual(int const handler, size_t const offset, size_t } void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1791,7 +1791,7 @@ void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, si break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1809,17 +1809,17 @@ void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, si } int allgather2_userbuff_residual_impl(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: return allgather2_userbuff_residual( handler, offset, elements, hidden_size, residual, comm, stream, force_enable); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: return allgather2_userbuff_residual<__nv_bfloat16>( handler, offset, elements, hidden_size, residual, comm, stream, force_enable); break; @@ -1830,11 +1830,11 @@ int allgather2_userbuff_residual_impl(int const handler, size_t const offset, si int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1849,7 +1849,7 @@ int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1870,12 +1870,12 @@ int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1890,7 +1890,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t con break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1914,11 +1914,11 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t con int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, void* residual_in, - void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream) + void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream) { switch (dataType) { - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: { if (kDISABLE_FP32_ACCUMULATION) { @@ -1935,7 +1935,7 @@ int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(int const handler, size_t break; } #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: { if (kDISABLE_FP32_ACCUMULATION) { diff --git a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h index 96f21b748282..d95a7e4f8499 100644 --- a/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h +++ b/cpp/tensorrt_llm/kernels/userbuffers/userbuffers.h @@ -120,25 +120,25 @@ namespace kernels::ub { using namespace ::tensorrt_llm::runtime::ub; void allreduce2_userbuff_inplace_impl(int const handler, size_t const offset, size_t const elements, - nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream = 0); + tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream = 0); // for TP-parallelism, only single node is implemented int allgather2_userbuff_residual_impl(int const handler, size_t const offset, size_t const elements, - int const hidden_size, void* residual, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream, + int const hidden_size, void* residual, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream, bool force_enable); int allreduce2_userbuff_rmsnorm_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, - float* scalefactor, void* residual_in, void* residual_out, nvinfer1::DataType dataType, communicator* comm, + float* scalefactor, void* residual_in, void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); int allreduce2_userbuff_inplace_rmsnorm_quant_fp4_impl(int const handler, size_t const offset, int const out_handler, size_t const out_offset, int const scale_handler, size_t const scale_offset, size_t const elements, int const hidden_size, void* beta, void* gamma, float eps, float* scalefactor, void* residual_in, - void* residual_out, nvinfer1::DataType dataType, communicator* comm, cudaStream_t stream); + void* residual_out, tensorrt_llm::DataType dataType, communicator* comm, cudaStream_t stream); } // namespace kernels::ub TRTLLM_NAMESPACE_END diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h index eb939b57c2db..7da7071a6c17 100644 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h @@ -24,7 +24,7 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" #include "tensorrt_llm/runtime/common.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include diff --git a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h index 616f9d25c2bf..3bc8e5bcf8fd 100644 --- a/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h +++ b/cpp/tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemmNVFP4.h @@ -24,7 +24,7 @@ #include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" #include "tensorrt_llm/runtime/common.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include diff --git a/cpp/tensorrt_llm/layers/decodingParams.h b/cpp/tensorrt_llm/layers/decodingParams.h index 1e77b8919ca1..426ed45de38c 100644 --- a/cpp/tensorrt_llm/layers/decodingParams.h +++ b/cpp/tensorrt_llm/layers/decodingParams.h @@ -194,7 +194,7 @@ class ExplicitDraftTokensSetupParams : public DecodingSetupParams // Hack to init some data for the context phase in the setup. TensorPtr randomDataSample; // [maxBatchSize], on gpu TensorPtr temperatures; // [maxBatchSize], on gpu - nvinfer1::DataType dtype; // [1] + tensorrt_llm::DataType dtype; // [1] }; class EagleSetupParams : public DecodingSetupParams @@ -204,7 +204,7 @@ class EagleSetupParams : public DecodingSetupParams // Hack to init some data for the context phase in the setup. TensorPtr randomDataSample; // [maxBatchSize], on gpu TensorPtr temperatures; // [maxBatchSize], on gpu - nvinfer1::DataType dtype; // [1] + tensorrt_llm::DataType dtype; // [1] }; class DynamicDecodeSetupParams : public BaseSetupParams diff --git a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp index e014ee4535e5..d7a789e58ae8 100644 --- a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp +++ b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.cpp @@ -93,15 +93,15 @@ void ExplicitDraftTokensLayer::setup(SizeType32 batchSize, SizeType32 beamWid batchSlots, getLimitsPenalty(DecodingPenaltyType::Temperature), "temperature penalty"); // Dispatch context buffer fill - if (mDecoderDtype == nvinfer1::DataType::kFLOAT) + if (mDecoderDtype == tensorrt_llm::DataType::kFLOAT) { fillContextBuffers(batchSize, batchSlots, *setupParams, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kHALF) + else if (mDecoderDtype == tensorrt_llm::DataType::kHALF) { fillContextBuffers(batchSize, batchSlots, *setupParams, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kBF16) + else if (mDecoderDtype == tensorrt_llm::DataType::kBF16) { fillContextBuffers<__nv_bfloat16>(batchSize, batchSlots, *setupParams, workspace); } @@ -126,15 +126,15 @@ void ExplicitDraftTokensLayer::forwardAsync(std::shared_ptr(*outputs, *inputs, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kHALF) + else if (mDecoderDtype == tensorrt_llm::DataType::kHALF) { splitInputDataToBatchSlots(*outputs, *inputs, workspace); } - else if (mDecoderDtype == nvinfer1::DataType::kBF16) + else if (mDecoderDtype == tensorrt_llm::DataType::kBF16) { splitInputDataToBatchSlots<__nv_bfloat16>(*outputs, *inputs, workspace); } diff --git a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h index 75883ded6e5a..66f68fa23193 100644 --- a/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h +++ b/cpp/tensorrt_llm/layers/explicitDraftTokensLayer.h @@ -83,7 +83,7 @@ class ExplicitDraftTokensLayer : public BaseLayer TensorPtr mTemperature; - std::optional mDecoderDtype{std::nullopt}; + std::optional mDecoderDtype{std::nullopt}; }; } // namespace tensorrt_llm::layers diff --git a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp b/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp index 09843fd7ce44..68ae1a727af7 100644 --- a/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp +++ b/cpp/tensorrt_llm/layers/lookaheadAlgorithm.cpp @@ -35,14 +35,14 @@ 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)}), nvinfer1::DataType::kINT32)) + runtime::ITensor::makeShape({(maxN <= 1 ? 0 : maxN - 2)}), tensorrt_llm::DataType::kINT32)) , mPastTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW * (maxN - 1)}), nvinfer1::DataType::kINT32)) - , mKeyTokensMax(runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxW}), nvinfer1::DataType::kINT32)) + 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}), nvinfer1::DataType::kINT32)) + runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxN * 2 - 1}), tensorrt_llm::DataType::kINT32)) , mGuessTokensMax( - runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxG * (maxN - 1)}), nvinfer1::DataType::kINT32)) + runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxG * (maxN - 1)}), tensorrt_llm::DataType::kINT32)) , mMaxW(maxW) , mMaxN(maxN) , mMaxG(maxG) @@ -52,12 +52,12 @@ LookaheadAlgorithm::LookaheadAlgorithm( std::tie(maxGeneratedLen, std::ignore, maxDraftLen, std::ignore) = executor::LookaheadDecodingConfig(maxW, maxN, maxG).calculateSpeculativeResource(); mAttentionMask = runtime::BufferManager::cpu( - runtime::ITensor::makeShape({maxDraftLen, maxDraftLen}), nvinfer1::DataType::kBOOL); + runtime::ITensor::makeShape({maxDraftLen, maxDraftLen}), tensorrt_llm::DataType::kBOOL); mDraftTokensMax - = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), nvinfer1::DataType::kINT32); + = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), tensorrt_llm::DataType::kINT32); mSampledTokensMax - = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxGeneratedLen}), nvinfer1::DataType::kINT32); - mEncodeMapMax = runtime::BufferManager::cpu(runtime::ITensor::makeShape({maxDraftLen}), nvinfer1::DataType::kINT32); + = 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) diff --git a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp b/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp index bf6e15080f3c..108dc8023f18 100644 --- a/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp +++ b/cpp/tensorrt_llm/layers/lookaheadDecodingLayer.cpp @@ -64,38 +64,38 @@ LookaheadDecodingLayer::CpuAlgorithmResources::CpuAlgorithmResources(DecoderD mPrompts.reserve(maxBatchSize); for (auto bi = 0; bi < maxBatchSize; bi++) { - mPrompts.emplace_back(BufferManager::cpu(ITensor::makeShape({0}), nvinfer1::DataType::kINT32)); + mPrompts.emplace_back(BufferManager::cpu(ITensor::makeShape({0}), tensorrt_llm::DataType::kINT32)); } auto const maxBatchShape1D = ITensor::makeShape({maxBatchSize}); - mBatchSlots = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + mBatchSlots = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); mTargetTokens - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), nvinfer1::DataType::kINT32); - mTokensPerStep = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); - mEndIds = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + = 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}), nvinfer1::DataType::kINT32); + mOutputIds = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxNumNewTokens}), tensorrt_llm::DataType::kINT32); mNewTokens = BufferManager::cpu( - ITensor::makeShape({maxTokensPerStep, maxBatchSize, beamWidth}), nvinfer1::DataType::kINT32); + ITensor::makeShape({maxTokensPerStep, maxBatchSize, beamWidth}), tensorrt_llm::DataType::kINT32); mPathsOffsets - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), nvinfer1::DataType::kINT32); + = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), tensorrt_llm::DataType::kINT32); mPathsOffsetsBatch - = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxAcceptedDraftLen}), nvinfer1::DataType::kINT32); - mNumNewTokens = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); - mNumNewTokensCumSum = BufferManager::cpu(ITensor::makeShape({maxBatchSize + 1}), nvinfer1::DataType::kINT32); - mNextDraftTokens = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), nvinfer1::DataType::kINT32); - mNextDraftPosIds = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxDraftLen}), nvinfer1::DataType::kINT32); - mGenerationLengths = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + = 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}), nvinfer1::DataType::kINT32); - mPositionIds = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep}), nvinfer1::DataType::kINT32); + = 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}), nvinfer1::DataType::kBOOL); + = BufferManager::cpu(ITensor::makeShape({maxTokensPerStep, maxTokensPerStep}), tensorrt_llm::DataType::kBOOL); mPackedMask = BufferManager::cpu(ITensor::makeShape({maxBatchSize, maxTokensPerStep, static_cast(divUp(maxTokensPerStep, 32))}), - nvinfer1::DataType::kINT32); - mNextDraftLengths = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); - mSequenceLengths = BufferManager::cpu(maxBatchShape1D, nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); + mNextDraftLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); + mSequenceLengths = BufferManager::cpu(maxBatchShape1D, tensorrt_llm::DataType::kINT32); } template @@ -117,12 +117,12 @@ LookaheadDecodingLayer::LookaheadDecodingLayer( auto const maxBatchShape2D = ITensor::makeShape({maxBatchSize, maxTokensPerStep}); mWorkspaceSize = getTopKWorkspaceSize(maxBatchSize, maxTokensPerStep, maxTopK, vocabSizePadded); - mTargetTokensDevice = mBufferManager->gpu(maxBatchShape2D, nvinfer1::DataType::kINT32); + mTargetTokensDevice = mBufferManager->gpu(maxBatchShape2D, tensorrt_llm::DataType::kINT32); mCurandStatesDevice - = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), nvinfer1::DataType::kINT8); + = mBufferManager->gpu(ITensor::makeShape({maxBatchSize, sizeof(curandState_t)}), tensorrt_llm::DataType::kINT8); mSetupWorkspaceSize = DecodingLayerWorkspace::calculateRequiredWorkspaceSize( - std::make_pair(maxBatchShape1D, nvinfer1::DataType::kINT64)); + std::make_pair(maxBatchShape1D, tensorrt_llm::DataType::kINT64)); TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } diff --git a/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h b/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h index 739cf65001ab..3731290953e7 100644 --- a/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h +++ b/cpp/tensorrt_llm/layers/lookaheadDecodingUtils.h @@ -318,12 +318,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return values(); - case nvinfer1::DataType::kFLOAT: return values(); - case nvinfer1::DataType::kINT8: return values(); - case nvinfer1::DataType::kINT32: return values(); - case nvinfer1::DataType::kINT64: return values(); - case nvinfer1::DataType::kUINT8: return values(); + 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"); } } @@ -376,12 +376,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return randomize(3); - case nvinfer1::DataType::kFLOAT: return randomize(3); - case nvinfer1::DataType::kINT8: return randomize(3); - case nvinfer1::DataType::kINT32: return randomize(3); - case nvinfer1::DataType::kINT64: return randomize(3); - case nvinfer1::DataType::kUINT8: return randomize(3); + 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; } } @@ -391,12 +391,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return randomize(0); - case nvinfer1::DataType::kFLOAT: return randomize(0); - case nvinfer1::DataType::kINT8: return randomize(0); - case nvinfer1::DataType::kINT32: return randomize(0); - case nvinfer1::DataType::kINT64: return randomize(0); - case nvinfer1::DataType::kUINT8: return randomize(0); + 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; } } @@ -405,12 +405,12 @@ class DebugTensor { switch (mTensor.getDataType()) { - case nvinfer1::DataType::kBOOL: return randomize(1); - case nvinfer1::DataType::kFLOAT: return randomize(1); - case nvinfer1::DataType::kINT8: return randomize(1); - case nvinfer1::DataType::kINT32: return randomize(1); - case nvinfer1::DataType::kINT64: return randomize(1); - case nvinfer1::DataType::kUINT8: return randomize(1); + 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; } } diff --git a/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp b/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp index 5954bc520ad0..b876d8cd8a4c 100644 --- a/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp +++ b/cpp/tensorrt_llm/layers/lookaheadPoolManager.cpp @@ -67,7 +67,7 @@ void LookaheadPoolManager::accept(TensorConstPtr const& prompt, SizeType32 level for (SizeType32 ti = 0; ti + level - 1 < length; ti++) { auto key = promptRange[ti]; - TensorPtr ngram = BufferManager::cpu(ITensor::makeShape({level - 1}), nvinfer1::DataType::kINT32); + 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()); @@ -107,7 +107,7 @@ void LookaheadPoolManager::update(TensorConstPtr const& keyTokens, TensorConstPt for (SizeType32 wi = 0; wi < window; wi++) { TensorConstPtr source = ITensor::at(ngramTokens, {wi}); - TensorPtr ngram = BufferManager::cpu(source->getShape(), nvinfer1::DataType::kINT32); + TensorPtr ngram = BufferManager::cpu(source->getShape(), tensorrt_llm::DataType::kINT32); BufferRange sourceRange(*source); BufferRange ngramRange(*ngram); std::copy(sourceRange.begin(), sourceRange.end(), ngramRange.begin()); diff --git a/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp b/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp index 40eff62c17d6..3870464b68f6 100644 --- a/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp +++ b/cpp/tensorrt_llm/layers/medusaDecodingLayer.cpp @@ -88,10 +88,10 @@ void MedusaDecodingLayer::allocateBuffer() mTiledBatchSlotsSetup = BufferManager::pinnedPool( ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mTiledBatchSlotsForward = BufferManager::pinnedPool( ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); mMedusaInputLogitsPtrs = BufferManager::pinnedPool( ITensor::makeShape({static_cast(mDecoderDomain.getBatchSize() * maxDraftPathLen)}), TRTDataType::value); diff --git a/cpp/tensorrt_llm/layers/penaltyLayer.cpp b/cpp/tensorrt_llm/layers/penaltyLayer.cpp index c6c57ca5034d..b79590073d1d 100644 --- a/cpp/tensorrt_llm/layers/penaltyLayer.cpp +++ b/cpp/tensorrt_llm/layers/penaltyLayer.cpp @@ -84,11 +84,11 @@ void PenaltyLayer::allocateWorkspace() auto const workspaceSize = mDecoderDomain.getBatchSize() * mDecoderDomain.getMaxDecodingTokens() * mConfiguredBeamWidth * mDecoderDomain.getVocabSize() * 2; - mPenaltyWorkspaceDevice = mBufferManager->gpu(workspaceSize, nvinfer1::DataType::kINT32); + mPenaltyWorkspaceDevice = mBufferManager->gpu(workspaceSize, tensorrt_llm::DataType::kINT32); if (mDecodingMode.isBeamSearch()) { - mPenaltyWorkspacePrevDevice = mBufferManager->gpu(workspaceSize, nvinfer1::DataType::kINT32); + mPenaltyWorkspacePrevDevice = mBufferManager->gpu(workspaceSize, tensorrt_llm::DataType::kINT32); } } @@ -111,27 +111,27 @@ void PenaltyLayer::allocateBuffer() if (mDecodingMode.isUseTemperature()) { - mTemperatureDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mTemperatureDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUseRepetitionPenalty()) { - mRepetitionPenaltyDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mRepetitionPenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUsePresencePenalty()) { - mPresencePenaltyDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mPresencePenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUseFrequencyPenalty()) { - mFrequencyPenaltyDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kFLOAT); + mFrequencyPenaltyDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kFLOAT); } if (mDecodingMode.isUseMinLength()) { - mMinLengthDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kINT32); + mMinLengthDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kINT32); } if (mDecodingMode.isUseOccurrencePenalty()) { - mPromptIgnoreLengthDevice = mBufferManager->gpu(batchSizeShape, nvinfer1::DataType::kINT32); + mPromptIgnoreLengthDevice = mBufferManager->gpu(batchSizeShape, tensorrt_llm::DataType::kINT32); } auto const logitsPtrDeviceDesc = std::make_pair(batchSizeShape, TRTDataType::value); diff --git a/cpp/tensorrt_llm/nanobind/CMakeLists.txt b/cpp/tensorrt_llm/nanobind/CMakeLists.txt index b523ae193871..729813f839ab 100755 --- a/cpp/tensorrt_llm/nanobind/CMakeLists.txt +++ b/cpp/tensorrt_llm/nanobind/CMakeLists.txt @@ -14,7 +14,6 @@ set(SRCS batch_manager/llmRequest.cpp common/tllmExceptions.cpp executor/bindings.cpp - executor/executor.cpp executor/executorConfig.cpp executor/request.cpp process_group/bindings.cpp diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp index 4ed8b1d1f0d4..5f4794e7916f 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/algorithms.cpp @@ -23,7 +23,6 @@ #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/logitsPostProcessor.h" #include "tensorrt_llm/batch_manager/medusaBuffers.h" #include "tensorrt_llm/batch_manager/microBatchScheduler.h" #include "tensorrt_llm/batch_manager/pauseRequests.h" @@ -127,12 +126,6 @@ void tensorrt_llm::nanobind::batch_manager::algorithms::initBindings(nb::module_ nb::call_guard()) .def("name", [](AllocateKvCache const&) { return AllocateKvCache::name; }); - nb::class_(m, LogitsPostProcessor::name) - .def(nb::init<>()) - .def("__call__", &LogitsPostProcessor::operator(), nb::arg("decoder_input_buffers"), - nb::arg("replicate_logits_post_processor"), nb::arg("world_config"), nb::arg("stream"), - nb::arg("logits_post_processor_batched") = std::nullopt) - .def("name", [](LogitsPostProcessor const&) { return LogitsPostProcessor::name; }); nb::class_(m, CreateNewDecoderRequests::name) .def(nb::init(), nb::arg("speculative_decoding_fast_logits"), @@ -141,7 +134,7 @@ void tensorrt_llm::nanobind::batch_manager::algorithms::initBindings(nb::module_ "__call__", [](CreateNewDecoderRequests& self, tr::ModelConfig const& modelConfig, tr::WorldConfig const& worldConfig, executor::DecodingConfig const& decodingConfig, RequestVector const& contextRequests, - nvinfer1::DataType logitsType, DecoderInputBuffers& inputBuffers, + 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) diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp index ee257713b85d..1a9dfb6f0123 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp @@ -483,7 +483,7 @@ void initBindings(nb::module_& m) nb::arg("max_num_sequences"), nb::arg("model_config"), nb::arg("world_config"), nb::arg("buffer_manager"), nb::call_guard()) .def(nb::init const&, tr::SizeType32>(), nb::arg("d_state"), nb::arg("d_conv"), nb::arg("num_heads"), nb::arg("n_groups"), nb::arg("head_dim"), nb::arg("max_batch_size"), nb::arg("world_config"), nb::arg("stream"), nb::arg("dtype"), diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp index 5257c7adb1c9..c0f92328ede1 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/cacheTransceiver.cpp @@ -118,7 +118,7 @@ void tb::CacheTransceiverBindings::initBindings(nb::module_& m) nb::class_(m, "CacheTransceiver") .def(nb::init, SizeType32, SizeType32, - runtime::WorldConfig, std::vector, nvinfer1::DataType, + runtime::WorldConfig, std::vector, tensorrt_llm::DataType, executor::kv_cache::CacheState::AttentionType, std::optional, tb::rnn_state_manager::RnnStateManager*, std::vector>(), nb::arg("cache_manager"), nb::arg("num_kv_heads_per_layer"), nb::arg("size_per_head"), diff --git a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp index 12b29d4981e2..5f22bc95c4fb 100644 --- a/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp +++ b/cpp/tensorrt_llm/nanobind/batch_manager/kvCacheManager.cpp @@ -348,7 +348,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::class_(m, "PoolConfiguration") .def(nb::init<>()) - .def(nb::init(), nb::arg("window_size"), nb::arg("size_per_head"), + .def(nb::init(), nb::arg("window_size"), nb::arg("size_per_head"), nb::arg("dtype")) .def_rw("window_size", &tbk::PoolConfiguration::windowSize) .def_rw("size_per_head", &tbk::PoolConfiguration::sizePerHead) @@ -651,7 +651,7 @@ void tb::kv_cache_manager::KVCacheManagerBindings::initBindings(nb::module_& m) nb::class_(m, "KVCacheManager") .def(nb::init const&, SizeType32, SizeType32, std::map> const&, SizeType32, SizeType32, - std::vector const&, nvinfer1::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool, + std::vector const&, tensorrt_llm::DataType, SizeType32, int64_t, SizeType32, SizeType32, bool, tbk::CacheType, std::optional, std::shared_ptr, bool, bool, std::shared_ptr, bool, SizeType32, SizeType32, bool, std::optional, diff --git a/cpp/tensorrt_llm/nanobind/bindings.cpp b/cpp/tensorrt_llm/nanobind/bindings.cpp index db263fa639f9..cd47fec28f83 100644 --- a/cpp/tensorrt_llm/nanobind/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/bindings.cpp @@ -168,17 +168,17 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def_rw("host_cache_size", &tb::PeftCacheManagerConfig::hostCacheSize) .def_rw("lora_prefetch_dir", &tb::PeftCacheManagerConfig::loraPrefetchDir); - nb::enum_(m, "DataType") - .value("FLOAT", nvinfer1::DataType::kFLOAT) - .value("HALF", nvinfer1::DataType::kHALF) - .value("INT8", nvinfer1::DataType::kINT8) - .value("INT32", nvinfer1::DataType::kINT32) - .value("BOOL", nvinfer1::DataType::kBOOL) - .value("UINT8", nvinfer1::DataType::kUINT8) - .value("FP8", nvinfer1::DataType::kFP8) - .value("BF16", nvinfer1::DataType::kBF16) - .value("INT64", nvinfer1::DataType::kINT64) - .value("NVFP4", nvinfer1::DataType::kFP4) + nb::enum_(m, "DataType") + .value("FLOAT", tensorrt_llm::DataType::kFLOAT) + .value("HALF", tensorrt_llm::DataType::kHALF) + .value("INT8", tensorrt_llm::DataType::kINT8) + .value("INT32", tensorrt_llm::DataType::kINT32) + .value("BOOL", tensorrt_llm::DataType::kBOOL) + .value("UINT8", tensorrt_llm::DataType::kUINT8) + .value("FP8", tensorrt_llm::DataType::kFP8) + .value("BF16", tensorrt_llm::DataType::kBF16) + .value("INT64", tensorrt_llm::DataType::kINT64) + .value("NVFP4", tensorrt_llm::DataType::kFP4) .export_values(); nb::enum_(m, "GptModelVariant") @@ -295,7 +295,7 @@ NB_MODULE(TRTLLM_NB_MODULE, m) .def(nb::self != nb::self); nb::class_(m, "ModelConfig") - .def(nb::init(), + .def(nb::init(), nb::arg("vocab_size"), nb::arg("num_layers"), nb::arg("num_attention_layers"), nb::arg("num_rnn_layers"), nb::arg("num_heads"), nb::arg("hidden_size"), nb::arg("data_type")) .def_prop_ro("vocab_size", &tr::ModelConfig::getVocabSize) diff --git a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp index b0ad31b7347e..a8d2301fa43d 100644 --- a/cpp/tensorrt_llm/nanobind/executor/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/executor/bindings.cpp @@ -16,7 +16,6 @@ */ #include "bindings.h" -#include "executor.h" #include "executorConfig.h" #include "request.h" #include "tensorrt_llm/executor/executor.h" @@ -287,7 +286,6 @@ void initBindings(nb::module_& m) tensorrt_llm::nanobind::executor::initRequestBindings(m); tensorrt_llm::nanobind::executor::initConfigBindings(m); - tensorrt_llm::nanobind::executor::Executor::initBindings(m); } } // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executor.cpp b/cpp/tensorrt_llm/nanobind/executor/executor.cpp deleted file mode 100644 index 34cc8182d1bb..000000000000 --- a/cpp/tensorrt_llm/nanobind/executor/executor.cpp +++ /dev/null @@ -1,225 +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. - */ - -#include "executor.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/executor/tensor.h" -#include "tensorrt_llm/nanobind/common/customCasters.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace nb = nanobind; -namespace tle = tensorrt_llm::executor; - -namespace nanobind::detail -{ - -template <> -struct dtype_traits -{ - static constexpr dlpack::dtype value{ - (uint8_t) dlpack::dtype_code::Float, // type code - 16, // size in bits - 1 // lanes (simd), usually set to 1 - }; - static constexpr auto name = const_name("float16"); -}; -} // namespace nanobind::detail - -namespace -{ -tle::Tensor numpyToTensor(nb::object const& object) -{ - std::string dtype_name = nb::cast(object.attr("dtype").attr("name")); - nb::object metadata = object.attr("dtype").attr("metadata"); - - tle::DataType dtype; - if (dtype_name == "float16") - { - dtype = tle::DataType::kFP16; - } - else if (dtype_name == "float32") - { - dtype = tle::DataType::kFP32; - } - else if (dtype_name == "int8") - { - dtype = tle::DataType::kINT8; - } - else if (dtype_name == "int32") - { - dtype = tle::DataType::kINT32; - } - else if (dtype_name == "int64") - { - dtype = tle::DataType::kINT64; - } - else if (dtype_name == "void8" && !metadata.is_none() && nb::cast(metadata["dtype"]) == "float8") - { - dtype = tle::DataType::kFP8; - } - else if (dtype_name == "void16" && !metadata.is_none() && nb::cast(metadata["dtype"]) == "bfloat16") - { - dtype = tle::DataType::kBF16; - } - else - { - TLLM_THROW("Unsupported numpy dtype."); - } - - nb::object array_interface = object.attr("__array_interface__"); - nb::object shape_obj = array_interface["shape"]; - std::vector dims; - dims.reserve(nb::len(shape_obj)); - - for (size_t i = 0; i < nb::len(shape_obj); ++i) - { - dims.push_back(nb::cast(shape_obj[i])); - } - - nb::object data_obj = array_interface["data"]; - uintptr_t addr = nb::cast(data_obj[0]); - void* data_ptr = reinterpret_cast(addr); - tle::Shape shape(dims.data(), dims.size()); - return tle::Tensor::of(dtype, data_ptr, shape); -} - -} // namespace - -namespace tensorrt_llm::nanobind::executor -{ - -Executor::Executor( - std::filesystem::path const& modelPath, tle::ModelType modelType, tle::ExecutorConfig const& executorConfig) -{ - mExecutor = std::make_unique(modelPath, modelType, executorConfig); -} - -Executor::Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, - tle::ModelType modelType, tle::ExecutorConfig const& executorConfig) -{ - mExecutor = std::make_unique(encoderModelPath, decoderModelPath, modelType, executorConfig); -} - -Executor::Executor(nb::bytes const& engineBuffer, std::string const& jsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig, std::optional managedWeights) -{ - uint8_t const* data = static_cast(engineBuffer.data()); - size_t size = engineBuffer.size(); - std::optional> managedWeightsMap = std::nullopt; - if (managedWeights.has_value() && !managedWeights.value().empty()) - { - managedWeightsMap = std::map(); - for (auto const& [rawName, rawArray] : managedWeights.value()) - { - std::string name = nb::cast(rawName); - nb::object array_obj = nb::cast(rawArray); - managedWeightsMap->emplace(name, numpyToTensor(array_obj)); - } - } - mExecutor = std::make_unique( - tle::BufferView(data, size), jsonConfigStr, modelType, executorConfig, managedWeightsMap); -} - -Executor::Executor(std::string const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, - std::string const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig) -{ - uint8_t const* encoderData = reinterpret_cast(encoderEngineBuffer.data()); - size_t encoderSize = encoderEngineBuffer.size(); - uint8_t const* decoderData = reinterpret_cast(decoderEngineBuffer.data()); - size_t decoderSize = decoderEngineBuffer.size(); - mExecutor = std::make_unique(tle::BufferView(encoderData, encoderSize), encoderJsonConfigStr, - tle::BufferView(decoderData, decoderSize), decoderJsonConfigStr, modelType, executorConfig); -} - -nb::object Executor::enter() -{ - TLLM_CHECK(static_cast(mExecutor)); - return nb::cast(this); -} - -void Executor::exit( - [[maybe_unused]] nb::handle type, [[maybe_unused]] nb::handle value, [[maybe_unused]] nb::handle traceback) -{ - shutdown(); - mExecutor = nullptr; -} - -void Executor::shutdown() -{ - // NOTE: we must release the GIL here. Executor has spawned a thread for the execution loop. That thread must be - // able to do forward progress for the shutdown process to succeed. It takes the GIL during its callbacks, so - // we release it now. Note that we shouldn't do anything related to python objects after that. - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - nb::gil_scoped_release release; - mExecutor->shutdown(); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void Executor::initBindings(nb::module_& m) -{ - nb::class_(m, "Executor") - .def(nb::init(), - nb::arg("model_path"), nb::arg("model_type"), nb::arg("executor_config")) - .def(nb::init(), - nb::arg("encoder_model_path"), nb::arg("decoder_model_path"), nb::arg("model_type"), - nb::arg("executor_config")) - .def(nb::init(), - nb::arg("engine_buffer"), nb::arg("json_config_str"), nb::arg("model_type"), nb::arg("executor_config"), - nb::arg("managed_weights") = nb::dict()) - .def(nb::init(), - nb::arg("encoder_engine_buffer"), nb::arg("encoder_json_config_str"), nb::arg("decoder_engine_buffer"), - nb::arg("decoder_json_config_str"), nb::arg("model_type"), nb::arg("executor_config")) - .def("shutdown", &Executor::shutdown) - .def("__enter__", &Executor::enter) - .def("__exit__", &Executor::exit, nb::arg("type").none(), nb::arg("value").none(), nb::arg("traceback").none()) - .def("enqueue_request", &Executor::enqueueRequest, nb::arg("request")) - .def("enqueue_requests", &Executor::enqueueRequests, nb::arg("requests")) - .def("await_responses", - nb::overload_cast const&>(&Executor::awaitResponses), - nb::arg("timeout") = nb::none()) - .def("await_responses", - nb::overload_cast const&>( - &Executor::awaitResponses), - nb::arg("id"), nb::arg("timeout") = nb::none()) - .def("await_responses", - nb::overload_cast const&, std::optional const&>( - &Executor::awaitResponses), - nb::arg("ids"), nb::arg("timeout") = nb::none()) - .def("get_num_responses_ready", &Executor::getNumResponsesReady, nb::arg("id") = nb::none()) - .def("cancel_request", &Executor::cancelRequest, nb::arg("id") = nb::none()) - .def("get_latest_iteration_stats", &Executor::getLatestIterationStats) - .def("get_latest_request_stats", &Executor::getLatestRequestStats) - .def("get_latest_debug_tensors", &Executor::getLatestDebugTensors) - .def("can_enqueue_requests", &Executor::canEnqueueRequests) - .def("get_kv_cache_event_manager", &Executor::getKVCacheEventManager); -} - -} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/executor/executor.h b/cpp/tensorrt_llm/nanobind/executor/executor.h deleted file mode 100644 index 22c24abb4bfd..000000000000 --- a/cpp/tensorrt_llm/nanobind/executor/executor.h +++ /dev/null @@ -1,129 +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/executor/executor.h" -#include "tensorrt_llm/executor/types.h" -#include - -namespace nb = nanobind; -namespace tle = tensorrt_llm::executor; - -namespace tensorrt_llm::nanobind::executor -{ - -class Executor -{ -public: - Executor( - std::filesystem::path const& modelPath, tle::ModelType modelType, tle::ExecutorConfig const& executorConfig); - - Executor(std::filesystem::path const& encoderModelPath, std::filesystem::path const& decoderModelPath, - tle::ModelType modelType, tle::ExecutorConfig const& executorConfig); - - Executor(nb::bytes const& engineBuffer, std::string const& jsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig, std::optional managedWeights); - - Executor(std::string const& encoderEngineBuffer, std::string const& encoderJsonConfigStr, - std::string const& decoderEngineBuffer, std::string const& decoderJsonConfigStr, tle::ModelType modelType, - tle::ExecutorConfig const& executorConfig); - - nb::object enter(); - void exit( - [[maybe_unused]] nb::handle type, [[maybe_unused]] nb::handle value, [[maybe_unused]] nb::handle traceback); - void shutdown(); - - [[nodiscard]] tle::IdType enqueueRequest(tle::Request const& request) - { - return mExecutor->enqueueRequest(request); - } - - [[nodiscard]] std::vector enqueueRequests(std::vector const& requests) - { - return mExecutor->enqueueRequests(requests); - } - - [[nodiscard]] std::vector awaitResponses( - std::optional const& timeout = std::nullopt) - { - // Await responses blocks until a response is received. Release GIL so that it can be ran in a background - // thread. - nb::gil_scoped_release release; - return mExecutor->awaitResponses(timeout); - } - - [[nodiscard]] std::vector awaitResponses( - tle::IdType const& requestId, std::optional const& timeout = std::nullopt) - { - // Await responses blocks until a response is received. Release GIL so that it can be ran in a background - // thread. - nb::gil_scoped_release release; - return mExecutor->awaitResponses(requestId, timeout); - } - - [[nodiscard]] std::vector> awaitResponses(std::vector const& requestIds, - std::optional const& timeout = std::nullopt) - { - // Await responses blocks until a response is received. Release GIL so that it can be ran in a background - // thread. - nb::gil_scoped_release release; - return mExecutor->awaitResponses(requestIds, timeout); - } - - [[nodiscard]] tle::SizeType32 getNumResponsesReady(std::optional const& requestId = std::nullopt) const - { - return mExecutor->getNumResponsesReady(requestId); - } - - void cancelRequest(tle::IdType requestId) - { - mExecutor->cancelRequest(requestId); - } - - std::deque getLatestIterationStats() - { - return mExecutor->getLatestIterationStats(); - } - - std::deque getLatestRequestStats() - { - return mExecutor->getLatestRequestStats(); - } - - std::deque getLatestDebugTensors() - { - return mExecutor->getLatestDebugTensors(); - } - - [[nodiscard]] bool canEnqueueRequests() const - { - return mExecutor->canEnqueueRequests(); - } - - [[nodiscard]] std::optional> getKVCacheEventManager() const - { - return mExecutor->getKVCacheEventManager(); - } - - static void initBindings(nb::module_& m); - -private: - std::unique_ptr mExecutor; -}; - -} // namespace tensorrt_llm::nanobind::executor diff --git a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp index 6d5d70aafb6b..52f0b1af6d16 100644 --- a/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp +++ b/cpp/tensorrt_llm/nanobind/runtime/bindings.cpp @@ -39,7 +39,6 @@ #include "tensorrt_llm/runtime/loraCache.h" #include "tensorrt_llm/runtime/mcastGPUBuffer.h" #include "tensorrt_llm/runtime/speculativeDecodingMode.h" -#include "tensorrt_llm/runtime/tllmRuntime.h" #include "tensorrt_llm/runtime/torchView.h" #include "tensorrt_llm/runtime/virtualMemory.h" @@ -68,7 +67,7 @@ class PyIGptDecoder : public tr::IGptDecoder 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 explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) override { @@ -125,46 +124,8 @@ void initBindings(nb::module_& m) .def("materialize_with_tag", &tr::CudaVirtualMemoryManager::materializeWithTag, nb::arg("tag"), nb::call_guard()); - nb::class_(m, "TllmRuntime") - .def( - "__init__", - [](tr::TllmRuntime* self, std::filesystem::path engine_path, float gpu_weights_percent = 1.0f, - bool use_shape_inference = true) - { - // Using default logger by passing nullptr - new (self) - tr::TllmRuntime(tr::RawEngine(engine_path), nullptr, gpu_weights_percent, use_shape_inference); - }, - nb::arg("engine_path"), nb::arg("gpu_weights_percent") = 1.0f, nb::arg("use_shape_inference") = true) - .def( - "__init__", - [](tr::TllmRuntime* self, nb::ndarray engine_buffer, float gpu_weights_percent = 1.0f, - bool use_shape_inference = true) - { - if (engine_buffer.ndim() != 1) - throw std::runtime_error("Expected 1-D array for engine buffer"); - new (self) tr::TllmRuntime(tr::RawEngine(engine_buffer.data(), engine_buffer.size()), nullptr, - gpu_weights_percent, use_shape_inference); - }, - nb::arg("engine_buffer"), nb::arg("gpu_weights_percent") = 1.0f, nb::arg("use_shape_inference") = true) - .def_prop_ro("num_contexts", &tr::TllmRuntime::getNbContexts) - .def_prop_ro("num_profiles", &tr::TllmRuntime::getNbProfiles) - .def("get_opt_profile_id", &tr::TllmRuntime::getOptProfileId, nb::arg("num_tokens"), nb::arg("split_points"), - nb::call_guard()) - .def("clear_contexts", &tr::TllmRuntime::clearContexts, nb::call_guard()) - .def("execute_context", &tr::TllmRuntime::executeContext, nb::arg("context_id"), - nb::call_guard()) - .def_prop_ro("stream_ptr", &tr::TllmRuntime::getStreamPtr) - .def_prop_ro("buffer_manager", - static_cast(&tr::TllmRuntime::getBufferManager)) - .def("set_layer_profiler", &tr::TllmRuntime::setLayerProfiler, nb::call_guard()) - .def("has_layer_profiler", &tr::TllmRuntime::hasLayerProfiler, nb::arg("context_id"), - nb::call_guard()) - .def_prop_ro("layer_profiler_info", &tr::TllmRuntime::getLayerProfileInfo) - .def("report_to_profiler", &tr::TllmRuntime::reportToProfiler, nb::arg("context_id"), - nb::call_guard()) - .def_prop_ro("logits_dtype_from_engine", - [](tr::TllmRuntime& self) { return self.getEngine().getTensorDataType("logits"); }); + // NOTE: The TensorRT-engine runtime (TllmRuntime) binding was removed along with the + // legacy TensorRT backend; the PyTorch backend does not use it. nb::class_(m, "LookaheadDecodingBuffers") .def(nb::init(), nb::arg("max_num_sequences"), @@ -204,7 +165,7 @@ void initBindings(nb::module_& m) "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 explicitDraftTokensDType = std::nullopt, std::optional> const& lookaheadPrompt = std::nullopt, std::optional> const& lookaheadAlgoConfigs = std::nullopt) { diff --git a/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp b/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp index caef94c5defd..2d2988392ae5 100644 --- a/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp +++ b/cpp/tensorrt_llm/nanobind/testing/modelSpecBinding.cpp @@ -44,7 +44,7 @@ void initBindings(nb::module_& m) .value("CUM_LOG_PROBS", OutputContentType::kCUM_LOG_PROBS, "Cumulative Log"); nb::class_(m, "ModelSpec") - .def(nb::init()) + .def(nb::init()) .def("use_gpt_plugin", &ModelSpec::useGptAttentionPlugin, nb::rv_policy::reference_internal) .def("use_packed_input", &ModelSpec::usePackedInput, nb::rv_policy::reference_internal) .def("set_kv_cache_type", &ModelSpec::setKVCacheType, nb::rv_policy::reference_internal) diff --git a/cpp/tensorrt_llm/plugins/CMakeLists.txt b/cpp/tensorrt_llm/plugins/CMakeLists.txt deleted file mode 100755 index 8b89cccdc813..000000000000 --- a/cpp/tensorrt_llm/plugins/CMakeLists.txt +++ /dev/null @@ -1,183 +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. -# - -set(PLUGIN_TARGET_NAME nvinfer_plugin_tensorrt_llm) -set(PLUGIN_SHARED_TARGET ${PLUGIN_TARGET_NAME}) - -set(TARGET_DIR ${CMAKE_CURRENT_SOURCE_DIR}) -set(PLUGIN_EXPORT_MAP ${TARGET_DIR}/exports.map) # Linux -set(PLUGIN_EXPORT_DEF ${TARGET_DIR}/exports.def) # Windows - -if(${CMAKE_BUILD_TYPE} MATCHES "Debug") - set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g") -endif() - -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --Wno-deprecated-declarations") -set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --diag-suppress 997") - -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 -Wno-overloaded-virtual") - 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() - -set(PLUGIN_SOURCES) -set(PLUGIN_CU_SOURCES) - -set(PLUGIN_LISTS - bertAttentionPlugin - cpSplitPlugin - fusedLayernormPlugin - gptAttentionCommon - gptAttentionPlugin - identityPlugin - gemmPlugin - gemmSwigluPlugin - fp8RowwiseGemmPlugin - smoothQuantGemmPlugin - fp4GemmPlugin - quantizePerTokenPlugin - quantizeTensorPlugin - quantizeToFP4Plugin - layernormQuantizationPlugin - rmsnormQuantizationPlugin - weightOnlyGroupwiseQuantMatmulPlugin - weightOnlyQuantMatmulPlugin - lookupPlugin - loraPlugin - doraPlugin - mixtureOfExperts - selectiveScanPlugin - mambaConv1dPlugin - lruPlugin - cumsumLastDimPlugin - topkLastDimPlugin - lowLatencyGemmPlugin - eaglePlugin - lowLatencyGemmSwigluPlugin - qserveGemmPlugin - cudaStreamPlugin - gemmAllReducePlugin) - -foreach(PLUGIN_ITER ${PLUGIN_LISTS}) - include_directories(${PLUGIN_ITER}) - add_subdirectory(${PLUGIN_ITER}) -endforeach(PLUGIN_ITER) - -if(ENABLE_MULTI_DEVICE) - include_directories(ncclPlugin) - add_subdirectory(ncclPlugin) -endif() -include_directories(common) -add_subdirectory(common) - -# Set gencodes -list(APPEND PLUGIN_SOURCES "${PLUGIN_CU_SOURCES}") - -list(APPEND PLUGIN_SOURCES "${CMAKE_CURRENT_SOURCE_DIR}/api/tllmPlugin.cpp") - -# ################################# SHARED LIBRARY -# ############################################################################## - -if(WIN32) - set(CMAKE_WINDOWS_EXPORT_ALL_SYMBOLS 1) -endif() - -add_library(${PLUGIN_SHARED_TARGET} SHARED ${PLUGIN_SOURCES}) -add_cuda_architectures(${PLUGIN_SHARED_TARGET} 89) - -target_include_directories( - ${PLUGIN_SHARED_TARGET} - PUBLIC ${CUDA_INSTALL_DIR}/include - PUBLIC - $ - PRIVATE ${TARGET_DIR}) - -if(USING_OSS_CUTLASS_FP4_GEMM) - target_compile_definitions(${PLUGIN_SHARED_TARGET} - PUBLIC USING_OSS_CUTLASS_FP4_GEMM) -endif() - -if(USING_OSS_CUTLASS_ALLREDUCE_GEMM) - target_compile_definitions(${PLUGIN_SHARED_TARGET} - PUBLIC USING_OSS_CUTLASS_ALLREDUCE_GEMM) -endif() - -if(USING_OSS_CUTLASS_MOE_GEMM) - target_compile_definitions(${PLUGIN_SHARED_TARGET} - PUBLIC USING_OSS_CUTLASS_MOE_GEMM) -endif() - -if(ENABLE_MULTI_DEVICE) - target_include_directories(${PLUGIN_SHARED_TARGET} - PUBLIC ${MPI_C_INCLUDE_DIRS}) -endif() - -if(CUDA_VERSION VERSION_LESS 11.0) - target_include_directories(${PLUGIN_SHARED_TARGET} PUBLIC ${CUB_ROOT_DIR}) -endif() - -set_target_properties( - ${PLUGIN_SHARED_TARGET} - PROPERTIES CXX_STANDARD "17" - CXX_STANDARD_REQUIRED "YES" - CXX_EXTENSIONS "NO" - ARCHIVE_OUTPUT_DIRECTORY "${TRT_OUT_DIR}" - LIBRARY_OUTPUT_DIRECTORY "${TRT_OUT_DIR}" - RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}") - -if(WIN32) - set_target_properties( - ${PLUGIN_SHARED_TARGET} - PROPERTIES LINK_FLAGS "/DEF:${PLUGIN_EXPORT_DEF} ${UNDEFINED_FLAG}") -else() - set_target_properties( - ${PLUGIN_SHARED_TARGET} - PROPERTIES - LINK_FLAGS - "-Wl,--exclude-libs,ALL -Wl,--version-script=${PLUGIN_EXPORT_MAP} -Wl,-rpath,'$ORIGIN' ${AS_NEEDED_FLAG} ${UNDEFINED_FLAG}" - ) -endif() - -set_property(TARGET ${PLUGIN_SHARED_TARGET} PROPERTY CUDA_STANDARD 17) - -target_link_libraries( - ${PLUGIN_SHARED_TARGET} - ${CUBLAS_LIB} - ${CUBLASLT_LIB} - ${TRT_LIB} - ${CUDA_DRV_LIB} - ${CUDA_RT_LIB} - ${CMAKE_DL_LIBS} - ${SHARED_TARGET}) - -if(WIN32) - target_link_libraries(${PLUGIN_SHARED_TARGET} context_attention_src) -endif() - -if(ENABLE_MULTI_DEVICE) - target_link_libraries(${PLUGIN_SHARED_TARGET} ${MPI_C_LIBRARIES} ${NCCL_LIB}) -endif() diff --git a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp b/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp deleted file mode 100644 index f0dceb2f4a99..000000000000 --- a/cpp/tensorrt_llm/plugins/api/tllmPlugin.cpp +++ /dev/null @@ -1,313 +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. - */ -#include "tensorrt_llm/plugins/api/tllmPlugin.h" - -#include "tensorrt_llm/common/stringUtils.h" -#include "tensorrt_llm/runtime/tllmLogger.h" - -#include "tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h" -#include "tensorrt_llm/plugins/doraPlugin/doraPlugin.h" -#include "tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h" -#include "tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" -#include "tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h" -#include "tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h" -#include "tensorrt_llm/plugins/identityPlugin/identityPlugin.h" -#include "tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h" -#include "tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h" -#include "tensorrt_llm/plugins/loraPlugin/loraPlugin.h" -#include "tensorrt_llm/plugins/lruPlugin/lruPlugin.h" -#include "tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h" -#include "tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h" -#include "tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h" -#if ENABLE_MULTI_DEVICE -#include "tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h" -#include "tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/recvPlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h" -#include "tensorrt_llm/plugins/ncclPlugin/sendPlugin.h" -#endif // ENABLE_MULTI_DEVICE -#include "tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h" -#include "tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h" -#include "tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h" -#include "tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h" -#include "tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h" -#include "tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h" -#include "tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h" -#include "tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h" -#include "tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h" -#include "tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h" -#include "tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h" -#include "tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h" -#include "tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h" -#include "tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h" -#include "tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h" - -#include -#include - -#include - -namespace tc = tensorrt_llm::common; - -namespace -{ - -nvinfer1::IPluginCreator* creatorPtr(nvinfer1::IPluginCreator& creator) -{ - return &creator; -} - -nvinfer1::IPluginCreatorInterface* creatorInterfacePtr(nvinfer1::IPluginCreatorInterface& creator) -{ - return &creator; -} - -auto tllmLogger = tensorrt_llm::runtime::TllmLogger(); - -nvinfer1::ILogger* gLogger{&tllmLogger}; - -class GlobalLoggerFinder : public nvinfer1::ILoggerFinder -{ -public: - nvinfer1::ILogger* findLogger() override - { - return gLogger; - } -}; - -GlobalLoggerFinder gGlobalLoggerFinder{}; - -#if !defined(_MSC_VER) -[[maybe_unused]] __attribute__((constructor)) -#endif -void initOnLoad() -{ - auto constexpr kLoadPlugins = "TRT_LLM_LOAD_PLUGINS"; - auto const loadPlugins = std::getenv(kLoadPlugins); - if (loadPlugins && loadPlugins[0] == '1') - { - initTrtLlmPlugins(gLogger); - } -} - -bool pluginsInitialized = false; - -} // namespace - -namespace tensorrt_llm::plugins::api -{ - -LoggerManager& tensorrt_llm::plugins::api::LoggerManager::getInstance() noexcept -{ - static LoggerManager instance; - return instance; -} - -void LoggerManager::setLoggerFinder(nvinfer1::ILoggerFinder* finder) -{ - std::lock_guard lk(mMutex); - if (mLoggerFinder == nullptr && finder != nullptr) - { - mLoggerFinder = finder; - } -} - -[[maybe_unused]] nvinfer1::ILogger* LoggerManager::logger() -{ - std::lock_guard lk(mMutex); - if (mLoggerFinder != nullptr) - { - return mLoggerFinder->findLogger(); - } - return nullptr; -} - -nvinfer1::ILogger* LoggerManager::defaultLogger() noexcept -{ - return gLogger; -} -} // namespace tensorrt_llm::plugins::api - -// New Plugin APIs - -extern "C" -{ - bool initTrtLlmPlugins(void* logger, char const* libNamespace) - { - if (pluginsInitialized) - { - return true; - } - - if (logger) - { - gLogger = static_cast(logger); - } - setLoggerFinder(&gGlobalLoggerFinder); - - auto registry = getPluginRegistry(); - - { - std::int32_t nbCreators; - auto creators = getPluginCreators(nbCreators); - - for (std::int32_t i = 0; i < nbCreators; ++i) - { - auto const creator = creators[i]; - creator->setPluginNamespace(libNamespace); - registry->registerCreator(*creator, libNamespace); - if (gLogger) - { - auto const msg = tc::fmtstr("Registered plugin creator %s version %s in namespace %s", - creator->getPluginName(), creator->getPluginVersion(), libNamespace); - gLogger->log(nvinfer1::ILogger::Severity::kVERBOSE, msg.c_str()); - } - } - } - - { - std::int32_t nbCreators; - auto creators = getCreators(nbCreators); - - for (std::int32_t i = 0; i < nbCreators; ++i) - { - auto const creator = creators[i]; - registry->registerCreator(*creator, libNamespace); - } - } - - pluginsInitialized = true; - return true; - } - - [[maybe_unused]] void setLoggerFinder([[maybe_unused]] nvinfer1::ILoggerFinder* finder) - { - tensorrt_llm::plugins::api::LoggerManager::getInstance().setLoggerFinder(finder); - } - - [[maybe_unused]] nvinfer1::IPluginCreator* const* getPluginCreators(std::int32_t& nbCreators) - { - static tensorrt_llm::plugins::IdentityPluginCreator identityPluginCreator; - static tensorrt_llm::plugins::BertAttentionPluginCreator bertAttentionPluginCreator; - static tensorrt_llm::plugins::FusedLayernormPluginCreator fusedLayernormPluginCreator; - static tensorrt_llm::plugins::GPTAttentionPluginCreator gptAttentionPluginCreator; - static tensorrt_llm::plugins::GemmPluginCreator gemmPluginCreator; - static tensorrt_llm::plugins::GemmSwigluPluginCreator gemmSwigluPluginCreator; - static tensorrt_llm::plugins::Fp8RowwiseGemmPluginCreator fp8RowwiseGemmPluginCreator; - static tensorrt_llm::plugins::MixtureOfExpertsPluginCreator moePluginCreator; -#if ENABLE_MULTI_DEVICE - static tensorrt_llm::plugins::SendPluginCreator sendPluginCreator; - static tensorrt_llm::plugins::RecvPluginCreator recvPluginCreator; - static tensorrt_llm::plugins::AllreducePluginCreator allreducePluginCreator; - static tensorrt_llm::plugins::AllgatherPluginCreator allgatherPluginCreator; - static tensorrt_llm::plugins::ReduceScatterPluginCreator reduceScatterPluginCreator; - static tensorrt_llm::plugins::GemmAllReducePluginCreator gemmAllReducePluginCreator; -#endif // ENABLE_MULTI_DEVICE - static tensorrt_llm::plugins::SmoothQuantGemmPluginCreator smoothQuantGemmPluginCreator; - static tensorrt_llm::plugins::QServeGemmPluginCreator qserveGemmPluginCreator; - static tensorrt_llm::plugins::LayernormQuantizationPluginCreator layernormQuantizationPluginCreator; - static tensorrt_llm::plugins::QuantizeToFP4PluginCreator quantizeToFP4PluginCreator; - static tensorrt_llm::plugins::QuantizePerTokenPluginCreator quantizePerTokenPluginCreator; - static tensorrt_llm::plugins::QuantizeTensorPluginCreator quantizeTensorPluginCreator; - static tensorrt_llm::plugins::RmsnormQuantizationPluginCreator rmsnormQuantizationPluginCreator; - static tensorrt_llm::plugins::WeightOnlyGroupwiseQuantMatmulPluginCreator - weightOnlyGroupwiseQuantMatmulPluginCreator; - static tensorrt_llm::plugins::WeightOnlyQuantMatmulPluginCreator weightOnlyQuantMatmulPluginCreator; - static tensorrt_llm::plugins::LookupPluginCreator lookupPluginCreator; - static tensorrt_llm::plugins::LoraPluginCreator loraPluginCreator; - static tensorrt_llm::plugins::SelectiveScanPluginCreator selectiveScanPluginCreator; - static tensorrt_llm::plugins::Fp4GemmPluginCreator fp4GemmPluginCreator; - static tensorrt_llm::plugins::MambaConv1dPluginCreator mambaConv1DPluginCreator; - static tensorrt_llm::plugins::lruPluginCreator lruPluginCreator; - static tensorrt_llm::plugins::CumsumLastDimPluginCreator cumsumLastDimPluginCreator; - static tensorrt_llm::plugins::TopkLastDimPluginCreator topkLastDimPluginCreator; - static tensorrt_llm::plugins::LowLatencyGemmPluginCreator lowLatencyGemmPluginCreator; - static tensorrt_llm::plugins::LowLatencyGemmSwigluPluginCreator lowLatencyGemmSwigluPluginCreator; - static tensorrt_llm::plugins::EagleDecodeDraftTokensPluginCreator eagleDecodeDraftTokensPluginCreator; - static tensorrt_llm::plugins::EagleSampleAndAcceptDraftTokensPluginCreator - eagleSampleAndAcceptDraftTokensPluginCreator; - static tensorrt_llm::plugins::CudaStreamPluginCreator cudaStreamPluginCreator; - - static std::array pluginCreators - = { creatorPtr(identityPluginCreator), - creatorPtr(bertAttentionPluginCreator), - creatorPtr(gptAttentionPluginCreator), - creatorPtr(gemmPluginCreator), - creatorPtr(gemmSwigluPluginCreator), - creatorPtr(fp8RowwiseGemmPluginCreator), - creatorPtr(moePluginCreator), -#if ENABLE_MULTI_DEVICE - creatorPtr(sendPluginCreator), - creatorPtr(recvPluginCreator), - creatorPtr(allreducePluginCreator), - creatorPtr(allgatherPluginCreator), - creatorPtr(reduceScatterPluginCreator), - creatorPtr(gemmAllReducePluginCreator), -#endif // ENABLE_MULTI_DEVICE - creatorPtr(fusedLayernormPluginCreator), - creatorPtr(smoothQuantGemmPluginCreator), - creatorPtr(qserveGemmPluginCreator), - creatorPtr(layernormQuantizationPluginCreator), - creatorPtr(quantizeToFP4PluginCreator), - creatorPtr(quantizePerTokenPluginCreator), - creatorPtr(quantizeTensorPluginCreator), - creatorPtr(rmsnormQuantizationPluginCreator), - creatorPtr(weightOnlyGroupwiseQuantMatmulPluginCreator), - creatorPtr(weightOnlyQuantMatmulPluginCreator), - creatorPtr(lookupPluginCreator), - creatorPtr(loraPluginCreator), - creatorPtr(selectiveScanPluginCreator), - creatorPtr(fp4GemmPluginCreator), - creatorPtr(mambaConv1DPluginCreator), - creatorPtr(lruPluginCreator), - creatorPtr(cumsumLastDimPluginCreator), - creatorPtr(topkLastDimPluginCreator), - creatorPtr(lowLatencyGemmPluginCreator), - creatorPtr(eagleDecodeDraftTokensPluginCreator), - creatorPtr(eagleSampleAndAcceptDraftTokensPluginCreator), - creatorPtr(lowLatencyGemmSwigluPluginCreator), - creatorPtr(cudaStreamPluginCreator), - }; - nbCreators = pluginCreators.size(); - return pluginCreators.data(); - } - - [[maybe_unused]] nvinfer1::IPluginCreatorInterface* const* getCreators(std::int32_t& nbCreators) - { - static tensorrt_llm::plugins::EaglePrepareDrafterInputsPluginCreator eaglePrepareDrafterInputsPluginCreator; -#if ENABLE_MULTI_DEVICE - static tensorrt_llm::plugins::CpSplitPluginCreator cpSplitPluginCreator; -#endif // ENABLE_MULTI_DEVICE - - static tensorrt_llm::plugins::DoraPluginCreator doraPluginCreator; - - static std::array creators - = { creatorInterfacePtr(eaglePrepareDrafterInputsPluginCreator), -#if ENABLE_MULTI_DEVICE - creatorInterfacePtr(cpSplitPluginCreator), -#endif // ENABLE_MULTI_DEVICE - creatorInterfacePtr(doraPluginCreator) }; - - nbCreators = creators.size(); - return creators.data(); - } -} // extern "C" diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp deleted file mode 100644 index 6acf0b3a9d25..000000000000 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.cpp +++ /dev/null @@ -1,1206 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "bertAttentionPlugin.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/kernels/recoverFromRingAtten.h" -#include "tensorrt_llm/kernels/sageAttentionKernels.h" -#include "tensorrt_llm/kernels/unfusedAttentionKernels.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -namespace tc = tensorrt_llm::common; - -using tensorrt_llm::plugins::BertAttentionPluginCreator; -using tensorrt_llm::plugins::BertAttentionPlugin; - -static char const* BERT_ATTENTION_PLUGIN_VERSION{"1"}; -static char const* BERT_ATTENTION_PLUGIN_NAME{"BertAttention"}; -PluginFieldCollection BertAttentionPluginCreator::mFC{}; -std::vector BertAttentionPluginCreator::mPluginAttributes; - -BertAttentionPlugin::BertAttentionPlugin(int num_heads, int head_size, float q_scaling, - ContextFMHAType context_fmha_type, nvinfer1::DataType type, bool do_relative_attention, int max_distance, - bool remove_padding, bool sage_attn, int sage_attn_q_block_size, int sage_attn_k_block_size, - int sage_attn_v_block_size, int cp_size, int cp_rank, std::set cp_group) - : mNumHeads(num_heads) - , mHeadSize(head_size) - , mQScaling(q_scaling) - , mType(type) - , mRelativeAttention(do_relative_attention) - , mMaxDistance(max_distance) - , mRemovePadding(remove_padding) - , mEnableContextFMHA(context_fmha_type != ContextFMHAType::DISABLED) - , mFMHAForceFP32Acc(context_fmha_type == ContextFMHAType::ENABLED_WITH_FP32_ACC) - , mSageAttn(sage_attn) - , mCpSize(cp_size) - , mCpRank(cp_rank) - , mCpGroup(std::move(cp_group)) -{ - // pre-check whether FMHA is supported in order to save memory allocation - if (mEnableContextFMHA) - { - mEnableContextFMHA = false; - if (!(mType == DataType::kHALF || mType == DataType::kBF16)) - { - TLLM_LOG_WARNING("Fall back to unfused MHA because of unsupported data type."); - } - else if (mRelativeAttention) - { - TLLM_LOG_WARNING("Fall back to unfused MHA because of relative position embedding."); - } - else - { - mEnableContextFMHA = true; - } - } - - if (mSageAttn) - { - mSageAttnQBlockSize = sage_attn_q_block_size; - mSageAttnKBlockSize = sage_attn_k_block_size; - mSageAttnVBlockSize = sage_attn_v_block_size; - std::vector blockSizeCombination - = {sage_attn_q_block_size, sage_attn_k_block_size, sage_attn_v_block_size}; - if (mSageAttnSupportedBlockSizes.find(blockSizeCombination) == mSageAttnSupportedBlockSizes.end() - || (head_size != 128 && head_size != 72 && head_size != 80)) - { - TLLM_LOG_WARNING(" Q, k ,v quant block size not support. disable sage attention"); - mSageAttn = false; - } - else - { - TLLM_LOG_INFO("SageAttnQBlockSize: %d, SageAttnKBlockSize: %d, SageAttnVBlockSize: %d", mSageAttnQBlockSize, - mSageAttnKBlockSize, mSageAttnVBlockSize); - } - } - - if (cp_group.size() > 1 && !mEnableContextFMHA) - { - TLLM_LOG_ERROR("Unfused MHA do not support context parallel now."); - } -} - -// Parameterized constructor -BertAttentionPlugin::BertAttentionPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mNumHeads); - read(d, mHeadSize); - read(d, mQScaling); - read(d, mQKHalfAccum); - read(d, mEnableContextFMHA); - read(d, mFMHAForceFP32Acc); - read(d, mType); - read(d, mRelativeAttention); - read(d, mMaxDistance); - read(d, mRemovePadding); - read(d, mSageAttn); - read(d, mSageAttnQBlockSize); - read(d, mSageAttnKBlockSize); - read(d, mSageAttnVBlockSize); - read(d, mCpSize); - read(d, mCpRank); - mCpGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mCpGroup.insert(groupItem); - } - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* BertAttentionPlugin::clone() const noexcept -{ - auto* plugin = new BertAttentionPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - plugin->initialize(); - return plugin; -} - -nvinfer1::DimsExprs BertAttentionPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(outputIndex == 0); - auto ret = inputs[0]; - ret.d[mRemovePadding ? 1 : 2] = exprBuilder.constant(ret.d[mRemovePadding ? 1 : 2]->getConstantValue() / 3); - return ret; -} - -bool BertAttentionPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - // inputs: [0] qkv, [1] input_lengths, [2] max_input_length (optional), [3] relative_attention_bias (optional) - // outputs: [X] hidden_states - if (nbInputs == 2) - { // BERT - if (pos == 1) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - if (nbInputs > 2) - { // Encoder in encoder-decoder - if (pos == 1 || pos == 2) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - - return false; -} - -void BertAttentionPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t BertAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - // if remove padding, inputs[0] "qkv_hidden_states" dim is [num_tokens, 3*hidden_dim] which doesn't have shape - // info should get max_batch_size and max_input_length from inputs[1] "input_lengths" and input[2] - // "max_input_length" - int const batch_size = mRemovePadding ? inputs[1].dims.d[0] : inputs[0].dims.d[0]; - int const input_seq_len = mRemovePadding ? inputs[2].dims.d[0] : inputs[0].dims.d[1]; - int const local_hidden_units_ = inputs[0].dims.d[mRemovePadding ? 1 : 2] / 3; - - auto const size = tensorrt_llm::runtime::BufferDataType(inputs[0].type).getSize(); - - size_t const attention_mask_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * input_seq_len; - size_t const cu_seqlens_size = sizeof(int) * (batch_size + 1); - size_t const q_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const k_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const v_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_size = mEnableContextFMHA ? 0 : size * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : size * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_float_size - = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * input_seq_len; - size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0; - int const paddedHeadSize = mSageAttn ? ((mHeadSize + 15) / 16) * 16 : mHeadSize; - const size_t quanted_qkv_size - = mSageAttn ? sizeof(__nv_fp8_e4m3) * batch_size * input_seq_len * mNumHeads * paddedHeadSize * 3 : 0; - const size_t q_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnQBlockSize - 1) / mSageAttnQBlockSize) * mNumHeads - : 0; - const size_t k_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnKBlockSize - 1) / mSageAttnKBlockSize) * mNumHeads - : 0; - const size_t v_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnVBlockSize - 1) / mSageAttnVBlockSize) * mNumHeads - : 0; - const size_t scale_bmm1_device_size = mSageAttn ? sizeof(float) * 2 : 0; - const size_t scale_bmm2_device_size = mSageAttn ? sizeof(float) : 0; - size_t sage_quant_space_size = mSageAttn ? sizeof(float) * batch_size * mNumHeads * mHeadSize : 0; - - if (paddedHeadSize != mHeadSize) - sage_quant_space_size - = sage_quant_space_size < (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - ? (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - : sage_quant_space_size; - - // workspace for RingAttention ping-pong buffer - bool const enableRingAttn = (mCpGroup.size() > 1); - const size_t ring_q_buf_size = enableRingAttn ? size * batch_size * input_seq_len * local_hidden_units_ : 0; - const size_t ring_kv_buf_size = enableRingAttn - ? 2 * size * batch_size * input_seq_len * local_hidden_units_ + sizeof(int) * (batch_size + 1) - : 0; - const size_t ring_softmax_stats_buf_size - = enableRingAttn ? 2 * sizeof(float) * batch_size * input_seq_len * mNumHeads : 0; - const size_t ring_softmax_stats_accu_buf_size - = enableRingAttn ? 2 * sizeof(float) * batch_size * input_seq_len * mNumHeads : 0; - const size_t ring_block_output_size = enableRingAttn ? size * batch_size * input_seq_len * local_hidden_units_ : 0; - - int const NUM_BUFFERS = 24; - - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = CUBLAS_WORKSPACE_SIZE; - workspaces[1] = attention_mask_size; - workspaces[2] = cu_seqlens_size; - workspaces[3] = q_buf_2_size; - workspaces[4] = k_buf_2_size; - workspaces[5] = v_buf_2_size; - workspaces[6] = qk_buf_size; - workspaces[7] = qkv_buf_2_size; - workspaces[8] = qk_buf_float_size; - workspaces[9] = padding_offset_size; - workspaces[10] = fmha_scheduler_counter; - workspaces[11] = quanted_qkv_size; - workspaces[12] = q_scale_size; - workspaces[13] = v_scale_size; - workspaces[14] = k_scale_size; - workspaces[15] = scale_bmm1_device_size; - workspaces[16] = scale_bmm2_device_size; - workspaces[17] = sage_quant_space_size; - workspaces[18] = ring_q_buf_size; - workspaces[19] = ring_kv_buf_size; // kv1 - workspaces[20] = ring_kv_buf_size; // kv2 - workspaces[21] = ring_softmax_stats_buf_size; - workspaces[22] = ring_softmax_stats_accu_buf_size; - workspaces[23] = ring_block_output_size; - - return tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); -} - -template -int BertAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - - // inputs - // input_tensor [batch_size, seq_len, local_hidden_size*3] or [num_tokens, local_hidden_size*3] - // input_lengths [batch_size] - // max_input_length [max_input_length] -- use shape dim to represent max value. If remove padding, this records - // the max input length among sequences; otherwise same as input_tensor's padded dim[1] relative_attention_bias - // [num_heads, num_buckets] (optional) - // outputs - // output_tensor [batch_size, seq_len, local_hidden_size] or [num_tokens, local_hidden_size] - - // if remove padding, inputs[0] dim is [num_tokens] which doesn't have workspace info - // should get max_batch_size from inputs[1] and max_input_length from plugin attribute - int const batch_size = mRemovePadding ? inputDesc[1].dims.d[0] : inputDesc[0].dims.d[0]; - int const input_seq_len = mRemovePadding ? inputDesc[2].dims.d[0] : inputDesc[0].dims.d[1]; - int const num_tokens = mRemovePadding ? inputDesc[0].dims.d[0] : batch_size * input_seq_len; - int const request_batch_size = batch_size; - int const request_seq_len = input_seq_len; - int const local_hidden_units_ = inputDesc[0].dims.d[mRemovePadding ? 1 : 2] / 3; - float const q_scaling = mQScaling; - - T const* attention_input = reinterpret_cast(inputs[0]); - int const* input_lengths = reinterpret_cast(inputs[1]); - T const* relative_attn_table = mRelativeAttention ? reinterpret_cast(inputs[3]) : nullptr; - T* context_buf_ = (T*) (outputs[0]); - - auto cublasHandle = mCublasWrapper->getCublasHandle(); - TLLM_CUDA_CHECK(cublasSetStream(cublasHandle, stream)); - mCublasWrapper->setStream(stream); - mCublasWrapper->setWorkspace(workspace); - if (inputDesc[0].type == DataType::kHALF) - { - mCublasWrapper->setFP16GemmConfig(); - } - else if (inputDesc[0].type == DataType::kFLOAT) - { - mCublasWrapper->setFP32GemmConfig(); - } -#ifdef ENABLE_BF16 - else if constexpr (std::is_same_v) - { - mCublasWrapper->setBF16GemmConfig(); - } -#endif - - size_t const attention_mask_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * input_seq_len; - size_t const cu_seqlens_size = sizeof(int) * (batch_size + 1); - size_t const q_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const k_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const v_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_size - = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const qkv_buf_2_size = mEnableContextFMHA ? 0 : sizeof(T) * batch_size * input_seq_len * local_hidden_units_; - size_t const qk_buf_float_size - = mEnableContextFMHA ? 0 : sizeof(float) * batch_size * mNumHeads * input_seq_len * input_seq_len; - size_t const padding_offset_size = mEnableContextFMHA ? 0 : sizeof(int) * batch_size * input_seq_len; - size_t const fmha_scheduler_counter = mEnableContextFMHA ? sizeof(uint32_t) : 0; - - int const paddedHeadSize = mSageAttn ? ((mHeadSize + 15) / 16) * 16 : mHeadSize; - const size_t quanted_qkv_size - = mSageAttn ? sizeof(__nv_fp8_e4m3) * batch_size * input_seq_len * mNumHeads * paddedHeadSize * 3 : 0; - const size_t q_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnQBlockSize - 1) / mSageAttnQBlockSize) * mNumHeads - : 0; - const size_t k_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnKBlockSize - 1) / mSageAttnKBlockSize) * mNumHeads - : 0; - const size_t v_scale_size = mSageAttn - ? sizeof(float) * batch_size * ((input_seq_len + mSageAttnVBlockSize - 1) / mSageAttnVBlockSize) * mNumHeads - : 0; - const size_t scale_bmm1_device_size = mSageAttn ? sizeof(float) * 2 : 0; - const size_t scale_bmm2_device_size = mSageAttn ? sizeof(float) : 0; - size_t sage_quant_space_size = mSageAttn ? sizeof(float) * batch_size * mNumHeads * mHeadSize : 0; - - if (paddedHeadSize != mHeadSize) - sage_quant_space_size - = sage_quant_space_size < (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - ? (batch_size * input_seq_len * mNumHeads * paddedHeadSize * sizeof(__nv_bfloat16)) - : sage_quant_space_size; - - bool const enableRingAttn = (mCpGroup.size() > 1); - const size_t ring_q_buf_size = enableRingAttn ? sizeof(T) * batch_size * input_seq_len * local_hidden_units_ : 0; - const size_t ring_kv_buf_size - = enableRingAttn ? 2 * sizeof(T) * batch_size * input_seq_len * local_hidden_units_ : 0; - const size_t ring_softmax_stats_buf_size - = enableRingAttn ? 2 * sizeof(float) * batch_size * input_seq_len * mNumHeads : 0; - const size_t ring_block_output_size - = enableRingAttn ? sizeof(T) * batch_size * input_seq_len * local_hidden_units_ : 0; - - // Workspace pointer shift - int8_t* workspace_byte_ptr = reinterpret_cast(workspace); - size_t offset = CUBLAS_WORKSPACE_SIZE; - - T* attention_mask = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, attention_mask_size)); - int* cu_seqlens = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, cu_seqlens_size)); - T* q_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, q_buf_2_size)); - T* k_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, k_buf_2_size)); - T* v_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, v_buf_2_size)); - T* qk_buf_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, qk_buf_size)); - T* qkv_buf_2_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, qkv_buf_2_size)); - float* qk_buf_float_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, qk_buf_float_size)); - int* padding_offset = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, padding_offset_size)); - uint32_t* fmha_tile_counter_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, fmha_scheduler_counter)); - - __nv_fp8_e4m3* quanted_qkv_ptr - = reinterpret_cast<__nv_fp8_e4m3*>(tc::nextWorkspacePtr(workspace_byte_ptr, offset, quanted_qkv_size)); - float* q_scale_ptr = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, q_scale_size)); - float* k_scale_ptr = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, k_scale_size)); - float* v_scale_ptr = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, v_scale_size)); - float* scale_bmm1_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, scale_bmm1_device_size)); - float* scale_bmm2_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, scale_bmm2_device_size)); - void* sage_quant_space_ptr - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, sage_quant_space_size)); - - T* ring_q_buf_ = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_q_buf_size)); - T* ring_kv_buf_1_ = reinterpret_cast( - tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_kv_buf_size + sizeof(int) * (batch_size + 1))); - T* ring_kv_buf_2_ = reinterpret_cast( - tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_kv_buf_size + sizeof(int) * (batch_size + 1))); - float* ring_softmax_stats_buf_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_softmax_stats_buf_size)); - float* ring_softmax_accu_stats_buf_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_softmax_stats_buf_size)); - T* ring_block_output_ - = reinterpret_cast(tc::nextWorkspacePtr(workspace_byte_ptr, offset, ring_block_output_size)); - - // build attention_mask, cu_seqlens, and padding_offset tensors - BuildDecoderInfoParams params{}; - params.seqQOffsets = cu_seqlens; - params.paddingOffsets = padding_offset; - params.attentionMask = attention_mask; - params.seqQLengths = input_lengths; - params.batchSize = batch_size; - params.maxQSeqLength = input_seq_len; - params.numTokens = num_tokens; - params.attentionMaskType = AttentionMaskType::PADDING; - params.fmhaTileCounter = fmha_tile_counter_ptr; - if (mSageAttn) - { - params.fmhaHostBmm1Scale = 1.0f / (sqrtf(mHeadSize * 1.0f) * q_scaling); - params.fmhaBmm1Scale = scale_bmm1_ptr; - params.fmhaBmm2Scale = scale_bmm2_ptr; - } - invokeBuildDecoderInfo(params, stream); - sync_check_cuda_error(stream); - - auto const gemm_data_type = tc::CudaDataType::value; - int const attention_seq_len_1 = request_seq_len; // q length - int const attention_seq_len_2 = request_seq_len; // kv length - - // If the model has relative attentiona bias, q scaling should be applied in QK gemm stage and use 1 in - // softamax stage (because to get softmax[scale(Q*K) + rel pos bias] here, q_scaling can't be applied during - // softmax phase by qk_scale); otherwise, use 1 in gemm stage and apply scaling in softmax stage - float const qk_scale - = 1.0f / (sqrtf(mHeadSize * 1.0f) * q_scaling); // q_scaling in denominator. by default q_scaling =1.0f - float const qk_scale_gemm = mRelativeAttention ? qk_scale : 1.0f; - T const qk_scale_softmax = static_cast(mRelativeAttention ? 1.0f : qk_scale); - - T* linear_bias_slopes = nullptr; - - // FMHA doesn't apply to MHA with relative attention bias, i.e. softmax(QK + bias) * V - // We update mEnableContextFMHA in constructor to check this condition - if (mEnableContextFMHA) - { - if (enableRingAttn) - { - // make sure the padding part of key/value buffer is 0 - cudaMemsetAsync(ring_kv_buf_1_, 0, - reinterpret_cast(ring_kv_buf_2_) - reinterpret_cast(ring_kv_buf_1_), stream); - - cudaMemcpyAsync(ring_q_buf_, attention_input, ring_q_buf_size, cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(ring_kv_buf_1_, - const_cast(reinterpret_cast(attention_input)) + ring_q_buf_size, ring_kv_buf_size, - cudaMemcpyDeviceToDevice, stream); - cudaMemcpyAsync(reinterpret_cast(ring_kv_buf_1_) + ring_kv_buf_size, cu_seqlens, - sizeof(int) * (batch_size + 1), cudaMemcpyDeviceToDevice, stream); - // init softmax_stats - cudaMemsetAsync(ring_softmax_accu_stats_buf_, 0, ring_softmax_stats_buf_size, stream); - -#if ENABLE_MULTI_DEVICE - // relative position of prev/next rank in cp group - int prev_rank = mCpRank > 0 ? mCpRank - 1 : mCpGroup.size() - 1; - int next_rank = (mCpRank == static_cast(mCpGroup.size() - 1)) ? 0 : mCpRank + 1; -#endif // ENABLE_MULTI_DEVICE - - common::check_cuda_error(cudaStreamCreate(&mNcclStream)); - common::check_cuda_error(cudaStreamSynchronize(stream)); - - uint32_t* fmha_scheduler_counter_h = (uint32_t*) malloc(sizeof(uint32_t)); - cudaMemcpyAsync( - fmha_scheduler_counter_h, fmha_tile_counter_ptr, sizeof(uint32_t), cudaMemcpyDeviceToHost, stream); - for (size_t iter = 0; iter < mCpGroup.size(); ++iter) - { - // KV buffer used by fmha - T* ring_fmha_kv_buf_ = (iter % 2 == 0) ? ring_kv_buf_1_ : ring_kv_buf_2_; -#if ENABLE_MULTI_DEVICE - T* ring_send_kv_buf_ = (iter % 2 == 0) ? ring_kv_buf_1_ : ring_kv_buf_2_; - T* ring_recv_kv_buf_ = (iter % 2 == 0) ? ring_kv_buf_2_ : ring_kv_buf_1_; - if (iter < mCpGroup.size() - 1) - { - NCCLCHECK(ncclGroupStart()); - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclSend(ring_send_kv_buf_, - ring_kv_buf_size / sizeof(T) + sizeof(int) / sizeof(T) * (batch_size + 1), - (*getDtypeMap())[inputDesc[0].type], next_rank, *mNcclComm, mNcclStream)); - NCCLCHECK(ncclRecv(ring_recv_kv_buf_, - ring_kv_buf_size / sizeof(T) + sizeof(int) / sizeof(T) * (batch_size + 1), - (*getDtypeMap())[inputDesc[0].type], prev_rank, *mNcclComm, mNcclStream)); - NCCLCHECK(ncclGroupEnd()); - } -#else - TLLM_LOG_ERROR("Please set ENABLE_MULTI_DEVICE to enable RingAttention"); - return 1; -#endif // ENABLE_MULTI_DEVICE - // Construct the fmha params for running kernels. - MHARunnerParams fmhaParams{}; - fmhaParams.b = request_batch_size; - fmhaParams.qSeqLen = request_seq_len; - fmhaParams.kvSeqLen = request_seq_len; - fmhaParams.totalQSeqLen = request_batch_size * request_seq_len; - // Device buffer pointers. - fmhaParams.qPtr = ring_q_buf_; - fmhaParams.kvPtr = ring_fmha_kv_buf_; - if (iter == 0) - { - fmhaParams.outputPtr = context_buf_; - fmhaParams.softmaxStatsPtr = ring_softmax_accu_stats_buf_; - } - else - { - cudaMemsetAsync(ring_softmax_stats_buf_, 0, ring_softmax_stats_buf_size, stream); - fmhaParams.outputPtr = ring_block_output_; - fmhaParams.softmaxStatsPtr = ring_softmax_stats_buf_; - } - fmhaParams.cuQSeqLenPtr = cu_seqlens; - fmhaParams.cuKvSeqLenPtr - = reinterpret_cast(reinterpret_cast(ring_fmha_kv_buf_) + ring_kv_buf_size); - - fmhaParams.tileCounterPtr = fmha_tile_counter_ptr; - fmhaParams.stream = stream; - // Run the fmha kernel. - cudaMemsetAsync(fmhaParams.outputPtr, 0, ring_block_output_size, stream); - cudaMemcpyAsync(fmhaParams.tileCounterPtr, fmha_scheduler_counter_h, sizeof(uint32_t), - cudaMemcpyHostToDevice, stream); - mFmhaDispatcher->run(fmhaParams); - if (iter != 0) - { - invokeRecoverFromRA((T*) context_buf_, (float*) ring_softmax_accu_stats_buf_, - (T*) ring_block_output_, (float*) ring_softmax_stats_buf_, fmhaParams.b, fmhaParams.qSeqLen, - mNumHeads, mHeadSize, cu_seqlens, stream); - } - cudaStreamSynchronize(stream); - cudaStreamSynchronize(mNcclStream); - } - common::check_cuda_error(cudaStreamDestroy(mNcclStream)); - free(fmha_scheduler_counter_h); - } - - else - { - if (mSageAttn && mHeadSize == 72 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 64 - && mSageAttnVBlockSize == 256) - { - sage_quant<72, 80, 64, 64, 256, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 80 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 64 - && mSageAttnVBlockSize == 256) - { - sage_quant<80, 80, 64, 64, 256, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 128 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 64 - && mSageAttnVBlockSize == 256) - { - sage_quant<128, 128, 64, 64, 256, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 128 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 32 - && mSageAttnVBlockSize == 32) - { - sage_quant<128, 128, 64, 32, 32, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 80 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 32 - && mSageAttnVBlockSize == 32) - { - sage_quant<80, 80, 64, 32, 32, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - if (mSageAttn && mHeadSize == 72 && mSageAttnQBlockSize == 64 && mSageAttnKBlockSize == 32 - && mSageAttnVBlockSize == 32) - { - sage_quant<72, 80, 64, 32, 32, __nv_bfloat16, __nv_fp8_e4m3, float>( - // host var - batch_size, mNumHeads, input_seq_len, true, true, - // device var - // q k v - attention_input, attention_input + mNumHeads * mHeadSize, - attention_input + 2 * mNumHeads * mHeadSize, - // stride - 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, 3 * mNumHeads * mHeadSize, cu_seqlens, - cu_seqlens, sage_quant_space_ptr, - // quant q k v - quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * paddedHeadSize, - quanted_qkv_ptr + 2 * mNumHeads * paddedHeadSize, - // quanted_qkv_ptr, quanted_qkv_ptr + mNumHeads * mHeadSize, context, - 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, 3 * mNumHeads * paddedHeadSize, - // scales - q_scale_ptr, k_scale_ptr, v_scale_ptr, stream); - - sync_check_cuda_error(stream); - } - - // Construct the fmha params for running kernels. - MHARunnerParams fmhaParams{}; - fmhaParams.b = request_batch_size; - fmhaParams.qSeqLen = request_seq_len; - fmhaParams.kvSeqLen = request_seq_len; - fmhaParams.totalQSeqLen = request_batch_size * request_seq_len; - // Device buffer pointers. - fmhaParams.qkvPtr = attention_input; - fmhaParams.outputPtr = context_buf_; - fmhaParams.cuQSeqLenPtr = cu_seqlens; - fmhaParams.cuKvSeqLenPtr = cu_seqlens; - fmhaParams.tileCounterPtr = fmha_tile_counter_ptr; - fmhaParams.stream = stream; - if (mSageAttn) - { - if (paddedHeadSize != mHeadSize) - fmhaParams.outputPtr = sage_quant_space_ptr; - fmhaParams.qkvPtr = quanted_qkv_ptr; - fmhaParams.scaleBmm1Ptr = scale_bmm1_ptr; - fmhaParams.scaleBmm2Ptr = scale_bmm2_ptr; - fmhaParams.qScalePtr = q_scale_ptr; - fmhaParams.kScalePtr = k_scale_ptr; - fmhaParams.vScalePtr = v_scale_ptr; - fmhaParams.qMaxNBlock = (input_seq_len + mSageAttnQBlockSize - 1) / mSageAttnQBlockSize; - fmhaParams.kMaxNBlock = (input_seq_len + mSageAttnKBlockSize - 1) / mSageAttnKBlockSize; - fmhaParams.vMaxNBlock = (input_seq_len + mSageAttnVBlockSize - 1) / mSageAttnVBlockSize; - } - - // Run the fmha kernel. - - // TODO: set it correctly for contiguous kv buffer (cross-attention). - fmhaParams.totalKvSeqLen = num_tokens; - - fmhaParams.cuKvSeqLenPtr = cu_seqlens; - fmhaParams.cuMaskRowsPtr = cu_seqlens; - fmhaParams.tileCounterPtr = fmha_tile_counter_ptr; - - fmhaParams.scaleBmm1Ptr = scale_bmm1_ptr; - fmhaParams.scaleBmm2Ptr = scale_bmm2_ptr; - fmhaParams.forceFp32Acc = mFMHAForceFP32Acc; - mFmhaDispatcher->run(fmhaParams); - sync_check_cuda_error(stream); - if (mSageAttn) - { - if (paddedHeadSize != mHeadSize && mHeadSize == 72) - { - unpadding<80, 72, __nv_bfloat16>(batch_size, mNumHeads, input_seq_len, sage_quant_space_ptr, - mNumHeads * 72, mNumHeads * 80, cu_seqlens, context_buf_, stream); - } - } - } - } - else - { - // FIXME: a temporary solution to make sure the padding part of key/value buffer is 0 - // NOTE: pointer subtraction is used below since there could be some extra gap due to alignment. - // Otherwise, we could do cudaMemsetAsync(k_buf_2_, 0, k_buf_2_size + v_buf_2_size, stream); - // cudaMemsetAsync(k_buf_2_, 0, reinterpret_cast(qk_buf_) - reinterpret_cast(k_buf_2_), - // stream); - // FIXME: the final solution is to change the add_fusedQKV_bias_transpose_kernel to map CTAs corresponding to - // the output shape, and set the padding part to 0. Without zero-initialize guarantee, these workspace buffers - // may contain random NaN values when IFB workload is high. - cudaMemsetAsync(k_buf_2_, 0, - reinterpret_cast(v_buf_2_) - reinterpret_cast(k_buf_2_) + v_buf_2_size, stream); - - // only non-FMHA path needs to split Q,K,V from QKV - invokeAddFusedQKVBiasTranspose(q_buf_2_, k_buf_2_, v_buf_2_, const_cast(attention_input), input_lengths, - mRemovePadding ? padding_offset : nullptr, batch_size, input_seq_len, num_tokens, mNumHeads, mNumHeads, - mHeadSize, 0, 0.0f, RotaryScalingType::kNONE, 0.0f, 0, PositionEmbeddingType::kLEARNED_ABSOLUTE, - (float*) nullptr, 0, stream); - - if (!mQKHalfAccum && gemm_data_type != CUDA_R_32F) - { - mCublasWrapper->stridedBatchedGemm(CUBLAS_OP_T, CUBLAS_OP_N, - attention_seq_len_2, // n - attention_seq_len_1, // m - mHeadSize, // k - qk_scale_gemm, k_buf_2_, gemm_data_type, - mHeadSize, // k - attention_seq_len_2 * mHeadSize, // n * k - q_buf_2_, gemm_data_type, - mHeadSize, // k - attention_seq_len_1 * mHeadSize, // m * k - 0.0f, qk_buf_float_, CUDA_R_32F, - attention_seq_len_2, // n - attention_seq_len_2 * attention_seq_len_1, - request_batch_size * mNumHeads, // global batch size - CUDA_R_32F); - - // add relative position bias - if (mRelativeAttention) - { - // add rel pos bias - // QK is (batch_size, local_head_num, q_length, k_length), rel pos bias is (1, local_head_num, - // max_output_len + 1, max_output_len + 1). broadcast along 1st dim. max_seq_len is already - // max_output_len + 1. In implicit mode, relative_attention_bias is rel attn table - // [num_heads, num_buckets], with necessary params (max_distance, num_buckets) passed at the end - invokeAddRelativeAttentionBiasUnaligned(qk_buf_float_, relative_attn_table, request_batch_size, - mNumHeads, attention_seq_len_1, attention_seq_len_2, stream, mMaxDistance > 0, - inputDesc[3].dims.d[1], mMaxDistance, true /* bidirectional */); - } - - MaskedSoftmaxParam param; - param.attention_score = qk_buf_; // (batch_size, head_num, q_length, k_length) - param.qk = qk_buf_float_; // (batch_size, head_num, q_length, k_length) - param.attention_mask = attention_mask; // (batch_size, q_length, k_length) - param.batch_size = request_batch_size; - param.q_length = attention_seq_len_1; - param.k_length = attention_seq_len_2; - param.num_heads = mNumHeads; - param.qk_scale = qk_scale_softmax; - param.linear_bias_slopes = const_cast(linear_bias_slopes); // (head_num,), optional - invokeMaskedSoftmax(param, stream); - } - else - { - mCublasWrapper->stridedBatchedGemm(CUBLAS_OP_T, CUBLAS_OP_N, attention_seq_len_2, attention_seq_len_1, - mHeadSize, k_buf_2_, mHeadSize, attention_seq_len_2 * mHeadSize, q_buf_2_, mHeadSize, - attention_seq_len_1 * mHeadSize, qk_buf_, attention_seq_len_2, - attention_seq_len_2 * attention_seq_len_1, request_batch_size * mNumHeads, qk_scale_gemm, - 0.0f); // alpha, beta - - // add relative position bias - if (mRelativeAttention) - { - // add rel pos bias - // QK is (batch_size, local_head_num, q_length, k_length), rel pos bias is (1, local_head_num, - // max_output_len + 1, max_output_len + 1). broadcast along 1st dim. max_seq_len is already - // max_output_len + 1. In implicit mode, relative_attention_bias is rel attn table - // [num_heads, num_buckets], with necessary params (max_distance, num_buckets) passed at the end - invokeAddRelativeAttentionBiasUnaligned(qk_buf_, relative_attn_table, request_batch_size, mNumHeads, - attention_seq_len_1, attention_seq_len_2, stream, mMaxDistance > 0, inputDesc[3].dims.d[1], - mMaxDistance, true /* bidirectional */); - } - - MaskedSoftmaxParam param; - param.attention_score = qk_buf_; // (batch_size, head_num, q_length, k_length) - param.qk = qk_buf_; // (batch_size, head_num, q_length, k_length) - param.attention_mask = attention_mask; // (batch_size, q_length, k_length) - param.batch_size = request_batch_size; - param.q_length = attention_seq_len_1; - param.k_length = attention_seq_len_2; - param.num_heads = mNumHeads; - param.qk_scale = qk_scale_softmax; - param.linear_bias_slopes = const_cast(linear_bias_slopes); // (head_num,), optional - invokeMaskedSoftmax(param, stream); - } - - mCublasWrapper->stridedBatchedGemm(CUBLAS_OP_N, CUBLAS_OP_N, mHeadSize, attention_seq_len_1, - attention_seq_len_2, v_buf_2_, mHeadSize, attention_seq_len_2 * mHeadSize, qk_buf_, attention_seq_len_2, - attention_seq_len_1 * attention_seq_len_2, qkv_buf_2_, mHeadSize, attention_seq_len_1 * mHeadSize, - request_batch_size * mNumHeads); - - if (!mRemovePadding) - { - invokeTransposeQKV(context_buf_, qkv_buf_2_, request_batch_size, attention_seq_len_1, mNumHeads, mHeadSize, - (float*) nullptr, 0, stream); - } - else - { - invokeTransposeAttentionOutRemovePadding(qkv_buf_2_, context_buf_, num_tokens, request_batch_size, - request_seq_len, mNumHeads, mHeadSize, padding_offset, (float*) nullptr, 0, stream); - } - } - sync_check_cuda_error(stream); - return 0; -} - -template int BertAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); - -template int BertAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); - -#ifdef ENABLE_BF16 -template int BertAttentionPlugin::enqueueImpl<__nv_bfloat16>(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); -#endif - -int BertAttentionPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType BertAttentionPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* BertAttentionPlugin::getPluginType() const noexcept -{ - return BERT_ATTENTION_PLUGIN_NAME; -} - -char const* BertAttentionPlugin::getPluginVersion() const noexcept -{ - return BERT_ATTENTION_PLUGIN_VERSION; -} - -int BertAttentionPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int BertAttentionPlugin::initialize() noexcept -{ - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - mCublasWrapper.reset(new tc::CublasMMWrapper(cublasHandle, cublasLtHandle, nullptr, nullptr)); - if (mEnableContextFMHA) - { - // Pre-checked during constructing. - Data_type data_type; - if (mType == DataType::kHALF) - { - data_type = DATA_TYPE_FP16; - } - else if (mType == DataType::kBF16) - { - data_type = DATA_TYPE_BF16; - } - else - { - TLLM_CHECK_WITH_INFO(false, "GPTAttentionPlugin received wrong data type."); - } - - // Construct the fmha runner. - MHARunnerFixedParams fmhaParams{}; - if (mSageAttn) - { - fmhaParams.dataType = DATA_TYPE_E4M3; - } - else - { - fmhaParams.dataType = data_type; - } - fmhaParams.dataTypeOut = data_type; - fmhaParams.forceFp32Acc = mFMHAForceFP32Acc; - fmhaParams.attentionMaskType = ContextAttentionMaskType::PADDING; - fmhaParams.isSPadded = !mRemovePadding; - fmhaParams.numQHeads = mNumHeads; - fmhaParams.numKvHeads = mNumHeads; - fmhaParams.headSize = mHeadSize; - fmhaParams.qScaling = mQScaling; - fmhaParams.sageBlockSizeQ = mSageAttnQBlockSize; - fmhaParams.sageBlockSizeK = mSageAttnKBlockSize; - fmhaParams.sageBlockSizeV = mSageAttnVBlockSize; - if (mSageAttn) - { - int const paddedHeadSize = ((mHeadSize + 15) / 16) * 16; - fmhaParams.headSize = paddedHeadSize; - } - - if (mCpGroup.size() > 1) - { - fmhaParams.attentionInputLayout = AttentionInputLayout::Q_CONTIGUOUS_KV; - fmhaParams.saveSoftmax = true; - } - - // Load kernels from the pre-compiled cubins. - // The KV input data type. The default is same as dataType. - fmhaParams.dataTypeKv = data_type; - fmhaParams.headSizeV = mHeadSize; - - // Load kernels from the pre-compiled cubins. - mFmhaDispatcher.reset(new FmhaDispatcher(fmhaParams)); - // Fall back to unfused MHA kernels if not supported. - mEnableContextFMHA = mFmhaDispatcher->isSupported(); - } - -#if ENABLE_MULTI_DEVICE - if (mCpGroup.size() > 1 && COMM_SESSION.getSize() > 1) - { - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - mNcclComm = getComm(mCpGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - } -#endif // ENABLE_MULTI_DEVICE - - return 0; -} - -void BertAttentionPlugin::destroy() noexcept -{ - delete this; -} - -size_t BertAttentionPlugin::getSerializationSize() const noexcept -{ - return sizeof(mNumHeads) + sizeof(mHeadSize) + sizeof(mQScaling) + sizeof(mQKHalfAccum) + sizeof(mEnableContextFMHA) - + sizeof(mFMHAForceFP32Acc) + sizeof(mType) + sizeof(mRelativeAttention) + sizeof(mMaxDistance) - + sizeof(mRemovePadding) + sizeof(mSageAttn) + sizeof(mSageAttnQBlockSize) + sizeof(mSageAttnKBlockSize) - + sizeof(mSageAttnVBlockSize) + sizeof(mCpSize) + sizeof(mCpRank) + sizeof(int32_t) * mCpGroup.size(); -} - -void BertAttentionPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mNumHeads); - write(d, mHeadSize); - write(d, mQScaling); - write(d, mQKHalfAccum); - write(d, mEnableContextFMHA); - write(d, mFMHAForceFP32Acc); - write(d, mType); - write(d, mRelativeAttention); - write(d, mMaxDistance); - write(d, mRemovePadding); - write(d, mSageAttn); - write(d, mSageAttnQBlockSize); - write(d, mSageAttnKBlockSize); - write(d, mSageAttnVBlockSize); - write(d, mCpSize); - write(d, mCpRank); - for (auto it = mCpGroup.begin(); it != mCpGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void BertAttentionPlugin::terminate() noexcept {} - -/////////////// - -BertAttentionPluginCreator::BertAttentionPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - - mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("head_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("q_scaling", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("context_fmha_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("do_relative_attention", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("max_distance", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("sage_attn", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("sage_attn_q_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sage_attn_k_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sage_attn_v_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_group", nullptr, PluginFieldType::kINT32)); - - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* BertAttentionPluginCreator::getPluginName() const noexcept -{ - return BERT_ATTENTION_PLUGIN_NAME; -} - -char const* BertAttentionPluginCreator::getPluginVersion() const noexcept -{ - return BERT_ATTENTION_PLUGIN_VERSION; -} - -PluginFieldCollection const* BertAttentionPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* BertAttentionPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int num_heads{}; - int head_size{}; - ContextFMHAType context_fmha_type{}; - float q_scaling{}; - nvinfer1::DataType type{}; - bool do_relative_attention{}; - int max_distance{}; - bool remove_padding{}; - bool sage_attn{}; - int sage_attn_q_block_size{}; - int sage_attn_k_block_size{}; - int sage_attn_v_block_size{}; - int cp_size{}; - int cp_rank{}; - std::set cp_group{}; - - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "num_heads")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - num_heads = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "head_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - head_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "q_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - q_scaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "context_fmha_type")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - context_fmha_type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "do_relative_attention")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - do_relative_attention = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "max_distance")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - max_distance = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - remove_padding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sage_attn")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - sage_attn = static_cast(*(static_cast(fields[i].data))); - if (sage_attn) - { - std::cout << "sage attn true!" << std::endl; - } - } - else if (!strcmp(attrName, "sage_attn_q_block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sage_attn_q_block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sage_attn_k_block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sage_attn_k_block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sage_attn_v_block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sage_attn_v_block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_rank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - cp_group.insert(*r); - ++r; - } - } - } - try - { - auto* obj = new BertAttentionPlugin(num_heads, head_size, q_scaling, context_fmha_type, type, - do_relative_attention, max_distance, remove_padding, sage_attn, sage_attn_q_block_size, - sage_attn_k_block_size, sage_attn_v_block_size, cp_size, cp_rank, cp_group); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* BertAttentionPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call BertAttentionPlugin::destroy() - try - { - auto* obj = new BertAttentionPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h b/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h deleted file mode 100644 index 2eb39086a005..000000000000 --- a/cpp/tensorrt_llm/plugins/bertAttentionPlugin/bertAttentionPlugin.h +++ /dev/null @@ -1,142 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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/common/cublasMMWrapper.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/fmhaDispatcher.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class BertAttentionPlugin : public BasePlugin -{ -public: - BertAttentionPlugin() = delete; - - BertAttentionPlugin(int num_heads, int head_size, float q_scaling, - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, nvinfer1::DataType type, - bool do_relative_attention = false, int max_distance = 0, bool remove_padding = false, bool sage_attn = false, - int sage_attn_q_block_size = 0, int sage_attn_k_block_size = 0, int sage_attn_v_block_size = 0, int cp_size = 1, - int cp_rank = 0, std::set cp_group = {}); - - BertAttentionPlugin(void const* data, size_t length); - - ~BertAttentionPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - - int mNumHeads; - int mHeadSize; - float mQScaling; - nvinfer1::DataType mType; - bool mRelativeAttention = false; - int mMaxDistance = 0; - bool mRemovePadding = false; - - // unfused mha - bool mQKHalfAccum = false; - - // fmha runner (disable by default) - bool mEnableContextFMHA = false; - bool mFMHAForceFP32Acc = false; - - // sage attention - bool mSageAttn = false; - int mSageAttnQBlockSize = 0; - int mSageAttnKBlockSize = 0; - int mSageAttnVBlockSize = 0; - std::set> mSageAttnSupportedBlockSizes{{64, 64, 256}, {64, 32, 32}}; - - int mSM = tensorrt_llm::common::getSMVersion(); - - // comm group for RingAttention - int mCpSize = 1; - int mCpRank = 0; - std::set mCpGroup = {}; -#if ENABLE_MULTI_DEVICE - std::shared_ptr mNcclComm; -#endif // ENABLE_MULTI_DEVICE - cudaStream_t mNcclStream; - - // The default copy constructor will leave them as nullptr. clone() shall initialize it. - UniqPtrWNullCopy mFmhaDispatcher; - UniqPtrWNullCopy mCublasWrapper; -}; - -class BertAttentionPluginCreator : public BaseCreator -{ -public: - BertAttentionPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/CMakeLists.txt b/cpp/tensorrt_llm/plugins/common/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/common/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.cpp b/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.cpp deleted file mode 100644 index 2aab6b3675d8..000000000000 --- a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.cpp +++ /dev/null @@ -1,35 +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. - */ - -#include "checkMacrosPlugin.h" - -#include "tensorrt_llm/common/logger.h" - -namespace tensorrt_llm::plugins -{ - -void caughtError(std::exception const& e) -{ - TLLM_LOG_EXCEPTION(e); -} - -void logError(char const* msg, char const* file, char const* fn, int line) -{ - TLLM_LOG_ERROR("Parameter check failed at: %s::%s::%d, condition: %s", file, fn, line, msg); -} - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.h b/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.h deleted file mode 100644 index d8d8af1ef220..000000000000 --- a/cpp/tensorrt_llm/plugins/common/checkMacrosPlugin.h +++ /dev/null @@ -1,29 +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. - */ -#pragma once - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" - -namespace tensorrt_llm::plugins -{ - -void logError(char const* msg, char const* file, char const* fn, int line); - -void caughtError(std::exception const& e); - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.cpp b/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.cpp deleted file mode 100644 index e5d6650648ab..000000000000 --- a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.cpp +++ /dev/null @@ -1,404 +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/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fused_gated_gemm/fused_gated_gemm.h" -#include "tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm.h" -#include "tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h" -#include "tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h" -#include "tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h" -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h" -#else -#include "fp4_gemm.h" -#endif -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -using GemmAllReduceImplInterface = tensorrt_llm::kernels::opened_cutlass_kernels::GemmAllReduceImplInterface; -#else -#include "allreduce_gemm_runner.h" -using GemmAllReduceImplInterface = tensorrt_llm::kernels::cutlass_kernels::GemmAllReduceImplInterface; -#endif - -#include - -namespace tensorrt_llm::plugins -{ - -template -GemmPluginProfiler::GemmPluginProfiler() -{ - mMNKProfileMap = std::make_shared(); - - // set SKIP_GEMM_PLUGIN_PROFILINGS=1 to avoid tactics profilings - auto const skipEnv = std::getenv("SKIP_GEMM_PLUGIN_PROFILINGS"); - mSkip = (skipEnv != NULL && std::stoi(skipEnv)); - if (mSkip) - { - TLLM_LOG_DEBUG( - "SKIP_GEMM_PLUGIN_PROFILINGS is set. Skipping GEMM plugin profilings. It could result in runtime error " - "if default tactic is not defined."); - } -} - -template -void GemmPluginProfiler::serialize( - char*& buffer, GemmIdType const& gemmId) const -{ - auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); - - // Save number of profiles for given GEMM ID - write(buffer, static_cast(mProfileMap->size())); - for (auto const& pair : *mProfileMap) - { - // Save pair of M to the best GEMM config - write(buffer, pair); - } -} - -template -void GemmPluginProfiler::deserialize( - char const*& data, GemmDims& dims, GemmIdType const& gemmId) -{ - // NOTE: this mutex is not needed since each thread owns its private map, but will put here for - // consistency - writer_lock lock(mMNKProfileMap->mutex); - - mDims = dims; - - // GemmId gemmId(dims.n, dims.k); - if (!mMNKProfileMap->existsMProfileMap(gemmId)) - { - // Create GEMM with GEMM ID if it does not exist - mMNKProfileMap->createMProfileMap(gemmId); - } - // Populate map with profiles of GEMM ID - auto profileMap = mMNKProfileMap->getMProfileMap(gemmId); - int selectedMapSize; - read(data, selectedMapSize); - for (int ii = 0; ii < selectedMapSize; ++ii) - { - std::pair> config; - read(data, config); - profileMap->insert(config); - } -} - -template -size_t GemmPluginProfiler::getSerializationSize( - GemmIdType const& gemmId) const -{ - reader_lock lock(mMNKProfileMap->mutex); - return sizeof(int) + // size of the tactics map - mMNKProfileMap->getMProfileMap(gemmId)->size() - * sizeof(std::pair>); // size of the tactics map -} - -template -int GemmPluginProfiler::getMaxProfileM() const -{ - return 8192; -} - -template -void GemmPluginProfiler::initTmpData( - int m, int n, int k, char* workspace, size_t size, cudaStream_t stream) -{ - /* Do nothing */ -} - -template -void GemmPluginProfiler::profileTactics(RunnerPtr const& runner, - nvinfer1::DataType const& type, GemmDims const& dims, GemmIdType const& gemmId, bool hasWeightOnlyCudaKernel) -{ - writer_lock lock(mMNKProfileMap->mutex); - - if (!dims.isInitialized()) - { - return; - } - - mRunner = runner; - mType = type; - - int const maxM = std::min(nextPowerOfTwo(dims.maxM), getMaxProfileM()); - computeTmpSize(maxM, dims.n, dims.k); - - if (!mMNKProfileMap->existsMProfileMap(gemmId)) - { - // Create map for GEMM ID - mMNKProfileMap->createMProfileMap(gemmId); - } - - if (mSkip) - { - return; - } - - auto mProfileMap = mMNKProfileMap->getMProfileMap(gemmId); - bool isAllocated{false}; - - auto profileTactics = [&mProfileMap, &isAllocated, this](int m, int n, int k) - { - if (mProfileMap->count(m) == 0) - { - if (!isAllocated) - { - // Allocate tmp data to run GEMMs - allocateTmpData(); - isAllocated = true; - } - initTmpData(m, n, k, mWorkspaceTmp, mTmpWorkspaceSizeInBytes, mStream); - auto tactics = this->getTactics(m, n, k); - - // Profile different tactics for particular m and insert best config to the map - mProfileMap->insert({m, this->profileTacticsForProblem(m, n, k, tactics)}); - } - }; - - common::check_cuda_error(cudaStreamCreate(&mStream)); - - int const startMinMRounded = nextPowerOfTwo(dims.minM); - - if (hasWeightOnlyCudaKernel) - { - // Profile tactics for finer granularity of M, - // if CUDA kernel is enabled for weight-only plugins - int minM = dims.minM; - for (int m = std::max(1, minM); m < std::min(16, maxM); m += 1) - { - profileTactics(m, dims.n, dims.k); - } - - for (int m = 16; m < maxM; m *= 2) - { - profileTactics(m, dims.n, dims.k); - } - } - else - { - // Profile tactics for CUTLASS kernel only - for (int m = std::max(1, startMinMRounded); m < maxM; m *= 2) - { - profileTactics(m, dims.n, dims.k); - } - } - - profileTactics(maxM, dims.n, dims.k); - - if (isAllocated) - { - // Free tmp data - freeTmpData(); - } - common::check_cuda_error(cudaStreamDestroy(mStream)); -} - -template -std::optional GemmPluginProfiler::getBestConfig( - int m, GemmIdType const& gemmId) const -{ - reader_lock lock(mMNKProfileMap->mutex); - - if (mSkip) - { - TLLM_LOG_TRACE("Skip is set, no best config is set for this instance"); - return std::nullopt; - } - - int const mRounded = std::min(std::max(1, nextPowerOfTwo(m)), getMaxProfileM()); - fflush(stdout); - - if (mMNKProfileMap->getMProfileMap(gemmId)->count(m) > 0) - { - return mMNKProfileMap->getMProfileMap(gemmId)->at(m); - } - else if (mMNKProfileMap->getMProfileMap(gemmId)->count(mRounded) > 0) - { - return mMNKProfileMap->getMProfileMap(gemmId)->at(mRounded); - } - else - { - std::ostringstream msg; - msg << "Cannot find best tactic for m=" << m << " and GEMM ID " << gemmId; - TLLM_LOG_WARNING(msg.str()); - return std::nullopt; - } -} - -template -void GemmPluginProfiler::allocateTmpData() -{ - TLLM_CHECK_WITH_INFO(mTmpWorkspaceSizeInBytes > 0, "tmpWorkspaceSizeInBytes must be larger than 0"); - auto const status = cudaMalloc(&mWorkspaceTmp, mTmpWorkspaceSizeInBytes); - TLLM_CHECK_WITH_INFO(status == cudaSuccess, "Can't allocate tmp workspace for GEMM tactics profiling."); -} - -template -void GemmPluginProfiler::freeTmpData() -{ - auto const status = cudaFree(mWorkspaceTmp); - TLLM_CHECK_WITH_INFO(status == cudaSuccess, "Can't free tmp workspace for GEMM tactics profiling."); -} - -template -std::optional GemmPluginProfiler::profileTacticsForProblem( - int m, int n, int k, std::vector const& tactics) -{ - TLLM_LOG_DEBUG(__PRETTY_FUNCTION__); - - float bestTime = std::numeric_limits::max(); - Config bestConfig; - bool foundOne = false; - - // Iterate over all tactics for given M, N and K - for (size_t ii = 0; ii < tactics.size(); ++ii) - { - Config const& candidateConfig = tactics[ii]; - float time = std::numeric_limits::max(); - try - { - if (!checkTactic(m, n, k, candidateConfig)) - { - continue; - } - // Profile particular tactic for given M, N and K - time = profileTacticForProblem(m, n, k, candidateConfig); - foundOne = true; - } - catch (std::exception const& e) - { - std::ostringstream msg; - msg << "Cannot profile configuration " << ii; - if constexpr (std::is_same_v) - { - msg << ": " << candidateConfig.toString(); - } - msg << "\n (for" - << " m=" << m << ", n=" << n << ", k=" << k << ")" - << ", reason: \"" << e.what() << "\". Skipped"; - TLLM_LOG_TRACE(msg.str()); - cudaGetLastError(); // Reset the last cudaError to cudaSuccess. - continue; - } - - // Choose the fastest tactic - if (time < bestTime) - { - bestConfig = candidateConfig; - bestTime = time; - } - } - - if (!foundOne) - { - std::ostringstream msg; - msg << "Have not found any valid GEMM config for shape (" - << "m=" << m << ", n=" << n << ", k=" << k << "). Will try to use default or fail at runtime"; - TLLM_LOG_WARNING(msg.str()); - return std::nullopt; - } - - return {bestConfig}; -} - -template -float GemmPluginProfiler::profileTacticForProblem( - int m, int n, int k, Config const& tactic) -{ - constexpr int warmup = 5; - constexpr int runs = 10; - - cudaStream_t stream = mStream; - - // Warmup the execution - for (int i = 0; i < warmup; ++i) - { - runTactic(m, n, k, tactic, mWorkspaceTmp, stream); - } - - cudaEvent_t start; - cudaEvent_t stop; - common::check_cuda_error(cudaEventCreate(&start)); - common::check_cuda_error(cudaEventCreate(&stop)); - common::check_cuda_error(cudaStreamSynchronize(stream)); - common::check_cuda_error(cudaEventRecord(start, stream)); - - // Profile GEMM - for (int i = 0; i < runs; ++i) - { - runTactic(m, n, k, tactic, mWorkspaceTmp, stream); - } - - common::check_cuda_error(cudaEventRecord(stop, stream)); - - common::check_cuda_error(cudaEventSynchronize(stop)); - - float elapsed; - common::check_cuda_error(cudaEventElapsedTime(&elapsed, start, stop)); - - common::check_cuda_error(cudaEventDestroy(start)); - common::check_cuda_error(cudaEventDestroy(stop)); - - return elapsed / runs; -} - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -template class GemmPluginProfiler, GemmIdCublas, GemmIdCublasHash>; - -// TODO I dont like the dependency on the MOE plugin here, but MOE needs the full context to run profiles -template class GemmPluginProfiler; - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; - -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -template class GemmPluginProfiler, GemmIdCore, GemmIdCoreHash>; -#else -template class GemmPluginProfiler, GemmIdCore, - GemmIdCoreHash>; -#endif - -template class GemmPluginProfiler; - -template class GemmPluginProfiler; - -template class GemmPluginProfiler, - GemmIdCore, GemmIdCoreHash>; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.h b/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.h deleted file mode 100644 index fe85b3b7e456..000000000000 --- a/cpp/tensorrt_llm/plugins/common/gemmPluginProfiler.h +++ /dev/null @@ -1,332 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 "pluginUtils.h" - -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -struct GemmDims -{ - using DimType64 = utils::DimType64; - - DimType64 minM; - DimType64 maxM; - DimType64 n; - DimType64 k; - - GemmDims() - : minM(-1) - , maxM(-1) - , n(-1) - , k(-1) - { - } - - GemmDims(DimType64 minM_, DimType64 maxM_, DimType64 n_, DimType64 k_) - : minM(minM_) - , maxM(maxM_) - , n(n_) - , k(k_) - { - } - - [[nodiscard]] bool isInitialized() const - { - return minM >= 0 && maxM >= 0 && n >= 0 && k >= 0; - } -}; - -// Unique ID of GEMM -// In our case GEMM is uniqly identified by N and K -class GemmIdCore -{ -public: - int n; - int k; - nvinfer1::DataType dtype; - - GemmIdCore(int n_, int k_, nvinfer1::DataType const& dtype_) - : n(n_) - , k(k_) - , dtype(dtype_) - { - } - - GemmIdCore() - : n(-1) - , k(-1) - , dtype(nvinfer1::DataType::kFLOAT) // dtype does not matter here - { - } - - bool operator==(GemmIdCore const& id) const - { - return isEqual(id); - } - - friend std::ostream& operator<<(std::ostream& out, GemmIdCore const& id) - { - out << "(N;K)=(" << id.n << ";" << id.k << "),"; - out << " type=" << static_cast(id.dtype); - return out; - } - -protected: - bool isEqual(GemmIdCore const& id) const - { - return n == id.n && k == id.k && dtype == id.dtype; - } -}; - -// Hash of GemmId -struct GemmIdCoreHash -{ - std::size_t operator()(GemmIdCore const& id) const - { - auto h1 = std::hash{}(id.n); - auto h2 = std::hash{}(id.k); - auto h3 = std::hash{}(static_cast(id.dtype)); - return h1 ^ h2 ^ h3; - } -}; - -class GemmIdCublas : public GemmIdCore -{ -public: - bool transA{}; - bool transB{}; - nvinfer1::DataType outputDtype; - - GemmIdCublas(int n_, int k_, nvinfer1::DataType const& dtype_, bool transA_, bool transB_, - nvinfer1::DataType const& output_dtype_) - : GemmIdCore(n_, k_, dtype_) - , transA(transA_) - , transB(transB_) - , outputDtype(output_dtype_) - { - } - - GemmIdCublas() {} - - bool operator==(GemmIdCublas const& id) const - { - return isEqual(id) && transA == id.transA && transB == id.transB && outputDtype == id.outputDtype; - } - - friend std::ostream& operator<<(std::ostream& out, GemmIdCublas const& id) - { - out << "(N;K)=(" << id.n << ";" << id.k << "),"; - out << " type=" << static_cast(id.dtype); - out << " transA=" << id.transA; - out << " transB=" << id.transB; - out << " outputDtype=" << static_cast(id.outputDtype); - return out; - } -}; - -// Hash of GemmIdCublas -struct GemmIdCublasHash -{ - std::size_t operator()(GemmIdCublas const& id) const - { - auto h1 = std::hash{}(id.n); - auto h2 = std::hash{}(id.k); - auto h3 = std::hash{}(static_cast(id.dtype)); - auto h4 = std::hash{}(id.transA); - auto h5 = std::hash{}(id.transB); - auto h6 = std::hash{}(static_cast(id.outputDtype)); - return h1 ^ h2 ^ h3 ^ h4 ^ h5 ^ h6; - } -}; - -template -class GemmPluginProfiler -{ -public: - // Map for single GEMM for different Ms (GEMM dimension) to the best config for particular M - using MProfileMap = std::unordered_map>; - using MProfileMapPtr = std::shared_ptr; - - // requires exclusive ownership to write to *this - using reader_lock = std::unique_lock; - // requires shared ownership to read from other - using writer_lock = std::shared_lock; - - // Struct of continuing map if GEMMs to the best profiles for different Ms - struct MNKProfileMap - { - // Mutex guarding map - std::shared_timed_mutex mutex; - // Map from GEMM Id to profile for particular GEMM - std::unordered_map profileMap; - - bool existsMProfileMap(GemmIdType const& id) - { - auto const iter = profileMap.find(id); - return iter != profileMap.end(); - } - - void createMProfileMap(GemmIdType const& id) - { - profileMap[id] = std::make_shared(); - } - - MProfileMapPtr getMProfileMap(GemmIdType const& id) - { - auto const iter = profileMap.find(id); - if (iter == profileMap.end()) - { - std::ostringstream msg; - msg << "Cannot find ID (" << id << ") in the profile map. Abort."; - TLLM_THROW(msg.str()); - } - return iter->second; - } - }; - - using MNKProfileMapPtr = std::shared_ptr; - - GemmPluginProfiler(); - - virtual ~GemmPluginProfiler() = default; - - void serialize(char*& buffer, GemmIdType const& gemmId) const; - - void deserialize(char const*& data, GemmDims& dims, GemmIdType const& gemmId); - size_t getSerializationSize(GemmIdType const& gemmId) const; - - void profileTactics(RunnerPtr const& runner, nvinfer1::DataType const& type, GemmDims const& dims, - GemmIdType const& gemmId, bool hasWeightOnlyCudaKernel = false); - - void setSelectionTactics(MNKProfileMapPtr const& map) - { - mMNKProfileMap = map; - } - - void setTmpWorkspaceSizeInBytes(size_t bytes) - { - mTmpWorkspaceSizeInBytes = bytes; - } - - void setSkip(bool skip) - { - mSkip = mSkip || skip; - } - - std::optional getBestConfig(int m, GemmIdType const& gemmId) const; - - virtual int getMaxProfileM() const; - -protected: - virtual void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) = 0; - - virtual void computeTmpSize(size_t maxM, size_t n, size_t k) = 0; - - virtual bool checkTactic(int m, int n, int k, Config const& tactic) const - { - return true; - } - - virtual std::vector getTactics(int m, int n, int k) const = 0; - - virtual void initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream); - -private: - void allocateTmpData(); - - void freeTmpData(); - - std::optional profileTacticsForProblem(int m, int n, int k, std::vector const& tactics); - - float profileTacticForProblem(int m, int n, int k, Config const& tactic); - - int nextPowerOfTwo(int v) const - { - --v; - v |= v >> 1; - v |= v >> 2; - v |= v >> 4; - v |= v >> 8; - v |= v >> 16; - return ++v; - } - -protected: - RunnerPtr mRunner{nullptr}; - - nvinfer1::DataType mType{}; - -private: - MNKProfileMapPtr mMNKProfileMap{}; - - size_t mTmpWorkspaceSizeInBytes{0}; - - char* mWorkspaceTmp{nullptr}; - - cudaStream_t mStream; - - GemmDims mDims{}; - - bool mSkip{false}; -}; - -template -class GemmPluginProfilerManager -{ -public: - using MNKProfileMap = typename GemmPluginProfilerType::MNKProfileMap; - using MNKProfileMapPtr = typename GemmPluginProfilerType::MNKProfileMapPtr; - using GemmPluginProfilerPtr = std::shared_ptr; - - GemmPluginProfilerManager() - { - mMNKProfileMap = std::make_shared(); - } - - GemmPluginProfilerPtr createGemmPluginProfiler(bool inference, bool skip = false) - { - auto profiler = std::make_shared(); - profiler->setSkip(skip); - // If the profiler is created during the engine build, - // mMNKProfileMap is shared between different profilers to minimize the time spent on the profiling - // and do not repeat profiling for the GEMMs of the same shape. - if (!inference) - { - profiler->setSelectionTactics(mMNKProfileMap); - } - return profiler; - } - -private: - MNKProfileMapPtr mMNKProfileMap{}; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/common/plugin.cpp b/cpp/tensorrt_llm/plugins/common/plugin.cpp deleted file mode 100644 index 82c8bf93b13c..000000000000 --- a/cpp/tensorrt_llm/plugins/common/plugin.cpp +++ /dev/null @@ -1,124 +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 "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include "checkMacrosPlugin.h" -#include -#include - -#ifdef _MSC_VER -#define FN_NAME __FUNCTION__ -#else -#define FN_NAME __func__ -#endif - -PluginFieldParser::PluginFieldParser(int32_t nbFields, nvinfer1::PluginField const* fields) - : mFields{fields} -{ - for (int32_t i = 0; i < nbFields; i++) - { - mMap.emplace(fields[i].name, PluginFieldParser::Record{i}); - } -} - -PluginFieldParser::~PluginFieldParser() -{ - for (auto const& [name, record] : mMap) - { - if (!record.retrieved) - { - std::stringstream ss; - ss << "unused plugin field with name: " << name; - tensorrt_llm::plugins::logError(ss.str().c_str(), __FILE__, FN_NAME, __LINE__); - } - } -} - -template -nvinfer1::PluginFieldType toFieldType(); -#define SPECIALIZE_TO_FIELD_TYPE(T, type) \ - template <> \ - nvinfer1::PluginFieldType toFieldType() \ - { \ - return nvinfer1::PluginFieldType::type; \ - } -SPECIALIZE_TO_FIELD_TYPE(half, kFLOAT16) -SPECIALIZE_TO_FIELD_TYPE(float, kFLOAT32) -SPECIALIZE_TO_FIELD_TYPE(double, kFLOAT64) -SPECIALIZE_TO_FIELD_TYPE(int8_t, kINT8) -SPECIALIZE_TO_FIELD_TYPE(int16_t, kINT16) -SPECIALIZE_TO_FIELD_TYPE(int32_t, kINT32) -SPECIALIZE_TO_FIELD_TYPE(char, kCHAR) -SPECIALIZE_TO_FIELD_TYPE(nvinfer1::Dims, kDIMS) -SPECIALIZE_TO_FIELD_TYPE(void, kUNKNOWN) -#undef SPECIALIZE_TO_FIELD_TYPE - -template -std::optional PluginFieldParser::getScalar(std::string_view const& name) -{ - auto const iter = mMap.find(name); - if (iter == mMap.end()) - { - return std::nullopt; - } - auto& record = mMap.at(name); - auto const& f = mFields[record.index]; - TLLM_CHECK(toFieldType() == f.type && f.length == 1); - record.retrieved = true; - return std::optional{*static_cast(f.data)}; -} - -#define INSTANTIATE_PluginFieldParser_getScalar(T) \ - template std::optional PluginFieldParser::getScalar(std::string_view const&) -INSTANTIATE_PluginFieldParser_getScalar(half); -INSTANTIATE_PluginFieldParser_getScalar(float); -INSTANTIATE_PluginFieldParser_getScalar(double); -INSTANTIATE_PluginFieldParser_getScalar(int8_t); -INSTANTIATE_PluginFieldParser_getScalar(int16_t); -INSTANTIATE_PluginFieldParser_getScalar(int32_t); -INSTANTIATE_PluginFieldParser_getScalar(char); -INSTANTIATE_PluginFieldParser_getScalar(nvinfer1::Dims); -#undef INSTANTIATE_PluginFieldParser_getScalar - -template -std::optional> PluginFieldParser::getSet(std::string_view const& name) -{ - auto const iter = mMap.find(name); - if (iter == mMap.end()) - { - return std::nullopt; - } - auto& record = mMap.at(name); - auto const& f = mFields[record.index]; - TLLM_CHECK(toFieldType() == f.type); - std::set group; - auto const* r = static_cast(f.data); - for (int j = 0; j < f.length; ++j) - { - group.insert(*r); - ++r; - } - - record.retrieved = true; - return std::optional{group}; -} - -#define INSTANTIATE_PluginFieldParser_getVector(T) \ - template std::optional> PluginFieldParser::getSet(std::string_view const&) -INSTANTIATE_PluginFieldParser_getVector(int32_t); -#undef INSTANTIATE_PluginFieldParser_getVector diff --git a/cpp/tensorrt_llm/plugins/common/plugin.h b/cpp/tensorrt_llm/plugins/common/plugin.h deleted file mode 100644 index a7febe4cc13d..000000000000 --- a/cpp/tensorrt_llm/plugins/common/plugin.h +++ /dev/null @@ -1,143 +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. - */ - -#pragma once - -#include "tensorrt_llm/common/opUtils.h" -#include "tensorrt_llm/plugins/api/tllmPlugin.h" -#include "tensorrt_llm/plugins/common/checkMacrosPlugin.h" - -#include - -#include -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using namespace tensorrt_llm::common::op; - -class BasePlugin : public nvinfer1::IPluginV2DynamicExt -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept override - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -class BasePluginV3 : public nvinfer1::IPluginV3, - public nvinfer1::IPluginV3OneCore, - public nvinfer1::IPluginV3OneBuild, - public nvinfer1::IPluginV3OneRuntime -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -class BaseCreator : public nvinfer1::IPluginCreator -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept override - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -class BaseCreatorV3 : public nvinfer1::IPluginCreatorV3One -{ -public: - void setPluginNamespace(char const* libNamespace) noexcept - { - mNamespace = libNamespace; - } - - [[nodiscard]] char const* getPluginNamespace() const noexcept override - { - return mNamespace.c_str(); - } - -protected: - std::string mNamespace{api::kDefaultNamespace}; -}; - -} // namespace tensorrt_llm::plugins - -// Init with O(n) and retrieve with O(1) -class PluginFieldParser -{ -public: - // field array must remain valid when calling getScalar() later. - PluginFieldParser(int32_t nbFields, nvinfer1::PluginField const* fields); - // delete to remind accidental mis-use (copy) which may result in false-alarm warnings about unused fields. - PluginFieldParser(PluginFieldParser const&) = delete; - PluginFieldParser& operator=(PluginFieldParser const&) = delete; - // check if all fields are retrieved and emit warning if some of them are not. - ~PluginFieldParser(); - template - std::optional getScalar(std::string_view const& name); - template - std::optional> getSet(std::string_view const& name); - -private: - nvinfer1::PluginField const* mFields; - - struct Record - { - Record(int32_t idx) - : index{idx} - { - } - - int32_t const index; - bool retrieved{false}; - }; - - std::unordered_map mMap; -}; diff --git a/cpp/tensorrt_llm/plugins/common/pluginUtils.h b/cpp/tensorrt_llm/plugins/common/pluginUtils.h deleted file mode 100644 index ee3e59d57c6d..000000000000 --- a/cpp/tensorrt_llm/plugins/common/pluginUtils.h +++ /dev/null @@ -1,78 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#pragma once - -#include - -#include "tensorrt_llm/common/logger.h" - -namespace tensorrt_llm::plugins::utils -{ -using DimType64 = int64_t; - -inline DimType64 computeMDimension(bool transA, nvinfer1::Dims const& dims) -{ - DimType64 M{1}; - if (transA) - { - for (int i = dims.nbDims - 1; i > 0; --i) - { - M *= dims.d[i]; - } - } - else - { - for (int i = 0; i < dims.nbDims - 1; ++i) - { - M *= dims.d[i]; - } - } - return M; -} - -inline DimType64 computeNDimension(bool transB, nvinfer1::Dims const& dims) -{ - DimType64 N{1}; - if (transB) - { - for (int32_t i = 0; i < dims.nbDims - 1; ++i) - { - N *= dims.d[i]; - } - } - else - { - for (int32_t i = dims.nbDims - 1; i > 0; --i) - { - N *= dims.d[i]; - } - } - return N; -} - -inline std::int32_t logErrorReturn0(char const* variable) -{ - TLLM_LOG_ERROR("Value of %s is out of range for int32_t", variable); - return 0; -} - -#define TLLM_INT32_CAST(value) \ - ((value > 0x7FFFFFFFLL || value < -0x80000000LL) ? tensorrt_llm::plugins::utils::logErrorReturn0(#value) \ - : static_cast(value)) - -} // namespace tensorrt_llm::plugins::utils diff --git a/cpp/tensorrt_llm/plugins/cpSplitPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/cpSplitPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/cpSplitPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp b/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp deleted file mode 100644 index 221d5dac2da1..000000000000 --- a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.cpp +++ /dev/null @@ -1,356 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 - -#include "cpSplitPlugin.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::CpSplitPluginCreator; -using tensorrt_llm::plugins::CpSplitPlugin; - -static char const* CPSPLIT_PLUGIN_VERSION{"1"}; -static char const* CPSPLIT_PLUGIN_NAME{"CpSplit"}; -PluginFieldCollection CpSplitPluginCreator::mFC{}; -std::vector CpSplitPluginCreator::mPluginAttributes; - -CpSplitPlugin::CpSplitPlugin() -{ - initFieldsToSerialize(); -} - -CpSplitPlugin::CpSplitPlugin(int cpSize, int cpRank) - : mCpSize(cpSize) - , mCpRank(cpRank) -{ - initFieldsToSerialize(); -} - -void CpSplitPlugin::initFieldsToSerialize() -{ - mDataToSerialize.clear(); - mDataToSerialize.emplace_back(PluginField("cp_size", &mCpSize, PluginFieldType::kINT32, 1)); - mDataToSerialize.emplace_back(PluginField("cp_rank", &mCpRank, PluginFieldType::kINT32, 1)); - mFCToSerialize.nbFields = mDataToSerialize.size(); - mFCToSerialize.fields = mDataToSerialize.data(); -} - -// IPluginV3 methods -nvinfer1::IPluginCapability* CpSplitPlugin::getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept -{ - switch (type) - { - case PluginCapabilityType::kBUILD: return static_cast(this); - case PluginCapabilityType::kRUNTIME: return static_cast(this); - case PluginCapabilityType::kCORE: return static_cast(this); - } - return nullptr; -} - -nvinfer1::IPluginV3* CpSplitPlugin::clone() noexcept -{ - std::unique_ptr plugin{std::make_unique(*this)}; - plugin->setPluginNamespace(mNamespace.c_str()); - plugin->initFieldsToSerialize(); - return plugin.release(); -} - -// IPluginV3OneCore methods -char const* CpSplitPlugin::getPluginName() const noexcept -{ - return CPSPLIT_PLUGIN_NAME; -} - -char const* CpSplitPlugin::getPluginVersion() const noexcept -{ - return CPSPLIT_PLUGIN_VERSION; -} - -// IPluginV3OneBuild methods -int32_t CpSplitPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::getOutputDataTypes( - DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept -{ - outputTypes[0] = inputTypes[0]; - outputTypes[1] = DataType::kINT32; - outputTypes[2] = DataType::kINT32; - return 0; -} - -int32_t CpSplitPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs, - int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept -{ - outputs[0].nbDims = 1; - - auto cpSize = exprBuilder.constant(mCpSize); - auto upper = inputs[0].d[0]; - auto opt = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *upper, *cpSize); - outputs[0].d[0] = exprBuilder.declareSizeTensor(1, *opt, *upper); - - // We must have such an output size tensor (with dim == 0) to notify the shape of output tensor above - outputs[1].nbDims = 0; - outputs[2].nbDims = 1; - outputs[2].d[0] = upper; - return 0; -} - -bool CpSplitPlugin::supportsFormatCombination( - int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept -{ - if (pos == IdxEntry::INPUT_IDS) - { - return ((inOut[pos].desc.type == DataType::kINT32) && (inOut[pos].desc.format == TensorFormat::kLINEAR)); - } - else if (pos == IdxEntry::REQUEST_TYPES || pos == IdxEntry::HOST_CONTEXT_LENGTH) - { - return inOut[pos].desc.type == DataType::kINT32; - } - else - { - return ((inOut[pos].desc.type == DataType::kINT32) && (inOut[pos].desc.format == TensorFormat::kLINEAR)); - } - return false; -} - -int32_t CpSplitPlugin::getNbOutputs() const noexcept -{ - return 3; -} - -size_t CpSplitPlugin::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::getNbTactics() noexcept -{ - return 0; -} - -char const* CpSplitPlugin::getTimingCacheID() noexcept -{ - return nullptr; -} - -int32_t CpSplitPlugin::getFormatCombinationLimit() noexcept -{ - return 1; -} - -char const* CpSplitPlugin::getMetadataString() noexcept -{ - return nullptr; -} - -// IPluginV3OneRuntime methods -int32_t CpSplitPlugin::setTactic(int32_t tactic) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::PluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t CpSplitPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // @param inputIds [tokenNum] - // @param host_request_types [batchSize]: Tensor = None (On CPU) - // The tensor on the host that indicates if a request is in context or - // generation phase. Its shape is [batch_size]. See Inflight Batching - // in docs/gpt_attention.md, - // @param host_context_lengths [batchSize]: Tensor = None (On CPU) - // A host tensor that contains the lengths of the different inputs - // outputs - // @param outputIds [tokenNum spiltted by cp] - // @param outputLength scalar - // @param joinIdx [tokenNum] - - int64_t tokenNum = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - tokenNum *= inputDesc[0].dims.d[i]; - } - - RequestType const* reqTypes = static_cast(inputs[IdxEntry::REQUEST_TYPES]); - int32_t const* hContextLengths = static_cast(inputs[IdxEntry::HOST_CONTEXT_LENGTH]); - int const* inputIds = reinterpret_cast(inputs[IdxEntry::INPUT_IDS]); - int* outputIds = reinterpret_cast(outputs[0]); - int32_t* outputLength = reinterpret_cast(outputs[1]); - int32_t* outputJoinIdx = reinterpret_cast(outputs[2]); - - int32_t const nbSeq = inputDesc[IdxEntry::HOST_CONTEXT_LENGTH].dims.d[0]; - - int32_t* hInputs = new int[inputDesc[IdxEntry::INPUT_IDS].dims.d[0]]; - int32_t* hOutputs = new int[inputDesc[IdxEntry::INPUT_IDS].dims.d[0]]; - int32_t* hOutputJoinIdx = new int[inputDesc[IdxEntry::INPUT_IDS].dims.d[0]]; - cudaMemcpyAsync( - hInputs, inputIds, sizeof(int32_t) * inputDesc[IdxEntry::INPUT_IDS].dims.d[0], cudaMemcpyDeviceToHost, stream); - sync_check_cuda_error(stream); - - int32_t inputIdx = 0; - int32_t outputIdx = 0; - for (int32_t seqIdx = 0; seqIdx < nbSeq; ++seqIdx) - { - if (reqTypes[seqIdx] == RequestType::kCONTEXT) - { - auto const& ctxLength = hContextLengths[seqIdx]; - int32_t partialAverageLength = (ctxLength + mCpSize - 1) / mCpSize; - int32_t partialLength - = mCpRank == mCpSize - 1 ? ctxLength - partialAverageLength * (mCpSize - 1) : partialAverageLength; - for (int i = 0; i < partialLength; i++) - { - hOutputs[outputIdx + i] = hInputs[inputIdx + partialAverageLength * mCpRank + i]; - } - inputIdx += ctxLength; - outputIdx += partialAverageLength; - } - else if (reqTypes[seqIdx] == RequestType::kGENERATION) - { - auto const& genLength = nbSeq - seqIdx; - int32_t partialAverageLength = (genLength + mCpSize - 1) / mCpSize; - int32_t partialLength - = mCpRank == mCpSize - 1 ? genLength - partialAverageLength * (mCpSize - 1) : partialAverageLength; - for (int i = 0; i < partialLength; i++) - { - hOutputs[outputIdx + i] = hInputs[inputIdx + partialAverageLength * mCpRank + i]; - } - outputIdx += partialAverageLength; - break; - } - } - int32_t hOutputLength = outputIdx; - inputIdx = 0; - outputIdx = 0; - for (int32_t seqIdx = 0; seqIdx < nbSeq; ++seqIdx) - { - if (reqTypes[seqIdx] == RequestType::kCONTEXT) - { - auto const& ctxLength = hContextLengths[seqIdx]; - int32_t partialAverageLength = (ctxLength + mCpSize - 1) / mCpSize; - for (int32_t idx = 0; idx < ctxLength; ++idx) - { - hOutputJoinIdx[inputIdx + idx] - = idx % partialAverageLength + idx / partialAverageLength * hOutputLength + outputIdx; - } - inputIdx += ctxLength; - outputIdx += partialAverageLength; - } - else if (reqTypes[seqIdx] == RequestType::kGENERATION) - { - auto const& genLength = nbSeq - seqIdx; - int32_t partialAverageLength = (genLength + mCpSize - 1) / mCpSize; - for (int32_t idx = 0; idx < genLength; ++idx) - { - hOutputJoinIdx[inputIdx + idx] - = idx % partialAverageLength + idx / partialAverageLength * hOutputLength + outputIdx; - } - break; - } - } - cudaMemcpyAsync(outputIds, hOutputs, sizeof(int32_t) * hOutputLength, cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(outputLength, &hOutputLength, sizeof(int32_t), cudaMemcpyHostToDevice, stream); - cudaMemcpyAsync(outputJoinIdx, hOutputJoinIdx, sizeof(int32_t) * tokenNum, cudaMemcpyHostToDevice, stream); - sync_check_cuda_error(stream); - return 0; -} - -nvinfer1::IPluginV3* CpSplitPlugin::attachToContext(nvinfer1::IPluginResourceContext* context) noexcept -{ - return clone(); -} - -nvinfer1::PluginFieldCollection const* CpSplitPlugin::getFieldsToSerialize() noexcept -{ - return &mFCToSerialize; -} - -CpSplitPluginCreator::CpSplitPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* CpSplitPluginCreator::getPluginName() const noexcept -{ - return CPSPLIT_PLUGIN_NAME; -} - -char const* CpSplitPluginCreator::getPluginVersion() const noexcept -{ - return CPSPLIT_PLUGIN_VERSION; -} - -nvinfer1::PluginFieldCollection const* CpSplitPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -nvinfer1::IPluginV3* CpSplitPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept -{ - PluginField const* fields = fc->fields; - int cp_size{}; - int cp_rank{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "cp_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "cp_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - cp_rank = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new CpSplitPlugin(cp_size, cp_rank); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h b/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h deleted file mode 100644 index 1dc8c15b355a..000000000000 --- a/cpp/tensorrt_llm/plugins/cpSplitPlugin/cpSplitPlugin.h +++ /dev/null @@ -1,112 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class CpSplitPlugin : public BasePluginV3 -{ -public: - CpSplitPlugin(); - CpSplitPlugin(int cpSize, int cpRank); - CpSplitPlugin(CpSplitPlugin const& p) = default; - void initFieldsToSerialize(); - - // IPluginV3 methods - nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override; - nvinfer1::IPluginV3* clone() noexcept override; - - // IPluginV3OneCore methods - char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; - - // IPluginV3OneBuild methods - int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override; // nochange - int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes, - int32_t nbInputs) const noexcept override; // fixed - int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs, - int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; // fixed - bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, - int32_t nbOutputs) noexcept override; // fixed - int32_t getNbOutputs() const noexcept override; // fixed - size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override; // fixed - int32_t getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept override; - int32_t getNbTactics() noexcept override; - char const* getTimingCacheID() noexcept override; - int32_t getFormatCombinationLimit() noexcept override; - char const* getMetadataString() noexcept override; - - // IPluginV3OneRuntime methods - int32_t setTactic(int32_t tactic) noexcept override; - int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept override; // fixed - nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override; - nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override; - -private: - int mCpSize; - int mCpRank; - std::vector mDataToSerialize; - nvinfer1::PluginFieldCollection mFCToSerialize; - - enum IdxEntry - { - INPUT_IDS, - REQUEST_TYPES, - HOST_CONTEXT_LENGTH, - }; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; -}; - -class CpSplitPluginCreator : public BaseCreatorV3 -{ -public: - CpSplitPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV3* createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/cudaStreamPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.cpp b/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.cpp deleted file mode 100644 index 802e828c9250..000000000000 --- a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.cpp +++ /dev/null @@ -1,293 +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. - */ - -#include "cudaStreamPlugin.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::CudaStreamPluginCreator; -using tensorrt_llm::plugins::CudaStreamPlugin; - -static char const* CUDA_STREAM_PLUGIN_VERSION{"1"}; -static char const* CUDA_STREAM_PLUGIN_NAME{"CudaStream"}; -PluginFieldCollection CudaStreamPluginCreator::mFC{}; -std::vector CudaStreamPluginCreator::mPluginAttributes; - -CudaStreamPlugin::CudaStreamPlugin(int sideStreamId, int nbInputs, nvinfer1::DataType type) - : mSideStreamId(sideStreamId) - , mNbInputs(nbInputs) - , mType(type) -{ - init(); -} - -CudaStreamPlugin::CudaStreamPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mSideStreamId); - read(d, mNbInputs); - read(d, mType); - - init(); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -CudaStreamPlugin::CudaStreamPlugin(CudaStreamPlugin const& other) - : mSideStreamId(other.mSideStreamId) - , mNbInputs(other.mNbInputs) - , mType(other.mType) -{ - init(); -} - -void CudaStreamPlugin::init() -{ - mSideStreamPtr = nullptr; -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* CudaStreamPlugin::clone() const noexcept -{ - auto* plugin = new CudaStreamPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs CudaStreamPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - assert(outputIndex == 0); - return inputs[outputIndex]; -} - -bool CudaStreamPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_CHECK_WITH_INFO(nbInputs == mNbInputs, "CudaStreamPlugin only accepts mNbInputs inputs"); - TLLM_CHECK_WITH_INFO(nbOutputs == 1, "CudaStreamPlugin only accepts 1 output"); - - auto const& desc = inOut[pos]; - if (desc.format != TensorFormat::kLINEAR) - { - return false; - } - - if (pos > 0 && pos < nbInputs) - { - return true; - } - return desc.type == mType; -} - -void CudaStreamPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t CudaStreamPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int CudaStreamPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (!mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - nvinfer1::pluginInternal::SideStream side_stream{}; - mSideStreamPtr = reinterpret_cast( - getPluginRegistry()->acquirePluginResource(resource_name.c_str(), &side_stream)); - } - mSideStreamPtr->waitSideStreamOnMainStream(stream); - size_t count = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - count *= inputDesc[0].dims.d[i]; - } - count *= tensorrt_llm::runtime::BufferDataType(inputDesc[0].type).getSize(); - TLLM_CUDA_CHECK(cudaMemcpyAsync(outputs[0], inputs[0], count, cudaMemcpyDeviceToDevice, stream)); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType CudaStreamPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* CudaStreamPlugin::getPluginType() const noexcept -{ - return CUDA_STREAM_PLUGIN_NAME; -} - -char const* CudaStreamPlugin::getPluginVersion() const noexcept -{ - return CUDA_STREAM_PLUGIN_VERSION; -} - -int CudaStreamPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int CudaStreamPlugin::initialize() noexcept -{ - return 0; -} - -void CudaStreamPlugin::terminate() noexcept -{ - if (mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - getPluginRegistry()->releasePluginResource(resource_name.c_str()); - mSideStreamPtr = nullptr; - } -} - -size_t CudaStreamPlugin::getSerializationSize() const noexcept -{ - return sizeof(mSideStreamId) + sizeof(mNbInputs) + sizeof(mType); -} - -void CudaStreamPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mSideStreamId); - write(d, mNbInputs); - write(d, mType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void CudaStreamPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -CudaStreamPluginCreator::CudaStreamPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("side_stream_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_inputs", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* CudaStreamPluginCreator::getPluginName() const noexcept -{ - return CUDA_STREAM_PLUGIN_NAME; -} - -char const* CudaStreamPluginCreator::getPluginVersion() const noexcept -{ - return CUDA_STREAM_PLUGIN_VERSION; -} - -PluginFieldCollection const* CudaStreamPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* CudaStreamPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int sideStreamId; - int nbInputs; - int type; - - // Read configurations from each fields - struct MapPair - { - char const* key; - int& field; - bool optional = false; - bool set = false; - }; - - std::array input_map{ - MapPair{"side_stream_id", std::ref(sideStreamId)}, - MapPair{"num_inputs", std::ref(nbInputs)}, - MapPair{"type_id", std::ref(type)}, - }; - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - for (auto& item : input_map) - { - if (!strcmp(item.key, attrName)) - { - TLLM_CHECK(fields[i].type == nvinfer1::PluginFieldType::kINT32); - TLLM_CHECK_WITH_INFO(!item.set, "Parameter %s was set twice", item.key); - item.field = static_cast(*(static_cast(fields[i].data))); - item.set = true; - } - } - } - - for (auto& item : input_map) - { - TLLM_CHECK_WITH_INFO(item.set || item.optional, "Parameter %s is required but not set", item.key); - } - - try - { - auto* obj = new CudaStreamPlugin(sideStreamId, nbInputs, static_cast(type)); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* CudaStreamPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call CudaStreamPlugin::destroy() - try - { - auto* obj = new CudaStreamPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h b/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h deleted file mode 100644 index 5b78c3b873bb..000000000000 --- a/cpp/tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h +++ /dev/null @@ -1,265 +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. - */ -#pragma once - -#include "NvInferPlugin.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/cudaMemPool.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" -#include -#include -#include - -namespace nvinfer1 -{ -namespace pluginInternal -{ -class SideWorkspace -{ -public: - SideWorkspace(cudaStream_t stream) - : mWorkspaceSize{0} - , mWorkspacePtr{nullptr} - , mStream{stream} - { - } - - ~SideWorkspace() - { - if (mWorkspacePtr) - { - TLLM_CUDA_CHECK(cudaFreeAsync(mWorkspacePtr, mStream)); - } - } - - void* get(size_t workspaceSize) - { - if (mWorkspacePtr && mWorkspaceSize < workspaceSize) - { - TLLM_CUDA_CHECK(cudaFreeAsync(mWorkspacePtr, mStream)); - mWorkspacePtr = nullptr; - } - if (!mWorkspacePtr) - { - mWorkspaceSize = workspaceSize; - auto pool_ptr - = tensorrt_llm::runtime::CudaMemPool::getPrimaryPoolForDevice(tensorrt_llm::common::getDevice()); - TLLM_CUDA_CHECK(cudaMallocFromPoolAsync(&mWorkspacePtr, mWorkspaceSize, pool_ptr->getPool(), mStream)); - } - return mWorkspacePtr; - } - -private: - size_t mWorkspaceSize; - void* mWorkspacePtr; - cudaStream_t mStream; -}; - -class SideStream : public IPluginResource -{ -public: - SideStream(bool init = false) - : mStream{} - , mMainEvent{} - , mSideEvent{} - , mWorkspace{} - , mInit{init} - { - // The object passed to acquirePluginResource should use the default value init=false - if (init) - { - TLLM_CUDA_CHECK(cudaStreamCreate(&mStream)); - TLLM_CUDA_CHECK(cudaEventCreateWithFlags(&mMainEvent, cudaEventDisableTiming)); - TLLM_CUDA_CHECK(cudaEventCreateWithFlags(&mSideEvent, cudaEventDisableTiming)); - mWorkspace = std::make_shared(mStream); - } - } - - void free() - { - if (mInit) - { - mWorkspace = nullptr; - TLLM_CUDA_CHECK(cudaStreamSynchronize(mStream)); - TLLM_CUDA_CHECK(cudaStreamDestroy(mStream)); - TLLM_CUDA_CHECK(cudaEventDestroy(mMainEvent)); - TLLM_CUDA_CHECK(cudaEventDestroy(mSideEvent)); - mInit = false; - } - } - - int32_t release() noexcept override - { - try - { - free(); - } - catch (std::exception const& e) - { - return -1; - } - return 0; - } - - IPluginResource* clone() noexcept override - { - // An object is cloned only when calling acquirePluginResource for the first time for each key - std::unique_ptr cloned{}; - try - { - if (!mInit) - { - cloned = std::make_unique(/* init */ true); - } - else - { - return nullptr; - } - } - catch (std::exception const& e) - { - return nullptr; - } - return cloned.release(); - } - - ~SideStream() override - { - free(); - } - - void* getWorkspacePtr(size_t workspaceSize) - { - return mWorkspace->get(workspaceSize); - } - - cudaStream_t getStream() const - { - return mStream; - } - - void waitMainStreamOnSideStream(cudaStream_t const stream) const - { - TLLM_CUDA_CHECK(cudaEventRecord(mMainEvent, stream)); - TLLM_CUDA_CHECK(cudaStreamWaitEvent(mStream, mMainEvent)); - } - - void waitSideStreamOnMainStream(cudaStream_t const stream) const - { - TLLM_CUDA_CHECK(cudaEventRecord(mSideEvent, mStream)); - TLLM_CUDA_CHECK(cudaStreamWaitEvent(stream, mSideEvent)); - } - - void stallMainStream(char const* name, cudaStream_t const stream, std::optional delay = std::nullopt) const - { - tensorrt_llm::runtime::utils::stallStream(name, stream, delay); - } - - void stallSideStream(char const* name, std::optional delay = std::nullopt) const - { - tensorrt_llm::runtime::utils::stallStream(name, mStream, delay); - } - - static std::string getResourceKey(int const stream_id) - { - return "side_stream_" + std::to_string(stream_id); - } - -private: - cudaStream_t mStream; - cudaEvent_t mMainEvent; - cudaEvent_t mSideEvent; - std::shared_ptr mWorkspace; - bool mInit; -}; - -} // namespace pluginInternal -} // namespace nvinfer1 - -namespace tensorrt_llm::plugins -{ - -class CudaStreamPlugin : public BasePlugin -{ -public: - CudaStreamPlugin(int sideStreamId, int nbInputs, nvinfer1::DataType type); - - CudaStreamPlugin(void const* data, size_t length); - - CudaStreamPlugin(CudaStreamPlugin const&); - - void init(); - - ~CudaStreamPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - int mSideStreamId; - int mNbInputs; - nvinfer1::DataType mType; - nvinfer1::pluginInternal::SideStream* mSideStreamPtr; -}; - -class CudaStreamPluginCreator : public BaseCreator -{ -public: - CudaStreamPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/CMakeLists.txt deleted file mode 100644 index ea25de075f34..000000000000 --- a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/CMakeLists.txt +++ /dev/null @@ -1,22 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.cpp b/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.cpp deleted file mode 100644 index 927a42ebac2f..000000000000 --- a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.cpp +++ /dev/null @@ -1,299 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "cumsumLastDimPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::CumsumLastDimPluginCreator; -using tensorrt_llm::plugins::CumsumLastDimPlugin; - -static char const* CUMSUM_LAST_DIM_PLUGIN_VERSION{"1"}; -static char const* CUMSUM_LAST_DIM_PLUGIN_NAME{"CumsumLastDim"}; -PluginFieldCollection CumsumLastDimPluginCreator::mFC{}; -std::vector CumsumLastDimPluginCreator::mPluginAttributes; - -static constexpr SizeType32 LENGTH_LIMIT_FOR_BLOCKSCAN = 4096; - -CumsumLastDimPlugin::CumsumLastDimPlugin(SizeType32 inputLength, nvinfer1::DataType type, size_t temp_storage_bytes) - : mInputLength(inputLength) - , mTempStorageBytes(temp_storage_bytes) - , mType(type) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); - if (mTempStorageBytes == 0) - { - mTempStorageBytes = getWorkspaceSizeNeeded(inputLength, type); - } -} - -// Parameterized constructor -CumsumLastDimPlugin::CumsumLastDimPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mInputLength); - read(d, mTempStorageBytes); - read(d, mType); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* CumsumLastDimPlugin::clone() const noexcept -{ - auto* plugin = new CumsumLastDimPlugin(mInputLength, mType, mTempStorageBytes); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, inputLength] -nvinfer1::DimsExprs CumsumLastDimPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK_WITH_INFO(outputIndex == 0, "Only one output."); - return inputs[getInputTensorIdx()]; -} - -bool CumsumLastDimPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void CumsumLastDimPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t CumsumLastDimPlugin::getWorkspaceSizeNeeded(SizeType32 inputLength, nvinfer1::DataType type) -{ - size_t tempStorageBytes{0}; - if (inputLength < LENGTH_LIMIT_FOR_BLOCKSCAN) // last dim unknown or small, use BlockScan - { - tempStorageBytes = 0; - } - else if (type == DataType::kINT32) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize(inputLength); - } - else if (type == DataType::kHALF) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize(inputLength); - } - else if (type == DataType::kFLOAT) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize(inputLength); - } -#ifdef ENABLE_BF16 - else if (type == DataType::kBF16) - { - tempStorageBytes = invokeComputeCumsumLastDimWorkspaceSize<__nv_bfloat16>(inputLength); - } -#endif - return tempStorageBytes; -} - -size_t CumsumLastDimPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mTempStorageBytes; -} - -template -int CumsumLastDimPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, inputLength] - // outputs - // 0. output_tensor [batch_size, inputLength] - auto const batchSize = inputDesc[getInputTensorIdx()].dims.d[0]; - auto const inputLength = inputDesc[getInputTensorIdx()].dims.d[1]; - /* - Two cases where we should use BlockScan: - 1. inputLength is small - 2. batchSize is large (since DeviceScan causes kernel launch per row) - */ - void* wp = inputLength < LENGTH_LIMIT_FOR_BLOCKSCAN || batchSize > 2 ? nullptr : workspace; - invokeCumsumLastDim( - batchSize, inputLength, inputs[getInputTensorIdx()], outputs[0], wp, mTempStorageBytes, stream); - - sync_check_cuda_error(stream); - return 0; -} - -int CumsumLastDimPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (mType == DataType::kINT32) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType CumsumLastDimPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index == 0, "Only one output."); - return inputTypes[getInputTensorIdx()]; -} - -// IPluginV2 Methods - -char const* CumsumLastDimPlugin::getPluginType() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_NAME; -} - -char const* CumsumLastDimPlugin::getPluginVersion() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_VERSION; -} - -int CumsumLastDimPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int CumsumLastDimPlugin::initialize() noexcept -{ - return 0; -} - -void CumsumLastDimPlugin::terminate() noexcept {} - -size_t CumsumLastDimPlugin::getSerializationSize() const noexcept -{ - return sizeof(mInputLength) + sizeof(mTempStorageBytes) + sizeof(mType); -} - -void CumsumLastDimPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mInputLength); - write(d, mTempStorageBytes); - write(d, mType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void CumsumLastDimPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -CumsumLastDimPluginCreator::CumsumLastDimPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("input_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* CumsumLastDimPluginCreator::getPluginName() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_NAME; -} - -char const* CumsumLastDimPluginCreator::getPluginVersion() const noexcept -{ - return CUMSUM_LAST_DIM_PLUGIN_VERSION; -} - -PluginFieldCollection const* CumsumLastDimPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* CumsumLastDimPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int inputLength{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "input_length")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - inputLength = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new CumsumLastDimPlugin(inputLength, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* CumsumLastDimPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call CumsumLastDimPlugin::destroy() - try - { - auto* obj = new CumsumLastDimPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h b/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h deleted file mode 100644 index 3cbf4e2356dd..000000000000 --- a/cpp/tensorrt_llm/plugins/cumsumLastDimPlugin/cumsumLastDimPlugin.h +++ /dev/null @@ -1,102 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_CUMSUM_LAST_DIM_PLUGIN_H -#define TRT_CUMSUM_LAST_DIM_PLUGIN_H - -#include "tensorrt_llm/kernels/cumsumLastDim.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -class CumsumLastDimPlugin : public BasePlugin -{ -public: - using SizeType32 = tensorrt_llm::kernels::SizeType32; - - CumsumLastDimPlugin(SizeType32 inputLength, nvinfer1::DataType type, size_t tempStorageBytes = 0); - CumsumLastDimPlugin(void const* data, size_t length); - ~CumsumLastDimPlugin() override = default; - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - size_t getWorkspaceSizeNeeded(SizeType32 inputLength, nvinfer1::DataType type); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - -private: - SizeType32 mInputLength; - size_t mTempStorageBytes; - nvinfer1::DataType mType; -}; - -class CumsumLastDimPluginCreator : public BaseCreator -{ -public: - CumsumLastDimPluginCreator(); - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif diff --git a/cpp/tensorrt_llm/plugins/doraPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/doraPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/doraPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.cpp b/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.cpp deleted file mode 100644 index 7c980f079ceb..000000000000 --- a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.cpp +++ /dev/null @@ -1,392 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "doraPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/runtime/iBuffer.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::DoraPlugin; -using tensorrt_llm::plugins::DoraPluginCreator; - -static char const* DORA_PLUGIN_VERSION{"1"}; -static char const* DORA_PLUGIN_NAME{"Dora"}; -PluginFieldCollection DoraPluginCreator::mFC{}; -std::vector DoraPluginCreator::mPluginAttributes; - -DoraPlugin::DoraPlugin(std::vector const& outHiddenSizes, nvinfer1::DataType type, bool removeInputPadding) - : mType(type) - , mRemoveInputPadding(removeInputPadding) - , mDoraImpl(outHiddenSizes, type) -{ - mOutHiddenSizes.resize(outHiddenSizes.size()); - mOutHiddenSizes.assign(outHiddenSizes.cbegin(), outHiddenSizes.cend()); - init(); -} - -void DoraPlugin::init() -{ - // initialize data to serialize - mDataToSerialize.clear(); - mDataToSerialize.emplace_back( - "out_hidden_sizes", mOutHiddenSizes.data(), PluginFieldType::kINT32, mOutHiddenSizes.size()); - mDataToSerialize.emplace_back("type", &mType, PluginFieldType::kINT32, 1); - mDataToSerialize.emplace_back("remove_input_padding", &mRemoveInputPadding, PluginFieldType::kINT8, 1); - mFieldsToSerialize.nbFields = static_cast(mDataToSerialize.size()); - mFieldsToSerialize.fields = mDataToSerialize.data(); -} - -// IPluginV3 methods -nvinfer1::IPluginCapability* DoraPlugin::getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept -{ - switch (type) - { - case PluginCapabilityType::kBUILD: return static_cast(this); - case PluginCapabilityType::kRUNTIME: return static_cast(this); - case PluginCapabilityType::kCORE: return static_cast(this); - } - return nullptr; -} - -nvinfer1::IPluginV3* DoraPlugin::clone() noexcept -{ - std::unique_ptr plugin{std::make_unique(mOutHiddenSizes, mType, mRemoveInputPadding)}; - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin.release(); -} - -// IPluginV3OneCore methods -char const* DoraPlugin::getPluginName() const noexcept -{ - return DORA_PLUGIN_NAME; -} - -char const* DoraPlugin::getPluginVersion() const noexcept -{ - return DORA_PLUGIN_VERSION; -} - -// IPluginV3OneBuild methods -int32_t DoraPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t DoraPlugin::getOutputDataTypes( - DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept -{ - try - { - TLLM_CHECK(nbOutputs == 1); - TLLM_CHECK(nbInputs == 2 + static_cast(mOutHiddenSizes.size()) + (mRemoveInputPadding ? 1 : 0)); - TLLM_CHECK(inputTypes[IdxEntry::kINPUT_TENSOR] == mType); - // output has the same dtype as the input, the plugin just applies scaling - outputTypes[0] = inputTypes[IdxEntry::kINPUT_TENSOR]; - } - catch (std::exception const& e) - { - caughtError(e); - } - return 0; -} - -int32_t DoraPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs, - int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbOutputs == 1); - TLLM_CHECK(nbShapeInputs == 0); - TLLM_CHECK(nbInputs == 2 + static_cast(mOutHiddenSizes.size()) + (mRemoveInputPadding ? 1 : 0)); - - auto const inputTensorDims = inputs[IdxEntry::kINPUT_TENSOR]; - TLLM_CHECK(inputTensorDims.nbDims == (mRemoveInputPadding ? 2 : 3)); - - auto const lastDim = inputTensorDims.d[inputTensorDims.nbDims - 1]; - TLLM_CHECK(lastDim->isConstant()); - TLLM_CHECK(lastDim->getConstantValue() == std::accumulate(mOutHiddenSizes.cbegin(), mOutHiddenSizes.cend(), 0)); - - outputs[0].nbDims = inputTensorDims.nbDims; - for (auto dim = 0; dim < inputTensorDims.nbDims; ++dim) - { - outputs[0].d[dim] = inputTensorDims.d[dim]; - } - } - catch (std::exception const& e) - { - caughtError(e); - } - return 0; -} - -bool DoraPlugin::supportsFormatCombination( - int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - auto const numModules = static_cast(mOutHiddenSizes.size()); - if (nbInputs != 2 + numModules + (mRemoveInputPadding ? 1 : 0)) - { - return false; - } - - bool const isInput = pos < nbInputs; - if (pos == IdxEntry::kHOST_REQUEST_TYPES) - { - return (inOut[pos].desc.type == nvinfer1::DataType::kINT32); - } - // optional host_context_lens after lora pointers - else if (pos == IdxEntry::kLORA_WEIGHTS_PTRS_START + numModules and isInput) - { - return (inOut[pos].desc.type == nvinfer1::DataType::kINT32 and mRemoveInputPadding); - } - // lora weight pointers - else if (pos >= IdxEntry::kLORA_WEIGHTS_PTRS_START and pos < IdxEntry::kLORA_WEIGHTS_PTRS_START + numModules) - { - return (inOut[pos].desc.type == nvinfer1::DataType::kINT64); - } - else if (pos != 0 and isInput) - { - TLLM_LOG_WARNING("%s: got an unexpected input at position %d", __PRETTY_FUNCTION__, pos); - return false; - } - - return (inOut[pos].desc.type == mType) and (inOut[pos].desc.format == TensorFormat::kLINEAR); -} - -int32_t DoraPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -size_t DoraPlugin::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept -{ - auto const inputTensorMax = inputs[IdxEntry::kINPUT_TENSOR].max; - auto const maxNumTokens = mRemoveInputPadding ? inputTensorMax.d[0] : inputTensorMax.d[0] * inputTensorMax.d[1]; - auto const size = mDoraImpl.getWorkspaceSize(maxNumTokens); - return size; -} - -int32_t DoraPlugin::getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept -{ - return 0; -} - -int32_t DoraPlugin::getNbTactics() noexcept -{ - return 0; -} - -char const* DoraPlugin::getTimingCacheID() noexcept -{ - return nullptr; -} - -int32_t DoraPlugin::getFormatCombinationLimit() noexcept -{ - return 1; -} - -char const* DoraPlugin::getMetadataString() noexcept -{ - return nullptr; -} - -// IPluginV3OneRuntime methods -int32_t DoraPlugin::setTactic(int32_t tactic) noexcept -{ - return 0; -} - -int32_t DoraPlugin::onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::PluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -int32_t DoraPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - - auto const numModules = static_cast(mOutHiddenSizes.size()); - auto const numReqs = inputDesc[IdxEntry::kHOST_REQUEST_TYPES].dims.d[0]; - - auto const inputTensorDesc = inputDesc[IdxEntry::kINPUT_TENSOR]; - auto const numTokens - = mRemoveInputPadding ? inputTensorDesc.dims.d[0] : inputTensorDesc.dims.d[0] * inputTensorDesc.dims.d[1]; - auto const seqLen = mRemoveInputPadding ? 0 : inputTensorDesc.dims.d[1]; - - void const* inputTensor = inputs[IdxEntry::kINPUT_TENSOR]; - auto const* hostRequestTypes = static_cast(inputs[IdxEntry::kHOST_REQUEST_TYPES]); - void const* const* loraWeightsPtrs = &inputs[IdxEntry::kLORA_WEIGHTS_PTRS_START]; - - int32_t const* hostContextLengths = mRemoveInputPadding - ? static_cast(inputs[IdxEntry::kLORA_WEIGHTS_PTRS_START + numModules]) - : nullptr; - - mExpandDoraWeightPtrs.clear(); - mExpandDoraWeightPtrs.reserve(numModules * numTokens); - - bool hasAnyDora = false; - - for (auto moduleIdx = 0; moduleIdx < numModules; moduleIdx++) - { - auto const loraWeightModulePtrs = static_cast(loraWeightsPtrs[moduleIdx]); - - int idx = 0; - for (int reqId = 0; reqId < numReqs; reqId++) - { - // loraWeightModulePtrs has 3 pointers for each module: A,B, and an optional DoRA magnitude - // the current DoRA plugin does not apply LoRA, so A and B are ignored. - RequestType const reqType = static_cast(hostRequestTypes[reqId]); - auto const* modulePtr = reinterpret_cast(loraWeightModulePtrs[reqId * 3 + 2]); - hasAnyDora = hasAnyDora or modulePtr != nullptr; - - if (reqType == RequestType::kGENERATION) - { - mExpandDoraWeightPtrs.push_back(modulePtr); - idx += 1; - } - else - { - int contextLen = (mRemoveInputPadding ? hostContextLengths[reqId] : seqLen); - - for (int contextId = 0; contextId < contextLen; contextId++) - { - mExpandDoraWeightPtrs.push_back(modulePtr); - idx += 1; - } - } - } - if (idx != numTokens) - { - TLLM_LOG_ERROR("LoraParams and input dims don't match, lora tokens %d input tokens %d", idx, numTokens); - return -1; - } - } - - if (hasAnyDora) - { - mDoraImpl.run(numTokens, inputTensor, mExpandDoraWeightPtrs.data(), outputs, workspace, stream); - } - else - { - // skip dora scaling if all requests are pure-lora - auto const inputRank = inputTensorDesc.dims.nbDims; - auto const numel - = std::accumulate(inputTensorDesc.dims.d, inputTensorDesc.dims.d + inputRank, 1, std::multiplies()); - auto const elemSize = tensorrt_llm::common::getDTypeSize(mType); - tensorrt_llm::common::cudaAutoCpy((int8_t*) outputs[0], (int8_t*) inputTensor, numel * elemSize, stream); - } - - sync_check_cuda_error(stream); - return 0; -} - -nvinfer1::IPluginV3* DoraPlugin::attachToContext(nvinfer1::IPluginResourceContext* context) noexcept -{ - return clone(); -} - -nvinfer1::PluginFieldCollection const* DoraPlugin::getFieldsToSerialize() noexcept -{ - return &mFieldsToSerialize; -} - -DoraPluginCreator::DoraPluginCreator() -{ - mPluginAttributes.clear(); - mPluginAttributes.emplace_back("num_modules", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("type", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("remove_input_padding", nullptr, PluginFieldType::kINT8, 1); - mFC.nbFields = static_cast(mPluginAttributes.size()); - mFC.fields = mPluginAttributes.data(); -} - -char const* DoraPluginCreator::getPluginName() const noexcept -{ - return DORA_PLUGIN_NAME; -} - -char const* DoraPluginCreator::getPluginVersion() const noexcept -{ - return DORA_PLUGIN_VERSION; -} - -nvinfer1::PluginFieldCollection const* DoraPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -nvinfer1::IPluginV3* DoraPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - bool removeInputPadding{}; - std::vector outHiddenSizes; - - // Read configurations from each field - for (int i = 0; i < fc->nbFields; ++i) - { - auto const field = fields[i]; - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type")) - { - TLLM_CHECK(field.type == PluginFieldType::kINT32 and field.length == 1); - type = *static_cast(field.data); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(field.type == PluginFieldType::kINT8 and field.length == 1); - removeInputPadding = *static_cast(field.data); - } - else if (!strcmp(attrName, "out_hidden_sizes")) - { - TLLM_CHECK(field.type == PluginFieldType::kINT32); - auto const* outHiddenSizesPtr = static_cast(field.data); - outHiddenSizes.resize(field.length); - outHiddenSizes.assign(outHiddenSizesPtr, outHiddenSizesPtr + field.length); - } - else - { - TLLM_LOG_WARNING("%s: got an unexpected attribute: %s", __PRETTY_FUNCTION__, attrName); - } - } - - try - { - auto* obj = new DoraPlugin(outHiddenSizes, type, removeInputPadding); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.h b/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.h deleted file mode 100644 index dfee11fdc90e..000000000000 --- a/cpp/tensorrt_llm/plugins/doraPlugin/doraPlugin.h +++ /dev/null @@ -1,114 +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. - */ -#include "tensorrt_llm/kernels/lora/dora.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -namespace tensorrt_llm::plugins -{ - -class DoraPlugin : public BasePluginV3 -{ -public: - DoraPlugin() = delete; - DoraPlugin(std::vector const& outHiddenSizes, nvinfer1::DataType type, bool removeInputPadding); - DoraPlugin(DoraPlugin const& p) = default; - - // IPluginV3 methods - nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override; - nvinfer1::IPluginV3* clone() noexcept override; - - // IPluginV3OneCore methods - char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; - - // IPluginV3OneBuild methods - int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override; - int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes, - int32_t nbInputs) const noexcept override; - int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs, - int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, - int32_t nbOutputs) noexcept override; - int32_t getNbOutputs() const noexcept override; - size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override; - int32_t getValidTactics(int32_t* tactics, int32_t nbTactics) noexcept override; - int32_t getNbTactics() noexcept override; - char const* getTimingCacheID() noexcept override; - int32_t getFormatCombinationLimit() noexcept override; - char const* getMetadataString() noexcept override; - - // IPluginV3OneRuntime methods - int32_t setTactic(int32_t tactic) noexcept override; - int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept override; // fixed - nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override; - nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override; - -private: - void init(); - - std::vector mDataToSerialize; - nvinfer1::PluginFieldCollection mFieldsToSerialize; - - enum IdxEntry - { - kINPUT_TENSOR = 0, - kHOST_REQUEST_TYPES = 1, - kLORA_WEIGHTS_PTRS_START = 2 - }; - - // TODO(oargov) this is shared with the LoRA plugin, put it somewhere else - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - - std::vector mOutHiddenSizes; - nvinfer1::DataType mType; - bool mRemoveInputPadding; - tensorrt_llm::kernels::DoraImpl mDoraImpl; - - std::vector mExpandDoraWeightPtrs{}; -}; - -class DoraPluginCreator : public BaseCreatorV3 -{ -public: - DoraPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV3* createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -}; // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/eaglePlugin/CMakeLists.txt deleted file mode 100644 index b6bd0439cc0c..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.cpp b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.cpp deleted file mode 100644 index 899c93855b9f..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.cpp +++ /dev/null @@ -1,945 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "eagleDecodeDraftTokensPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::EagleDecodeDraftTokensPluginCreator; -using tensorrt_llm::plugins::EagleDecodeDraftTokensPlugin; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -static char const* EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_VERSION{"1"}; -static char const* EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_NAME{"EagleDecodeDraftTokens"}; -PluginFieldCollection EagleDecodeDraftTokensPluginCreator::mFC{}; -std::vector EagleDecodeDraftTokensPluginCreator::mPluginAttributes; - -EagleDecodeDraftTokensPlugin::EagleDecodeDraftTokensPlugin( - nvinfer1::DataType type, int32_t layerIdx, int32_t numEagleLayers, bool topKSampling) - : mDtype(type) - , mLayerIdx(layerIdx) - , mNumEagleLayers(numEagleLayers) - , mTopKSampling(topKSampling) -{ - TLLM_CHECK_WITH_INFO(mTopKSampling, "Multinomial sampling is not supported yet."); -} - -// Parameterized constructor -EagleDecodeDraftTokensPlugin::EagleDecodeDraftTokensPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDtype); - read(d, mLayerIdx); - read(d, mNumEagleLayers); - read(d, mTopKSampling); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - static_cast(length), static_cast(d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* EagleDecodeDraftTokensPlugin::clone() const noexcept -{ - auto* plugin = new EagleDecodeDraftTokensPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs EagleDecodeDraftTokensPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(outputIndex < getNbOutputs()); - TLLM_CHECK(nbInputs == 12); - auto const batchSizeExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[0]; - auto const maxDecodingTokensExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[1]; - auto const maxPathLengthExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[2]; - auto const maxDecodingDraftTokensExpr - = exprBuilder.operation(DimensionOperation::kSUB, *maxDecodingTokensExpr, *exprBuilder.constant(1)); - - auto const numEagleLayersExpr - = exprBuilder.operation(DimensionOperation::kSUB, *maxPathLengthExpr, *exprBuilder.constant(1)); - auto const maxDecodingDraftTokensSquareExpr - = exprBuilder.operation(DimensionOperation::kPROD, *maxDecodingDraftTokensExpr, - *maxDecodingDraftTokensExpr); // maxDecodingDraftTokensExpr * maxDecodingDraftTokensExpr - - nvinfer1::DimsExprs ret; - if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_DRAFT_TOKEN_IDS)) - { - // output_draft_token_ids: [batch_size, max_decoding_draft_tokens] - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_DRAFT_LENS)) - { - // output_draft_lens: [batch_size] - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_PATHS)) - { - // output_path: [batch_size, max_decoding_tokens, max_path_len] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingTokensExpr; - ret.d[2] = maxPathLengthExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)) - { - // output_current_scores: [batch_size, max_decoding_draft_tokens] - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_NEXT_EXPAND_INDICES)) - { - // output_next_expand_index - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES)) - { - // output_all_layers_scores: - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = numEagleLayersExpr; - ret.d[2] = maxDecodingDraftTokensSquareExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS)) - { - // output_all_layers_draft_token_ids: - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = numEagleLayersExpr; - ret.d[2] = maxDecodingDraftTokensSquareExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR)) - { - // output_all_layers_draft_token_ids_predecessor - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = numEagleLayersExpr; - ret.d[2] = maxDecodingDraftTokensSquareExpr; - } - else - { - TLLM_CHECK_WITH_INFO( - false, "Wrong outputIndex %d in EagleDecodeDraftTokensPlugin::getOutputDimensions", outputIndex); - } - return ret; -} - -bool EagleDecodeDraftTokensPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_CHECK(nbInputs == 12 && nbOutputs == getNbOutputs()); - TLLM_CHECK(pos < nbInputs + nbOutputs); - - if (pos == getIdx(InputIdxEntry::LOGITS)) - { - // input: logits - // output: output_all_layers_scores - return (inOut[pos].type == mDtype) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES) || pos == getIdx(InputIdxEntry::INPUT_PREV_SCORES) - || pos == nbInputs + getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES) - || pos == nbInputs + getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)) - { - // input: rand_sample, input_all_layers_scores, input_prev_scores - // output: output_all_layers_scores, output_current_scores - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else - { - // input: path, num_valid_logits, use_dynamic_tree, dynamic_tree_max_topK, input_draft_token_ids, - // input_draft_lens, input_current_expand_index, input_all_layers_draft_token_ids - // output: output_draft_token_ids, output_draft_lens, output_path, output_next_expand_index - // output_all_layers_draft_token_ids, output_all_alyers_draft_token_predecessor - return (inOut[pos].type == nvinfer1::DataType::kINT32) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void EagleDecodeDraftTokensPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -template -size_t EagleDecodeDraftTokensPlugin::getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - size_t workspaceSize{0}; - auto const numInputLogits = inputs[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const batchSize = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const vocabSizePadded = inputs[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const maxDecodingTokens = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - auto const maxDecodingDraftTokens = maxDecodingTokens - 1; - auto const maxTopK = maxDecodingDraftTokens; - auto const mNumEagleLayers = inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES)].dims.d[1]; - - // Greedy sampling - if (mTopKSampling) - { - // 0. The first topK sampling workspace - auto const draftTokenSamplingWorkspaceSize - = getTopKWorkspaceSize(numInputLogits, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, vocabSizePadded); - - // 1. The first TopKs [numInputLogits] - auto const topKsSize = numInputLogits * sizeof(SizeType32); - - // 2. Topks offset [batchSize] - // Each request will have different number of logits that need to be sampled - // This tensor will record the start offset of the topK for each request - auto const topKOffsetSize = batchSize * sizeof(SizeType32); - - // 3. Logits ptrs [numInputLogits] - auto const logitsPtrsSize = numInputLogits * sizeof(T*); - - // 4. The first topK sampling's output ids ptrs [numInputLogits][maxDecodingDraftTokens] - auto const firstTopKOutputIdsPtrsSize = numInputLogits * sizeof(TokenIdType*); - - // 5. The first topK sampling's output ids (temporary buffer) [numInputLogits * maxDecodingDraftTokens] - auto const firstTopKOutputIdsSize = numInputLogits * maxDecodingDraftTokens * sizeof(TokenIdType); - - // 6. Number of successors for each nodes, extract from the paths and layerId - // [batchSize * maxDecodingTokens] - auto const numSuccessorsForEachNodeSize = batchSize * maxDecodingTokens * sizeof(SizeType32); - - // 7. Flag whether to do decoding or not. SamplingTopK is done for numInputLogits tokens. - // But only sum(numValidLogitsPerRequest[:]) of them are valid. - // [batchSize * maxDecodingTokens] - auto const skipDecodeSize = numInputLogits * sizeof(bool); - - // 8. The first topK sampling's logprobs [batchSize * maxDecodingDraftTokens] - auto const firstTopKOutputLogProbsSize = numInputLogits * maxDecodingDraftTokens * sizeof(float); - - // 9. Eagle-2, the second topK sampling workspace - // Sampling from [batchSize, maxTopK * maxTopK] to [batchSize, maxTopK] - auto const secondTopKSamplingWorkspaceSize = getTopKWorkspaceSize( - batchSize, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, maxTopK * maxTopK); - - // 10. Eagle-2, the outputIds of the second topK sampling, shape [batchSize, maxDecodingTokens] - auto const secondTopKOutputIdsSize = batchSize * maxDecodingTokens * sizeof(TokenIdType); - // 11. Eagle-2, the outputIdsPtr of the second topK sampling, shape [batchSize] - auto const secondTopKOutputIdsPtrSize = batchSize * sizeof(TokenIdType*); - // 12. Eagle-2, the inputScoresPtrs of the second topK sampling, shape [batchSize] - auto const secondTopKInputScoresPtrsSize = batchSize * sizeof(float*); - // 13. Eagle-2, the outpuLogProbs of the second topK samplig, shape [batchSize, maxDecodingDraftTokens] - auto const secondTopKOutputLogProbsSize = batchSize * maxDecodingDraftTokens * sizeof(float); - - // 14. Eagle-2, the input scores pointers of the third topK sampling, shape [batchSize] - // Each points to a vocabSize = '(mNumEagleLayers - 1) * dynamicTreeMaxTopK * dynamicTreeMaxTopK + - // dynamicTreeMaxTopK' - auto const thirdTopKInputScoresPtrsSize = batchSize * sizeof(float*); - // 15. Eagle-2, the output of the third topK sampling, shape [batchSize, maxDecodingDraftTokens] - auto const thirdTopKOutputIdsSize = batchSize * maxDecodingDraftTokens * sizeof(TokenIdType); - // 16. Eagle-2, the output pointers of the third topK sampling, shape [batchSize] - auto const thirdTopKOutputIdsPtrsSize = batchSize * sizeof(TokenIdType*); - // 17. Eagle-2, the workspace of the third topK sampling - // Sampling from [batchSize, '(mNumEagleLayers - 1) * dynamicTreeMaxTopK * dynamicTreeMaxTopK + - // dynamicTreeMaxTopK'] to [batchSize, maxDecodingDraftTokens] We over-set the vocabsize here. - auto const thridTopKSamplingWorkspaceSize = getTopKWorkspaceSize(batchSize, /* maxTokensPerStep */ 1, - /* maxTopK */ maxDecodingDraftTokens, mNumEagleLayers * maxDecodingDraftTokens * maxDecodingDraftTokens); - - // 18. Eagle-2, the topKs for each request in the third topK sampling - // The real topK value is min(maxDecodingDraftTokens, totalNumDraftTokensForAllLayers) - auto const thirdTopKsSize = batchSize * sizeof(SizeType32); - - SizeType32 constexpr NUM_BUFFERS{19}; - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = draftTokenSamplingWorkspaceSize; - workspaces[1] = topKsSize; - workspaces[2] = topKOffsetSize; - workspaces[3] = logitsPtrsSize; - workspaces[4] = firstTopKOutputIdsPtrsSize; - workspaces[5] = firstTopKOutputIdsSize; - workspaces[6] = numSuccessorsForEachNodeSize; - workspaces[7] = skipDecodeSize; - workspaces[8] = firstTopKOutputLogProbsSize; - workspaces[9] = secondTopKSamplingWorkspaceSize; - workspaces[10] = secondTopKOutputIdsSize; - workspaces[11] = secondTopKOutputIdsPtrSize; - workspaces[12] = secondTopKInputScoresPtrsSize; - workspaces[13] = secondTopKOutputLogProbsSize; - workspaces[14] = thirdTopKInputScoresPtrsSize; - workspaces[15] = thirdTopKOutputIdsSize; - workspaces[16] = thirdTopKOutputIdsPtrsSize; - workspaces[17] = thridTopKSamplingWorkspaceSize; - workspaces[18] = thirdTopKsSize; - workspaceSize = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); - } - else - { - // TODO fill me - // Multinomial sampling - TLLM_CHECK_WITH_INFO(false, "Multinomial sampling is not supported yet."); - } - - return workspaceSize; -} - -size_t EagleDecodeDraftTokensPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - auto const logitsType = inputs[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - return getWorkspaceSizeType(inputs, nbInputs, outputs, nbOutputs); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - return getWorkspaceSizeType<__half>(inputs, nbInputs, outputs, nbOutputs); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - return 0; -} - -template -void EagleDecodeDraftTokensPlugin::doTopKSampling(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // We allocate many buffers with 'numInputLogits' size, but the input logits will include some padding logits. - // So only 'batchSize' or 'numValidLogits' size will be actually used. - auto const numInputLogits = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[2]; - auto const maxDecodingDraftTokens = maxDecodingTokens - 1; - auto const maxTopK = maxDecodingDraftTokens; - - ////////////////////////////////////////// Get plugin inputs ////////////////////////////////////////// - // Plugin inputs - // Input logits for sampling, shape: [numInputLogits, vocabSizePadded] - auto pluginInputLogits = static_cast(inputs[getIdx(InputIdxEntry::LOGITS)]); - // Input paths, shape: [batchSize, maxDecodingTokens, maxPathLen] - auto pluginInputPaths = static_cast(inputs[getIdx(InputIdxEntry::PATHS)]); - auto numValidLogits = static_cast(inputs[getIdx(InputIdxEntry::NUM_VALID_LOGITS)]); - // For Eagle-2 - // Whether to use dynamic tree (i.e., Eagle-2) - auto useDynamicTree = *(static_cast(inputs[getIdx(InputIdxEntry::USE_DYNAMIC_TREE)])); - // The max topK for dynamic tree. All the requests have the same expand topK. - // In Eagle-2, dynamicTreeMaxTopK is equal to maxNonLeavesPerLayer in the internal EagleNets. - auto dynamicTreeMaxTopK = *(static_cast(inputs[getIdx(InputIdxEntry::DYNAMIC_TREE_MAX_TOPK)])); - // All layer's draft tokenIds, shape: [batchSize, maxDecodingDraftTokens] - auto pluginInputDraftTokenIds - = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_DRAFT_TOKEN_IDS)]); - // The number of all layer's draft tokenIds, shape: [batchSize] - auto pluginInputDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_DRAFT_LENS)]); - // The previous EagleNet's scores, shape: [batchSize, maxDecodingDraftTokens] - auto pluginInputPrevScores = static_cast(inputs[getIdx(InputIdxEntry::INPUT_PREV_SCORES)]); - // The indices of the nodes that will be expand in this layer, shape: [batchSize, maxDecodingDraftTokens] - // The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - auto pluginInputCurrentExpandIndices - = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_CURRENT_EXPAND_INDICES)]); - // The scores from all previous EagleNets, - // shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginInputAllLayersScores = static_cast(inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES)]); - // The draft tokens from all previous EagleNets, - // shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginInputAllLayersDraftTokenIds - = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS)]); - // The predecessor of all the draft tokens, - // shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginInputAllLayersDraftTokenIdsPredecessor = reinterpret_cast( - inputs[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR)]); - - ////////////////////////////////////////// Get plugin outputs ////////////////////////////////////////// - // Plugin outputs - // All layer's draft tokenIds, shape: [batchSize, maxDecodingDraftTokens] - auto pluginOutputDraftTokenIds - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_DRAFT_TOKEN_IDS)]); - // The number of all layer's draft tokenIds, shape: [batchSize] - auto pluginOutputDraftLens = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_DRAFT_LENS)]); - // For Eagle-2 - // Updated paths base on this layer's sampling result, shape: [batchSize, maxDecodingTokens, maxPathLen] - auto pluginOutputPaths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_PATHS)]); - // This layer's scores, which will be used in next layers [batchSize, maxDecodingDraftTokens] - auto pluginOutputCurrentScores = static_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)]); - // The indices of the nodes that will be expand in next layer, shape: [batchSize, maxDecodingDraftTokens] - // The index is related to the final output tree, which has max_decoding_draft_tokens draft tokens. - auto pluginOutputNextExpandIndices - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_NEXT_EXPAND_INDICES)]); - // Updated scores, shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginOutputAllLayersScores = static_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES)]); - // Updated draft tokens, shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x maxDecodingDraftTokens] - auto pluginOutputAllLayersDraftTokenIds - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS)]); - // Update the predecessor of the draft tokens, shape: [batchSize, mNumEagleLayers, maxDecodingDraftTokens x - // maxDecodingDraftTokens] - auto pluginOutputAllLayersDraftTokenIdsPredecessor - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR)]); - - ////////////////////////////////////////// Get workspaces ////////////////////////////////////////// - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - // Workspace 0: Sampling workspace. - // Treat numInputLogits as batchSize - auto const samplingWorkspaceSize - = getTopKWorkspaceSize(numInputLogits, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, vocabSizePadded); - void* workspaceSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, samplingWorkspaceSize)); - - // Workspace 1: Topks tensor: shape [numInputLogits] - SizeType32* topKs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(SizeType32))); - - // Workspace 2: topKOffset tensor: shape: [batchSize], number of nodes that have successors for each requests - SizeType32* topKOffset - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - // Workspace 3: logits pointers tensor: shape: [numInputLogits] - T const** logitsPtrs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(T*))); - - // Workspace 4: outputIds pointers tensor: shape [numInputLogits], each points to a [maxDecodingDraftTokens] buffer - TokenIdType** firstTopKOutputIdsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(TokenIdType*))); - - // Workspace 5: outputIds tensor: flatten outputIds, shape [numInputLogits * maxDecodingDraftTokens] - TokenIdType* firstTopKOutputIdsFlatten = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * maxDecodingDraftTokens * sizeof(TokenIdType))); - - // Workspace 6: number of successors for each nodes tensor: shape [batchSize * maxDecodingTokens] - SizeType32* numSuccessorsForEachNode = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - - // Workspace 7: skip decoding mask [numInputLogits] - bool* skipDecode - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * sizeof(bool))); - - // In Eagle-1, we do not need to return logProbs - float* firstTopKOutputLogProbs = nullptr; - if (useDynamicTree) - { - // Workspace 8. The output logProbs of the first topK sampling. - // Which will be updated with the previous layer's scores (i.e., pluginInputPrevScores), and will be treat as - // the input of the second topK sampling. For mLayerIdx == 0, shape: [numInputLogits(batchSize), - // maxDecodingDraftTokens] For mLayerIdx > 0, shape: [numInputLogits(batchSize * dynamicTreeMaxTopK), - // maxDecodingDraftTokens] - firstTopKOutputLogProbs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, numInputLogits * maxDecodingDraftTokens * sizeof(float))); - } - - SizeType32 const secondTopKVocabSize = dynamicTreeMaxTopK * maxDecodingDraftTokens; - // Workspace 9: Sampling from [batchSize, dynamicTreeMaxTopK * maxDecodingDraftTokens] to [batchSize, - // dynamicTreeMaxTopK] - auto const secondTopKSamplingWorkspaceSize - = getTopKWorkspaceSize(batchSize, /* maxTokensPerStep */ 1, /* maxTopK */ maxTopK, secondTopKVocabSize); - void* workspaceScoresSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, secondTopKSamplingWorkspaceSize)); - - // Workspace 10: the second (scores) sampling's outputIds, shape: [batchSize, maxDecodingDraftTokens] - TokenIdType* secondTopKOutputIdsFlatten = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingDraftTokens * sizeof(TokenIdType))); - - // Workspace 11: the second (scores) sampling's outputIdsPtrs - TokenIdType** secondTopKOutputIdsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(TokenIdType*))); - - // Workspace 12: input scores pointers - float** secondTopKInputScoresPtrs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(float*))); - - // Workspace 13: the second sampling's outputLogProbs - float* secondTopKOutputLogProbs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingDraftTokens * sizeof(float))); - - // Workspace 14: The input scores pointers of the third topK sampling, shape [batchSize] - float** thirdTopKInputScoresPtrs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(float*))); - - // Workspace 15: The output of the third topK sampling, shape [batchSize, maxDecodingDraftTokens] - TokenIdType* thirdTopKOutputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingDraftTokens * sizeof(TokenIdType))); - - // Workspace 16: The output pointers of the third topK sampling, shape [batchSize] - TokenIdType** thirdTopKOutputIdsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(TokenIdType*))); - - // The number of draft tokens among all layers - long const totalNumDraftTokensForAllLayers - = (mNumEagleLayers - 1) * dynamicTreeMaxTopK * dynamicTreeMaxTopK + dynamicTreeMaxTopK; - - auto const thridTopKSamplingWorkspaceSize = getTopKWorkspaceSize( - batchSize, /* maxTokensPerStep */ 1, /* maxTopK */ maxDecodingDraftTokens, totalNumDraftTokensForAllLayers); - // Workspace 17: The workspace of the third topK sampling - void* workspaceThirdTopKSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, thridTopKSamplingWorkspaceSize)); - - // Workspace 18. Eagle-2, the topKs for each request in the third topK sampling, shape [batchSize] - // The real topK value is min(maxDecodingDraftTokens, totalNumDraftTokensForAllLayers) - SizeType32* thirdTopKs - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - ////////////////////////////////////////// Main logic ////////////////////////////////////////// - // Fill logitsPtrs from plugin input logits - // And fill firstTopKOutputIdsPtrs from firstTopKOutputIdsFlatten - invokeAssembleDraftLogitsOffsets(logitsPtrs, pluginInputLogits, firstTopKOutputIdsPtrs, firstTopKOutputIdsFlatten, - skipDecode, numValidLogits, numInputLogits, batchSize, maxDecodingDraftTokens, vocabSizePadded, stream); - sync_check_cuda_error(stream); - - if (useDynamicTree) - { - // For Eagle-2, the topK value between different requests are the same, all set to 'dynamicTreeMaxTopK'. - invokeSetTopKsFromDyanmicTreeMaxTopK( - mLayerIdx, batchSize, numInputLogits, topKs, topKOffset, dynamicTreeMaxTopK, numValidLogits, stream); - sync_check_cuda_error(stream); - - // Do softmax for the input logits - // We set the 'batchSize' and 'maxBatchSize' to 'numInputLogits', while 'numInputLogits' logits may contain - // some padding logits, which do not need to be calculated. - // We use 'skipDecode' list to skip these padding logits. This could avoid redundant calculations. - BiasSoftmaxParams biasSoftmaxParams; - biasSoftmaxParams.logits = const_cast(pluginInputLogits); - biasSoftmaxParams.logitsPtrs = nullptr; - biasSoftmaxParams.probs = const_cast(pluginInputLogits); - biasSoftmaxParams.maxBeamWidth = 1; - biasSoftmaxParams.batchSlots = nullptr; - biasSoftmaxParams.batchSize = numInputLogits; - biasSoftmaxParams.maxBatchSize = numInputLogits; - biasSoftmaxParams.vocabSize = vocabSizePadded; - biasSoftmaxParams.vocabSizePadded = vocabSizePadded; - biasSoftmaxParams.skipSoftMax = false; - biasSoftmaxParams.batchSlotsLogits = false; - biasSoftmaxParams.skipDecode = skipDecode; - biasSoftmaxParams.checkParams(); - - invokeAddBiasSoftMax(biasSoftmaxParams, stream); - sync_check_cuda_error(stream); - } - else - { - // For Eagle-1, extract topK value from input path. - invokeExtractTopKsFromPath(pluginInputPaths, topKs, topKOffset, numSuccessorsForEachNode, mLayerIdx, batchSize, - maxDecodingTokens, maxPathLen, stream); - sync_check_cuda_error(stream); - } - - TopKSamplingKernelParams params{}; - params.logProbsPtrs = logitsPtrs; // [numInputLogits][vocabSizePadded] - params.outputIdsPtrs = firstTopKOutputIdsPtrs; // [numInputLogits][maxDecodingDraftTokens] - params.workspace = workspaceSampling; - params.maxTopK = maxTopK; - params.topKs = topKs; // [numInputLogits] - params.batchSize = numInputLogits; - params.maxBatchSize = numInputLogits; - params.maxTokensPerStep = 1; - params.vocabSizePadded = vocabSizePadded; - params.returnAllSelectedTokens = true; - params.strictTopPBoundary = false; - params.skipDecode = skipDecode; - params.outputLogProbs = firstTopKOutputLogProbs; // [numInputLogits * maxDecodingDraftTokens] - params.logitsHasProbs = true; - - invokeBatchTopKSampling(params, stream); - sync_check_cuda_error(stream); - - if (useDynamicTree) - { - // When mLayerIdx == 0, we do not need to update scores. - // We take the outputLogProbs of the first topK sampling as the scores directly. - if (mLayerIdx != 0) - { - // Update firstTopKOutputLogProbs with pluginInputPrevScores, which is the scores from the previous layer - invokeUpdateScores(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, firstTopKOutputLogProbs, - pluginInputPrevScores, stream); - sync_check_cuda_error(stream); - - // Do the second top-dynamicTreeMaxTopK sampling among this dynamicTreeMaxTopK x dynamicTreeMaxTopK draft - // tokens. Through the second topK sampling, we obtain the dynamicTreeMaxTopK output draft tokens of this - // layer. - - // Although theoretically we only need to select 'dynamicTreeMaxTopK' draft tokens from 'dynamicTreeMaxTopK - // * dynamicTreeMaxTopK' draft tokens, we over-set vocabSize here. This is because when we write the scores - // into firstTopKOutputLogProbs, we store it in the form of [batchSize * dynamicTreeMaxTopK, - // maxDecodingDraftTokens]. For each request, these 'dynamicTreeMaxTopK * dynamicTreeMaxTopK' scores are not - // saved continuously, but in the format of [dynamicTreeMaxTopK, maxDecodingDraftTokens]. For unused - // positions, we set '-inf' to ensure that they will not be sampled. Examples: For a request, - // dynamicTreeMaxTopK == 3, the scores in its buffer ([dynamicTreeMaxTopK, maxDecodingDraftTokens]) are as - // follow: - // [[1.1, 2.2, 3.3, -inf, -inf, ...], - // [4.4, 5.5, 6.6, -inf, -inf, ...], - // [7.7, 8.8, 9.9, -inf, -inf, ...]] - - // Prepare the input of the second topK sampling. - invokeAssembleSecondTopKSamplingInputs(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, - firstTopKOutputLogProbs, secondTopKInputScoresPtrs, secondTopKOutputIdsFlatten, secondTopKOutputIdsPtrs, - stream); - sync_check_cuda_error(stream); - - TopKSamplingKernelParams params{}; - params.logProbsPtrs = secondTopKInputScoresPtrs; - params.outputIdsPtrs = secondTopKOutputIdsPtrs; - params.workspace = workspaceScoresSampling; - params.maxTopK = maxTopK; // Same to maxDecodingTokens - params.topKs = topKs; // [batchSize], all set to dynamicTreeMaxTopK - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.maxTokensPerStep = 1; - params.vocabSizePadded = secondTopKVocabSize; - params.returnAllSelectedTokens = true; - params.strictTopPBoundary = false; - - invokeBatchTopKSampling(params, stream); - sync_check_cuda_error(stream); - } - - // Copy this layer's scores and draft tokensId: - // 1) Copy this layer's scores to pluginOutputAllLayersScores - // 2) Copy dynamicTreeMaxTopK (or dynamicTreeMaxTopK * dynamicTreeMaxTopK) draft tokens to - // pluginOutputAllLayersDraftTokenIds 3) Set the predecessors of these draft tokens and save to - // pluginOutputAllLayersDraftTokenIdsPredecessor, - // which will be used to reconstruct the final output tree at the last layer - invokeCopyScoresAndDraftTokenIds(mLayerIdx, mNumEagleLayers, maxDecodingDraftTokens, batchSize, - dynamicTreeMaxTopK, - pluginInputCurrentExpandIndices, // The indices of the nodes that expand in this layer (i.e., the input - // logits). The index is related to the final tree. - pluginInputAllLayersScores, pluginInputAllLayersDraftTokenIds, pluginInputAllLayersDraftTokenIdsPredecessor, - pluginOutputAllLayersScores, pluginOutputAllLayersDraftTokenIds, - pluginOutputAllLayersDraftTokenIdsPredecessor, - firstTopKOutputLogProbs, // This layer's scores - firstTopKOutputIdsFlatten, // This layer's draft tokens - stream); - sync_check_cuda_error(stream); - - // Update Path - // For mLayerIdx == 0, the output of the first topK sampling are the output draft tokens of this layers. The - // update logic is simple. For mLayerIdx > 0, the output of the second topK sampling are the output draft tokens - // of this layers. 'secondTopKOutputIdsPtrs' contains the top-dynamicTreeMaxTopK selected from the second topK - // sampling. 'pluginOutputNextExpandIndices' record the selected the top-dynamicTreeMaxTopK draft token's Id of - // this layer, - // which will be used in the next layer to compute the predecessors. - // The last layer will completely reconstruct the paths, so there is no need to update the paths here. - if (mLayerIdx != mNumEagleLayers - 1) - { - invokeUpdatePath(mLayerIdx, batchSize, dynamicTreeMaxTopK, maxDecodingTokens, maxPathLen, pluginInputPaths, - pluginOutputPaths, - secondTopKOutputIdsPtrs, // if mLayerIdx == 0, secondTopKOutputIdsPtrs == nullptr, and it's useless - // during update paths - pluginOutputNextExpandIndices, stream); - sync_check_cuda_error(stream); - } - - if (mLayerIdx != 0) - { - // We will extract the real draft tokenIds and scores from 'firstTopKOutputIdsFlatten' and - // 'secondTopKInputScoresPtrs' according to the 'secondTopKOutputIdsPtrs'. And store them into - // 'secondTopKOutputIdsPtrs' and 'secondTopKOutputLogProbs' (reuse these buffers). - // secondTopKInputScoresPtrs: shape [batchSize * dynamicTreeMaxTopK, maxDecodingDraftTokens] - // The original scores, which were used to do the second TopK sampling - // secondTopKOutputIdsPtrs: shape [batchSize], each points to a [maxDecodingDraftTokens] buffer - // The output of the second TopK sampling, which are the indices of the top-dynamicTreeMaxTopK among - // 'dynamicTreeMaxTopK * dynamicTreeMaxTopK'. We need to figure out what these top-dynamicTreeMaxTopK - // draft tokens' real tokenIds. - // firstTopKOutputIdsFlatten: shape [batchSize * dynamicTreeMaxTopK, maxDecodingDraftTokens] - // The value are related to the vocabSize, which is the real tokenIds. - invokeExtractScoresAndRealDraftTokensIds(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, - secondTopKInputScoresPtrs, secondTopKOutputIdsPtrs, firstTopKOutputIdsFlatten, secondTopKOutputLogProbs, - stream); - sync_check_cuda_error(stream); - } - - // Copy this layer's output draft tokens and scores. - // This layer's output scores is next layer's previous scores. - // if mLayerIdx == 0, directly use the first topK's outputIds / logProbs as this layer's output draft tokens / - // scores if mLayerIdx > 0, we use the second topK's outputIds / logProbs, - // which is updated with the real draft tokenIds / logprobs in 'invokeExtractScoresAndRealDraftTokensIds' - invokeUpdateDraftTokensAndLensAndCurScores(mLayerIdx, batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, - mLayerIdx == 0 ? firstTopKOutputIdsPtrs : secondTopKOutputIdsPtrs, pluginInputDraftTokenIds, - pluginInputDraftLens, pluginOutputDraftTokenIds, pluginOutputDraftLens, - mLayerIdx == 0 ? firstTopKOutputLogProbs : secondTopKOutputLogProbs, pluginOutputCurrentScores, stream); - sync_check_cuda_error(stream); - - if (mLayerIdx == mNumEagleLayers - 1) - { - // The maximum number of nodes on the final tree (exclude the root node) - auto const maxNodesOnFinalTree = std::min(maxDecodingDraftTokens, totalNumDraftTokensForAllLayers); - - // When reach the last EagleNet, we need to do the third sampling, which take all layers' draft tokens and - // scores as input, and then select top-maxDecodingDraftTokens draft tokens among them. We need to - // reconstruct the path/tree after the third topK sampling. - invokeAssembleThridTopKSamplingInputs(batchSize, maxDecodingDraftTokens, mNumEagleLayers, - maxNodesOnFinalTree, thirdTopKs, pluginOutputAllLayersScores, thirdTopKInputScoresPtrs, - thirdTopKOutputIds, thirdTopKOutputIdsPtrs, stream); - sync_check_cuda_error(stream); - - // 1) Do topK sampling among all previous draft tokens - TopKSamplingKernelParams params{}; - params.logProbsPtrs = thirdTopKInputScoresPtrs; - params.outputIdsPtrs = thirdTopKOutputIdsPtrs; - params.workspace = workspaceThirdTopKSampling; - params.topKs = thirdTopKs; // All set to 'maxNodesOnFinalTree' - params.maxTopK = maxDecodingDraftTokens; // We set maxTopK to 'maxDecodingDraftTokens' to align the - // outputIdsPtrs offsets when written back. - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.maxTokensPerStep = 1; - params.vocabSizePadded = totalNumDraftTokensForAllLayers; - params.returnAllSelectedTokens = true; - params.strictTopPBoundary = false; // Make sure to select topK tokens. - - invokeBatchTopKSampling(params, stream); - sync_check_cuda_error(stream); - - // 2) Reconstruct the Path - invokeReconstructFinalPath(batchSize, dynamicTreeMaxTopK, maxDecodingDraftTokens, maxDecodingTokens, - maxPathLen, mNumEagleLayers, maxNodesOnFinalTree, thirdTopKOutputIdsPtrs, - pluginOutputAllLayersDraftTokenIdsPredecessor, pluginOutputPaths, stream); - sync_check_cuda_error(stream); - - // 3) Copy this layer's outputIds to outputDraftTokenIds - invokeCopyFinalDraftTokens(batchSize, maxDecodingDraftTokens, mNumEagleLayers, maxNodesOnFinalTree, - thirdTopKOutputIdsPtrs, pluginOutputAllLayersDraftTokenIds, pluginOutputDraftTokenIds, - pluginOutputDraftLens, stream); - sync_check_cuda_error(stream); - } - } - else - { - // Eagle-1: Copy output token id from outputIdsPtrs to the plugin output buffer - invokeCopyOutputTokensIds(firstTopKOutputIdsPtrs, topKs, topKOffset, pluginInputDraftTokenIds, - pluginInputDraftLens, numValidLogits, pluginOutputDraftTokenIds, pluginOutputDraftLens, mLayerIdx, - batchSize, maxDecodingDraftTokens, pluginInputPaths, pluginOutputPaths, maxPathLen, stream); - sync_check_cuda_error(stream); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleDecodeDraftTokensPlugin::enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // TODO split batch into greedy and non-greedy and execute both paths - if (mTopKSampling) - { - doTopKSampling(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - // TODO fill me - TLLM_CHECK_WITH_INFO(false, "Multinomial sampling is not supported yet"); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -int EagleDecodeDraftTokensPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - auto const logitsType = inputDesc[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - enqueueType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - enqueueType<__half>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType EagleDecodeDraftTokensPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index < getNbOutputs()); - TLLM_CHECK(index < getNbOutputs()); - if (index == getIdx(OutputIdxEntry::OUTPUT_ALL_LAYERS_SCORES) - || index == getIdx(OutputIdxEntry::OUTPUT_CURRENT_SCORES)) - { - // Only output_prev_socres are float - return inputTypes[getIdx(InputIdxEntry::INPUT_ALL_LAYERS_SCORES)]; - } - else - { - // output_draft_token_ids, output_draft_lens, output_paths, output_next_expand_index, - // output_all_layers_draft_token_ids, output_all_layers_draft_token_ids_predecessor - // are all int32 type, same as path - return inputTypes[getIdx(InputIdxEntry::PATHS)]; - } -} - -// IPluginV2 Methods - -char const* EagleDecodeDraftTokensPlugin::getPluginType() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleDecodeDraftTokensPlugin::getPluginVersion() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_VERSION; -} - -int EagleDecodeDraftTokensPlugin::getNbOutputs() const noexcept -{ - return 8; -} - -int EagleDecodeDraftTokensPlugin::initialize() noexcept -{ - return 0; -} - -void EagleDecodeDraftTokensPlugin::terminate() noexcept {} - -size_t EagleDecodeDraftTokensPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDtype) + sizeof(mLayerIdx) + sizeof(mNumEagleLayers) + sizeof(mTopKSampling); -} - -void EagleDecodeDraftTokensPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDtype); - write(d, mLayerIdx); - write(d, mNumEagleLayers); - write(d, mTopKSampling); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void EagleDecodeDraftTokensPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -EagleDecodeDraftTokensPluginCreator::EagleDecodeDraftTokensPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("layer_idx", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_eagle_layers", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("top_k_sampling", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* EagleDecodeDraftTokensPluginCreator::getPluginName() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleDecodeDraftTokensPluginCreator::getPluginVersion() const noexcept -{ - return EAGLE_DECODE_DRAFT_TOKENS_PLUGIN_VERSION; -} - -PluginFieldCollection const* EagleDecodeDraftTokensPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* EagleDecodeDraftTokensPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int32_t layerIdx{}; - int32_t numEagleLayers{}; - nvinfer1::DataType type{}; - bool topKSampling{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "layer_idx")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - layerIdx = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "num_eagle_layers")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - numEagleLayers = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "top_k_sampling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - topKSampling = static_cast(*static_cast(fields[i].data)); - } - } - - try - { - auto* obj = new EagleDecodeDraftTokensPlugin(type, layerIdx, numEagleLayers, topKSampling); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* EagleDecodeDraftTokensPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call EagleDecodeDraftTokensPlugin::destroy() - try - { - auto* obj = new EagleDecodeDraftTokensPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h deleted file mode 100644 index 8c144a1bc073..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleDecodeDraftTokensPlugin.h +++ /dev/null @@ -1,174 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class EagleDecodeDraftTokensPlugin : public BasePlugin -{ -public: - EagleDecodeDraftTokensPlugin(nvinfer1::DataType type, int32_t layerIdx, int32_t numEagleLayers, bool topKSampling); - - EagleDecodeDraftTokensPlugin(void const* data, size_t length); - - ~EagleDecodeDraftTokensPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - enum class InputIdxEntry : int32_t - { - // 12 inputs - // [num_input_logits, vocab_size_padded] - LOGITS = 0, - // [batch_size, max_decoding_tokens, max_path_len] - PATHS, - // [1] - NUM_VALID_LOGITS, - // [1] - USE_DYNAMIC_TREE, - // [1] - DYNAMIC_TREE_MAX_TOPK, - - // [batch_size, max_decoding_draft_tokens] - INPUT_DRAFT_TOKEN_IDS, - // [batch_size] - INPUT_DRAFT_LENS, - - // [batch_size, max_decoding_draft_tokens] - INPUT_PREV_SCORES, - - // [batch_size, max_decoding_draft_tokens] - INPUT_CURRENT_EXPAND_INDICES, - - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - INPUT_ALL_LAYERS_SCORES, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - INPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR - }; - - enum class OutputIdxEntry : int32_t - { - // 8 outputs - // [batch_size, max_decoding_draft_tokens] - OUTPUT_DRAFT_TOKEN_IDS = 0, - // [batch_size] - OUTPUT_DRAFT_LENS, - - // [batch_size, max_decoding_tokens, max_path_len] - OUTPUT_PATHS, - - // [batch_size, max_decoding_draft_tokens] - OUTPUT_CURRENT_SCORES, - - // [batch_size, max_decoding_draft_tokens] - OUTPUT_NEXT_EXPAND_INDICES, - - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - OUTPUT_ALL_LAYERS_SCORES, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS, - // [batch_size, num_eagle_layers, max_decoding_draft_tokens x max_decoding_draft_tokens] - OUTPUT_ALL_LAYERS_DRAFT_TOKEN_IDS_PREDECESSOR - }; - - int32_t getIdx(InputIdxEntry idx) const - { - return static_cast(idx); - } - - int32_t getIdx(OutputIdxEntry idx) const - { - return static_cast(idx); - } - -private: - template - size_t getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept; - - template - void enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - - template - void doTopKSampling(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - -private: - nvinfer1::DataType mDtype; // Logit datatype - int32_t mLayerIdx{-1}; // Index of eagle layer - int32_t mNumEagleLayers{-1}; // Number of eagle layers - bool mTopKSampling; // Use TopK sampling or multinomial sampling -}; - -class EagleDecodeDraftTokensPluginCreator : public BaseCreator -{ -public: - EagleDecodeDraftTokensPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.cpp b/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.cpp deleted file mode 100644 index 2cd8c695e296..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.cpp +++ /dev/null @@ -1,548 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "eaglePrepareDrafterInputsPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::EaglePrepareDrafterInputsPluginCreator; -using tensorrt_llm::plugins::EaglePrepareDrafterInputsPlugin; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -static char const* EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_VERSION{"1"}; -static char const* EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_NAME{"EaglePrepareDrafterInputs"}; -PluginFieldCollection EaglePrepareDrafterInputsPluginCreator::mFC{}; -std::vector EaglePrepareDrafterInputsPluginCreator::mPluginAttributes; - -EaglePrepareDrafterInputsPlugin::EaglePrepareDrafterInputsPlugin( - int32_t layerIdx, int32_t numLayers, int32_t maxNonLeavesPerLayer) - : mLayerIdx(layerIdx) - , mNumLayers(numLayers) - , mMaxNonLeavesPerLayer(maxNonLeavesPerLayer) -{ -} - -void EaglePrepareDrafterInputsPlugin::initFieldsToSerialize() -{ - mDataToSerialize.clear(); - mDataToSerialize.emplace_back(PluginField("layer_idx", &mLayerIdx, PluginFieldType::kINT32, 1)); - mDataToSerialize.emplace_back(PluginField("num_layers", &mNumLayers, PluginFieldType::kINT32, 1)); - mDataToSerialize.emplace_back( - PluginField("max_non_leaves_per_layer", &mMaxNonLeavesPerLayer, PluginFieldType::kINT32, 1)); - mFCToSerialize.nbFields = mDataToSerialize.size(); - mFCToSerialize.fields = mDataToSerialize.data(); -} - -nvinfer1::IPluginCapability* EaglePrepareDrafterInputsPlugin::getCapabilityInterface( - nvinfer1::PluginCapabilityType type) noexcept -{ - try - { - if (type == nvinfer1::PluginCapabilityType::kBUILD) - { - return static_cast(this); - } - if (type == nvinfer1::PluginCapabilityType::kRUNTIME) - { - return static_cast(this); - } - TLLM_CHECK(type == nvinfer1::PluginCapabilityType::kCORE); - return static_cast(this); - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -// IPluginV3 methods -nvinfer1::IPluginV3* EaglePrepareDrafterInputsPlugin::clone() noexcept -{ - auto clone = std::make_unique(*this); - clone->initFieldsToSerialize(); - return clone.release(); -} - -// IPluginV3OneCore methods -char const* EaglePrepareDrafterInputsPlugin::getPluginName() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_NAME; -} - -char const* EaglePrepareDrafterInputsPlugin::getPluginVersion() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_VERSION; -} - -char const* EaglePrepareDrafterInputsPlugin::getPluginNamespace() const noexcept -{ - return tensorrt_llm::plugins::api::kDefaultNamespace; -} - -// IPluginV3OneBuild methods -int32_t EaglePrepareDrafterInputsPlugin::getNbOutputs() const noexcept -{ - return 11; -} - -int32_t EaglePrepareDrafterInputsPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -bool EaglePrepareDrafterInputsPlugin::supportsFormatCombination( - int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept -{ - return (inOut[pos].desc.type == nvinfer1::DataType::kINT32) && (inOut[pos].desc.format == TensorFormat::kLINEAR); -} - -int32_t EaglePrepareDrafterInputsPlugin::getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, - nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept -{ - outputTypes[0] = nvinfer1::DataType::kINT32; - outputTypes[1] = nvinfer1::DataType::kINT32; - outputTypes[2] = nvinfer1::DataType::kINT32; - outputTypes[3] = nvinfer1::DataType::kINT32; - outputTypes[4] = nvinfer1::DataType::kINT32; - outputTypes[5] = nvinfer1::DataType::kINT32; - outputTypes[6] = nvinfer1::DataType::kINT32; - outputTypes[7] = nvinfer1::DataType::kINT32; - outputTypes[8] = nvinfer1::DataType::kINT32; - outputTypes[9] = nvinfer1::DataType::kINT32; - outputTypes[10] = nvinfer1::DataType::kINT32; - outputTypes[11] = nvinfer1::DataType::kINT32; - return 0; -} - -int32_t EaglePrepareDrafterInputsPlugin::getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, - nvinfer1::DimsExprs const* shapeInputs, int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(nbOutputs == 11); - TLLM_CHECK(nbInputs == 15); - TLLM_CHECK(nbShapeInputs == 0); - auto const numTokens = inputs[getIdx(InputIdxEntry::INPUT_IDS)].d[0]; - auto const batchSizeExpr = inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)].d[0]; - auto const numGenRequestsExpr = inputs[getIdx(InputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)].d[0]; - auto const numInputGenTokensExpr = inputs[getIdx(InputIdxEntry::INPUT_GEN_TOKENS)].d[0]; - auto const maxDecodingLenExpr = inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)].d[1]; - auto const maxPathLenExpr = inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)].d[2]; - - for (SizeType32 outputIndex = 0; outputIndex < nbOutputs; ++outputIndex) - { - if (outputIndex == getIdx(OutputIdxEntry::SEQUENCE_LENGTHS) - || outputIndex == getIdx(OutputIdxEntry::CONTEXT_LENGTHS) - || outputIndex == getIdx(OutputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)) - { - outputs[outputIndex] = inputs[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)]; - } - else if (outputIndex == getIdx(OutputIdxEntry::SPEC_DECODING_PACKED_MASK)) - { - outputs[outputIndex].nbDims = 3; - outputs[outputIndex].d[0] = batchSizeExpr; - outputs[outputIndex].d[1] = maxDecodingLenExpr; - outputs[outputIndex].d[2] - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *maxDecodingLenExpr, *exprBuilder.constant(32)); - } - else if (outputIndex == getIdx(OutputIdxEntry::SPEC_DECODING_POSITION_OFFSETS)) - { - outputs[outputIndex].nbDims = 2; - outputs[outputIndex].d[0] = batchSizeExpr; - outputs[outputIndex].d[1] = maxDecodingLenExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::OUTPUT_IDS) - || outputIndex == getIdx(OutputIdxEntry::HIDDEN_STATES_INDICES) - || (mLayerIdx == 0 && outputIndex == getIdx(OutputIdxEntry::POSITION_IDS))) - { - if (mLayerIdx == 0) - { - // We have at most numGenRequests * (mNumLayers + 1) accepted tokens per step for gen requests and - // input_ids - numGenTokens tokens for context requests. - auto numOutputGenTokensExpr = exprBuilder.operation( - DimensionOperation::kPROD, *numGenRequestsExpr, *exprBuilder.constant(mNumLayers + 1)); - auto numInputCtxTokensExpr - = exprBuilder.operation(DimensionOperation::kSUB, *numTokens, *numInputGenTokensExpr); - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation(DimensionOperation::kMAX, *exprBuilder.constant(1), - *exprBuilder.operation(DimensionOperation::kSUM, *numOutputGenTokensExpr, *numInputCtxTokensExpr)); - } - else - { - // At most we have mMaxNonLeavesPerLayer non-leaves at this layer. - // And in total we pass all non-leaves + all their preceding nodes. - // batchSize * mMaxNonLeavesPerLayer * layerIdx - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation(DimensionOperation::kPROD, - *exprBuilder.operation(DimensionOperation::kPROD, *exprBuilder.constant(mLayerIdx), - *exprBuilder.constant(mMaxNonLeavesPerLayer)), - *batchSizeExpr); - } - } - else if (mLayerIdx > 0 && outputIndex == getIdx(OutputIdxEntry::POSITION_IDS)) - { - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::LAST_TOKEN_INDICES)) - { - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation( - DimensionOperation::kPROD, *exprBuilder.constant(mMaxNonLeavesPerLayer), *batchSizeExpr); - } - else if (outputIndex == getIdx(OutputIdxEntry::NUM_LAST_TOKEN_INDICES)) - { - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.constant(1); - } - else if (outputIndex == getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)) - { - // batchSize * (maxPathLen - 1) + 1 - outputs[outputIndex].nbDims = 1; - outputs[outputIndex].d[0] = exprBuilder.operation(DimensionOperation::kSUM, *exprBuilder.constant(1), - *exprBuilder.operation(DimensionOperation::kPROD, *batchSizeExpr, - *exprBuilder.operation(DimensionOperation::kSUB, *maxPathLenExpr, *exprBuilder.constant(1)))); - } - } - return 0; -} - -int32_t EaglePrepareDrafterInputsPlugin::onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::PluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - return 0; -} - -nvinfer1::IPluginV3* EaglePrepareDrafterInputsPlugin::attachToContext( - nvinfer1::IPluginResourceContext* context) noexcept -{ - return clone(); -} - -PluginFieldCollection const* EaglePrepareDrafterInputsPlugin::getFieldsToSerialize() noexcept -{ - return &mFCToSerialize; -} - -size_t EaglePrepareDrafterInputsPlugin::getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - size_t workspaceSize{0}; - - auto const batchSize = inputs[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].max.d[0]; - auto const maxDecodingTokens = inputs[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].max.d[1]; - - if (mLayerIdx > 0) - { - SizeType32 constexpr NUM_BUFFERS{9}; - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = batchSize * maxDecodingTokens * sizeof(int8_t); // isLeafMask - workspaces[1] = batchSize * maxDecodingTokens * sizeof(SizeType32); // selectedDraftIndices - workspaces[2] = batchSize * maxDecodingTokens * sizeof(SizeType32); // selectedDraftPosOffsets - workspaces[3] = batchSize * sizeof(SizeType32); // numSelectedDraftIndices - workspaces[4] = batchSize * maxDecodingTokens * maxDecodingTokens * sizeof(int8_t); // selectedMasks - workspaces[5] = (batchSize + 1) * sizeof(SizeType32); // cumSumGenerationLengths - workspaces[6] = batchSize * maxDecodingTokens * sizeof(SizeType32); // nonLeavesInLevelOffsets - workspaces[7] = batchSize * maxDecodingTokens * sizeof(SizeType32); // parentNonLeafInLevelOffset - workspaces[8] = 1 * sizeof(SizeType32); // maxGenerationLength - workspaceSize = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); - } - - return workspaceSize; -} - -void EaglePrepareDrafterInputsPlugin::prepareCtxEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)].dims.d[0]; - - auto const numTokens = inputDesc[getIdx(InputIdxEntry::INPUT_IDS)].dims.d[0]; - auto const numGenRequests = inputDesc[getIdx(InputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)].dims.d[0]; - auto const numInputGenTokens = inputDesc[getIdx(InputIdxEntry::INPUT_GEN_TOKENS)].dims.d[0]; - - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::ACCEPTED_TOKENS)].dims.d[1]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].dims.d[1]; - - auto eagleNetSequenceLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SEQUENCE_LENGTHS)]); - auto eagleNetContextLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::CONTEXT_LENGTHS)]); - auto outputIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_IDS)]); - auto positionIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::POSITION_IDS)]); - auto hiddenStatesIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_STATES_INDICES)]); - auto lastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::LAST_TOKEN_INDICES)]); - auto numLastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::NUM_LAST_TOKEN_INDICES)]); - auto hiddenSizeBatchLevelStarts - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)]); - - auto inputIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::INPUT_IDS)]); - auto chunkedContextNextTokens - = reinterpret_cast(inputs[getIdx(InputIdxEntry::CHUNKED_CONTEXT_NEXT_TOKENS)]); - auto baseNetSequenceLengths = reinterpret_cast(inputs[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)]); - auto baseNetContextLengths = reinterpret_cast(inputs[getIdx(InputIdxEntry::CONTEXT_LENGTHS)]); - auto acceptedTokens = reinterpret_cast(inputs[getIdx(InputIdxEntry::ACCEPTED_TOKENS)]); - auto acceptedLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::ACCEPTED_LENS)]); - auto prevDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::PREV_DRAFT_LENS)]); - auto prevPaths = reinterpret_cast(inputs[getIdx(InputIdxEntry::PREV_DRAFT_PATHS)]); - auto bestPathIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::ACCEPTED_PATHS)]); - - auto const numOutputTokens = (numTokens - numInputGenTokens) + (numGenRequests * (mNumLayers + 1)); - cudaMemsetAsync(positionIds, 0, numOutputTokens * sizeof(SizeType32), stream); - cudaMemsetAsync(hiddenStatesIndices, 0, numOutputTokens * sizeof(SizeType32), stream); - - invokePrepareCtxEagleNetInputs(eagleNetSequenceLengths, eagleNetContextLengths, outputIds, positionIds, - hiddenStatesIndices, lastTokenIndices, numLastTokenIndices, hiddenSizeBatchLevelStarts, inputIds, - chunkedContextNextTokens, baseNetSequenceLengths, baseNetContextLengths, acceptedTokens, acceptedLens, - prevDraftLens, prevPaths, bestPathIds, batchSize, maxPathLen, maxDecodingTokens, mMaxNonLeavesPerLayer, stream); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void EaglePrepareDrafterInputsPlugin::prepareGenEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].dims.d[1]; - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)].dims.d[2]; - - auto eagleNetSequenceLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SEQUENCE_LENGTHS)]); - auto eagleNetContextLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::CONTEXT_LENGTHS)]); - auto outputIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::OUTPUT_IDS)]); - auto positionIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::POSITION_IDS)]); - auto specDecodingGenLengths - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SPEC_DECODING_GENERATION_LENGTHS)]); - auto specDecodingPositionOffsets - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SPEC_DECODING_POSITION_OFFSETS)]); - auto specDecodingPackedMasks - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::SPEC_DECODING_PACKED_MASK)]); - auto hiddenStatesIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_STATES_INDICES)]); - auto lastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::LAST_TOKEN_INDICES)]); - auto numLastTokenIndices = reinterpret_cast(outputs[getIdx(OutputIdxEntry::NUM_LAST_TOKEN_INDICES)]); - auto outputHiddenSizeBatchStartsPerLevel - = reinterpret_cast(outputs[getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)]); - - auto eagleNet0SequenceLengths - = reinterpret_cast(inputs[getIdx(InputIdxEntry::SEQUENCE_LENGTHS)]); - auto eagleNet0ContextLength = reinterpret_cast(inputs[getIdx(InputIdxEntry::CONTEXT_LENGTHS)]); - auto nextDraftPaths = reinterpret_cast(inputs[getIdx(InputIdxEntry::NEXT_DRAFT_PATHS)]); - auto nextDraftIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::NEXT_DRAFT_TOKENS)]); - auto inputHiddenSizeBatchStartsPerLevel - = reinterpret_cast(inputs[getIdx(InputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)]); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - int8_t* isLeafMask = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(int8_t))); - TokenIdType* selectedDraftIndices = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - SizeType32* selectedDraftPosOffsets = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - SizeType32* numSelectedDraftIndices - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - bool* selectedMasks = reinterpret_cast(tc::nextWorkspacePtr( - workspaceBytePtr, offset, batchSize * maxDecodingTokens * maxDecodingTokens * sizeof(int8_t))); - SizeType32* cumSumGenerationLengths = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, (batchSize + 1) * sizeof(SizeType32))); - SizeType32* nonLeavesInLevelOffsets = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - SizeType32* parentNonLeafInLevelOffset = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(SizeType32))); - SizeType32* maxGenerationLength - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, 1 * sizeof(SizeType32))); - - cudaMemsetAsync(hiddenStatesIndices, 0, batchSize * mMaxNonLeavesPerLayer * mLayerIdx * sizeof(SizeType32), stream); - cudaMemsetAsync(selectedMasks, 0, batchSize * maxDecodingTokens * maxDecodingTokens * sizeof(int8_t), stream); - // Prefill mask setting all to leaves. - cudaMemsetAsync(isLeafMask, 1, batchSize * maxDecodingTokens * sizeof(int8_t), stream); - - PrepareGenEagleNetInputsParams params; - params.nextSequenceLengths = eagleNetSequenceLengths; - params.nextContextLengths = eagleNetContextLengths; - params.outputIds = outputIds; - params.positionIds = positionIds; - params.specDecodingGenLengths = specDecodingGenLengths; - params.specDecodingPositionOffsets = specDecodingPositionOffsets; - params.specDecodingPackedMasks = specDecodingPackedMasks; - params.hiddenStatesIndices = hiddenStatesIndices; - params.lastTokenIndices = lastTokenIndices; - params.numLastTokenIndices = numLastTokenIndices; - params.outputHiddenSizeBatchStartsPerLevel = outputHiddenSizeBatchStartsPerLevel; - - // tmp data - params.isLeafMask = isLeafMask; - params.selectedDraftIndices = selectedDraftIndices; - params.selectedDraftPosOffsets = selectedDraftPosOffsets; - params.numSelectedDraftIndices = numSelectedDraftIndices; - params.selectedMasks = selectedMasks; - params.cumSumGenerationLengths = cumSumGenerationLengths; - params.maxGenerationLength = maxGenerationLength; - params.nonLeavesInLevelOffsets = nonLeavesInLevelOffsets; - params.parentNonLeafInLevelOffset = parentNonLeafInLevelOffset; - - params.nextDraftIds = nextDraftIds; - params.eagleNet0SequenceLengths = eagleNet0SequenceLengths; - params.prevContextLengths = eagleNet0ContextLength; - params.nextPaths = nextDraftPaths; - params.inputHiddenSizeBatchStartsPerLevel = inputHiddenSizeBatchStartsPerLevel; - params.levelIdx = mLayerIdx; - params.batchSize = batchSize; - params.maxPathLen = maxPathLen; - params.maxDecodingTokens = maxDecodingTokens; - params.maxNonLeavesPerLayer = mMaxNonLeavesPerLayer; - params.stream = stream; - - params.checkParams(); - - invokePrepareGenEagleNetInputs(params); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -int EaglePrepareDrafterInputsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // First EagleNet instance (EagleNet0) is always chunked context attn, - // where we process either context tokens or newly accepted tokens and append them to EagleNet KV cache. - - // For all following EagleNetX (X > 0) instances there is need for masked spec decoding attn. - // Ideally with mask for context. - // Let's say we have prompt ABCD and two variants of tokens spec decoding tokens E and F - // predicted by EagleNet0. If we draw full attn mask, it becomes: - // |A|B|C|D|E|F - // E|1|1|1|1|1|0 - // F|1|1|1|1|0|1 - // - // In the next step we predict token G from ABCDE branch and token H from ABCDF branch -- like beam search. - // And we'd need spec decoding mask that includes kv cache: - // |A|B|C|D|E|F|G|H - // G|1|1|1|1|1|0|1|0 - // H|1|1|1|1|0|1|0|1 - // - // But TRT-LLM does not support such mask for now. We can only provide - // |G|H - // G|1|0 - // H|0|1 - // , which is wrong mask. - // - // For now we WAR this by passing EFGH for the EagleNet1 with right mask - // and using only G and H logits for sampling, but that's redundant compute: - // |E|F|G|H - // E|1|0|0|0 - // F|0|1|0|0 - // G|1|0|1|0 - // H|0|1|0|1 - - if (mLayerIdx == 0) - { - prepareCtxEagleNetData(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - prepareGenEagleNetData(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - - return 0; -} - -/////////////// - -EaglePrepareDrafterInputsPluginCreator::EaglePrepareDrafterInputsPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("layer_idx", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_layers", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("max_non_leaves_per_layer", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* EaglePrepareDrafterInputsPluginCreator::getPluginName() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_NAME; -} - -char const* EaglePrepareDrafterInputsPluginCreator::getPluginVersion() const noexcept -{ - return EAGLE_PREPARE_DRAFTER_INPUTS_PLUGIN_VERSION; -} - -PluginFieldCollection const* EaglePrepareDrafterInputsPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -nvinfer1::IPluginV3* EaglePrepareDrafterInputsPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept -{ - try - { - int32_t layerIdx{0}; - int32_t numLayers{0}; - int32_t maxNonLeavesPerLayer{0}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fc->fields[i].name; - if (!strcmp(attrName, "layer_idx")) - { - TLLM_CHECK(fc->fields[i].type == PluginFieldType::kINT32); - layerIdx = *static_cast(fc->fields[i].data); - } - else if (!strcmp(attrName, "num_layers")) - { - TLLM_CHECK(fc->fields[i].type == PluginFieldType::kINT32); - numLayers = *static_cast(fc->fields[i].data); - } - else if (!strcmp(attrName, "max_non_leaves_per_layer")) - { - TLLM_CHECK(fc->fields[i].type == PluginFieldType::kINT32); - maxNonLeavesPerLayer = *static_cast(fc->fields[i].data); - } - } - return new EaglePrepareDrafterInputsPlugin(layerIdx, numLayers, maxNonLeavesPerLayer); - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -char const* EaglePrepareDrafterInputsPluginCreator::getPluginNamespace() const noexcept -{ - return tensorrt_llm::plugins::api::kDefaultNamespace; -} diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h b/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h deleted file mode 100644 index 0059c46f6c8d..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eaglePrepareDrafterInputsPlugin.h +++ /dev/null @@ -1,186 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class EaglePrepareDrafterInputsPlugin : public nvinfer1::IPluginV3, - public nvinfer1::IPluginV3OneCore, - public nvinfer1::IPluginV3OneBuild, - public nvinfer1::IPluginV3OneRuntime -{ -public: - EaglePrepareDrafterInputsPlugin(EaglePrepareDrafterInputsPlugin const& p) = default; - - EaglePrepareDrafterInputsPlugin(int32_t layerIdx, int32_t numLayers, int32_t maxNonLeavesPerLayer); - - nvinfer1::IPluginV3* clone() noexcept override; - - nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override; - - void initFieldsToSerialize(); - - char const* getPluginName() const noexcept override; - char const* getPluginVersion() const noexcept override; - char const* getPluginNamespace() const noexcept override; - - int32_t getNbOutputs() const noexcept override; - - bool supportsFormatCombination( - int pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override; - int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override; - - int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes, - int32_t nbInputs) const noexcept override; - - int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs, - int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - - int32_t onShapeChange(nvinfer1::PluginTensorDesc const* in, int32_t nbInputs, nvinfer1::PluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - - nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override; - - size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - -private: - enum class InputIdxEntry : int32_t - { - //! [batch_size] - SEQUENCE_LENGTHS = 0, - //! [batch_size] - CONTEXT_LENGTHS, - //! [num_tokens] - INPUT_IDS, - //! [batch_size] - CHUNKED_CONTEXT_NEXT_TOKENS, - //! [batch_size, max_path_len] - ACCEPTED_TOKENS, - //! [batch_size] - ACCEPTED_LENS, - //! [batch_size] - ACCEPTED_PATHS, - //! [batch_size, max_decoding_draft_tokens] - NEXT_DRAFT_TOKENS, - //! [batch_size] - NEXT_DRAFT_LENS, - //! [batch_size, max_decoding_tokens, max_path_len] - NEXT_DRAFT_PATHS, - //! [batch_size] - PREV_DRAFT_LENS, - //! [batch_size, max_decoding_tokens, max_path_len] - PREV_DRAFT_PATHS, - //! [(max_path_len - 1) * batch_size + 1] - HIDDEN_SIZE_BATCH_LEVEL_STARTS, - //! [num_gen_tokens] - INPUT_GEN_TOKENS, - //! [num_gen_requests] - SPEC_DECODING_GENERATION_LENGTHS, - }; - - enum class OutputIdxEntry : int32_t - { - //! [batch_size] - SEQUENCE_LENGTHS = 0, - //! [batch_size] - CONTEXT_LENGTHS, - //! [batch_size] - SPEC_DECODING_GENERATION_LENGTHS, - //! [batch_size, max_decoding_tokens] - SPEC_DECODING_POSITION_OFFSETS, - //! [batchSize, maxDecodingTokens, ceil(maxDecodingTokens / 32)] - SPEC_DECODING_PACKED_MASK, - //! [batchSize * mMaxNonLeavesPerLayer * layerIdx] for layerIdx > 0 - //! [num_tokens - numGenTokens + numGenRequests * (mNumLayers + 1)] for layerIdx == 0 - OUTPUT_IDS, - //! [batchSize] for layerIdx > 0 - //! [num_tokens - numGenTokens + numGenRequests * (mNumLayers + 1)] for layerIdx == 0 - POSITION_IDS, - //! [batchSize * mMaxNonLeavesPerLayer * layerIdx] for layerIdx > 0 - //! [num_tokens - numGenTokens + numGenRequests * (mNumLayers + 1)] for layerIdx == 0 - HIDDEN_STATES_INDICES, - //! [batchSize * mMaxNonLeavesPerLayer] - LAST_TOKEN_INDICES, - //! [1] - NUM_LAST_TOKEN_INDICES, - //! [(max_path_len - 1) * batch_size + 1] - HIDDEN_SIZE_BATCH_LEVEL_STARTS, - }; - - int32_t getIdx(InputIdxEntry idx) const - { - return static_cast(idx); - } - - int32_t getIdx(OutputIdxEntry idx) const - { - return static_cast(idx); - } - -private: - void prepareCtxEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept; - - void prepareGenEagleNetData(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept; - -private: - int32_t mLayerIdx{0}; - int32_t mNumLayers{0}; - int32_t mMaxNonLeavesPerLayer{0}; - std::vector mDataToSerialize; - nvinfer1::PluginFieldCollection mFCToSerialize; -}; - -class EaglePrepareDrafterInputsPluginCreator : public nvinfer1::IPluginCreatorV3One -{ -public: - EaglePrepareDrafterInputsPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - char const* getPluginNamespace() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV3* createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.cpp b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.cpp deleted file mode 100644 index 5fb30f583712..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.cpp +++ /dev/null @@ -1,565 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "eagleSampleAndAcceptDraftTokensPlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/memoryUtils.h" -#include "tensorrt_llm/kernels/samplingTopKKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/common.h" -#include "tensorrt_llm/kernels/speculativeDecoding/eagleDecodingKernels.h" -#include "tensorrt_llm/kernels/speculativeDecoding/medusaDecodingKernels.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::EagleSampleAndAcceptDraftTokensPluginCreator; -using tensorrt_llm::plugins::EagleSampleAndAcceptDraftTokensPlugin; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::kernels::speculative_decoding; -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -static char const* EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_VERSION{"1"}; -static char const* EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_NAME{"EagleSampleAndAcceptDraftTokens"}; -PluginFieldCollection EagleSampleAndAcceptDraftTokensPluginCreator::mFC{}; -std::vector EagleSampleAndAcceptDraftTokensPluginCreator::mPluginAttributes; - -EagleSampleAndAcceptDraftTokensPlugin::EagleSampleAndAcceptDraftTokensPlugin(nvinfer1::DataType type) - : mDtype(type) -{ -} - -// Parameterized constructor -EagleSampleAndAcceptDraftTokensPlugin::EagleSampleAndAcceptDraftTokensPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDtype); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* EagleSampleAndAcceptDraftTokensPlugin::clone() const noexcept -{ - auto* plugin = new EagleSampleAndAcceptDraftTokensPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs EagleSampleAndAcceptDraftTokensPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK(nbInputs == 10); - TLLM_CHECK(outputIndex < 7); - auto const batchSizeExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[0]; - auto const maxDecodingDraftTokensExpr = inputs[getIdx(InputIdxEntry::DRAFT_TOKEN_IDS)].d[1]; - auto const maxDecodingTokensExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[1]; - auto const maxPathLenExpr = inputs[getIdx(InputIdxEntry::PATHS)].d[2]; - - nvinfer1::DimsExprs ret; - if (outputIndex == getIdx(OutputIdxEntry::ACCEPTED_TOKENS)) - { - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxPathLenExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::ACCEPTED_LENS)) - { - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::BEST_ACCEPTED_PATHS)) - { - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::NEXT_DRAFT_TOKEN_IDS)) - { - ret.nbDims = 2; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingDraftTokensExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::NEXT_DRAFT_LENS)) - { - ret.nbDims = 1; - ret.d[0] = batchSizeExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::NEXT_DRAFT_PATHS)) - { - ret.nbDims = 3; - ret.d[0] = batchSizeExpr; - ret.d[1] = maxDecodingTokensExpr; - ret.d[2] = maxPathLenExpr; - } - else if (outputIndex == getIdx(OutputIdxEntry::HIDDEN_SIZE_BATCH_LEVEL_STARTS)) - { - ret.nbDims = 1; - ret.d[0] = exprBuilder.operation(DimensionOperation::kSUM, *exprBuilder.constant(1), - *exprBuilder.operation(DimensionOperation::kPROD, - *exprBuilder.operation(DimensionOperation::kSUB, *maxPathLenExpr, *exprBuilder.constant(1)), - *batchSizeExpr)); - } - return ret; -} - -bool EagleSampleAndAcceptDraftTokensPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getIdx(InputIdxEntry::LOGITS)) // logits - { - return (inOut[pos].type == mDtype) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == getIdx(InputIdxEntry::TEMPERATURE) || pos == getIdx(InputIdxEntry::RAND_VALIDATION) - || pos == getIdx(InputIdxEntry::POSTERIOR_ALPHA) - || pos == getIdx(InputIdxEntry::POSTERIOR_THRESHOLD)) // temperature, rand_validation - { - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else // everything else - { - return (inOut[pos].type == nvinfer1::DataType::kINT32) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void EagleSampleAndAcceptDraftTokensPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -template -size_t EagleSampleAndAcceptDraftTokensPlugin::getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, - int nbInputs, nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - size_t workspaceSize{0}; - - auto const vocabSizePadded = inputs[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const batchSize = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputs[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - - // Greedy sampling - // Top1 sampling workspace - auto const greedySamplingWorkspaceSize - = getTopKWorkspaceSize(batchSize, maxDecodingTokens, /* maxTopK */ 1, vocabSizePadded); - - // Multinomial sampling - auto const typicalSamplingWorkspaceSize - = getTypicalAcceptanceWorkspaceSize(batchSize, maxDecodingTokens, vocabSizePadded); - - auto const primarySamplingWorkspaceSize = std::max(greedySamplingWorkspaceSize, typicalSamplingWorkspaceSize); - - // Target output ids - auto const targetOutputIdsSize = batchSize * maxDecodingTokens * sizeof(TokenIdType); - // Logits ptrs - auto const logitsPtrsSize = batchSize * maxDecodingTokens * sizeof(T*); - SizeType32 constexpr NUM_BUFFERS{4}; - size_t workspaces[NUM_BUFFERS]; - workspaces[0] = targetOutputIdsSize; - workspaces[1] = primarySamplingWorkspaceSize; - workspaces[2] = logitsPtrsSize; - workspaces[3] = batchSize * sizeof(SizeType32); - workspaceSize = tc::calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); - - return workspaceSize; -} - -size_t EagleSampleAndAcceptDraftTokensPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - auto const logitsType = inputs[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - return getWorkspaceSizeType(inputs, nbInputs, outputs, nbOutputs); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - return getWorkspaceSizeType<__half>(inputs, nbInputs, outputs, nbOutputs); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - return 0; -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::samplePrimeHeadTokens(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // auto const maxNumTokens = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - - auto logits = static_cast(inputs[getIdx(InputIdxEntry::LOGITS)]); - auto prevDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::DRAFT_LENS)]); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - auto const samplingWorkspaceSize - = getTopKWorkspaceSize(batchSize, maxDecodingTokens, /* maxTopK */ 1, vocabSizePadded); - - TokenIdType* outputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - void* workspaceSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, samplingWorkspaceSize)); - T const** logitsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(T*))); - SizeType32* decodingTokens - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - // Assemble pointers to logits - invokeAssembleTargetLogitsOffsets( - logitsPtrs, decodingTokens, logits, prevDraftLens, batchSize, maxDecodingTokens, vocabSizePadded, stream); - - sync_check_cuda_error(stream); - - TopKSamplingKernelParams params; - params.logProbsPtrs = logitsPtrs; - params.outputIds = outputIds; - params.workspace = workspaceSampling; - params.maxTopK = 1; - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.tokensPerStep = decodingTokens; - params.maxTokensPerStep = maxDecodingTokens; - params.maxSeqLen = maxDecodingTokens; - params.vocabSizePadded = vocabSizePadded; - - invokeBatchTopKSampling(params, stream); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::doTypicalAcceptance(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // auto const maxNumTokens = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - // auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[2]; - // auto const maxDraftPathLen = maxPathLen - 1; - - auto logits = static_cast(inputs[getIdx(InputIdxEntry::LOGITS)]); - auto prevDraftLens = reinterpret_cast(inputs[getIdx(InputIdxEntry::DRAFT_LENS)]); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - // Multinomial sampling - auto const primarySamplingWorkspaceSize - = getTypicalAcceptanceWorkspaceSize(batchSize, maxDecodingTokens, vocabSizePadded); - - TokenIdType* outputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - void* workspaceSampling - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, primarySamplingWorkspaceSize)); - T** logitsPtrs = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(T*))); - SizeType32* decodingTokens - = reinterpret_cast(tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * sizeof(SizeType32))); - - // Assemble pointers to logits - invokeAssembleTargetLogitsOffsets(const_cast(logitsPtrs), decodingTokens, logits, prevDraftLens, - batchSize, maxDecodingTokens, vocabSizePadded, stream); - - sync_check_cuda_error(stream); - - TypicalAcceptanceSampling params; - params.logitsPtrs = logitsPtrs; - params.generationLengths = decodingTokens; - params.temperatures = reinterpret_cast(inputs[getIdx(InputIdxEntry::TEMPERATURE)]); - params.posteriorThresholds = reinterpret_cast(inputs[getIdx(InputIdxEntry::POSTERIOR_THRESHOLD)]); - params.posteriorAlphas = reinterpret_cast(inputs[getIdx(InputIdxEntry::POSTERIOR_ALPHA)]); - params.outputIds = outputIds; - params.workspace = reinterpret_cast(workspaceSampling); - params.randomVals = reinterpret_cast(inputs[getIdx(InputIdxEntry::RAND_VALIDATION)]); - - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.maxDecodingTokens = maxDecodingTokens; - params.vocabSize = vocabSizePadded; - - if (mSmCnt <= 0) - { - auto const deviceId = tensorrt_llm::common::getDevice(); - cudaDeviceProp prop{}; - TLLM_CUDA_CHECK(cudaGetDeviceProperties(&prop, deviceId)); - mSmCnt = prop.multiProcessorCount; - } - params.smCnt = mSmCnt; - - params.checkParams(); - - typicalAcceptanceSampling(params, stream); - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::acceptDraftTokens(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - // auto const maxNumTokens = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[0]; - auto const vocabSizePadded = inputDesc[getIdx(InputIdxEntry::LOGITS)].dims.d[1]; - - auto const batchSize = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[0]; - auto const maxDecodingTokens = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[1]; - auto const maxPathLen = inputDesc[getIdx(InputIdxEntry::PATHS)].dims.d[2]; - auto const maxDraftPathLen = maxPathLen - 1; - - auto const useDynamicTree = *(reinterpret_cast(inputs[getIdx(InputIdxEntry::USE_DYNAMIC_TREE)])); - - int8_t* workspaceBytePtr = reinterpret_cast(workspace); - size_t offset{0}; - - // auto const samplingWorkspaceSize - // = getTopKWorkspaceSize(batchSize, maxDecodingTokens, /* maxTopK */ 1, vocabSizePadded); - - TokenIdType* outputIds = reinterpret_cast( - tc::nextWorkspacePtr(workspaceBytePtr, offset, batchSize * maxDecodingTokens * sizeof(TokenIdType))); - - AcceptDraftTokensByIdsWithPathsParams params; - params.outputIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::ACCEPTED_TOKENS)]); - params.draftIds = reinterpret_cast(inputs[getIdx(InputIdxEntry::DRAFT_TOKEN_IDS)]); - params.targetIds = outputIds; - params.acceptedLengths = reinterpret_cast(outputs[getIdx(OutputIdxEntry::ACCEPTED_LENS)]); - params.paths = reinterpret_cast(inputs[getIdx(InputIdxEntry::PATHS)]); - params.bestPathIds = reinterpret_cast(outputs[getIdx(OutputIdxEntry::BEST_ACCEPTED_PATHS)]); - params.batchSize = batchSize; - params.maxBatchSize = batchSize; - params.vocabSize = vocabSizePadded; - params.maxSeqLen = maxPathLen; - params.maxDraftPathLen = maxDraftPathLen; - params.maxDecodingTokens = maxDecodingTokens; - params.stream = stream; - - params.checkParams(); - - acceptDraftTokensByIdsWithPaths(params); - - if (useDynamicTree) - { - // For Eagle-2, after verification and acceptance, the original path becomes useless. - // All set to '-1' - cudaMemsetAsync(outputs[getIdx(OutputIdxEntry::NEXT_DRAFT_PATHS)], -1, - batchSize * maxDecodingTokens * maxPathLen * sizeof(SizeType32), stream); - } - else - { - // For Eagle-1 - // Copy input paths to the output - cudaMemcpyAsync(outputs[getIdx(OutputIdxEntry::NEXT_DRAFT_PATHS)], inputs[getIdx(InputIdxEntry::PATHS)], - batchSize * maxDecodingTokens * maxPathLen * sizeof(SizeType32), cudaMemcpyDeviceToDevice, stream); - } - - sync_check_cuda_error(stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -template -void EagleSampleAndAcceptDraftTokensPlugin::enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const greedySampling = reinterpret_cast(inputs[getIdx(InputIdxEntry::GREEDY_SAMPLING)])[0]; - // TODO split batch into greedy and non-greedy and execute both paths - if (greedySampling) - { - // Sample all main head tokens with Top-1. - samplePrimeHeadTokens(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - // Typical sampling for typical acceptance. - doTypicalAcceptance(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - - // Accept tokens based on token ids, write the best path and best token id. - acceptDraftTokens(inputDesc, outputDesc, inputs, outputs, workspace, stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -int EagleSampleAndAcceptDraftTokensPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - auto const logitsType = inputDesc[getIdx(InputIdxEntry::LOGITS)].type; - if (logitsType == nvinfer1::DataType::kFLOAT) - { - enqueueType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (logitsType == nvinfer1::DataType::kHALF) - { - enqueueType<__half>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported logits type"); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType EagleSampleAndAcceptDraftTokensPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index < 7); - // input 1 is draft tokens now of int32 type. All outputs are int32_t as well. - return inputTypes[getIdx(InputIdxEntry::DRAFT_TOKEN_IDS)]; -} - -// IPluginV2 Methods - -char const* EagleSampleAndAcceptDraftTokensPlugin::getPluginType() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleSampleAndAcceptDraftTokensPlugin::getPluginVersion() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_VERSION; -} - -int EagleSampleAndAcceptDraftTokensPlugin::getNbOutputs() const noexcept -{ - return 7; -} - -int EagleSampleAndAcceptDraftTokensPlugin::initialize() noexcept -{ - return 0; -} - -void EagleSampleAndAcceptDraftTokensPlugin::terminate() noexcept {} - -size_t EagleSampleAndAcceptDraftTokensPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDtype); -} - -void EagleSampleAndAcceptDraftTokensPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDtype); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void EagleSampleAndAcceptDraftTokensPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -EagleSampleAndAcceptDraftTokensPluginCreator::EagleSampleAndAcceptDraftTokensPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* EagleSampleAndAcceptDraftTokensPluginCreator::getPluginName() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_NAME; -} - -char const* EagleSampleAndAcceptDraftTokensPluginCreator::getPluginVersion() const noexcept -{ - return EAGLE_SAMPLE_AND_ACCEPT_DRAFT_TOKENS_PLUGIN_VERSION; -} - -PluginFieldCollection const* EagleSampleAndAcceptDraftTokensPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* EagleSampleAndAcceptDraftTokensPluginCreator::createPlugin( - char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new EagleSampleAndAcceptDraftTokensPlugin(type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* EagleSampleAndAcceptDraftTokensPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call EagleSampleAndAcceptDraftTokensPlugin::destroy() - try - { - auto* obj = new EagleSampleAndAcceptDraftTokensPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h b/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h deleted file mode 100644 index 3b14bab83170..000000000000 --- a/cpp/tensorrt_llm/plugins/eaglePlugin/eagleSampleAndAcceptDraftTokensPlugin.h +++ /dev/null @@ -1,167 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class EagleSampleAndAcceptDraftTokensPlugin : public BasePlugin -{ -public: - EagleSampleAndAcceptDraftTokensPlugin(nvinfer1::DataType type); - - EagleSampleAndAcceptDraftTokensPlugin(void const* data, size_t length); - - ~EagleSampleAndAcceptDraftTokensPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - enum class InputIdxEntry : int32_t - { - //! [num_tokens, vocab_size_padded] - LOGITS = 0, - //! [batch_size, max_decoding_draft_tokens] - DRAFT_TOKEN_IDS, - //! [batch_size] - DRAFT_LENS, - //! [batch_size] - TEMPERATURE, - //! [batch_size, max_decoding_tokens] - RAND_VALIDATION, - //! [batch_size] - POSTERIOR_ALPHA, - //! [batch_size] - POSTERIOR_THRESHOLD, - //! [batch_size, max_decoding_tokens, max_path_len] - PATHS, - //! [1] - GREEDY_SAMPLING, - //! [1] - USE_DYNAMIC_TREE - }; - - enum class OutputIdxEntry : int32_t - { - //! [batch_size, max_path_len] - ACCEPTED_TOKENS = 0, - //! [batch_size] - ACCEPTED_LENS, - //! [batch_size] - BEST_ACCEPTED_PATHS, - //! [batch_size, max_decoding_draft_tokens] - NEXT_DRAFT_TOKEN_IDS, - //! [batch_size] - NEXT_DRAFT_LENS, - //! [batch_size, max_decoding_tokens, max_path_len] - NEXT_DRAFT_PATHS, - //! [max_draft_path_len * batch_size] - HIDDEN_SIZE_BATCH_LEVEL_STARTS, - }; - - int32_t getIdx(InputIdxEntry idx) const - { - return static_cast(idx); - } - - int32_t getIdx(OutputIdxEntry idx) const - { - return static_cast(idx); - } - -private: - template - size_t getWorkspaceSizeType(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept; - - template - void samplePrimeHeadTokens(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept; - - template - void doTypicalAcceptance(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - - template - void acceptDraftTokens(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - - template - void enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept; - -private: - nvinfer1::DataType mDtype; - int32_t mSmCnt{0}; -}; - -class EagleSampleAndAcceptDraftTokensPluginCreator : public BaseCreator -{ -public: - EagleSampleAndAcceptDraftTokensPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/exports.def b/cpp/tensorrt_llm/plugins/exports.def deleted file mode 100644 index 5d4ac9e3e793..000000000000 --- a/cpp/tensorrt_llm/plugins/exports.def +++ /dev/null @@ -1,19 +0,0 @@ -; SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. - -LIBRARY nvinfer_plugin_tensorrt_llm -EXPORTS -getPluginRegistry -initLibNvInferPlugins diff --git a/cpp/tensorrt_llm/plugins/exports.map b/cpp/tensorrt_llm/plugins/exports.map deleted file mode 100644 index c6c949775079..000000000000 --- a/cpp/tensorrt_llm/plugins/exports.map +++ /dev/null @@ -1,34 +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. - */ - -/* Hides all symbols except those specified in the global section */ -{ - global: - initTrtLlmPlugins; - setLoggerFinder; - getPluginCreators; - getCreators; - extern "C++" { - nvinfer1::IPluginCreator::*; - nvinfer1::IPluginV2Ext::*; - nvinfer1::IPluginV2IOExt::*; - nvinfer1::PluginRegistrar*; - tensorrt_llm::plugins::api::*; - tensorrt_llm::plugins::*; - }; - local: *; -}; diff --git a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fp4GemmPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.cpp b/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.cpp deleted file mode 100644 index 05f06ae38feb..000000000000 --- a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.cpp +++ /dev/null @@ -1,434 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 - -#include "fp4GemmPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::Fp4GemmPluginCreator; -using tensorrt_llm::plugins::Fp4GemmPlugin; -using tensorrt_llm::plugins::Fp4GemmPluginProfiler; -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -using namespace tensorrt_llm::kernels::cutlass_kernels; -#else -using namespace tensorrt_llm::kernels::internal_cutlass_kernels; -#endif - -constexpr nvinfer1::DataType FP4_DTYPE = nvinfer1::DataType::kFP4; -constexpr nvinfer1::DataType FP8_DTYPE = nvinfer1::DataType::kFP8; - -static char const* FP4_GEMM_PLUGIN_VERSION{"1"}; -static char const* FP4_GEMM_PLUGIN_NAME{"Fp4Gemm"}; -PluginFieldCollection Fp4GemmPluginCreator::mFC{}; -std::vector Fp4GemmPluginCreator::mPluginAttributes; - -void Fp4GemmPluginProfiler::runTactic( - int m, int n, int k, Fp4GemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - // Workspace size required by gemm runner - // NB: this function will throw exception when selected tactic exceeds SMEM, which is then - // caught by gemmPluginProfiler and it will register this tactic as invalid - size_t wsSizeRunner = mRunner->getWorkspaceSize(m, n, k, /* batch_count */ 1); - - // Workspace size required by profiling - size_t wsByteOffset = 0; - int8_t* wsBytePointer = reinterpret_cast(workspace); - void* aTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, (m * k) / 2)); - void* bTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, (n * k) / 2)); - void* dTmp = reinterpret_cast( - nextWorkspacePtr(wsBytePointer, wsByteOffset, m * n * (mType == nvinfer1::DataType::kFLOAT ? 4u : 2u))); - // SF M/N is padded along 128 and K is padded along 4. - int vector_size = 16; - int sf_round_m = ((m + 127) / 128) * 128; - int sf_round_n = ((n + 127) / 128) * 128; - int sf_round_k = ((k / vector_size + 3) / 4) * 4; - float* a_sf = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, sf_round_m * sf_round_k)); - float* b_sf = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, sf_round_n * sf_round_k)); - float* global_sf = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, sizeof(float))); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, wsSizeRunner)); - - // Run profiling - mRunner->gemm(dTmp, aTmp, bTmp, a_sf, b_sf, global_sf, m, n, k, /* batch_count */ 1, tactic, workspaceTmp, - wsSizeRunner, stream); - sync_check_cuda_error(stream); -} - -void Fp4GemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - size_t vector_size = 16; - size_t sf_round_m = ((maxM + 127) / 128) * 128; - size_t sf_round_n = ((n + 127) / 128) * 128; - size_t sf_round_k = ((k / vector_size + 3) / 4) * 4; - std::vector workspaces = { - (size_t) (maxM * k / 2), // A - (size_t) (n * k / 2), // B - maxM * n * (mType == nvinfer1::DataType::kFLOAT ? 4u : 2u), // D - (size_t) (sf_round_m * sf_round_k), // A_SF - (size_t) (sf_round_n * sf_round_k), // B_SF - sizeof(float), // Global_SF - mRunner->getWorkspaceSize(maxM, n, k, /* batch_count */ 1) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector Fp4GemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -Fp4GemmPlugin::Fp4GemmPlugin( - int sfVecSize, nvinfer1::DataType OutputType, Fp4GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) - , mSfVecSize(sfVecSize) - , mOutputType(OutputType) -{ - init(OutputType); -} - -Fp4GemmPlugin::Fp4GemmPlugin(void const* data, size_t length, Fp4GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mSfVecSize); - read(d, mOutputType); - read(d, mDims); - - init(mOutputType); - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK(d == a + length); -} - -void Fp4GemmPlugin::init(nvinfer1::DataType type) -{ - TLLM_CHECK_WITH_INFO((getSMVersion() >= 100), "FP4 Gemm not supported before Blackwell"); - TLLM_CHECK_WITH_INFO( - (mOutputType == DataType::kBF16) || (mOutputType == DataType::kFLOAT) || (mOutputType == DataType::kHALF), - "Only support float, half, bfloat16, got %d.", (int) mOutputType); - mOutputType = type; - if (mOutputType == nvinfer1::DataType::kHALF) - { - mGemmRunner = std::make_shared>(); - } - else if (mOutputType == nvinfer1::DataType::kFLOAT) - { - mGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - else if (mOutputType == nvinfer1::DataType::kBF16) - { - mGemmRunner = std::make_shared>(); - } -#endif - - mGemmId = GemmIdCore(mDims.n, mDims.k, mOutputType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* Fp4GemmPlugin::clone() const noexcept -{ - auto* plugin = new Fp4GemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs Fp4GemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK_WITH_INFO(outputIndex == 0, "Only support one output"); - auto const& dimsInput = inputs[getInputTensorIdx()]; - auto const& dimsWeights = inputs[getWeightsTensorIdx()]; - TLLM_CHECK_WITH_INFO(dimsInput.nbDims >= 2 && dimsWeights.nbDims == 2, "Fp4GemmPlugin input dim=%d, weights dim=%d", - dimsInput.nbDims, dimsWeights.nbDims); - nvinfer1::DimsExprs ret; - if (outputIndex == 0) - { - ret.nbDims = dimsInput.nbDims; - for (int i = 0; i < dimsInput.nbDims - 1; ++i) - { - ret.d[i] = dimsInput.d[i]; - } - ret.d[dimsInput.nbDims - 1] = dimsWeights.d[0]; - } - else - { - TLLM_CHECK_WITH_INFO(outputIndex == 0, "output fp4 not supported now."); - ret.nbDims = 1; - auto vecCount = dimsInput.d[0]; - int numDim = dimsInput.nbDims; - for (int idx = 1; idx < numDim - 1; ++idx) - { - vecCount = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *vecCount, *dimsInput.d[idx]); - } - auto constant128 = exprBuilder.constant(128); - auto alignedRowCount = exprBuilder.operation(nvinfer1::DimensionOperation::kCEIL_DIV, *vecCount, *constant128); - alignedRowCount = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *alignedRowCount, *constant128); - auto constant4 = exprBuilder.constant(4); - auto constantSFSize = exprBuilder.constant(mSfVecSize); - auto sfColumn - = exprBuilder.operation(nvinfer1::DimensionOperation::kCEIL_DIV, *dimsInput.d[numDim - 1], *constantSFSize); - auto alignedColumnCount = exprBuilder.operation(nvinfer1::DimensionOperation::kCEIL_DIV, *sfColumn, *constant4); - alignedColumnCount - = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *alignedColumnCount, *constant4); - auto totalSize - = exprBuilder.operation(nvinfer1::DimensionOperation::kPROD, *alignedColumnCount, *alignedRowCount); - ret.d[0] = totalSize; - } - return ret; -} - -bool Fp4GemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (inOut[pos].format != TensorFormat::kLINEAR) - { - return false; - } - if (pos == getInputTensorIdx()) - { - return (inOut[pos].type == FP4_DTYPE); - } - else if (pos == getWeightsTensorIdx()) - { - return (inOut[pos].type == FP4_DTYPE); - } - else if (pos == getInputSFTensorIdx() || pos == getWeightsSFTensorIdx()) - { - return (inOut[pos].type == FP8_DTYPE); - } - else if (pos == getGlobalSFTensorIdx()) - { - return (inOut[pos].type == DataType::kFLOAT); - } - else if (pos == nbInputs) - { - // Output - return (inOut[pos].type == DataType::kFLOAT || inOut[pos].type == DataType::kBF16 - || inOut[pos].type == DataType::kHALF); - } - return false; -} - -void Fp4GemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[2].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[2].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mOutputType}; - m_workspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK, /* batch_count */ 1); -} - -size_t Fp4GemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int Fp4GemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // 0. input_tensor [num_tokens, dim] - // 1. input_block_scale [num_tokens, dim / SFVecSize] (padded) - // 2. weights_tensor [out_dim, dim] - // 3. weights_block_scale [out_dim, dim / SFVecSize] (padded) - // 4. alpha (global scaling factor) [1] - // outputs - // 0. output_tensor [num_tokens, out_dim] - int64_t m = 1; - for (int i = 0; i < inputDesc[getInputTensorIdx()].dims.nbDims - 1; ++i) - { - m *= inputDesc[getInputTensorIdx()].dims.d[i]; - } - int const n = inputDesc[getWeightsTensorIdx()].dims.d[0]; - int const k = inputDesc[getWeightsTensorIdx()].dims.d[1]; - TLLM_CHECK_WITH_INFO(k % 32 == 0, "K dim should be aligned to 16 Bytes"); - int N_align = mOutputType == nvinfer1::DataType::kFLOAT ? 4u : 8u; - TLLM_CHECK_WITH_INFO(n % N_align == 0, "N dim should be aligned to 16 Bytes"); - size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k, /* batch_count */ 1); - auto const bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid FP4 GEMM tactic"); - if (m >= 1) - { - mGemmRunner->gemm(outputs[0], inputs[0], inputs[2], inputs[1], inputs[3], - reinterpret_cast(inputs[4]), m, n, k, /* batch_count */ 1, *bestTactic, - reinterpret_cast(workspace), wsSize, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType Fp4GemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index == 0, "Only support one output"); - return mOutputType; -} - -// IPluginV2 Methods - -char const* Fp4GemmPlugin::getPluginType() const noexcept -{ - return FP4_GEMM_PLUGIN_NAME; -} - -char const* Fp4GemmPlugin::getPluginVersion() const noexcept -{ - return FP4_GEMM_PLUGIN_VERSION; -} - -int Fp4GemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int Fp4GemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void Fp4GemmPlugin::terminate() noexcept {} - -size_t Fp4GemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(mSfVecSize) + // mSfVecSize - sizeof(nvinfer1::DataType) + // dtype - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void Fp4GemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mSfVecSize); - write(d, mOutputType); - write(d, mDims); - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void Fp4GemmPlugin::destroy() noexcept -{ - delete this; -} - -void Fp4GemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mGemmRunner, mOutputType, mDims, mGemmId); -} - -/////////////// - -Fp4GemmPluginCreator::Fp4GemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("sv_vec_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("output_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* Fp4GemmPluginCreator::getPluginName() const noexcept -{ - return FP4_GEMM_PLUGIN_NAME; -} - -char const* Fp4GemmPluginCreator::getPluginVersion() const noexcept -{ - return FP4_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* Fp4GemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* Fp4GemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 2); - int sf_vec_size{}; - nvinfer1::DataType output_type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "sf_vec_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sf_vec_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "output_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - output_type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // Fp4GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - auto* obj = new Fp4GemmPlugin(sf_vec_size, output_type, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* Fp4GemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call CumsumLastDimPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new Fp4GemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h b/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h deleted file mode 100644 index 9947e849d84f..000000000000 --- a/cpp/tensorrt_llm/plugins/fp4GemmPlugin/fp4GemmPlugin.h +++ /dev/null @@ -1,162 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/fp4_gemm.h" -#else -#include "fp4_gemm.h" -#endif - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -#if defined(USING_OSS_CUTLASS_FP4_GEMM) -using Fp4GemmRunnerPtr = std::shared_ptr; -#else -using Fp4GemmRunnerPtr - = std::shared_ptr; -#endif - -class Fp4GemmPluginProfiler : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class Fp4GemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - Fp4GemmPlugin() = delete; - - Fp4GemmPlugin(int sfVecSize, nvinfer1::DataType OutputType, PluginProfilerPtr const& pluginProfiler); - - Fp4GemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - - ~Fp4GemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - - IndexType getInputSFTensorIdx() const - { - return 1; - }; - - IndexType getWeightsTensorIdx() const - { - return 2; - }; - - IndexType getWeightsSFTensorIdx() const - { - return 3; - }; - - IndexType getGlobalSFTensorIdx() const - { - return 4; - } - - void init(nvinfer1::DataType type); - void configGemm(); - - Fp4GemmRunnerPtr mGemmRunner; - PluginProfilerPtr mPluginProfiler; - - int mSfVecSize; - nvinfer1::DataType mOutputType; - size_t m_workspaceMaxSize; - GemmDims mDims{}; - GemmIdCore mGemmId{}; -}; - -class Fp4GemmPluginCreator : public BaseCreator -{ -public: - Fp4GemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager mGemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/CMakeLists.txt deleted file mode 100644 index 3b714a3928fb..000000000000 --- a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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 *.cu) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.cpp deleted file mode 100644 index 84963df50a21..000000000000 --- a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.cpp +++ /dev/null @@ -1,422 +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. - */ - -#include "fp8RowwiseGemmPlugin.h" -#include "cutlass_extensions/gemm_configs.h" - -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::Fp8RowwiseGemmPluginCreator; -using tensorrt_llm::plugins::Fp8RowwiseGemmPlugin; -using tensorrt_llm::plugins::Fp8RowwiseGemmPluginProfiler; - -static char const* FP8_ROWWISE_GEMM_PLUGIN_VERSION{"1"}; -static char const* FP8_ROWWISE_GEMM_PLUGIN_NAME{"Fp8RowwiseGemm"}; -PluginFieldCollection Fp8RowwiseGemmPluginCreator::mFC{}; -std::vector Fp8RowwiseGemmPluginCreator::mPluginAttributes; - -size_t Fp8RowwiseGemmPluginProfiler::getBytePerElement(nvinfer1::DataType type) -{ - size_t bpe; - if (type == nvinfer1::DataType::kHALF || type == nvinfer1::DataType::kBF16) - { - bpe = 2; - } - else if (type == nvinfer1::DataType::kINT8 || type == nvinfer1::DataType::kFP8) - { - bpe = 1; - } - else - { - TLLM_THROW("Not recognized/implemented"); - } - return bpe; -} - -void Fp8RowwiseGemmPluginProfiler::setQuantMode(tensorrt_llm::common::QuantMode const& quantMode) -{ - mQuantMode = quantMode; -} - -void Fp8RowwiseGemmPluginProfiler::runTactic(int m, int n, int k, Fp8RowwiseGemmPluginProfiler::Config const& tactic, - char* workspace, cudaStream_t const& stream) -{ - size_t bpeIn = getBytePerElement(nvinfer1::DataType::kFP8); - size_t bpeOut = getBytePerElement(mType); - - // Workspace size required by gemm runner - // NB: this function will throw exception when selected tactic exceeds SMEM, which is then - // caught by gemmPluginProfiler and it will register this tactic as invalid - size_t wsSizeRunner = mRunner->getWorkspaceSize(m, n, k); - - // Workspace size required by profiling - size_t wsByteOffset = 0; - int8_t* wsBytePointer = reinterpret_cast(workspace); - void* aTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * k * bpeIn)); - void* bTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * k * bpeIn)); - // void* cTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * bpeOut)); - void* dTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * n * bpeOut)); - float* scaleD0Tmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * sizeof(float))); - float* scaleD1Tmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * sizeof(float))); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, wsSizeRunner)); - - // Run profiling - mRunner->gemm(dTmp, aTmp, bTmp, nullptr, mQuantMode, m, n, k, scaleD0Tmp, scaleD1Tmp, tactic, workspaceTmp, - wsSizeRunner, stream); - sync_check_cuda_error(stream); -} - -int Fp8RowwiseGemmPluginProfiler::getMaxProfileM() const -{ - // Max_num_tokens are not suggested to be set larger than 16k. - return 16384; -} - -void Fp8RowwiseGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - std::vector workspaces = { - maxM * k * getBytePerElement(nvinfer1::DataType::kFP8), // A - n * k * getBytePerElement(nvinfer1::DataType::kFP8), // B - // n * getBytePerElement(mType), // C_bias - maxM * n * getBytePerElement(mType), // D - maxM * sizeof(float), // alphaRow - n * sizeof(float), // alphaCol - maxM * sizeof(float), // alphaOutput - mRunner->getWorkspaceSize(maxM, n, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector Fp8RowwiseGemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -Fp8RowwiseGemmPlugin::Fp8RowwiseGemmPlugin( - QuantMode quantMode, nvinfer1::DataType type, Fp8RowwiseGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mQuantMode(quantMode) - , mPluginProfiler(pluginProfiler) -{ - init(type); -} - -// Parameterized constructor -Fp8RowwiseGemmPlugin::Fp8RowwiseGemmPlugin( - void const* data, size_t length, Fp8RowwiseGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - unsigned int quantMode; - read(d, quantMode); - read(d, type); - read(d, mDims); - - mQuantMode = QuantMode(quantMode); - - init(type); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK(d == a + length); -} - -void Fp8RowwiseGemmPlugin::init(nvinfer1::DataType type) -{ - mType = type; - if (mType == nvinfer1::DataType::kHALF) - { - mGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - mGemmRunner = std::make_shared>(); - } -#endif - else - { - TLLM_THROW("Fp8 Rowwise Gemm plugin doesn't support this type now"); - } - - mPluginProfiler->setQuantMode(mQuantMode); - - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* Fp8RowwiseGemmPlugin::clone() const noexcept -{ - auto* plugin = new Fp8RowwiseGemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs Fp8RowwiseGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 4); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool Fp8RowwiseGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have fp8 type - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // scales channels - case 3: - // scales tokens - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // All other format combinations are unsupported. - return false; - } -} - -void Fp8RowwiseGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - mWorkspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t Fp8RowwiseGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mWorkspaceMaxSize; -} - -int Fp8RowwiseGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M(*), K] - // mat2 [N, K] - // scale_tokens [M, 1] if has_per_token_scaling else [1, 1] - // scale_channels [1, N] if has_per_channel_scaling else [1, 1] - // outputs - // mat [M(*), N] - int m = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m *= inputDesc[0].dims.d[ii]; - } - int const n = inputDesc[1].dims.d[0]; - int const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k); - - auto const bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid GEMM tactic"); - mGemmRunner->gemm(outputs[0], inputs[0], inputs[1], nullptr, mQuantMode, m, n, k, - reinterpret_cast(inputs[2]), reinterpret_cast(inputs[3]), *bestTactic, - reinterpret_cast(workspace), wsSize, stream); - sync_check_cuda_error(stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType Fp8RowwiseGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* Fp8RowwiseGemmPlugin::getPluginType() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_NAME; -} - -char const* Fp8RowwiseGemmPlugin::getPluginVersion() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_VERSION; -} - -int Fp8RowwiseGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int Fp8RowwiseGemmPlugin::initialize() noexcept -{ - configGemm(); // gemm profiler in action - return 0; -} - -void Fp8RowwiseGemmPlugin::terminate() noexcept {} - -size_t Fp8RowwiseGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(unsigned int) + // QuantMode - sizeof(nvinfer1::DataType) + // dtype - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void Fp8RowwiseGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void Fp8RowwiseGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void Fp8RowwiseGemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mGemmRunner, mType, mDims, mGemmId); -} - -Fp8RowwiseGemmPluginCreator::Fp8RowwiseGemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("has_per_channel_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("has_per_token_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* Fp8RowwiseGemmPluginCreator::getPluginName() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_NAME; -} - -char const* Fp8RowwiseGemmPluginCreator::getPluginVersion() const noexcept -{ - return FP8_ROWWISE_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* Fp8RowwiseGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* Fp8RowwiseGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 3); - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // Fp8RowwiseGemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - QuantMode quantMode = QuantMode{}; - auto* obj = new Fp8RowwiseGemmPlugin(quantMode, type, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* Fp8RowwiseGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call Fp8RowwiseGemmPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new Fp8RowwiseGemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h b/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h deleted file mode 100644 index 36f22ad5885d..000000000000 --- a/cpp/tensorrt_llm/plugins/fp8RowwiseGemmPlugin/fp8RowwiseGemmPlugin.h +++ /dev/null @@ -1,140 +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. - */ -#pragma once - -#include "tensorrt_llm/kernels/cutlass_kernels/fp8_rowwise_gemm/fp8_rowwise_gemm.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using Fp8RowwiseGemmRunnerPtr - = std::shared_ptr; - -class Fp8RowwiseGemmPluginProfiler : public GemmPluginProfiler - -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantMode(tensorrt_llm::common::QuantMode const& quantMode); - - virtual int getMaxProfileM() const override; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - size_t getBytePerElement(nvinfer1::DataType type); - - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class Fp8RowwiseGemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - Fp8RowwiseGemmPlugin() = delete; - - Fp8RowwiseGemmPlugin( - tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, PluginProfilerPtr const& pluginProfiler); - - Fp8RowwiseGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~Fp8RowwiseGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - - void configGemm(); - -private: - const std::string mLayerName; - - Fp8RowwiseGemmRunnerPtr mGemmRunner; - tensorrt_llm::common::QuantMode mQuantMode; // not configurable yet - size_t mWorkspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; -}; - -class Fp8RowwiseGemmPluginCreator : public BaseCreator -{ -public: - Fp8RowwiseGemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager mGemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/CMakeLists.txt deleted file mode 100755 index 7cc985b60b7a..000000000000 --- a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.cpp b/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.cpp deleted file mode 100644 index 541afdadc4c8..000000000000 --- a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.cpp +++ /dev/null @@ -1,388 +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. - */ -#include "fusedLayernormPlugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::FusedLayernormPluginCreator; -using tensorrt_llm::plugins::FusedLayernormPlugin; - -static char const* FUSED_LAYERNORM_PLUGIN_VERSION{"1"}; -static char const* FUSED_LAYERNORM_PLUGIN_NAME{"FusedLayernorm"}; -PluginFieldCollection FusedLayernormPluginCreator::mFC{}; -std::vector FusedLayernormPluginCreator::mPluginAttributes; - -FusedLayernormPlugin::FusedLayernormPlugin(float eps, bool needFP32Output, bool needQuantize, nvinfer1::DataType type) - : mEps(eps) - , mNeedFP32Output(needFP32Output) - , mNeedQuantize(needQuantize) - , mType(type) -{ -} - -// Parameterized constructor -FusedLayernormPlugin::FusedLayernormPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mEps); - read(d, mNeedFP32Output); - read(d, mNeedQuantize); - read(d, mType); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* FusedLayernormPlugin::clone() const noexcept -{ - auto* plugin = new FusedLayernormPlugin(mEps, mNeedFP32Output, mNeedQuantize, mType); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs FusedLayernormPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - // Dim should be the same as input hidden states - if (!mNeedQuantize) - { - return inputs[0]; - } - - if (outputIndex == 1) // un-normed output fp16 - { - return inputs[0]; - } - if (outputIndex == 0) // quantized normed output - { - // Quantized output with int64_t data type (16 FP4 values per element). - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - return ret; - } - - // Scaling Factors. - try - { - TLLM_CHECK(outputIndex == 2); - - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - // Sequence dimension or token dimension. - // Pad to multiple of 128. - auto dimM - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 2], *exprBuilder.constant(128)); - ret.d[ret.nbDims - 2] = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - // Hidden size dimension. - ret.d[ret.nbDims - 1] - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16)); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool FusedLayernormPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int const totalPoses = 5 + 2 * static_cast(mNeedQuantize); - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 3 + static_cast(mNeedQuantize)); - if (pos < nbInputs) - { - switch (pos) - { - case 0: - case 1: - case 2: return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - case 3: return (inOut[pos].type == nvinfer1::DataType::kFLOAT); - } - } - if (pos == nbInputs) // Normed output - { - if (mNeedQuantize) - { - // fp4 quantized output -- fp4 padded tp int64 - return (inOut[pos].type == nvinfer1::DataType::kFP4) && (inOut[pos].format == TensorFormat::kLINEAR); - } - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == nbInputs + 1) // Un-normed output - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - // fp4 act_per_block_scale -- fp8 padded to int32 - return (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void FusedLayernormPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t FusedLayernormPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return sizeof(WarpSpecializedCounters); -} - -int FusedLayernormPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] - // residual [M(*), N] - // weight [N, ] - // scale [1, ] - if needQuantize - // outputs - // output [M(*), N] - fp4 padded to int64 / fp16 - // un-normed output [M(*), N] - fp16 - // act_per_block_scale - fp8 padded to int32 - if needQuantize - -#define SETUP_PARAM \ - Param param; \ - int64_t m64 = 1; \ - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) \ - { \ - m64 *= inputDesc[0].dims.d[i]; \ - } \ - int const m = TLLM_INT32_CAST(m64); \ - int const n = TLLM_INT32_CAST(inputDesc[2].dims.d[0]); \ - param.m = m; \ - param.n = n; \ - param.layernorm_eps = mEps; \ - param.input = const_cast(reinterpret_cast(inputs[0])); \ - param.residual = const_cast(reinterpret_cast(inputs[1])); \ - param.gamma = const_cast(reinterpret_cast(inputs[2])); \ - if (mNeedQuantize) \ - { \ - param.sf_scale = const_cast(reinterpret_cast(inputs[3])); \ - } \ - param.counters = reinterpret_cast(workspace); \ - param.stream = stream; \ - param.normed_output = reinterpret_cast(outputs[0]); \ - param.output = reinterpret_cast(outputs[1]); \ - param.sf_out = reinterpret_cast(outputs[2]); - -#define CLEANUP_AND_INVOKE \ - TLLM_CUDA_CHECK(cudaMemsetAsync(workspace, 0, sizeof(WarpSpecializedCounters), stream)); \ - invokeWSLayerNorm(param, true, num_sms); - - int num_sms = tensorrt_llm::common::getMultiProcessorCount(); - - if (mType == DataType::kHALF) - { - using Input = half; - using Param = WarpSpecializedParam>; - SETUP_PARAM - CLEANUP_AND_INVOKE - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - using Input = __nv_bfloat16; - using Param = WarpSpecializedParam>; - SETUP_PARAM - CLEANUP_AND_INVOKE - } -#endif - else - { - TLLM_LOG_ERROR("Unsupported data type"); - return 1; - } - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType FusedLayernormPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - // assert((mNeedFP32Output && index < 3) || (!mNeedFP32Output && index < 2)); - assert((mNeedQuantize && index < 3) || (!mNeedQuantize && index < 2)); - if (index == 0) - { - // Output 0 quantized output of layernorm - fp4 padded to int64 - if (mNeedQuantize) - { - return nvinfer1::DataType::kFP4; - } - return mType; - } - else if (index == 1) - { - // Output 1 un-normed output - return mType; - } - // Output 2 act_per_block_scale - fp8 padded to int32 - return nvinfer1::DataType::kFP8; -} - -// IPluginV2 Methods - -char const* FusedLayernormPlugin::getPluginType() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_NAME; -} - -char const* FusedLayernormPlugin::getPluginVersion() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_VERSION; -} - -int FusedLayernormPlugin::getNbOutputs() const noexcept -{ - return 2 + static_cast(mNeedQuantize); -} - -int FusedLayernormPlugin::initialize() noexcept -{ - return 0; -} - -void FusedLayernormPlugin::terminate() noexcept {} - -size_t FusedLayernormPlugin::getSerializationSize() const noexcept -{ - return sizeof(mEps) + sizeof(mNeedFP32Output) + sizeof(mNeedQuantize) + sizeof(mType); -} - -void FusedLayernormPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mEps); - write(d, mNeedFP32Output); - write(d, mNeedQuantize); - write(d, mType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void FusedLayernormPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -FusedLayernormPluginCreator::FusedLayernormPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("need_fp32_output", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("need_quantize", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* FusedLayernormPluginCreator::getPluginName() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_NAME; -} - -char const* FusedLayernormPluginCreator::getPluginVersion() const noexcept -{ - return FUSED_LAYERNORM_PLUGIN_VERSION; -} - -PluginFieldCollection const* FusedLayernormPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* FusedLayernormPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - float eps{}; - nvinfer1::DataType type{}; - bool needFP32Output{}; - bool needQuantize{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "need_fp32_output")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - needFP32Output = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "need_quantize")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - needQuantize = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new FusedLayernormPlugin(eps, needFP32Output, needQuantize, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* FusedLayernormPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call FusedLayernormPlugin::destroy() - try - { - auto* obj = new FusedLayernormPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h b/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h deleted file mode 100755 index c6c899950fdc..000000000000 --- a/cpp/tensorrt_llm/plugins/fusedLayernormPlugin/fusedLayernormPlugin.h +++ /dev/null @@ -1,98 +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. - */ -#pragma once - -#include "tensorrt_llm/kernels/fusedLayernormKernels/layernorm_param.h" -#include "tensorrt_llm/kernels/fusedLayernormKernels/ws_layernorm.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class FusedLayernormPlugin : public BasePlugin -{ -public: - FusedLayernormPlugin() = delete; - - FusedLayernormPlugin(float eps, bool needFP32Output, bool needQuantize, nvinfer1::DataType type); - - FusedLayernormPlugin(void const* data, size_t length); - - ~FusedLayernormPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - float mEps; - bool mNeedFP32Output; - bool mNeedQuantize; - nvinfer1::DataType mType; - - const std::string mLayerName; -}; - -class FusedLayernormPluginCreator : public BaseCreator -{ -public: - FusedLayernormPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/CMakeLists.txt deleted file mode 100644 index 1d1fa98f4132..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/CMakeLists.txt +++ /dev/null @@ -1,19 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp deleted file mode 100644 index 08ee2af55406..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.cpp +++ /dev/null @@ -1,721 +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 "gemmAllReducePlugin.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" -#include "tensorrt_llm/plugins/common/pluginUtils.h" - -#include - -static char const* GEMM_ALLREDUCE_PLUGIN_VERSION = "1"; -static char const* GEMM_ALLREDUCE_PLUGIN_NAME = "GemmAllReduce"; -template -using CutlassType = ::tensorrt_llm::kernels::cutlass_kernels::CutlassType; - -namespace tensorrt_llm::plugins -{ -template -static std::pair makeEntry() -{ - return {std::make_tuple(ElementA, ElementB, ElementD), - [&]() - { - using GemmTraits - = cutlass_kernels::GemmTypes::type, typename CutlassType::type, - typename CutlassType::type, // C, unused - typename CutlassType::type, - std::conditional_t, // SFA - std::conditional_t, // SFB - cutlass::layout::RowMajor, cutlass::layout::ColumnMajor, - cutlass::layout::RowMajor, // C, unused - cutlass::layout::RowMajor>; - return new cutlass_kernels::GemmAllReduceImplRunner(); - }}; -} - -template -static std::map getTypedInstantiators() -{ - return std::map({makeEntry(), - makeEntry(), - makeEntry(), - makeEntry(), - makeEntry(), - makeEntry()}); -} - -//////////////////////////////////////////////////////////// -// GemmAllReducePlugin Methods -//////////////////////////////////////////////////////////// -GemmAllReducePlugin::GemmAllReducePlugin(GemmAllReducePluginOptions const& options) - : mOptions(options) - , mGemmId(GemmIdCore(options.maxProblemShape.n, options.maxProblemShape.k, options.typeD)) - , mProfiler(mGemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/options.deserialize)) -{ - // construct mapping of input/output pos to argument - int argIdx = 0; - // inputs - mArgMap[argIdx++] = TensorArg::IN_ACTIVATION; - mArgMap[argIdx++] = TensorArg::IN_WEIGHT; - if (mOptions.hasSFA) - { - mArgMap[argIdx++] = TensorArg::IN_ACTIVATION_SF; - } - if (mOptions.hasSFB) - { - mArgMap[argIdx++] = TensorArg::IN_WEIGHT_SF; - } - if (mOptions.alphaIsPtr) - { - mArgMap[argIdx++] = TensorArg::IN_ALPHA; - } - mNbInputs = argIdx; - // outputs - mArgMap[argIdx++] = TensorArg::OUT_D_UC; - mArgMap[argIdx++] = TensorArg::OUT_D_MC; - mArgMap[argIdx++] = TensorArg::OUT_D_IPC; - mNbOutputs = argIdx - mNbInputs; - - // Create mapping of argument to tensor pos - for (auto const& pair : mArgMap) - { - mArgInvMap[pair.second] = pair.first; - } - - // Use map instead of huge switch case - mTypedInstantiators = getTypedInstantiators(); - - auto key = std::make_tuple(mOptions.typeA, mOptions.typeB, mOptions.typeD); - - TLLM_CHECK_WITH_INFO(mTypedInstantiators.count(key) > 0, "No cutlass gemm for impl."); - mGemm = std::shared_ptr(mTypedInstantiators[key]()); -} - -void GemmAllReducePlugin::allocatePersistentWorkspace() -{ - TLLM_CHECK(mOptions.maxProblemShape.isInitialized()); - - mWorkspaceKey = "gemm_allreduce_workspace_m" + std::to_string(mOptions.maxProblemShape.maxM); - - cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig smallest_tile_config - = mGemm->getSupportedLaunchConfigs()[0]; - cutlass_kernels::GemmAllReduceImplInterface::ProblemArgs args; - args.argProblemShape(mOptions.maxProblemShape.maxM, mOptions.maxProblemShape.n, mOptions.maxProblemShape.k, 1) - .argRanks(mRank, mOptions.group) - .argLaunchConfig(smallest_tile_config); - - TLLM_CHECK(mWorkspace == nullptr); - - // Wrap persistent workspace in IPluginResource type - // so that clone() can be called to allocate memory - GemmAllReducePersistentWorkspace unallocated_resource(mGemm->getPersistentWorkspace(args)); - - // Register and allocate workspace - mWorkspace = static_cast( - getPluginRegistry()->acquirePluginResource(mWorkspaceKey.c_str(), &unallocated_resource)); - TLLM_CHECK(mWorkspace != nullptr); -} - -LaunchConfig GemmAllReducePlugin::getStaticHeuristicLaunchConfig(int M) const -{ - using namespace tensorrt_llm::cutlass_extensions; - // This is only applicable when we swap and transpose A & B. - // When M is small we want to select tile that best fits it to maximize MMA efficiency. - auto filterByM = [&](std::vector candidateConfigs) - { - std::vector result; - if (M <= 16) - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x16x128 and config.transposed; }); - } - else if (M <= 32) - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x32x128 and config.transposed; }); - } - else if (M <= 64) - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x64x128 and config.transposed; }); - } - else - { - std::copy_if(candidateConfigs.begin(), candidateConfigs.end(), std::back_inserter(result), - [](const LaunchConfig& config) - { return config.tile_shape == TileShape::TileShape_128x128x128 and config.transposed; }); - } - // If result empty then use any. - if (result.empty()) - { - result = candidateConfigs; - } - return result; - }; - - auto bestLaunchConfigs = mGemm->getSupportedLaunchConfigs(); - bestLaunchConfigs = filterByM(bestLaunchConfigs); - TLLM_CHECK(!bestLaunchConfigs.empty()); - // Return first one, because who knows which is best. - return bestLaunchConfigs.front(); -} - -static GemmAllReducePluginOptions deserializeOptions(void const*& data, size_t length) -{ - char const* begin = reinterpret_cast(data); - char const*& end = reinterpret_cast(data); - GemmAllReducePluginOptions options; - options.deserialize = true; - - read(end, options.typeA); - read(end, options.typeB); - read(end, options.typeD); - read(end, options.transA); - read(end, options.transB); - read(end, options.alpha); - read(end, options.maxProblemShape); - read(end, options.groupSize); - for (int i = 0; i < options.groupSize; ++i) - { - int rank = -1; - read(end, rank); - options.group.insert(rank); - } - read(end, options.hasSFA); - read(end, options.hasSFB); - read(end, options.alphaIsPtr); - - TLLM_CHECK_WITH_INFO(end == begin + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (end - begin)); - - return options; -} - -GemmAllReducePlugin::GemmAllReducePlugin(void const* data, size_t length) - : GemmAllReducePlugin(deserializeOptions(std::ref(data), length)) -{ - if (mProfiler->useProfiler()) - { - mProfiler->deserializeFromOwnFile(mGemmId, mOptions.maxProblemShape); - } -} - -////////////////////////////////// -// IPluginV2DynamicExt Methods -////////////////////////////////// -IPluginV2DynamicExt* GemmAllReducePlugin::clone() const noexcept -{ - return new GemmAllReducePlugin(*this); -} - -DimsExprs GemmAllReducePlugin::getOutputDimensions( - int outputIndex, DimsExprs const* inputs, int nbInputs, IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == mNbInputs); // number of input tensors - TLLM_CHECK(inputs[0].nbDims == inputs[1].nbDims); - TLLM_CHECK(outputIndex < getNbOutputs()); - - // List of pointers to D on each rank - if ((nbInputs + outputIndex) == TensorArg::OUT_D_IPC) - { - DimsExprs out_dims; - out_dims.nbDims = 1; - out_dims.d[0] = exprBuilder.constant(mOptions.groupSize); - return out_dims; - } - - TLLM_CHECK(mOptions.transA == false); - TLLM_CHECK(mOptions.transB == true); - - int const nbDimsA = inputs[0].nbDims; // number of dims - int const nbDimsB = inputs[1].nbDims; - - DimsExprs out_dims; - // subtract 2 -> K from each input - out_dims.nbDims = nbDimsA + nbDimsB - 2; - - if (mOptions.transA) - { - for (int i = 1; i < nbDimsA; ++i) - { - out_dims.d[i - 1] = inputs[0].d[i]; - } - } - else - { - for (int i = 0; i < nbDimsA - 1; ++i) - { - out_dims.d[i] = inputs[0].d[i]; - } - } - if (mOptions.transB) - { - for (int i = 0; i < nbDimsB - 1; ++i) - { - out_dims.d[nbDimsA - 1 + i] = inputs[1].d[i]; - } - } - else - { - for (int i = 1; i < nbDimsB; ++i) - { - out_dims.d[nbDimsA - 2 + i] = inputs[1].d[i]; - } - } - return out_dims; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool GemmAllReducePlugin::supportsFormatCombination( - int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept -{ - // inOut[0] -> activation - // inOut[1] -> weight - // inOut[1+hasInputSF] -> activation_sf - // inOut[1+hasInputSF*2] -> weight_sf - // inOut[2+hasInputSF*2] -> output[0] = D_uc - // inOut[3+hasInputSF*2] -> output[1] = D_mc - - TLLM_CHECK_WITH_INFO(pos < mNbInputs + mNbOutputs, "Unexpected pos: %d", pos); - auto const& desc = inOut[pos]; - - TLLM_CHECK_WITH_INFO(mArgMap.count(pos) > 0, "pos %d not found in mArgMap.", pos); - TensorArg arg = mArgMap[pos]; - - auto typeExists = [&](DataType dtype, auto idx) -> bool - { - for (const auto& [key, value] : mTypedInstantiators) - { - // key format: - if (std::get(key) == dtype) - { - return true; - } - } - return false; - }; - - switch (arg) - { - case TensorArg::IN_ACTIVATION: return typeExists(desc.type, std::integral_constant{}); - case TensorArg::IN_WEIGHT: return typeExists(desc.type, std::integral_constant{}); - case TensorArg::IN_ACTIVATION_SF: - case TensorArg::IN_WEIGHT_SF: - // Assumed SF for only FP4 at the moment - return desc.type == DataType::kFP8; - case TensorArg::IN_ALPHA: return desc.type == DataType::kFLOAT; - case TensorArg::OUT_D_UC: - case TensorArg::OUT_D_MC: - case TensorArg::OUT_D_IPC: return typeExists(desc.type, std::integral_constant{}); - default: return false; - } -} - -void GemmAllReducePlugin::configurePlugin( - DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept -{ - // Get problem shape - int const nbDimsA = in[0].max.nbDims; - int const minM = utils::computeMDimension(mOptions.transA, in[0].min); - int const maxM = utils::computeMDimension(mOptions.transA, in[0].max); - int const N = utils::computeNDimension(mOptions.transB, in[1].max); - int const K = mOptions.transA ? in[0].max.d[0] : in[0].max.d[nbDimsA - 1]; - - TLLM_CHECK_WITH_INFO(out[0].desc.type == mOptions.typeD, "Output type mismatch."); - - // Ensure call from execution phase does - // not override call from build phase - if (!mOptions.maxProblemShape.isInitialized()) - { - mOptions.maxProblemShape = {minM, maxM, N, K}; - mGemmId = {N, K, mOptions.typeD}; - } - - // Build phase doesn't have COMM_SESSION (i.e built on single rank) - // so do not allocate persistent workspace - if (!isBuilding()) - { - auto getTPRank = [&]() - { - int rank = COMM_SESSION.getRank(); - auto it = std::find(mOptions.group.begin(), mOptions.group.end(), rank); - TLLM_CHECK_WITH_INFO(it != mOptions.group.end(), - "Incorrect group specified - rank " + std::to_string(rank) + " not found in group"); - return std::distance(mOptions.group.begin(), it); - }; - - mRank = getTPRank(); - - if (mWorkspace == nullptr) - { - allocatePersistentWorkspace(); - } - } -} - -size_t GemmAllReducePlugin::getWorkspaceSize( - PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept -{ - return 0; -} - -int GemmAllReducePlugin::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs[0] -> [M(*), K] - // inputs[1] -> [K, N] - // outputs[0] -> [M(*), N] unicast ptr - // outputs[1] -> [M(*), N] multicast ptr - auto const nbDimsA = inputDesc[0].dims.nbDims; - auto const M = utils::computeMDimension(mOptions.transA, inputDesc[0].dims); - auto const N = utils::computeNDimension(mOptions.transB, inputDesc[1].dims); - auto const K = mOptions.transA ? inputDesc[0].dims.d[0] : inputDesc[0].dims.d[nbDimsA - 1]; - - TLLM_CHECK_WITH_INFO(M <= mOptions.maxProblemShape.maxM, "GemmAllReducePlugin M > maxM."); - TLLM_CHECK_WITH_INFO(M > 0, "GemmAllReducePlugin M is 0."); - TLLM_CHECK_WITH_INFO(N > 0, "GemmAllReducePlugin N is 0."); - TLLM_CHECK_WITH_INFO(K > 0, "GemmAllReducePlugin K is 0."); - TLLM_CHECK_WITH_INFO(mWorkspace != nullptr, "GemmAllReducePlugin workspace is null."); - - LaunchConfig bestLaunchConfig; - if (mProfiler->useProfiler()) - { - bestLaunchConfig = mProfiler->getBestConfig(M, mGemmId).value(); - } - else - { - bestLaunchConfig = getStaticHeuristicLaunchConfig(M); - } - - void const* activation = inputs[mArgInvMap[TensorArg::IN_ACTIVATION]]; - void const* weight = inputs[mArgInvMap[TensorArg::IN_WEIGHT]]; - void* D_out_uc = outputs[mArgInvMap[TensorArg::OUT_D_UC] - mNbInputs]; - void* D_out_mc = outputs[mArgInvMap[TensorArg::OUT_D_MC] - mNbInputs]; - void* D_out_ipc = outputs[mArgInvMap[TensorArg::OUT_D_IPC] - mNbInputs]; - - TLLM_CHECK_WITH_INFO(activation != nullptr, "GemmAllReducePlugin activation is NULL"); - TLLM_CHECK_WITH_INFO(weight != nullptr, "GemmAllReducePlugin weight is NULL"); - TLLM_CHECK_WITH_INFO(D_out_uc != nullptr, "GemmAllReducePlugin out_uc is NULL"); - TLLM_CHECK_WITH_INFO(D_out_mc != nullptr, "GemmAllReducePlugin out_mc is NULL"); - TLLM_CHECK_WITH_INFO(D_out_ipc != nullptr, "GemmAllReducePlugin out_ipc is NULL"); - - cutlass_kernels::GemmAllReduceImplInterface::ProblemArgs args; - args.argProblemShape(M, N, K, 1) - .argA(activation) - .argB(weight) - .argC(nullptr) - .argD(D_out_uc, D_out_mc, (void**) D_out_ipc) - .argRanks(mRank, mOptions.group) - .argBeta(0.f) // no bias - .argLaunchConfig(bestLaunchConfig) - .argWorkspace(mWorkspace->mWorkspace.get()); - // tensor for scaling input A - if (mOptions.hasSFA) - { - void const* activation_sf = inputs[mArgInvMap[TensorArg::IN_ACTIVATION_SF]]; - TLLM_CHECK_WITH_INFO(activation_sf != nullptr, "GemmAllReducePlugin activation_sf is NULL"); - args.argAScale(activation_sf); - } - // tensor for scaling input B - if (mOptions.hasSFB) - { - void const* weight_sf = inputs[mArgInvMap[TensorArg::IN_WEIGHT_SF]]; - TLLM_CHECK_WITH_INFO(weight_sf != nullptr, "GemmAllReducePlugin weight_sf is NULL"); - args.argBScale(weight_sf); - } - // tensor for scaling output D - if (mOptions.alphaIsPtr) - { - void const* alpha_vec = inputs[mArgInvMap[TensorArg::IN_ALPHA]]; - TLLM_CHECK_WITH_INFO(alpha_vec != nullptr, "GemmAllReducePlugin alpha_vec is NULL"); - args.argAlphaPtr(reinterpret_cast(alpha_vec)); - } - else - { - args.argAlpha(mOptions.alpha); - } - - mGemm->run(args, stream); - - return 0; -} - -////////////////////////////////// -// IPluginV2Ext Methods -////////////////////////////////// -DataType GemmAllReducePlugin::getOutputDataType(int index, DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index < getNbOutputs(), "Output index out of bounds: %d", index); - return mOptions.typeD; -} - -////////////////////////////////// -// IPluginV2 Methods -////////////////////////////////// -char const* GemmAllReducePlugin::getPluginType() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_NAME; -} - -char const* GemmAllReducePlugin::getPluginVersion() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_VERSION; -} - -int GemmAllReducePlugin::getNbOutputs() const noexcept -{ - return mNbOutputs; -} - -int GemmAllReducePlugin::initialize() noexcept -{ - if (isBuilding() && mProfiler->useProfiler()) - { - // TODO (xsimmons): interfaces between GemmPluginProfiler and Plugin - // needs to be relooked at - current interface implicitly assigns runner to profiler - // object in profileTactics() - assert(mOptions.maxProblemShape.isInitialized()); - mProfiler->profileTactics(mGemm, mOptions.typeD, mOptions.maxProblemShape, mGemmId); - } - return 0; -} - -void GemmAllReducePlugin::terminate() noexcept -{ - if (isBuilding()) // need this otherwise getComm will crash during build phase - { - return; - } - - // free mWorkspace - if (mWorkspace) - { - getPluginRegistry()->releasePluginResource(mWorkspaceKey.c_str()); - mWorkspace = nullptr; - } -} - -size_t GemmAllReducePlugin::getSerializationSize() const noexcept -{ - // cannot use sizeof(GemmAllReducePluginOptions) - // becaused need packed attribute which doesn't work on enum - // without making the enum also packed - size_t size = 0; - size += sizeof(mOptions.typeA); - size += sizeof(mOptions.typeB); - size += sizeof(mOptions.typeD); - size += sizeof(mOptions.transA); - size += sizeof(mOptions.transB); - size += sizeof(mOptions.alpha); - size += sizeof(mOptions.maxProblemShape); - size += sizeof(mOptions.groupSize); - size += mOptions.group.size() * sizeof(int); - size += sizeof(mOptions.hasSFA); - size += sizeof(mOptions.hasSFB); - size += sizeof(mOptions.alphaIsPtr); - return size; -} - -void GemmAllReducePlugin::serialize(void* buffer) const noexcept -{ - char* begin = reinterpret_cast(buffer); - char* end = reinterpret_cast(buffer); - - write(end, mOptions.typeA); - write(end, mOptions.typeB); - write(end, mOptions.typeD); - write(end, mOptions.transA); - write(end, mOptions.transB); - write(end, mOptions.alpha); - write(end, mOptions.maxProblemShape); - write(end, mOptions.groupSize); - for (auto const& rank : mOptions.group) - { - write(end, rank); - } - write(end, mOptions.hasSFA); - write(end, mOptions.hasSFB); - write(end, mOptions.alphaIsPtr); - TLLM_CHECK(end == begin + getSerializationSize()); - - // Profiler MNK->kernel mappings need to be deterministic and consistent across ranks - // to ensure correct functionality (unlike standalone GEMMs). - // Since by default each rank will generate and serialize its own profiler mapping - // this can lead to different mappings between ranks which will result in fatal - // error. Therefore only generate and use profiler mapping for single rank. - if (mProfiler->useProfiler() && COMM_SESSION.getRank() == 0) - { - mProfiler->serializeToOwnFile(mGemmId); - } -} - -void GemmAllReducePlugin::destroy() noexcept -{ - delete this; -} - -//////////////////////////////////////////////////////////// -// GemmAllReducePluginCreator Methods -//////////////////////////////////////////////////////////// -PluginFieldCollection GemmAllReducePluginCreator::mFC; -std::vector GemmAllReducePluginCreator::mPluginAttributes; - -GemmAllReducePluginCreator::GemmAllReducePluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back("type_a", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("type_b", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("type_d", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("transa", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("transb", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("alpha", nullptr, PluginFieldType::kFLOAT32, 1); - mPluginAttributes.emplace_back("group", nullptr, PluginFieldType::kINT32, 1); - mPluginAttributes.emplace_back("has_sfa", nullptr, PluginFieldType::kINT8, 1); - mPluginAttributes.emplace_back("has_sfb", nullptr, PluginFieldType::kINT8, 1); - mPluginAttributes.emplace_back("alpha_is_ptr", nullptr, PluginFieldType::kINT8, 1); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* GemmAllReducePluginCreator::getPluginName() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_NAME; -} - -char const* GemmAllReducePluginCreator::getPluginVersion() const noexcept -{ - return GEMM_ALLREDUCE_PLUGIN_VERSION; -} - -PluginFieldCollection const* GemmAllReducePluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GemmAllReducePluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - GemmAllReducePluginOptions options; - options.deserialize = false; - - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_a")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.typeA = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "type_b")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.typeB = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "type_d")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.typeD = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "transa")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.transA = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "transb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - options.transB = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "alpha")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - options.alpha = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* ranks = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - options.group.insert(ranks[j]); - } - options.groupSize = options.group.size(); - } - else if (!strcmp(attrName, "has_sfa")) // passed in as input tensor - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - options.hasSFA = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "has_sfb")) // passed in as input tensor - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - options.hasSFB = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "alpha_is_ptr")) // passed in as input tensor - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - options.alphaIsPtr = *static_cast(fields[i].data); - } - } - - try - { - // GemmAllReducePluginCreator is unique and shared for an engine generation - auto* obj = new GemmAllReducePlugin(options); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - return nullptr; - } -} - -IPluginV2* GemmAllReducePluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GemmAllReducePlugin::destroy() - try - { - auto* obj = new GemmAllReducePlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h deleted file mode 100644 index 457926246002..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePlugin.h +++ /dev/null @@ -1,189 +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 - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -#else -#include "allreduce_gemm_runner.h" -#endif - -#include "gemmAllReducePluginProfiler.h" -#include "gemmAllReducePluginResource.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -using namespace nvinfer1; - -using nvinfer1::DataType; -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -namespace cutlass_kernels = ::tensorrt_llm::kernels::opened_cutlass_kernels; -#else -namespace cutlass_kernels = ::tensorrt_llm::kernels::cutlass_kernels; -#endif - -using LaunchConfig = typename cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig; - -namespace tensorrt_llm::plugins -{ -struct GemmAllReducePluginOptions -{ - // Don't need to specify problem shape, this - // is specified in configurePlugin - DataType typeA; - DataType typeB; - DataType typeD; - int transA; - int transB; - float alpha; - // ranks participating in collective - std::set group; - int groupSize; - // Set in configurePlugin during build phase - GemmDims maxProblemShape; - bool deserialize; // used for profiler instantiation - int8_t hasSFA = 0; - int8_t hasSFB = 0; - int8_t alphaIsPtr = 0; -}; - -class GemmAllReducePlugin : public BasePlugin -{ - friend class GemmAllReducePluginCreator; - -public: - ~GemmAllReducePlugin() override = default; - - ////////////////////////////////// - // IPluginV2DynamicExt Methods - ////////////////////////////////// - IPluginV2DynamicExt* clone() const noexcept override; - - DimsExprs getOutputDimensions( - int outputIndex, DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override; - - // inOut[0] -> activation - // inOut[1] -> weight - // inOut[2] -> result - bool supportsFormatCombination( - int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override; - - // in[0] -> activation - // in[1] -> weight - // no bias needed - void configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, - int32_t nbOutputs) noexcept override; - - size_t getWorkspaceSize(PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, - int32_t nbOutputs) const noexcept override; - - // in[0] -> activation - // in[1] -> weight - // out[0] -> result_uc - // out[1] -> result_mc - int enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs, - void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - ////////////////////////////////// - // IPluginV2Ext Methods - ////////////////////////////////// - DataType getOutputDataType(int index, DataType const* inputTypes, int nbInputs) const noexcept override; - - ////////////////////////////////// - // IPluginV2 Methods - ////////////////////////////////// - char const* getPluginType() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - int getNbOutputs() const noexcept override; - - int initialize() noexcept override; - - void terminate() noexcept override; - - size_t getSerializationSize() const noexcept override; - - void serialize(void* buffer) const noexcept override; - - void destroy() noexcept override; - -private: - explicit GemmAllReducePlugin(GemmAllReducePluginOptions const& options); - // Parameterized constructor - explicit GemmAllReducePlugin(void const* data, size_t length); - - void allocatePersistentWorkspace(); - - LaunchConfig getStaticHeuristicLaunchConfig(int M) const; - - // Params that are initialized during constructor - using KeyType = std::tuple; - using ValueType = std::function; - GemmAllReducePluginOptions mOptions; - int mRank = 0; - - enum TensorArg - { - IN_ACTIVATION, - IN_ACTIVATION_SF, - IN_WEIGHT, - IN_WEIGHT_SF, - IN_ALPHA, - OUT_D_UC, - OUT_D_MC, - OUT_D_IPC - }; - - std::unordered_map mArgMap; - std::unordered_map mArgInvMap; - int mNbInputs = 0; - int mNbOutputs = 0; - - std::map mTypedInstantiators; - std::string mWorkspaceKey; - std::shared_ptr mGemm; - // Params that are initialized during configurePlugin() - GemmAllReducePersistentWorkspace* mWorkspace = nullptr; - - // Used for selecting best GEMM for given problem shapes - GemmIdCore mGemmId{}; - GemmPluginProfilerManager mGemmPluginProfileManager; - std::shared_ptr mProfiler; -}; - -class GemmAllReducePluginCreator : public BaseCreator -{ -public: - GemmAllReducePluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp deleted file mode 100644 index a6f7ca2615df..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.cpp +++ /dev/null @@ -1,140 +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 "gemmAllReducePlugin.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/kernels/cutlass_kernels/cutlass_type_conversion.h" -#include "tensorrt_llm/plugins/common/pluginUtils.h" - -namespace tc = tensorrt_llm::common; - -namespace tensorrt_llm::plugins -{ -void GemmAllReducePluginProfiler::serializeToOwnFile(GemmIdCore gemmId) -{ - std::vector file_buf(getSerializationSize(gemmId)); - char* begin = file_buf.data(); - char* end = file_buf.data(); - serialize(end, gemmId); - assert(end == begin + file_buf.size()); - - auto fileName = getCacheFileName(gemmId); - std::ofstream file(fileName, std::ios::binary); - TLLM_CHECK(file.is_open()); - file.write(begin, file_buf.size()); - file.flush(); - file.close(); -} - -void GemmAllReducePluginProfiler::deserializeFromOwnFile(GemmIdCore gemmId, GemmDims problemShape) -{ - auto fileName = getCacheFileName(gemmId); - std::ifstream file(fileName, std::ios::binary); - TLLM_CHECK(file.is_open()); - file.seekg(0, std::ios::end); - std::streamsize size = file.tellg(); - TLLM_CHECK(size > 0); - file.seekg(0, std::ios::beg); - - std::vector file_buf(size); - file.read(file_buf.data(), size); - file.close(); - - char const* begin = const_cast(file_buf.data()); - char const* end = begin; - deserialize(end, problemShape, gemmId); - assert(end == begin + size); -} - -bool GemmAllReducePluginProfiler::useProfiler() -{ - // char const* envDir = getenv("GEMM_AR_PLUGIN_PROFILE_DIR"); - // return envDir != nullptr; - // TODO(xsimmons): currently the profiler does not add any perf gain - // due to static heuristics being sufficient. We can re-enable this - // when we need more configurations. - return false; -} - -std::string GemmAllReducePluginProfiler::getCacheFileName(GemmIdCore gemmId) -{ - std::stringstream fileName; - char const* envDir = getenv("GEMM_AR_PLUGIN_PROFILE_DIR"); - std::string directory = envDir ? std::string(envDir) : "/tmp/"; - fileName << directory + "/gemm-AR"; - fileName << "-n" << std::to_string(gemmId.n); - fileName << "-k" << std::to_string(gemmId.k); - fileName << "-" << tc::getDtypeString(gemmId.dtype); - fileName << ".prof_cache"; - return fileName.str(); -} - -void GemmAllReducePluginProfiler::runTactic(int m, int n, int k, - cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig const& tactic, char* workspace, - cudaStream_t const& stream) -{ - const size_t dtype_size = tc::getDTypeSize(mType); - char* inputA = workspace; - char* inputB = inputA + m * k * dtype_size; - char* outputD = inputB + n * k * dtype_size; - char* inputSFA = outputD + m * n * dtype_size; - char* inputSFB = inputSFA + m * k * dtype_size; - std::set tpGroup = {0}; - - // Run on single-GPU - cutlass_kernels::GemmAllReduceImplInterface::ProblemArgs args; - args.argProblemShape(m, n, k, 1) - .argA((void*) inputA) - .argB((void*) inputB) - .argD((void*) outputD, /*output_mc=*/nullptr) - .argAScale((void*) inputSFA) - .argBScale((void*) inputSFB) - .argRanks(0, tpGroup) - .argAlpha(1.f) - .argBeta(0.f) // no bias - .argLaunchConfig(tactic); - - TLLM_CHECK(mRunner != nullptr); - mRunner->run(args, stream); -} - -void GemmAllReducePluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - TLLM_CHECK(maxM != 0); - TLLM_CHECK(n != 0); - TLLM_CHECK(k != 0); - // mType refers to the output data type - // WARNING: This code assumes that the output precision is >= to input precision - const size_t dtype_size = tc::getDTypeSize(mType); - size_t bytes = 0; - bytes += maxM * k * dtype_size; // A - bytes += n * k * dtype_size; // B - // No C - // Note that D is typically IPC, however, when tuning GEMM we need it to run on single GPU - bytes += maxM * n * dtype_size; // D - // scale tensors for A & B - will at most be same size as A/B - bytes += maxM * k * dtype_size; // A - bytes += n * k * dtype_size; // B - - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector GemmAllReducePluginProfiler::getTactics( - int m, int n, int k) const -{ - TLLM_CHECK(mRunner != nullptr); - return mRunner->getSupportedLaunchConfigs(); -} -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.h b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.h deleted file mode 100644 index faacbb3b8c0f..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginProfiler.h +++ /dev/null @@ -1,68 +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 - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -#else -#include "allreduce_gemm_runner.h" -#endif -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -namespace tensorrt_llm::plugins -{ -/* - * Used for tuning to find best GEMM configs for different problem shapes. - * WARNING: Tuning GEMM+AR kernel may not be fully representable of real - * multi-GPU workloads as tuning only runs on single-GPU. - * IMPORTANT: TRT-LLM does not support deterministic tuning across ranks. - * Because of this, we have to serialize/deserialize our own configuration file. - */ - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -namespace cutlass_kernels = ::tensorrt_llm::kernels::opened_cutlass_kernels; -#else -namespace cutlass_kernels = ::tensorrt_llm::kernels::cutlass_kernels; -#endif -class GemmAllReducePluginProfiler - : public GemmPluginProfiler, GemmIdCore, GemmIdCoreHash> -{ -public: - void serializeToOwnFile(GemmIdCore gemmId); - - void deserializeFromOwnFile(GemmIdCore gemmId, GemmDims problemShape); - - bool useProfiler(); - -protected: - //////////////////////////////////// - // GemmPluginProfiler methods - //////////////////////////////////// - void runTactic(int m, int n, int k, cutlass_kernels::GemmAllReduceImplInterface::LaunchConfig const& tactic, - char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics( - int m, int n, int k) const override; - -private: - static std::string getCacheFileName(GemmIdCore gemmId); -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginResource.h b/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginResource.h deleted file mode 100644 index 8136bd363bd7..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmAllReducePlugin/gemmAllReducePluginResource.h +++ /dev/null @@ -1,82 +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 "NvInferPlugin.h" - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/allreduce_gemm_runner.h" -#else -#include "allreduce_gemm_runner.h" -#endif -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -using namespace nvinfer1; - -namespace tensorrt_llm::plugins -{ - -#if defined(USING_OSS_CUTLASS_ALLREDUCE_GEMM) -namespace cutlass_kernels = ::tensorrt_llm::kernels::opened_cutlass_kernels; -#else -namespace cutlass_kernels = ::tensorrt_llm::kernels::cutlass_kernels; -#endif -class GemmAllReducePersistentWorkspace : public IPluginResource -{ -public: - GemmAllReducePersistentWorkspace(std::shared_ptr workspace) - : mWorkspace(workspace) - { - } - - ////////////////////////////////// - // IPluginResource Methods - ////////////////////////////////// - IPluginResource* clone() noexcept override - { - auto copy = new GemmAllReducePersistentWorkspace(mWorkspace); - // Resource initialization (if any) may be skipped for non-cloned objects - // since only clones will be registered by TensorRT. - try - { - copy->mWorkspace->allocate(); - return copy; - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR(e.what()); - return nullptr; - } - } - - int32_t release() noexcept override - { - try - { - return mWorkspace->free(); - } - catch (std::exception const& e) - { - TLLM_LOG_ERROR(e.what()); - return -1; - } - } - - std::shared_ptr mWorkspace; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gemmPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.cpp b/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.cpp deleted file mode 100644 index 9e06ad01d10f..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.cpp +++ /dev/null @@ -1,614 +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. - */ - -#include "gemmPlugin.h" - -#include "gemmPluginProfiler.h" -#include "plugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/cudaCoreGemm.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -#include - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::GemmDims; -using tensorrt_llm::plugins::GemmPluginCreator; -using tensorrt_llm::plugins::GemmPlugin; -using tensorrt_llm::plugins::CublasLtGemmPluginProfiler; -using tensorrt_llm::plugins::CublasGemmWrapperPtr; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* GEMM_PLUGIN_VERSION{"1"}; -static char const* GEMM_PLUGIN_NAME{"Gemm"}; -PluginFieldCollection GemmPluginCreator::mFC{}; -std::vector GemmPluginCreator::mPluginAttributes; - -void getProblemParams(cublasOperation_t& transa, cublasOperation_t& transb, int& m, int& n, int& k, int& lda, int& ldb, - int& ldc, bool transA, bool transB, int M, int N, int K, int padLda, int padLdb, int padLdc) -{ - transa = transB ? CUBLAS_OP_T : CUBLAS_OP_N; - transb = transA ? CUBLAS_OP_T : CUBLAS_OP_N; - m = N; - n = M; - k = K; - lda = transB ? K + padLdb : N + padLdb; - ldb = transA ? M + padLda : K + padLda; - ldc = N + padLdc; -} - -void runGemm(int const M, int const N, int const K, bool const transA, bool const transB, int const padLda, - int const padLdb, int const padLdc, nvinfer1::DataType const type, CublasGemmWrapperPtr const& cublasWrapperPtr, - void const* act, void const* weight, float const alpha, void* output, - std::optional const& heuristic, void* workspace, cudaStream_t stream) -{ - if (M == 0 || N == 0 || K == 0) - return; - - cublasWrapperPtr->setStream(stream); - cublasWrapperPtr->setWorkspace(workspace); - - cublasOperation_t transa, transb; - int m, n, k; - int lda, ldb, ldc; - getProblemParams(transa, transb, m, n, k, lda, ldb, ldc, transA, transB, M, N, K, padLda, padLdb, padLdc); - - cublasWrapperPtr->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - cublasWrapperPtr->Gemm(transa, transb, m, n, k, weight, lda, act, ldb, output, ldc, alpha, 0.0f, heuristic); - cublasWrapperPtr->destroyDescriptors(); -} - -void CublasLtGemmPluginProfiler::runTactic( - int m, int n, int k, CublasLtGemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - size_t dataSize = sizeof(half); - if (mType == nvinfer1::DataType::kFLOAT) - { - dataSize = sizeof(float); - } - - void* actPtr = reinterpret_cast(workspace); - void* weightPtr = reinterpret_cast( - nextWorkspacePtrWithAlignment(reinterpret_cast(actPtr), m * k * dataSize, ALIGNMENT)); - void* outputPtr = reinterpret_cast( - nextWorkspacePtrWithAlignment(reinterpret_cast(weightPtr), n * k * dataSize, ALIGNMENT)); - char* workspacePtr = reinterpret_cast( - nextWorkspacePtrWithAlignment(reinterpret_cast(outputPtr), m * (n + mPadLdc) * dataSize, ALIGNMENT)); - runGemm(m, n, k, mTransA, mTransB, mPadLda, mPadLdb, mPadLdc, mType, mRunner, actPtr, weightPtr, 1.0f, outputPtr, - {tactic}, workspacePtr, stream); -} - -bool CublasLtGemmPluginProfiler::checkTactic(int m, int n, int k, Config const& tactic) const -{ - cublasOperation_t transa, transb; - int M = m, N = n, K = k; - int lda, ldb, ldc; - getProblemParams(transa, transb, m, n, k, lda, ldb, ldc, mTransA, mTransB, M, N, K, mPadLda, mPadLdb, mPadLdc); - - mRunner->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - - auto const checkResult = mRunner->checkTactic(transa, transb, m, n, k, lda, ldb, ldc, tactic.algo); - - mRunner->destroyDescriptors(); - - return checkResult; -} - -void CublasLtGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - size_t dataSize = getDTypeSize(mType); - size_t outputDataSize = getDTypeSize(mOutputType); - - std::vector workspaces = { - maxM * k * dataSize, // A - n * k * dataSize, // B - maxM * (n + mPadLdc) * outputDataSize, // C - CUBLAS_WORKSPACE_SIZE // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size(), ALIGNMENT); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector CublasLtGemmPluginProfiler::getTactics(int M, int N, int K) const -{ - cublasOperation_t transa, transb; - int m, n, k; - int lda, ldb, ldc; - getProblemParams(transa, transb, m, n, k, lda, ldb, ldc, mTransA, mTransB, M, N, K, mPadLda, mPadLdb, mPadLdc); - - mRunner->createDescriptors(transa, transb, m, n, k, lda, ldb, ldc); - auto const heruistics = mRunner->getTactics(transa, transb, m, n, k, lda, ldb, ldc); - mRunner->destroyDescriptors(); - - return heruistics; -} - -GemmPlugin::GemmPlugin(int transA, int transB, int padLda, int padLdb, int padLdc, nvinfer1::DataType type, bool useFp8, - float alpha, GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mTransA(transA) - , mTransB(transB) - , mPadLda(padLda) - , mPadLdb(padLdb) - , mPadLdc(padLdc) - , mType(type) - , mOutputType(type) - , mUseFp8(useFp8) - , mAlpha(alpha) - , mPluginProfiler(pluginProfiler) -{ - init(); -} - -// Parameterized constructor -GemmPlugin::GemmPlugin(void const* data, size_t length, GemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mTransA); - read(d, mTransB); - read(d, mPadLda); - read(d, mPadLdb); - read(d, mPadLdc); - read(d, mType); - read(d, mUseFp8); - read(d, mAlpha); - read(d, mDims); - read(d, mOutputType); - - init(); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -thread_local CublasGemmWrapperPtr GemmPlugin::mCublasWrapper = nullptr; - -void GemmPlugin::init() -{ - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - mCublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - - mPluginProfiler->setTranspose(mTransA, mTransB); - mPluginProfiler->setOutputType(mOutputType); - mPluginProfiler->setPadLd(mPadLda, mPadLdb, mPadLdc); - - mGemmId = GemmIdCublas(mDims.n, mDims.k, mType, mTransA, mTransB, mOutputType); - - mArch = tensorrt_llm::common::getSMVersion(); -} - -void GemmPlugin::setGemmConfig() -{ - if (mType == nvinfer1::DataType::kHALF) - { - mCublasWrapper->setFP16GemmConfig(trtToCublasDtype(mOutputType)); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - mCublasWrapper->setFP32GemmConfig(); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - mCublasWrapper->setBF16GemmConfig(trtToCublasDtype(mOutputType)); - } -#endif - -#ifdef ENABLE_FP8 - if (mUseFp8) - { - mCublasWrapper->setFP8GemmConfig(trtToCublasDtype(mOutputType)); - } -#endif -} - -void GemmPlugin::configGemm() -{ - if (!mDims.isInitialized()) - { - return; - } - - setGemmConfig(); - - mPluginProfiler->profileTactics(mCublasWrapper, mType, mDims, mGemmId); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* GemmPlugin::clone() const noexcept -{ - auto* plugin = new GemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs GemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - int const nbDimsB = inputs[1].nbDims; - DimsExprs ret; - ret.nbDims = nbDimsA + nbDimsB - 2; - - if (mTransA) - { - for (int i = 1; i < nbDimsA; ++i) - { - ret.d[i - 1] = inputs[0].d[i]; - } - } - else - { - for (int i = 0; i < nbDimsA - 1; ++i) - { - ret.d[i] = inputs[0].d[i]; - } - } - if (mTransB) - { - for (int i = 0; i < nbDimsB - 1; ++i) - { - ret.d[nbDimsA - 1 + i] = exprBuilder.constant(inputs[1].d[i]->getConstantValue() + mPadLdc); - } - } - else - { - for (int i = 1; i < nbDimsB; ++i) - { - ret.d[nbDimsA - 2 + i] = exprBuilder.constant(inputs[1].d[i]->getConstantValue() + mPadLdc); - } - } - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool GemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - auto const& desc = inOut[pos]; - if (desc.format != TensorFormat::kLINEAR) - { - return false; - } - - if (pos < nbInputs) - { - // If use FP8, act/weight dtype should be kFP8 - if (mUseFp8) - { - return desc.type == nvinfer1::DataType::kFP8; - } - else - { - return desc.type == mType; - } - } - - return desc.type == mType || desc.type == nvinfer1::DataType::kFLOAT; -} - -void GemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const nbDimsA = in[0].max.nbDims; - - auto const minM = utils::computeMDimension(mTransA, in[0].min); - auto const maxM = utils::computeMDimension(mTransA, in[0].max); - auto const N = utils::computeNDimension(mTransB, in[1].max); - auto const K = static_cast(mTransA ? in[0].max.d[0] : in[0].max.d[nbDimsA - 1]); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - mGemmId.n = N; - mGemmId.k = K; - - mOutputType = out[0].desc.type; -} - -size_t GemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return CUBLAS_WORKSPACE_SIZE; -} - -int GemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M, K] (mTransA = False) - // mat2 [K, N] (mTransB = False) - // outputs - // mat [M, N] - if (mCublasWrapper == nullptr) - { - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - mCublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - } - setGemmConfig(); - - int const nbDimsA = inputDesc[0].dims.nbDims; - int const padM = mTransA ? mPadLda : 0; - int const padN = mTransB ? 0 : mPadLdb; - int const padK = mTransA ? 0 : mPadLda; - auto const M = utils::computeMDimension(mTransA, inputDesc[0].dims) - padM; - auto const N = utils::computeNDimension(mTransB, inputDesc[1].dims) - padN; - int const K = static_cast( - mTransA ? inputDesc[0].dims.d[0] - padK : inputDesc[0].dims.d[nbDimsA - 1] - padK); - - bool noPadDim = padM == 0 && padN == 0 && padK == 0 && mPadLdc == 0; - bool cudaKernelSupportType = mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kFLOAT - || mType == nvinfer1::DataType::kBF16; - - // skip computation for a TRT empty tensor - if (M == 0) - { - return 0; - } - - std::string mnkStr = "MNK={" + std::to_string(M) + ", " + std::to_string(N) + ", " + std::to_string(K) + "}"; - { - std::string const activationStr = "GEMM layer's activation before GEMM with " + mnkStr; - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(M, K, mType, inputs[0], stream, activationStr) == false, - "Found invalid number (NaN or Inf) in " + activationStr); - } - - bool cudaKernelFinished = false; - bool isArch90or100 = mArch >= 90 && mArch < 120; - // TODO: sub tensor matmul is not supported in fp8 gemm cuda kernel - if (!isArch90or100 && M <= 4 && N <= 128000 && mUseFp8 && noPadDim && cudaKernelSupportType) - { - tensorrt_llm::kernels::cuda_core_gemm::Params params(reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1]), mAlpha, reinterpret_cast(outputs[0]), M, N, K, - CUDA_R_8F_E4M3, trtToCublasDtype(mOutputType)); - cudaKernelFinished = tensorrt_llm::kernels::cuda_core_gemm::cudaCoreGemmDispatcher(params, stream); - } - else if (!isArch90or100 && ((mArch < 90 && M <= 6) || (isArch90or100 && M <= 2)) && N <= 128000 && !mUseFp8 - && noPadDim && cudaKernelSupportType) - { - tensorrt_llm::kernels::cuda_core_gemm::Params params(reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1]), mAlpha, reinterpret_cast(outputs[0]), M, N, K, - trtToCublasDtype(mType), trtToCublasDtype(mOutputType)); - cudaKernelFinished = tensorrt_llm::kernels::cuda_core_gemm::cudaCoreGemmDispatcher(params, stream); - } - - if (!cudaKernelFinished) - { - auto bestTactic = mPluginProfiler->getBestConfig(M, mGemmId); - runGemm(M, N, K, mTransA, mTransB, mPadLda, mPadLdb, mPadLdc, mType, mCublasWrapper, inputs[0], inputs[1], - mAlpha, outputs[0], bestTactic, workspace, stream); - } - - { - std::string const outputStr = "GEMM layer's output after GEMM with " + mnkStr; - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(M, N + mPadLdc, mType, outputs[0], stream, outputStr) - == false, - "Found invalid number (NaN or Inf) in " + outputStr); - } - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType GemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* GemmPlugin::getPluginType() const noexcept -{ - return GEMM_PLUGIN_NAME; -} - -char const* GemmPlugin::getPluginVersion() const noexcept -{ - return GEMM_PLUGIN_VERSION; -} - -int GemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int GemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void GemmPlugin::destroy() noexcept -{ - delete this; -} - -size_t GemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(mTransA) + sizeof(mTransB) + sizeof(mPadLda) + sizeof(mPadLdb) + sizeof(mPadLdc) + sizeof(mType) - + sizeof(mDims) + sizeof(mUseFp8) + sizeof(mAlpha) + mPluginProfiler->getSerializationSize(mGemmId) - + sizeof(mOutputType); // selected tactics container size -} - -void GemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mTransA); - write(d, mTransB); - write(d, mPadLda); - write(d, mPadLdb); - write(d, mPadLdc); - write(d, mType); - write(d, mUseFp8); - write(d, mAlpha); - write(d, mDims); - write(d, mOutputType); - mPluginProfiler->serialize(d, mGemmId); - - TLLM_CHECK(d == a + getSerializationSize()); -} - -void GemmPlugin::terminate() noexcept {} - -/////////////// - -GemmPluginCreator::GemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("transA", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("transB", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("padLda", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("padLdb", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("padLdc", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("use_fp8", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* GemmPluginCreator::getPluginName() const noexcept -{ - return GEMM_PLUGIN_NAME; -} - -char const* GemmPluginCreator::getPluginVersion() const noexcept -{ - return GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* GemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int transA{}; - int transB{}; - int padLda{}; - int padLdb{}; - int padLdc{}; - nvinfer1::DataType type{}; - int useFp8{}; - float alpha = 1.F; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "transa")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transA = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "transb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transB = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pad_lda")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - padLda = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pad_ldb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - padLdb = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pad_ldc")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - padLdc = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "use_fp8")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - useFp8 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "alpha")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - alpha = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - auto* obj = new GemmPlugin(transA, transB, padLda, padLdb, padLdc, type, useFp8, alpha, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* GemmPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GemmPlugin::destroy() - try - { - // GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true, /* skip */ true); - auto* obj = new GemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h b/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h deleted file mode 100644 index 1ba553c23d4b..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h +++ /dev/null @@ -1,169 +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. - */ -#ifndef TRT_GEMM_PLUGIN_H -#define TRT_GEMM_PLUGIN_H - -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include - -namespace tensorrt_llm::plugins -{ - -using CublasGemmWrapper = tensorrt_llm::common::CublasMMWrapper; -using CublasGemmWrapperPtr = std::shared_ptr; - -class CublasLtGemmPluginProfiler - : public GemmPluginProfiler -{ -public: - using Config = cublasLtMatmulHeuristicResult_t; - - void setTranspose(bool transposeA, bool transposeB) - { - mTransA = transposeA; - mTransB = transposeB; - } - - void setPadLd(int padLda, int padLdb, int padLdc) - { - mPadLda = padLda; - mPadLdb = padLdb; - mPadLdc = padLdc; - } - - void setOutputType(nvinfer1::DataType type) - { - mOutputType = type; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - bool checkTactic(int m, int n, int k, Config const& tactic) const override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - bool mTransA; - bool mTransB; - int mPadLda; - int mPadLdb; - int mPadLdc; - nvinfer1::DataType mOutputType; - - static constexpr size_t ALIGNMENT = 256; -}; - -class GemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - GemmPlugin() = delete; - - GemmPlugin(int transA, int transB, int padLda, int padLdb, int padLdc, nvinfer1::DataType type, bool useFp8, - float alpha, PluginProfilerPtr const& profiler); - - GemmPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~GemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(); - void configGemm(); - void setGemmConfig(); - -private: - const std::string mLayerName; - - int mTransA; - int mTransB; - int mPadLda; - int mPadLdb; - int mPadLdc; - int mArch; - nvinfer1::DataType mType; - nvinfer1::DataType mOutputType; - - static thread_local CublasGemmWrapperPtr mCublasWrapper; - - GemmDims mDims{}; - GemmIdCublas mGemmId{}; - bool mUseFp8{false}; - float mAlpha{1.f}; - - PluginProfilerPtr mPluginProfiler; -}; - -class GemmPluginCreator : public BaseCreator -{ -public: - GemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_GEMM_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/CMakeLists.txt deleted file mode 100644 index 3b714a3928fb..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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 *.cu) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cpp b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cpp deleted file mode 100644 index ed964ace695f..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cpp +++ /dev/null @@ -1,446 +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. - */ - -#include "gemmSwigluPlugin.h" -#include "cutlass_extensions/gemm_configs.h" - -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::GemmSwigluPluginCreator; -using tensorrt_llm::plugins::GemmSwigluPlugin; -using tensorrt_llm::plugins::GemmSwigluPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* GEMM_SWIGLU_PLUGIN_VERSION{"1"}; -static char const* GEMM_SWIGLU_PLUGIN_NAME{"GemmSwiglu"}; -PluginFieldCollection GemmSwigluPluginCreator::mFC{}; -std::vector GemmSwigluPluginCreator::mPluginAttributes; - -size_t GemmSwigluPluginProfiler::getBytePerElement(nvinfer1::DataType type) -{ - size_t bpe; - if (type == nvinfer1::DataType::kHALF || type == nvinfer1::DataType::kBF16) - { - bpe = 2; - } - else if (type == nvinfer1::DataType::kINT8 || type == nvinfer1::DataType::kFP8) - { - bpe = 1; - } - else - { - TLLM_THROW("Not recognized/implemented"); - } - return bpe; -} - -void GemmSwigluPluginProfiler::setQuantMode(tensorrt_llm::common::QuantMode const& quantMode) -{ - mQuantMode = quantMode; -} - -void GemmSwigluPluginProfiler::runTactic( - int m, int n, int k, GemmSwigluPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - size_t bpe = getBytePerElement(mType); - - // Workspace size required by gemm runner - // NB: this function will throw exception when selected tactic exceeds SMEM, which is then - // caught by gemmPluginProfiler and it will register this tactic as invalid - size_t wsSizeRunner = mRunner->getWorkspaceSize(m, n, k); - - // Workspace size required by profiling - size_t wsByteOffset = 0; - int8_t* wsBytePointer = reinterpret_cast(workspace); - void* aTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * k * bpe)); - void* bTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, n * k * bpe)); - void* cTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, 1 * n * bpe)); - void* dTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, m * (n / 2) * bpe)); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr(wsBytePointer, wsByteOffset, wsSizeRunner)); - - // Run profiling - mRunner->gemm( - dTmp, aTmp, bTmp, cTmp, mQuantMode, m, n, k, 1.0, 1.0, 1.0, tactic, workspaceTmp, wsSizeRunner, stream); -} - -int GemmSwigluPluginProfiler::getMaxProfileM() const -{ - return 32768; -} - -void GemmSwigluPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - std::vector workspaces = { - maxM * k * getBytePerElement(mType), // A - n * k * getBytePerElement(mType), // B - 1 * n * getBytePerElement(mType), // C_bias - maxM * (n / 2) * getBytePerElement(mType), // D - mRunner->getWorkspaceSize(maxM, n, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector GemmSwigluPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -GemmSwigluPlugin::GemmSwigluPlugin(QuantMode quantMode, nvinfer1::DataType type, bool hasBias, float scale_d0, - float scale_d1, float scale_output, GemmSwigluPlugin::PluginProfilerPtr const& pluginProfiler) - : mQuantMode(quantMode) - , mPluginProfiler(pluginProfiler) - , mHasBias(hasBias) - , mScaleD0(scale_d0) - , mScaleD1(scale_d1) - , mScaleOutput(scale_output) -{ - init(type); -} - -// Parameterized constructor -GemmSwigluPlugin::GemmSwigluPlugin( - void const* data, size_t length, GemmSwigluPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - unsigned int quantMode; - read(d, quantMode); - read(d, type); - read(d, mHasBias); - read(d, mScaleD0); - read(d, mScaleD1); - read(d, mScaleOutput); - read(d, mDims); - - mQuantMode = QuantMode(quantMode); - - init(type); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK(d == a + length); -} - -void GemmSwigluPlugin::init(nvinfer1::DataType type) -{ - mType = type; - if (mType == nvinfer1::DataType::kFP8) - { - mGemmRunner = std::make_shared>(); - } - else - { - TLLM_THROW("Gemm Swiglu plugin only supports fp8 now"); - } - - mPluginProfiler->setQuantMode(mQuantMode); - - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* GemmSwigluPlugin::clone() const noexcept -{ - auto* plugin = new GemmSwigluPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs GemmSwigluPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 3); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue() / 2); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool GemmSwigluPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // bias - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - case 3: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - TLLM_CHECK(false); - return false; - } -} - -void GemmSwigluPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[1]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[1]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - mWorkspaceMaxSize = mGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t GemmSwigluPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mWorkspaceMaxSize; -} - -int GemmSwigluPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M(*), K] - // mat2 [K, N] - // bias [1, N] - // outputs - // mat [M(*), N / 2] - int m = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m *= inputDesc[0].dims.d[ii]; - } - int const n = inputDesc[1].dims.d[1]; - int const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - size_t const wsSize = mGemmRunner->getWorkspaceSize(m, n, k); - - auto const bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid GEMM tactic"); - mGemmRunner->gemm(outputs[0], inputs[0], inputs[1], inputs[2], mQuantMode, m, n, k, mScaleD0, mScaleD1, - mScaleOutput, *bestTactic, reinterpret_cast(workspace), wsSize, stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType GemmSwigluPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* GemmSwigluPlugin::getPluginType() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* GemmSwigluPlugin::getPluginVersion() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_VERSION; -} - -int GemmSwigluPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int GemmSwigluPlugin::initialize() noexcept -{ - configGemm(); // gemm profiler in action - return 0; -} - -void GemmSwigluPlugin::terminate() noexcept {} - -size_t GemmSwigluPlugin::getSerializationSize() const noexcept -{ - return sizeof(unsigned int) + // QuantMode - sizeof(nvinfer1::DataType) + // dtype - sizeof(bool) + // hasBias - sizeof(float) * 3 + // scales - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void GemmSwigluPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mHasBias); - write(d, mScaleD0); - write(d, mScaleD1); - write(d, mScaleOutput); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void GemmSwigluPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void GemmSwigluPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mGemmRunner, mType, mDims, mGemmId); -} - -/////////////// - -GemmSwigluPluginCreator::GemmSwigluPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("has_bias", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("scale_d0", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_d1", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_output", nullptr, PluginFieldType::kFLOAT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* GemmSwigluPluginCreator::getPluginName() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* GemmSwigluPluginCreator::getPluginVersion() const noexcept -{ - return GEMM_SWIGLU_PLUGIN_VERSION; -} - -PluginFieldCollection const* GemmSwigluPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GemmSwigluPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 5); - nvinfer1::DataType type{}; - bool hasBias{}; - float scale_d0{}; - float scale_d1{}; - float scale_output{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "has_bias")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - hasBias = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d0")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d0 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d1")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d1 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_output")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_output = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // GemmSwigluPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - QuantMode quantMode = QuantMode{}; - auto* obj = new GemmSwigluPlugin(quantMode, type, hasBias, scale_d0, scale_d1, scale_output, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* GemmSwigluPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GemmSwigluPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = mGemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new GemmSwigluPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cu b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cu deleted file mode 100644 index 339c432b1113..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.cu +++ /dev/null @@ -1,41 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 "gemmSwigluPlugin.h" - -#include "cutlass/util/reference/device/tensor_fill.h" -#include "cutlass_extensions/gemm_configs.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::GemmSwigluPluginCreator; -using tensorrt_llm::plugins::GemmSwigluPlugin; -using tensorrt_llm::plugins::GemmSwigluPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -void GemmSwigluPluginProfiler::initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream) -{ - size_t bpe = getBytePerElement(mType); - - if (mType == nvinfer1::DataType::kFP8) - { - cutlass::reference::device::BlockFillRandomUniform(reinterpret_cast(workspace), - m * k + n * k + 1 * n, 42, cutlass::float_e4m3_t{128}, -cutlass::float_e4m3_t{128}, -1, 0, stream); - } -} diff --git a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h b/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h deleted file mode 100644 index 766e59aad258..000000000000 --- a/cpp/tensorrt_llm/plugins/gemmSwigluPlugin/gemmSwigluPlugin.h +++ /dev/null @@ -1,150 +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. - */ -#pragma once - -#include "tensorrt_llm/kernels/cutlass_kernels/fused_gated_gemm/fused_gated_gemm.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using GemmSwigluRunnerPtr - = std::shared_ptr; - -class GemmSwigluPluginProfiler : public GemmPluginProfiler - -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantMode(tensorrt_llm::common::QuantMode const& quantMode); - - virtual int getMaxProfileM() const override; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - // TODO(anchengc) implement checkTactic - // bool checkTactic(int m, int n, int k, const Config& tactic) const override; - - std::vector getTactics(int m, int n, int k) const override; - - void initTmpData(int m, int n, int k, char* workspace, size_t size, cudaStream_t stream) override; - -private: - size_t getBytePerElement(nvinfer1::DataType type); - - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class GemmSwigluPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - GemmSwigluPlugin() = delete; - - GemmSwigluPlugin(tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, bool hasBias, float scale_d0, - float scale_d1, float scale_output, PluginProfilerPtr const& pluginProfiler); - - GemmSwigluPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~GemmSwigluPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - - void configGemm(); - // void setGemmConfig(); - -private: - const std::string mLayerName; - - GemmSwigluRunnerPtr mGemmRunner; - tensorrt_llm::common::QuantMode mQuantMode; // not configurable yet - size_t mWorkspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; - bool mHasBias; - float mScaleD0; - float mScaleD1; - float mScaleOutput; -}; - -class GemmSwigluPluginCreator : public BaseCreator -{ -public: - GemmSwigluPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager mGemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gptAttentionCommon/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp deleted file mode 100644 index 717ab3083e5f..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.cpp +++ /dev/null @@ -1,380 +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. - */ -#include "gptAttentionCommon.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -namespace tc = tensorrt_llm::common; -using tensorrt_llm::plugins::GPTAttentionPluginCreatorCommon; -using tensorrt_llm::plugins::GPTAttentionPluginCommon; - -GPTAttentionPluginCommon::GPTAttentionPluginCommon(int layer_idx, int num_heads, int vision_start, int vision_length, - int num_kv_heads, int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, - float attn_logit_softcapping_scale, tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. Use 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, float rotary_embedding_long_m_scale, - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, tensorrt_llm::kernels::BlockSparseParams block_sparse_params, - bool paged_kv_cache, int tokens_per_block, nvinfer1::DataType type, int32_t max_context_length, - bool qkv_bias_enabled, bool cross_attention, int max_distance, bool pos_shift_enabled, bool dense_context_fmha, - bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, - bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, - int32_t spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool fuse_fp4_quant, bool skip_attn, int cp_size, - int cp_rank, std::set cp_group) - : mResource{DecoderXQARunner::getResourceGlobal()} -{ - mLayerIdx = layer_idx; - mNumHeads = num_heads; - mVisionStart = vision_start; - mVisionLength = vision_length; - mNumKVHeads = num_kv_heads; - mNumKVHeadsOrigin = num_kv_heads_origin; - mHeadSize = head_size; - mUnidirectional = unidirectional; - mQScaling = q_scaling; - mAttnLogitSoftcappingScale = attn_logit_softcapping_scale; - mRotaryEmbeddingDim = rotary_embedding_dim; - mRotaryEmbeddingBase = rotary_embedding_base; - mRotaryEmbeddingScaleType = rotary_embedding_scale_type; - mRotaryEmbeddingScale = rotary_embedding_scale; - mRotaryEmbeddingShortMscale = rotary_embedding_short_m_scale; - mRotaryEmbeddingLongMscale = rotary_embedding_long_m_scale; - mRotaryEmbeddingMaxPositions = rotary_embedding_max_positions; - mRotaryEmbeddingOriginalMaxPositions = rotary_embedding_original_max_positions; - mPositionEmbeddingType = position_embedding_type; - mEnableContextFMHA = context_fmha_type != ContextFMHAType::DISABLED; - mFMHAForceFP32Acc = type == nvinfer1::DataType::kBF16; - mMaskType = mask_type; - mBlockSparseParams = block_sparse_params; - mType = type; - mMultiBlockMode = true; - mEnableXQA = true; - mKVCacheQuantMode = tc::QuantMode(kv_cache_quant_mode); - mRemovePadding = remove_input_padding; - mPagedKVCache = paged_kv_cache; - mTokensPerBlock = tokens_per_block; - mTpSize = tp_size; - mTpRank = tp_rank; - mUnfuseQkvGemm = unfuse_qkv_gemm; - mUseLognScaling = use_logn_scaling; - mMaxContextLength = max_context_length; - mQKVBiasEnabled = qkv_bias_enabled; - mCrossAttention = cross_attention; - mMaxDistance = max_distance; - mPosShiftEnabled = pos_shift_enabled; - mDenseContextFMHA = dense_context_fmha; - mPagedContextFMHA = use_paged_context_fmha; - mFP8ContextFMHA = use_fp8_context_fmha; - mFP8AttenOutput = use_fp8_context_fmha; - mHasFullAttentionMask = has_full_attention_mask; - mUseKVCache = use_cache; - mIsSpecDecodingEnabled = is_spec_decoding_enabled; - mSpecDecodingIsGenerationLengthVariable = spec_decoding_is_generation_length_variable; - mSpecDecodingMaxGenerationLength = spec_decoding_max_generation_length; - mIsMLAEnabled = is_mla_enabled; - mMLAParams = {q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim}; - mCpSize = cp_size; - mCpRank = cp_rank; - mCpGroup = std::move(cp_group); - mFuseFp4Quant = fuse_fp4_quant; - mSkipAttn = skip_attn; -} - -// Parameterized constructor -GPTAttentionPluginCommon::GPTAttentionPluginCommon(void const* data, size_t length) - : mResource{DecoderXQARunner::getResourceGlobal()} -{ - char const *d = reinterpret_cast(data), *a = d; - unsigned int kvCacheQuantMode; - - read(d, mLayerIdx); - read(d, mNumHeads); - read(d, mVisionStart); - read(d, mVisionLength); - read(d, mNumKVHeads); - read(d, mNumKVHeadsOrigin); - read(d, mHeadSize); - read(d, mUnidirectional); - read(d, mQScaling); - read(d, mAttnLogitSoftcappingScale); - read(d, mPositionEmbeddingType); - read(d, mRotaryEmbeddingDim); - read(d, mRotaryEmbeddingBase); - read(d, mRotaryEmbeddingScaleType); - read(d, mRotaryEmbeddingScale); - read(d, mRotaryEmbeddingShortMscale); - read(d, mRotaryEmbeddingLongMscale); - read(d, mRotaryEmbeddingMaxPositions); - read(d, mRotaryEmbeddingOriginalMaxPositions); - read(d, mTpSize); - read(d, mTpRank); - read(d, mUnfuseQkvGemm); - read(d, mUseLognScaling); - read(d, mEnableContextFMHA); - read(d, mFMHAForceFP32Acc); - read(d, mMultiBlockMode); - read(d, mEnableXQA); - read(d, kvCacheQuantMode); - read(d, mRemovePadding); - read(d, mMaskType); - read(d, mBlockSparseParams); - read(d, mPagedKVCache); - read(d, mTokensPerBlock); - read(d, mType); - read(d, mMaxContextLength); - read(d, mQKVBiasEnabled); - read(d, mCrossAttention); - read(d, mMaxDistance); - read(d, mPosShiftEnabled); - read(d, mDenseContextFMHA); - read(d, mPagedContextFMHA); - read(d, mFP8ContextFMHA); - read(d, mFP8AttenOutput); - read(d, mHasFullAttentionMask); - read(d, mUseKVCache); - read(d, mIsSpecDecodingEnabled); - read(d, mUseSpecDecoding); - read(d, mSpecDecodingIsGenerationLengthVariable); - read(d, mSpecDecodingMaxGenerationLength); - read(d, mIsMLAEnabled); - read(d, mMLAParams); - read(d, mNbMultiBlockSemaphores); - read(d, mFuseFp4Quant); - read(d, mSkipAttn); - read(d, mCpSize); - read(d, mCpRank); - - mKVCacheQuantMode = tc::QuantMode(kvCacheQuantMode); - - uint32_t decoderXQARunnerResourceSerializedSize; - read(d, decoderXQARunnerResourceSerializedSize); - mResource->merge(DecoderXQARunnerResource(d, decoderXQARunnerResourceSerializedSize), /*initialize=*/true); - d += decoderXQARunnerResourceSerializedSize; - - mCpGroup.clear(); - int32_t groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mCpGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); - TLLM_CHECK_WITH_INFO((smVersion() >= 80) || (mType != nvinfer1::DataType::kBF16), - "Unsupported data type, pre SM 80 GPUs do not support bfloat16"); -} - -int GPTAttentionPluginCommon::initialize() noexcept -{ - return AttentionOp::initialize(); -} - -void GPTAttentionPluginCommon::destroy() noexcept -{ - delete this; -} - -size_t GPTAttentionPluginCommon::getCommonSerializationSize() const noexcept -{ - return sizeof(mLayerIdx) + sizeof(mNumHeads) + +sizeof(mVisionStart) + sizeof(mVisionLength) + sizeof(mNumKVHeads) - + sizeof(mNumKVHeadsOrigin) + sizeof(mHeadSize) + sizeof(mUnidirectional) + sizeof(mQScaling) - + sizeof(mAttnLogitSoftcappingScale) + sizeof(mPositionEmbeddingType) + sizeof(mRotaryEmbeddingDim) - + sizeof(mRotaryEmbeddingBase) + sizeof(mRotaryEmbeddingScaleType) + sizeof(mRotaryEmbeddingScale) - + sizeof(mRotaryEmbeddingShortMscale) + sizeof(mRotaryEmbeddingLongMscale) - + sizeof(mRotaryEmbeddingMaxPositions) + sizeof(mRotaryEmbeddingOriginalMaxPositions) + sizeof(mTpSize) - + sizeof(mTpRank) + sizeof(mEnableContextFMHA) + sizeof(mFMHAForceFP32Acc) + sizeof(mMultiBlockMode) - + sizeof(mEnableXQA) + sizeof(unsigned int) // mKVCacheQuantMode - + sizeof(mRemovePadding) + sizeof(mMaskType) + sizeof(mBlockSparseParams) + sizeof(mPagedKVCache) - + sizeof(mTokensPerBlock) + sizeof(mType) + sizeof(mMaxContextLength) + sizeof(mQKVBiasEnabled) - + sizeof(mCrossAttention) + sizeof(mMaxDistance) + sizeof(mPosShiftEnabled) + sizeof(mDenseContextFMHA) - + sizeof(mPagedContextFMHA) + sizeof(mFP8ContextFMHA) + sizeof(mFP8AttenOutput) + sizeof(mHasFullAttentionMask) - + sizeof(mUseKVCache) + sizeof(mUnfuseQkvGemm) + sizeof(mUseLognScaling) + sizeof(mIsSpecDecodingEnabled) - + sizeof(mUseSpecDecoding) + sizeof(mSpecDecodingIsGenerationLengthVariable) - + sizeof(mSpecDecodingMaxGenerationLength) + sizeof(mNbMultiBlockSemaphores) + sizeof(mIsMLAEnabled) - + sizeof(mMLAParams) + sizeof(mFuseFp4Quant) + sizeof(mSkipAttn) - + sizeof(uint32_t) // size of DecoderXQARunnerResource buffer. - + sizeof(mCpSize) + sizeof(mCpRank) + sizeof(int32_t) * mCpGroup.size() + mResource->getSerializationSize(); -} - -void GPTAttentionPluginCommon::serializeCommon(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mLayerIdx); - write(d, mNumHeads); - write(d, mVisionStart); - write(d, mVisionLength); - write(d, mNumKVHeads); - write(d, mNumKVHeadsOrigin); - write(d, mHeadSize); - write(d, mUnidirectional); - write(d, mQScaling); - write(d, mAttnLogitSoftcappingScale); - write(d, mPositionEmbeddingType); - write(d, mRotaryEmbeddingDim); - write(d, mRotaryEmbeddingBase); - write(d, mRotaryEmbeddingScaleType); - write(d, mRotaryEmbeddingScale); - write(d, mRotaryEmbeddingShortMscale); - write(d, mRotaryEmbeddingLongMscale); - write(d, mRotaryEmbeddingMaxPositions); - write(d, mRotaryEmbeddingOriginalMaxPositions); - write(d, mTpSize); - write(d, mTpRank); - write(d, mUnfuseQkvGemm); - write(d, mUseLognScaling); - write(d, mEnableContextFMHA); - write(d, mFMHAForceFP32Acc); - write(d, mMultiBlockMode); - write(d, mEnableXQA); - write(d, mKVCacheQuantMode.value()); - write(d, mRemovePadding); - write(d, mMaskType); - write(d, mBlockSparseParams); - write(d, mPagedKVCache); - write(d, mTokensPerBlock); - write(d, mType); - write(d, mMaxContextLength); - write(d, mQKVBiasEnabled); - write(d, mCrossAttention); - write(d, mMaxDistance); - write(d, mPosShiftEnabled); - write(d, mDenseContextFMHA); - write(d, mPagedContextFMHA); - write(d, mFP8ContextFMHA); - write(d, mFP8AttenOutput); - write(d, mHasFullAttentionMask); - write(d, mUseKVCache); - write(d, mIsSpecDecodingEnabled); - write(d, mUseSpecDecoding); - write(d, mSpecDecodingIsGenerationLengthVariable); - write(d, mSpecDecodingMaxGenerationLength); - write(d, mIsMLAEnabled); - write(d, mMLAParams); - write(d, mNbMultiBlockSemaphores); - write(d, mFuseFp4Quant); - write(d, mSkipAttn); - write(d, mCpSize); - write(d, mCpRank); - - // An uint32_t that specifies the size of the serialized buffer, followed by the actual content. - uint32_t decoderXQARunnerResourceSerializedSize = mResource->getSerializationSize(); - write(d, decoderXQARunnerResourceSerializedSize); - mResource->serialize(d, decoderXQARunnerResourceSerializedSize); - d += decoderXQARunnerResourceSerializedSize; - - for (auto it = mCpGroup.begin(); it != mCpGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getCommonSerializationSize()); -} - -void GPTAttentionPluginCommon::terminate() noexcept -{ - // Do nothing, destroy will always be called, so release the resources there. -} - -/////////////// - -GPTAttentionPluginCreatorCommon::GPTAttentionPluginCreatorCommon() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("layer_idx", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("vision_start", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("vision_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_kv_heads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_kv_heads_origin", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("layer_idx_in_cache_pool", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("head_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("unidirectional", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("q_scaling", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("attn_logit_softcapping_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("position_embedding_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_base", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_scale_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_short_m_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_long_m_scale", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("rotary_embedding_max_positions", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back( - PluginField("rotary_embedding_original_max_positions", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("tp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("tp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("unfuse_qkv_gemm", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_logn_scaling", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("context_fmha_type", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("kv_cache_quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("mask_type", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_homo_head_pattern", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_num_local_blocks", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_sparse_vertical_stride", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("paged_kv_cache", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("tokens_per_block", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("max_context_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("qkv_bias_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("do_cross_attention", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("max_distance", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("pos_shift_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("dense_context_fmha", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_paged_context_fmha", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_fp8_context_fmha", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("has_full_attention_mask", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("use_cache", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("is_spec_decoding_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back( - PluginField("spec_decoding_is_generation_length_variable", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back( - PluginField("spec_decoding_max_generation_length", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("is_mla_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("q_lora_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("kv_lora_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("qk_nope_head_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("qk_rope_head_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("v_head_dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("fuse_fp4_quant", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("skip_attn", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("cp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("cp_group", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -PluginFieldCollection const* GPTAttentionPluginCreatorCommon::getFieldNames() noexcept -{ - return &mFC; -} diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h deleted file mode 100644 index dd87d67aab9e..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h +++ /dev/null @@ -1,112 +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. - */ -#pragma once - -#include "tensorrt_llm/common/attentionOp.h" -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::kernels -{ -class DecoderXQARunnerResource; -} - -namespace tensorrt_llm::plugins -{ - -class GPTAttentionPluginCommon : public BasePlugin, public tensorrt_llm::common::op::AttentionOp -{ -public: - GPTAttentionPluginCommon() = delete; - - GPTAttentionPluginCommon(int layer_idx, int num_heads, int vision_start, int vision_length, int num_kv_heads, - int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, float attn_logit_softcapping_scale, - tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. Use 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, float rotary_embedding_long_m_scale, - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, - tensorrt_llm::kernels::BlockSparseParams block_sparse_params, bool paged_kv_cache, int tokens_per_block, - nvinfer1::DataType type, int32_t max_context_length, bool qkv_bias_enabled, bool cross_attention = false, - int max_distance = 0, bool pos_shift_enabled = false, bool dense_context_fmha = false, - bool use_paged_context_fmha = true, bool use_fp8_context_fmha = true, bool has_full_attention_mask = false, - bool use_cache = true, bool is_spec_decoding_enabled = false, - bool spec_decoding_is_generation_length_variable = false, int32_t spec_decoding_max_generation_length = 1, - bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool fuse_fp4_quant = false, bool skip_attn = false, - int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); - - GPTAttentionPluginCommon(void const* data, size_t length); - - ~GPTAttentionPluginCommon() override = default; - - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - //! This is called on every trt Engine creation - int initialize() noexcept override; - //! This is called on every trt Engine destroy - void terminate() noexcept override; - - //! This is called on every trt ExecutionContext creation by TRT - //! Note TRT does not call the initialize on cloned plugin, so clone internally should do initialization. - template - T* cloneImpl() const noexcept; - - //! This is called on evert trt Engine or ExecutionContext destroy. - //! None-cloned plugins will call terminate and then call destroy, while the cloned plugins will call destroy only - //! So plugin should put the resource release inside destroy. - void destroy() noexcept override; - - size_t getCommonSerializationSize() const noexcept; - void serializeCommon(void* buffer) const noexcept; - -protected: - std::string const mLayerName; - -private: - std::shared_ptr mResource; -}; - -class GPTAttentionPluginCreatorCommon : public BaseCreator -{ -public: - GPTAttentionPluginCreatorCommon(); - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - template - T* deserializePluginImpl(char const* name, void const* serialData, size_t serialLength) noexcept; - -protected: - std::vector mPluginAttributes; - nvinfer1::PluginFieldCollection mFC{}; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h b/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h deleted file mode 100644 index 51462cee6f40..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h +++ /dev/null @@ -1,54 +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. - */ - -#pragma once - -#include "gptAttentionCommon.h" - -namespace tensorrt_llm::plugins -{ -template -T* GPTAttentionPluginCommon::cloneImpl() const noexcept -{ - static_assert(std::is_base_of_v); - auto* plugin = new T(static_cast(*this)); - plugin->setPluginNamespace(mNamespace.c_str()); - - // Cloned plugins should be in initialized state with correct resources ready to be enqueued. - plugin->initialize(); - return plugin; -} - -template -T* GPTAttentionPluginCreatorCommon::deserializePluginImpl( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GPTAttentionPluginCommon::destroy() - try - { - auto* obj = new T(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp deleted file mode 100644 index 6f8c41c94131..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.cpp +++ /dev/null @@ -1,1387 +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. - */ - -#include "gptAttentionPlugin.h" - -#include "tensorrt_llm/batch_manager/contextProgress.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/kernels/decoderMaskedMultiheadAttention.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/kernels/unfusedAttentionKernels.h" -#include "tensorrt_llm/plugins/common/checkMacrosPlugin.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h" -#include "tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommonImpl.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" - -#include -#include -#include -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::GPTAttentionPluginCreator; -using tensorrt_llm::plugins::GPTAttentionPlugin; - -static char const* GPT_ATTENTION_PLUGIN_VERSION{"1"}; -static char const* GPT_ATTENTION_PLUGIN_NAME{"GPTAttention"}; - -GPTAttentionPlugin::GPTAttentionPlugin(int layer_idx, int num_heads, int vision_start, int vision_length, - int num_kv_heads, int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, - float attn_logit_softcapping_scale, tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, - float rotary_embedding_long_m_scale, // magnitude scaling factors for Phi-3 long RoPE - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, tensorrt_llm::kernels::BlockSparseParams block_sparse_params, - bool paged_kv_cache, int tokens_per_block, nvinfer1::DataType type, int32_t max_context_length, - bool qkv_bias_enabled, bool cross_attention, int max_distance, bool pos_shift_enabled, bool dense_context_fmha, - bool use_paged_context_fmha, bool use_fp8_context_fmha, bool has_full_attention_mask, bool use_cache, - bool is_spec_decoding_enabled, bool spec_decoding_is_generation_length_variable, - int spec_decoding_max_generation_length, bool is_mla_enabled, int q_lora_rank, int kv_lora_rank, - int qk_nope_head_dim, int qk_rope_head_dim, int v_head_dim, bool fuse_fp4_quant, bool skip_attn, int cp_size, - int cp_rank, std::set cp_group) - : GPTAttentionPluginCommon(layer_idx, num_heads, vision_start, vision_length, num_kv_heads, num_kv_heads_origin, - head_size, unidirectional, q_scaling, attn_logit_softcapping_scale, position_embedding_type, - rotary_embedding_dim, rotary_embedding_base, rotary_embedding_scale_type, rotary_embedding_scale, - rotary_embedding_short_m_scale, rotary_embedding_long_m_scale, rotary_embedding_max_positions, - rotary_embedding_original_max_positions, tp_size, tp_rank, unfuse_qkv_gemm, use_logn_scaling, context_fmha_type, - kv_cache_quant_mode, remove_input_padding, mask_type, block_sparse_params, paged_kv_cache, tokens_per_block, - type, max_context_length, qkv_bias_enabled, cross_attention, max_distance, pos_shift_enabled, - dense_context_fmha, use_paged_context_fmha, use_fp8_context_fmha, has_full_attention_mask, use_cache, - is_spec_decoding_enabled, spec_decoding_is_generation_length_variable, spec_decoding_max_generation_length, - is_mla_enabled, q_lora_rank, kv_lora_rank, qk_nope_head_dim, qk_rope_head_dim, v_head_dim, fuse_fp4_quant, - skip_attn, cp_size, cp_rank, cp_group) -{ - TLLM_CHECK_WITH_INFO( - !is_mla_enabled, "GPTAttentionPlugin no longer supports MLA. Please use the PyTorch workflow instead."); - initEntryIdx(); -} - -GPTAttentionPlugin::GPTAttentionPlugin(void const* data, size_t length) - : GPTAttentionPluginCommon(data, length) -{ - initEntryIdx(); -} - -std::string GPTAttentionPlugin::toString(IdxEntry const& entry) const -{ -#define TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(name) \ - case IdxEntry::name: return #name - - switch (entry) - { - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(QKV_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(K_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(V_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_MASK); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_PACKED_MASK); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SEQUENCE_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_PAST_KEY_VALUE_LENGTHS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_MAX_ATTENTION_WINDOW); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_SINK_TOKEN_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CONTEXT_LENGTHS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CACHE_INDIR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(REQUEST_TYPES); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(KV_CACHE_BLOCK_OFFSETS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_KV_CACHE_BLOCK_OFFSETS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_KV_CACHE_POOL_POINTERS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_KV_CACHE_POOL_MAPPING); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(PAST_KEY_VALUE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(KV_CACHE_QUANTIZATION_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(KV_CACHE_DEQUANTIZATION_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_OUTPUT_QUANTIZATION_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ATTENTION_OUTPUT_SF_SCALE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ROTARY_INV_FREQ); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ROTARY_COS_SIN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ALIBI_SLOPES); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(RELATIVE_ATTENTION_BIAS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CROSS_KV); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(CROSS_KV_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ENCODER_INPUT_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_CONTEXT_LENGTH); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(QKV_BIAS_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_GENERATION_LENGTHS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_PACKED_MASK); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_POSITION_OFFSETS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SPEC_DECODING_USE); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(LONG_ROPE_ROTARY_INV_FREQ); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(LONG_ROPE_ROTARY_COS_SIN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MROPE_ROTARY_COS_SIN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MROPE_POSITION_DELTAS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_RUNTIME_PERF_KNOBS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(HOST_CONTEXT_PROGRESS); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MLA_Q_B_PROJ_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MLA_KV_B_PROJ_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(MLA_K_B_PROJ_TRANS_TENSOR); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(SKIP_ATTN); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(LOGN_SCALING); - TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING(ENUM_SIZE); - } -#undef TLLM_GPT_ATTN_IDX_ENTRY_TO_STRING - TLLM_LOG_TRACE(common::fmtstr("Missing string description for IdxEntry enum %lu.\n", static_cast(entry))); - return ""; -} - -bool GPTAttentionPlugin::isEntryUsed(IdxEntry const& entry) const -{ - switch (entry) - { - case IdxEntry::QKV_TENSOR: return true; - case IdxEntry::K_TENSOR: return mUnfuseQkvGemm; - case IdxEntry::V_TENSOR: return mUnfuseQkvGemm; - case IdxEntry::ATTENTION_MASK: return useFullCustomMask(); - case IdxEntry::ATTENTION_PACKED_MASK: return useCustomMask(); - case IdxEntry::SEQUENCE_LENGTH: return useKVCache(); - case IdxEntry::HOST_PAST_KEY_VALUE_LENGTHS: return useKVCache(); - case IdxEntry::HOST_MAX_ATTENTION_WINDOW: return true; - case IdxEntry::HOST_SINK_TOKEN_LENGTH: return true; - case IdxEntry::CONTEXT_LENGTHS: return true; - case IdxEntry::CACHE_INDIR: return useKVCache(); - case IdxEntry::REQUEST_TYPES: return true; - case IdxEntry::KV_CACHE_BLOCK_OFFSETS: return useKVCache() && mPagedKVCache; - case IdxEntry::HOST_KV_CACHE_BLOCK_OFFSETS: return useKVCache() && mPagedKVCache; - case IdxEntry::HOST_KV_CACHE_POOL_POINTERS: return useKVCache() && mPagedKVCache; - case IdxEntry::HOST_KV_CACHE_POOL_MAPPING: return useKVCache() && mPagedKVCache; - case IdxEntry::PAST_KEY_VALUE: return useKVCache() && !mPagedKVCache; - case IdxEntry::KV_CACHE_QUANTIZATION_SCALE: return useKVCache() && mKVCacheQuantMode.hasKvCacheQuant(); - case IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE: return useKVCache() && mKVCacheQuantMode.hasKvCacheQuant(); - case IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE: return mFP8ContextFMHA; - case IdxEntry::ATTENTION_OUTPUT_SF_SCALE: return mFuseFp4Quant; - case IdxEntry::ROTARY_INV_FREQ: return isRoPE(); - case IdxEntry::ROTARY_COS_SIN: return isRoPE(); - case IdxEntry::ALIBI_SLOPES: return isALiBi(); - case IdxEntry::RELATIVE_ATTENTION_BIAS: return isRelativePosition(); - case IdxEntry::CROSS_KV: return isCrossAttention(); - case IdxEntry::CROSS_KV_LENGTH: return isCrossAttention(); - case IdxEntry::LOGN_SCALING: return isLognScaling(); - case IdxEntry::ENCODER_INPUT_LENGTH: return isCrossAttention(); - case IdxEntry::HOST_CONTEXT_LENGTH: return mRemovePadding; - case IdxEntry::QKV_BIAS_TENSOR: return mQKVBiasEnabled; - case IdxEntry::SPEC_DECODING_GENERATION_LENGTHS: return mIsSpecDecodingEnabled; - case IdxEntry::SPEC_DECODING_PACKED_MASK: return mIsSpecDecodingEnabled; - case IdxEntry::SPEC_DECODING_POSITION_OFFSETS: return mIsSpecDecodingEnabled; - case IdxEntry::SPEC_DECODING_USE: return mIsSpecDecodingEnabled; - case IdxEntry::LONG_ROPE_ROTARY_INV_FREQ: return isLongRoPE(); - case IdxEntry::LONG_ROPE_ROTARY_COS_SIN: return isLongRoPE(); - case IdxEntry::MROPE_ROTARY_COS_SIN: return isMRoPE(); - case IdxEntry::MROPE_POSITION_DELTAS: return isMRoPE(); - case IdxEntry::HOST_RUNTIME_PERF_KNOBS: return true; - case IdxEntry::HOST_CONTEXT_PROGRESS: return true; - case IdxEntry::MLA_Q_B_PROJ_TENSOR: return mIsMLAEnabled; - case IdxEntry::MLA_KV_B_PROJ_TENSOR: return mIsMLAEnabled; - case IdxEntry::MLA_K_B_PROJ_TRANS_TENSOR: return mIsMLAEnabled; - case IdxEntry::SKIP_ATTN: return mSkipAttn; - default: return false; - } -} - -void GPTAttentionPlugin::initEntryIdx() -{ - mEntryIdx.resize(static_cast(IdxEntry::ENUM_SIZE)); - size_t entryIdx = 0; - for (size_t i = 0; i < static_cast(IdxEntry::ENUM_SIZE); i++) - { - mEntryIdx[i] = entryIdx; - entryIdx += isEntryUsed(static_cast(i)); - } -} - -GPTAttentionPlugin::IndexType GPTAttentionPlugin::getIdx(IdxEntry const& entry) const -{ - TLLM_CHECK_WITH_INFO( - isEntryUsed(entry), common::fmtstr("getIdx() should not be used with entry %s.\n", toString(entry).data())); - return mEntryIdx[static_cast(entry)]; -} - -// IPluginV2DynamicExt Methods -GPTAttentionPlugin* GPTAttentionPlugin::clone() const noexcept -{ - return dynamic_cast(this->cloneImpl()); -} - -static int getPackedTensorHiddenDimIndex(bool removePadding) -{ - return removePadding ? 1 : 2; -} - -// NOTE: generation input length might be larger than one in the spec decoding mode. -int GPTAttentionPlugin::getGenerationInputSequenceLength( - nvinfer1::PluginTensorDesc const* inputDesc, int32_t localNbSeq, int32_t localNbTokens) const -{ - if (mRemovePadding) - { - // Speculative decoding mode might need variable generation input sequence length. - if (mIsSpecDecodingEnabled && mUseSpecDecoding) - { - TLLM_CHECK_WITH_INFO(mCpSize <= 1, "Context Parallel does not support speculative decoding mode for now"); - // SPEC_DECODING_POSITION_OFFSETS: [batch_size, max_generation_input_length]. - return inputDesc[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)].dims.d[1]; - } - else - { - if (mCpSize > 1) - { - // Given that localNbTokens == (beamSize * localNbSeq + mCpSize - 1) / mCpSize, but when mCpSize - 1 > - // localNbSeq, there are multiple choices for beamSize. Assume beamSize == 1 here. - TLLM_CHECK_WITH_INFO(localNbTokens == (localNbSeq + mCpSize - 1) / mCpSize, - "Context Parallel does not support beamSize > 1 for non-speculative decoding mode, " - "localNbTokens=%d, localNbSeq=%d", - localNbTokens, localNbSeq); - return 1; - } - // [num_tokens, local_hidden_size] where num_tokens = batch_size * generation_input_length - TLLM_CHECK_WITH_INFO(localNbTokens % localNbSeq == 0, - "seq_len should be same for all generation requests, localNbTokens=%d, localNbSeq=%d", localNbTokens, - localNbSeq); - return localNbTokens / localNbSeq; - } - } - else - { - // We don't have IFB without mRemovePadding, so just take it out from inputDesc - // [batch_size, seq_len, local_hidden_size] - return inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]; - } -} - -// outputs -// output_tensor [batch_size, seq_len, local_hidden_size] or [num_tokens, local_hidden_size] -// present_key_value_pool (optional if mPagedKVCache is false) [batch_size, 2, local_num_kv_heads, max_seq_len, -// head_size] -nvinfer1::DimsExprs GPTAttentionPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (mFuseFp4Quant) - { - TLLM_CHECK(outputIndex == 0 || outputIndex == 1 || (!mPagedKVCache && useKVCache() && outputIndex == 2)); - // Compute the output dimension for FP4 quantized tensor. Consistent with QuantizeToFP4Plugin. - if (outputIndex == 0) - { - auto ret = inputs[getIdx(IdxEntry::QKV_TENSOR)]; - return ret; - } - // Compute the output dimension for output scaling factor tensor. Consistent with QuantizeToFP4Plugin. - if (outputIndex == 1) - { - auto ret = inputs[getIdx(IdxEntry::QKV_TENSOR)]; - // Sequence dimension or token dimension. - // Pad to multiple of 128. - auto dimM = exprBuilder.operation(DimensionOperation::kCEIL_DIV, - *ret.d[getPackedTensorHiddenDimIndex(mRemovePadding) - 1], *exprBuilder.constant(128)); - ret.d[getPackedTensorHiddenDimIndex(mRemovePadding) - 1] - = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - // Hidden size dimension. - // Div (rounding up) by 16 since 16 elements share one SF and SF padded to k%4==0. - ret.d[getPackedTensorHiddenDimIndex(mRemovePadding)] = exprBuilder.operation(DimensionOperation::kCEIL_DIV, - *ret.d[getPackedTensorHiddenDimIndex(mRemovePadding)], *exprBuilder.constant(16)); - return ret; - } - } - else - { - TLLM_CHECK(outputIndex == 0 || (!mPagedKVCache && useKVCache() && outputIndex == 1)); - if (outputIndex == 0) - { - auto ret = inputs[getIdx(IdxEntry::QKV_TENSOR)]; - // In MLA, the output dim is v_head_dim - auto const head_size = mHeadSize; - ret.d[getPackedTensorHiddenDimIndex(mRemovePadding)] = exprBuilder.operation( - DimensionOperation::kPROD, *exprBuilder.constant(head_size), *exprBuilder.constant(mNumHeads)); - return ret; - } - } - return inputs[getIdx(IdxEntry::PAST_KEY_VALUE)]; -} - -bool GPTAttentionPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - bool result = false; - int posCaseLine = -1; - if (pos == getIdx(IdxEntry::CONTEXT_LENGTHS) || pos == getIdx(IdxEntry::REQUEST_TYPES) - || pos == getIdx(IdxEntry::HOST_MAX_ATTENTION_WINDOW) || pos == getIdx(IdxEntry::HOST_SINK_TOKEN_LENGTH) - || (isEntryUsed(IdxEntry::SPEC_DECODING_PACKED_MASK) && pos == getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)) - || (isEntryUsed(IdxEntry::SPEC_DECODING_POSITION_OFFSETS) - && pos == getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)) - || (isEntryUsed(IdxEntry::SPEC_DECODING_GENERATION_LENGTHS) - && pos == getIdx(IdxEntry::SPEC_DECODING_GENERATION_LENGTHS)) - || (isEntryUsed(IdxEntry::SPEC_DECODING_USE) && pos == getIdx(IdxEntry::SPEC_DECODING_USE))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (isMRoPE() && (pos == getIdx(IdxEntry::MROPE_ROTARY_COS_SIN))) - { - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (isMRoPE() && (pos == getIdx(IdxEntry::MROPE_POSITION_DELTAS))) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos == getIdx(IdxEntry::HOST_RUNTIME_PERF_KNOBS) || pos == getIdx(IdxEntry::HOST_CONTEXT_PROGRESS)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (useKVCache() - && (pos == getIdx(IdxEntry::SEQUENCE_LENGTH) || pos == getIdx(IdxEntry::HOST_PAST_KEY_VALUE_LENGTHS) - || pos == getIdx(IdxEntry::CACHE_INDIR))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (isRoPE() && (pos == getIdx(IdxEntry::ROTARY_INV_FREQ) || pos == getIdx(IdxEntry::ROTARY_COS_SIN))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (isLongRoPE() - && (pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_INV_FREQ) || pos == getIdx(IdxEntry::LONG_ROPE_ROTARY_COS_SIN))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (useKVCache() && mKVCacheQuantMode.hasKvCacheQuant() - && (pos == getIdx(IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE) - || pos == getIdx(IdxEntry::KV_CACHE_QUANTIZATION_SCALE))) - { - // kv_scale for mType->int8/fp8 and int8/fp8->mType conversion - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (mFP8ContextFMHA && pos == getIdx(IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (mFuseFp4Quant && pos == getIdx(IdxEntry::ATTENTION_OUTPUT_SF_SCALE)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useFullCustomMask() && pos == getIdx(IdxEntry::ATTENTION_MASK)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kBOOL && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useCustomMask() && pos == getIdx(IdxEntry::ATTENTION_PACKED_MASK)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mPagedKVCache - && (pos == getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS) || pos == getIdx(IdxEntry::HOST_KV_CACHE_BLOCK_OFFSETS))) - { - // kv cache block offsets - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mPagedKVCache && (pos == getIdx(IdxEntry::HOST_KV_CACHE_POOL_POINTERS))) - { - // kv cache pool pointers - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT64 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mPagedKVCache && (pos == getIdx(IdxEntry::HOST_KV_CACHE_POOL_MAPPING))) - { - // kv cache pool mapping - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (useKVCache() && mKVCacheQuantMode.hasInt8KvCache() - && (!mPagedKVCache && (pos == getIdx(IdxEntry::PAST_KEY_VALUE) || pos == nbInputs + 1))) - { - // If use Int8 K/V cache we require I/O KV values to int8 - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kINT8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (useKVCache() && mKVCacheQuantMode.hasFp8KvCache() - && (!mPagedKVCache && (pos == getIdx(IdxEntry::PAST_KEY_VALUE) || pos == nbInputs + 1))) - { - // If use FP8 K/V cache we require I/O KV values to FP8 - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (mRemovePadding && (pos == getIdx(IdxEntry::HOST_CONTEXT_LENGTH))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32 && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (mCrossAttention - && (pos == getIdx(IdxEntry::CROSS_KV_LENGTH) || pos == getIdx(IdxEntry::ENCODER_INPUT_LENGTH))) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (isLognScaling() && pos == getIdx(IdxEntry::LOGN_SCALING)) - { - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (pos == nbInputs && mFuseFp4Quant) - { - // Set dtype for output FP4 quantized tensor. - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP4) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == nbInputs + 1 && mFuseFp4Quant) - { - // Set dtype for output scaling factor tensor. Use kINT32 as storage type (same as QuantizeToFP4Plugin). - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == nbInputs && mFP8ContextFMHA) - { - // Output tensor now supports fp8 data type. - posCaseLine = __LINE__; - result = (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (mSkipAttn && pos == getIdx(IdxEntry::SKIP_ATTN)) - { - posCaseLine = __LINE__; - result = inOut[pos].type == nvinfer1::DataType::kBOOL && inOut[pos].format == TensorFormat::kLINEAR; - } - else - { - posCaseLine = __LINE__; - result = (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - TLLM_LOG_DEBUG( - "%s: pos: %d, result: %d, posCaseLine: %d", __PRETTY_FUNCTION__, pos, static_cast(result), posCaseLine); - return result; -} - -template -void GPTAttentionPlugin::configurePluginImpl(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - TLLM_CHECK(mHeadSize > 0); - - int beamWidth = -1; - if (!isCrossAttention() && useKVCache()) - { - // desc_val == -1 means beam_width is not static, we should look at min/max/opt. - // - // In prepareEnqueueGeneration, we'll prepare for all cases where beam_width doesn't exceed max. - // TODO: pass min AND max to prepareEnqueueGeneration instead of max only. - int desc_val = in[getIdx(IdxEntry::CACHE_INDIR)].desc.dims.d[1]; - int max_val = in[getIdx(IdxEntry::CACHE_INDIR)].max.d[1]; - beamWidth = desc_val == -1 ? max_val : desc_val; - } - else - { - beamWidth = 1; - } - TLLM_CHECK(beamWidth != -1); - - // Commonly, cyclic_attention_window_size, and max_attention_window_size will be the same - // unless each layer has different attention window sizes. - // the kv_cache capacity. - int max_encoder_context_len = isCrossAttention() ? in[getIdx(IdxEntry::CROSS_KV_LENGTH)].desc.dims.d[0] : 0; - int const max_attention_window_size = isCrossAttention() - ? max_encoder_context_len - : (useKVCache() ? in[getIdx(IdxEntry::CACHE_INDIR)].desc.dims.d[2] : 0); - int const cyclic_attention_window_size = max_attention_window_size; - - int const num_requests = 256; - int const sink_token_length = 0; - - EnqueueGenerationParams enqueueParams; - enqueueParams.max_attention_window_size = max_attention_window_size; - enqueueParams.cyclic_attention_window_size = cyclic_attention_window_size; - enqueueParams.max_cyclic_attention_window_size = cyclic_attention_window_size; - enqueueParams.sink_token_length = sink_token_length; - enqueueParams.beam_width = beamWidth; - enqueueParams.num_requests = num_requests; - - prepareEnqueueGeneration(enqueueParams); - - // Always reserve SemaphoreArray (for multi-block mode) as MMHA may enable multi-block mode when shared memory is - // not enough. - auto const& ctxLenTensor = in[getIdx(IdxEntry::CONTEXT_LENGTHS)]; - TLLM_CHECK_DEBUG(ctxLenTensor.max.nbDims == 1); - int32_t const max_batch_beam = in[getIdx(IdxEntry::CONTEXT_LENGTHS)].max.d[0]; - reserveSemaphoreArray(mNumHeads * max_batch_beam); -} - -template -void GPTAttentionPlugin::configurePluginDispatchKVCacheType(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - if (mPagedKVCache) - { - configurePluginImpl(in, nbInputs, out, nbOutputs); - } - else - { - configurePluginImpl(in, nbInputs, out, nbOutputs); - } -} - -void GPTAttentionPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - if (mType == nvinfer1::DataType::kHALF) - { - configurePluginDispatchKVCacheType(in, nbInputs, out, nbOutputs); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - configurePluginDispatchKVCacheType(in, nbInputs, out, nbOutputs); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - configurePluginDispatchKVCacheType<__nv_bfloat16>(in, nbInputs, out, nbOutputs); - } -#endif -} - -size_t GPTAttentionPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - int const max_context_length = mMaxContextLength; - int const cross_kv_length = isCrossAttention() ? inputs[getIdx(IdxEntry::CROSS_KV_LENGTH)].dims.d[0] : 0; - int const max_num_seq = inputs[getIdx(IdxEntry::CONTEXT_LENGTHS)].dims.d[0]; - auto const type = inputs[getIdx(IdxEntry::QKV_TENSOR)].type; - int const max_kv_cache_length - = isCrossAttention() ? cross_kv_length : (useKVCache() ? inputs[getIdx(IdxEntry::CACHE_INDIR)].dims.d[2] : 0); - int const max_num_tokens - = mRemovePadding ? inputs[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0] : max_num_seq * max_context_length; - int const max_blocks_per_sequence - = (useKVCache() && mPagedKVCache) ? inputs[getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS)].dims.d[3] : 0; - - size_t const context_workspace_size - = getWorkspaceSizeForContext(type, max_num_seq, max_context_length, cross_kv_length, max_num_tokens); - - size_t const generation_workspace_size = getWorkspaceSizeForGeneration( - type, max_num_seq, max_kv_cache_length, max_num_tokens, max_blocks_per_sequence); - - size_t attention_input_workspace_size = 0; - - if (mUnfuseQkvGemm) - { - int const local_hidden_units_q - = inputs[getIdx(IdxEntry::QKV_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - int const local_hidden_units_kv - = inputs[getIdx(IdxEntry::K_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - size_t const size = tensorrt_llm::runtime::BufferDataType(type).getSize(); - size_t const attention_input_size = size * max_num_tokens * (local_hidden_units_q + 2 * local_hidden_units_kv); - size_t workspaces[1]; - workspaces[0] = attention_input_size; - attention_input_workspace_size = tensorrt_llm::common::calculateTotalWorkspaceSize(workspaces, 1); - } - - return std::max(context_workspace_size, generation_workspace_size) + attention_input_workspace_size; -} - -static size_t getStride(nvinfer1::Dims const& dims, int n) -{ - TLLM_CHECK(n >= 0 && n < dims.nbDims); - return std::accumulate(dims.d + n + 1, dims.d + dims.nbDims, 1, std::multiplies{}); -} - -template -int GPTAttentionPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - TLLM_LOG_TRACE("Attention plugin start at layer %d", mLayerIdx); - - using runtime::RequestType; - - int32_t const nbSeq = inputDesc[getIdx(IdxEntry::CONTEXT_LENGTHS)].dims.d[0]; - RequestType const* reqTypes = static_cast(inputs[getIdx(IdxEntry::REQUEST_TYPES)]); - - int32_t nbContextRequests = 0; - int32_t contextTokenIdxEnd = 0; - int32_t contextTokenIdxEndForCp = 0; - // count context requests - for (int32_t seqIdx = 0; seqIdx < nbSeq; seqIdx++) - { - if (reqTypes[seqIdx] != RequestType::kCONTEXT) - { - break; - } - ++nbContextRequests; - contextTokenIdxEnd += mRemovePadding - ? static_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)])[seqIdx] - : inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]; - contextTokenIdxEndForCp += mRemovePadding - ? (static_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)])[seqIdx] + mCpSize - 1) - / mCpSize - : (inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1] + mCpSize - 1) / mCpSize; - } - - for (int32_t seqIdx = nbContextRequests; seqIdx < nbSeq; seqIdx++) - { - TLLM_CHECK(reqTypes[seqIdx] == RequestType::kGENERATION); - } - - // mixed requests require mRemovePadding and mPagedKVCache - if (nbContextRequests != 0 && nbContextRequests != nbSeq) - { - TLLM_CHECK(mRemovePadding && mPagedKVCache); - } - - if (nbContextRequests > 0) - { - auto seqIdxBeg = 0; - auto tokenIdxBeg = 0; - auto localNbTokens = contextTokenIdxEnd; - enqueueSome(seqIdxBeg, nbContextRequests, tokenIdxBeg, localNbTokens, - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - - if (auto nbGenerationSeq = nbSeq - nbContextRequests; nbGenerationSeq > 0) - { - auto seqIdxBeg = nbContextRequests; - auto tokenIdxBeg = mCpSize > 1 ? contextTokenIdxEndForCp : contextTokenIdxEnd; - // if mRemovePadding is true, we may have IFB, and need to remove context tokens. - // if mRemovePadding is false, it is only generation requests, so just multiply batch_beam and seq_len (May not - // 1 for Parallel Decoding) - auto localNbTokens = mRemovePadding - ? inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0] - tokenIdxBeg - : inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0] * inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]; - enqueueSome(seqIdxBeg, nbGenerationSeq, tokenIdxBeg, localNbTokens, inputDesc, - outputDesc, inputs, outputs, workspace, stream); - } - - sync_check_cuda_error(stream); - TLLM_LOG_TRACE("Attention plugin stop at layer %d", mLayerIdx); - - return 0; -} - -template -int GPTAttentionPlugin::enqueueSome(int32_t seqIdxBeg, int32_t localNbSeq, int32_t tokenIdxBeg, int32_t localNbTokens, - nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) -{ - // relative_attention_bias [head_num, max_seq_len, max_seq_len] (optional in relative position) - // or [head_num, num_buckets] (optional in implicit relative attention) - // cross_kv [batch_size, seq_len, 2 * local_hidden_size] or [num_tokens, 2 * local_hidden_size] - // when enable remove_input_padding (optional in cross attention mode) - // cross_kv_length [int] max encoder input context length (optional in cross attention mode) - // encoder_input_lengths [batch_size] raw sequence lengths (optional in cross attention mode) - - using runtime::RequestType; - - auto const* const reqTypeInBatchPtr - = static_cast(inputs[getIdx(IdxEntry::REQUEST_TYPES)]) + seqIdxBeg; - bool const is_context = (reqTypeInBatchPtr[0] == RequestType::kCONTEXT); - - T const* attention_input = static_cast(inputs[getIdx(IdxEntry::QKV_TENSOR)]) - + inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)] - * size_t(tokenIdxBeg); - - bool changeSpecDecodingMode = false; - if (mIsSpecDecodingEnabled) - { - bool useSpecDecoding - = static_cast(reinterpret_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_USE)])[0]); - changeSpecDecodingMode = mUseSpecDecoding != useSpecDecoding; - mUseSpecDecoding = useSpecDecoding; - } - - [[maybe_unused]] MlaParams mla_params; - - T const* qkv_bias = nullptr; - if (mQKVBiasEnabled) - { - qkv_bias = reinterpret_cast(inputs[getIdx(IdxEntry::QKV_BIAS_TENSOR)]); - } - - // Note we still need context length during generation for MMHA optimization. - int32_t const max_context_q_len = [&]() - { - if (!mRemovePadding) - { - return static_cast(inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[1]); - } - auto const host_context_lengths - = static_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)]) + seqIdxBeg; - return *std::max_element(host_context_lengths, host_context_lengths + localNbSeq); - }(); - - // Rotary inv_freq, cos_sin cache to avoid re-computing. - float const* rotary_inv_freq = nullptr; - float2 const* rotary_cos_sin = nullptr; - - bool const useLongRoPECache = isLongRoPE() && max_context_q_len > mRotaryEmbeddingOriginalMaxPositions; - if (isRoPE()) - { - auto inputName = useLongRoPECache ? IdxEntry::LONG_ROPE_ROTARY_INV_FREQ : IdxEntry::ROTARY_INV_FREQ; - rotary_inv_freq = reinterpret_cast(inputs[getIdx(inputName)]); - } - if (isRoPE()) - { - auto inputName = useLongRoPECache ? IdxEntry::LONG_ROPE_ROTARY_COS_SIN : IdxEntry::ROTARY_COS_SIN; - rotary_cos_sin = reinterpret_cast(inputs[getIdx(inputName)]); - } - - auto const mrope_rotary_cos_sin - = isMRoPE() ? reinterpret_cast(inputs[getIdx(IdxEntry::MROPE_ROTARY_COS_SIN)]) : nullptr; - - auto const mrope_position_deltas - = isMRoPE() ? reinterpret_cast(inputs[getIdx(IdxEntry::MROPE_POSITION_DELTAS)]) : nullptr; - - if (mUnfuseQkvGemm) - { - int const max_seqlen = inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[mRemovePadding ? 0 : 1]; - int const batch_size = mRemovePadding ? 1 : inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[0]; - - T const* attention_input_q = static_cast(inputs[getIdx(IdxEntry::QKV_TENSOR)]); - T const* attention_input_k = static_cast(inputs[getIdx(IdxEntry::K_TENSOR)]); - T const* attention_input_v = static_cast(inputs[getIdx(IdxEntry::V_TENSOR)]); - size_t const hidden_units_q - = inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - size_t const hidden_units_kv - = inputDesc[getIdx(IdxEntry::K_TENSOR)].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)]; - size_t const hidden_units = hidden_units_q + 2 * hidden_units_kv; - size_t const size_qkv = sizeof(T) * hidden_units; - size_t const size_q = sizeof(T) * hidden_units_q; - size_t const size_kv = sizeof(T) * hidden_units_kv; - size_t const total_size = size_qkv * batch_size * max_seqlen; - int8_t* workspace_byte_ptr = reinterpret_cast(workspace); - size_t offset = 0; - T* attention_input_qkv = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, total_size)); - workspace = reinterpret_cast(workspace_byte_ptr + offset); - - cudaMemcpy2DAsync(attention_input_qkv, size_qkv, attention_input_q, size_q, size_q, batch_size * max_seqlen, - cudaMemcpyDeviceToDevice, stream); - cudaMemcpy2DAsync(attention_input_qkv + hidden_units_q, size_qkv, attention_input_k, size_kv, size_kv, - batch_size * max_seqlen, cudaMemcpyDeviceToDevice, stream); - cudaMemcpy2DAsync(attention_input_qkv + hidden_units_q + hidden_units_kv, size_qkv, attention_input_v, size_kv, - size_kv, batch_size * max_seqlen, cudaMemcpyDeviceToDevice, stream); - - attention_input = attention_input_qkv + hidden_units * tokenIdxBeg; - } - - int const* context_q_lengths = reinterpret_cast(inputs[getIdx(IdxEntry::CONTEXT_LENGTHS)]) + seqIdxBeg; - int const* sequence_kv_length = useKVCache() - ? static_cast(inputs[getIdx(IdxEntry::SEQUENCE_LENGTH)]) + seqIdxBeg - : context_q_lengths; - - int max_encoder_context_len = isCrossAttention() ? inputDesc[getIdx(IdxEntry::CROSS_KV_LENGTH)].dims.d[0] : 0; - // for enc-dec model, since decoder_input_ids could be longer than 1, - // such model has an encoder context (for cross attn) and an decoder context (for self attn) - // clarify 3 lens: - // -- max_context_q_len: len of decoder input. No "max" concept, it's what it is given. - // Also called (decoder_)input_seq_length, normally 1 for encoder-decoder start token - // -- max_seq_len: max allowed len of decoder output, i.e. final results - // -- max_encoder_context_len: len of encoder input (in cross attn). Also called encoder_input_seq_length - - int const beamWidth - = isCrossAttention() ? 1 : (useKVCache() ? inputDesc[getIdx(IdxEntry::CACHE_INDIR)].dims.d[1] : 1); - - // Commonly, cyclic_attention_window_size, and max_attention_window_size will be the same - // unless each layer has different attention window sizes. - // the kv_cache capacity. - int const max_attention_window_size = isCrossAttention() - ? max_encoder_context_len - : (useKVCache() ? inputDesc[getIdx(IdxEntry::CACHE_INDIR)].dims.d[2] : 0); - // The cyclic_attention_window_size will determine the cyclic kv cache position of new tokens. - // Note that this cyclic_attention_window_size might be smaller than the actual kv cache capactity. - int const* cyclic_attention_window_sizes - = reinterpret_cast(inputs[getIdx(IdxEntry::HOST_MAX_ATTENTION_WINDOW)]); - int const cyclic_attention_window_size - = isCrossAttention() ? max_encoder_context_len : cyclic_attention_window_sizes[mLayerIdx]; - int const sink_token_length = reinterpret_cast(inputs[getIdx(IdxEntry::HOST_SINK_TOKEN_LENGTH)])[0]; - int const num_attn_layer = inputDesc[getIdx(IdxEntry::HOST_MAX_ATTENTION_WINDOW)].dims.d[0]; - int const max_cyclic_attention_window_size = isCrossAttention() - ? max_encoder_context_len - : *std::max_element(cyclic_attention_window_sizes, cyclic_attention_window_sizes + num_attn_layer); - bool const can_use_one_more_block = beamWidth > 1; - - float const* kv_scale_orig_quant = nullptr; - float const* kv_scale_quant_orig = nullptr; - if (useKVCache() && mKVCacheQuantMode.hasKvCacheQuant()) - { - assert(inputDesc[getIdx(IdxEntry::KV_CACHE_QUANTIZATION_SCALE)].type == nvinfer1::DataType::kFLOAT); - assert(inputDesc[getIdx(IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE)].type == nvinfer1::DataType::kFLOAT); - kv_scale_orig_quant = reinterpret_cast(inputs[getIdx(IdxEntry::KV_CACHE_QUANTIZATION_SCALE)]); - kv_scale_quant_orig = reinterpret_cast(inputs[getIdx(IdxEntry::KV_CACHE_DEQUANTIZATION_SCALE)]); - } - - float const* attention_output_orig_quant = nullptr; - if (mFP8ContextFMHA) - { - assert(inputDesc[getIdx(IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE)].type == nvinfer1::DataType::kFLOAT); - attention_output_orig_quant - = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_OUTPUT_QUANTIZATION_SCALE)]); - } - float const* attention_output_sf_scale = nullptr; - if (mFuseFp4Quant) - { - assert(inputDesc[getIdx(IdxEntry::ATTENTION_OUTPUT_SF_SCALE)].type == nvinfer1::DataType::kFLOAT); - attention_output_sf_scale = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_OUTPUT_SF_SCALE)]); - } - uint32_t const* attention_packed_mask = nullptr; - if (useCustomMask()) - { - assert(inputDesc[getIdx(IdxEntry::ATTENTION_PACKED_MASK)].type == nvinfer1::DataType::kINT32); - attention_packed_mask = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_PACKED_MASK)]); - } - bool const* attention_mask = nullptr; - int attention_mask_stride = 0; - if (useFullCustomMask()) - { - attention_mask_stride = static_cast(inputDesc[getIdx(IdxEntry::ATTENTION_MASK)].dims.d[1]); - attention_mask = reinterpret_cast(inputs[getIdx(IdxEntry::ATTENTION_MASK)]) - + attention_mask_stride * static_cast(tokenIdxBeg); - } - - int max_blocks_per_sequence = 0; - kernels::KVBlockArray::DataType* block_offsets = nullptr; - void* host_primary_pool_pointer = nullptr; - void* host_secondary_pool_pointer = nullptr; - if (useKVCache() && mPagedKVCache) - { - auto const& kvCacheBlockOffsetsShape = inputDesc[getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS)].dims; - max_blocks_per_sequence = kvCacheBlockOffsetsShape.d[kvCacheBlockOffsetsShape.nbDims - 1]; - - std::int32_t const* host_pool_mapping - = static_cast(inputs[getIdx(IdxEntry::HOST_KV_CACHE_POOL_MAPPING)]); - - int32_t const layerToPool = host_pool_mapping[mLayerIdx * 2]; - int32_t const layerIdxInCachePool = host_pool_mapping[mLayerIdx * 2 + 1]; - TLLM_LOG_TRACE("Layer%d: LayerCachePoolLocator{.indexOfPool=%d, .layerIdxInCachePool=%d}", mLayerIdx, - layerToPool, layerIdxInCachePool); - auto const seqStride = getStride(kvCacheBlockOffsetsShape, 1); - auto const poolStride = getStride(kvCacheBlockOffsetsShape, 0); - auto const seqOffset = seqIdxBeg * seqStride; - auto const poolOffset = layerToPool * poolStride; - - block_offsets - = reinterpret_cast(inputs[getIdx(IdxEntry::KV_CACHE_BLOCK_OFFSETS)]) - + poolOffset + seqOffset; - - auto const* const typed_host_pool_pointers - = static_cast(inputs[getIdx(IdxEntry::HOST_KV_CACHE_POOL_POINTERS)]); - - auto const cacheElemSize = (mKVCacheQuantMode.hasKvCacheQuant() ? 1 : sizeof(T)); - - auto const kv_cache_head_num = (mNumKVHeads + mCpSize - 1) / mCpSize; - auto const blockSize = mTokensPerBlock * kv_cache_head_num * mHeadSize; - auto const bytesPerBlock = blockSize * cacheElemSize; - auto const layerOffset = layerIdxInCachePool * 2 * bytesPerBlock; - - host_primary_pool_pointer = reinterpret_cast(typed_host_pool_pointers[layerToPool * 2] + layerOffset); - host_secondary_pool_pointer - = reinterpret_cast(typed_host_pool_pointers[layerToPool * 2 + 1] + layerOffset); - } - - // The index of kv cache tensor in outputs. If fuse FP4 quant, an additional scaling factor output is added before - // the kv cache tensor. - int const kvCacheIdxInOutputs = mFuseFp4Quant ? 2 : 1; - // The number of elements per storage type. For FP4 output, storage type is uint8_t. - int const numEltsPerStorageType = mFuseFp4Quant ? 2 : 1; - - AttentionOutT* context_buf_ = static_cast(outputs[0]) - + outputDesc[0].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)] * tokenIdxBeg / numEltsPerStorageType; - - __nv_fp8_e4m3* context_buf_sf_ = nullptr; - if (mFuseFp4Quant) - { - // The output address for FP4 scaling factor. - context_buf_sf_ = static_cast<__nv_fp8_e4m3*>(outputs[1]); - } - - void* key_value_cache = nullptr; - if (useKVCache() && !mPagedKVCache) - { - auto const cacheElemSize = (mKVCacheQuantMode.hasKvCacheQuant() ? 1 : sizeof(T)); - key_value_cache = static_cast(outputs[kvCacheIdxInOutputs]) - + cacheElemSize * getStride(outputDesc[kvCacheIdxInOutputs].dims, 0) * seqIdxBeg; - void const* past_key_value_cache = inputs[getIdx(IdxEntry::PAST_KEY_VALUE)]; - if (past_key_value_cache != outputs[kvCacheIdxInOutputs]) - { - auto shape = outputDesc[kvCacheIdxInOutputs].dims; - auto const size - = cacheElemSize * std::accumulate(shape.d, shape.d + shape.nbDims, 1, std::multiplies{}); - cudaMemcpyAsync(outputs[kvCacheIdxInOutputs], past_key_value_cache, size, cudaMemcpyDeviceToDevice, stream); - } - } - - T const* alibi_slopes = isALiBi() ? static_cast(inputs[getIdx(IdxEntry::ALIBI_SLOPES)]) : nullptr; - - int const* spec_decoding_packed_mask = nullptr; - int const* spec_decoding_position_offsets = nullptr; - int const* spec_decoding_generation_lengths = nullptr; - int num_decoding_draft_tokens = 0; - if (mIsSpecDecodingEnabled && mUseSpecDecoding) - { - // Second dimension of spec_decoding_position_offsets is num_decoding_draft_tokens + 1. - // [batch_size, num_decoding_draft_tokens + 1] - num_decoding_draft_tokens = inputDesc[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)].dims.d[1] - 1; - if (num_decoding_draft_tokens > 0) - { - // spec_decoding_* tensors are not filled for context requests. Hence, always strting from 0th index - int32_t constexpr genSeqIdx = 0; - spec_decoding_packed_mask = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)]) - + genSeqIdx * getStride(inputDesc[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)].dims, 0); - // Packed as [num_tokens, packed_mask_size] - // Use seqIdxBeg * (num_decoding_draft_tokens + 1) here as only generation tokens have the packed_mask - // buffer. - // TODO: support variable sequence length based on generationTokenIdxBeg. - spec_decoding_packed_mask = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)]) - + genSeqIdx * (num_decoding_draft_tokens + 1) - * getStride(inputDesc[getIdx(IdxEntry::SPEC_DECODING_PACKED_MASK)].dims, 0); - spec_decoding_position_offsets - = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)]) - + genSeqIdx * getStride(inputDesc[getIdx(IdxEntry::SPEC_DECODING_POSITION_OFFSETS)].dims, 0); - spec_decoding_generation_lengths - = static_cast(inputs[getIdx(IdxEntry::SPEC_DECODING_GENERATION_LENGTHS)]) + genSeqIdx; - } - } - - int32_t const* host_past_kv_len_list = useKVCache() - ? static_cast(inputs[getIdx(IdxEntry::HOST_PAST_KEY_VALUE_LENGTHS)]) + seqIdxBeg - : nullptr; - int32_t const max_context_kv_len = useKVCache() - ? *std::max_element(host_past_kv_len_list, host_past_kv_len_list + localNbSeq) - : max_context_q_len; - - int const* host_context_lengths - = mRemovePadding ? reinterpret_cast(inputs[getIdx(IdxEntry::HOST_CONTEXT_LENGTH)]) : nullptr; - - int64_t const* runtime_perf_knobs = static_cast(inputs[getIdx(IdxEntry::HOST_RUNTIME_PERF_KNOBS)]); - - EnqueueParams common_enqueue_params; - common_enqueue_params.attention_input = attention_input; - common_enqueue_params.qkv_bias = qkv_bias; - common_enqueue_params.attention_mask = attention_mask; - common_enqueue_params.rotary_inv_freq = rotary_inv_freq; - common_enqueue_params.rotary_cos_sin = rotary_cos_sin; - common_enqueue_params.max_attention_window_size = max_attention_window_size; - common_enqueue_params.cyclic_attention_window_size = cyclic_attention_window_size; - common_enqueue_params.max_cyclic_attention_window_size = max_cyclic_attention_window_size; - common_enqueue_params.can_use_one_more_block = can_use_one_more_block; - common_enqueue_params.sink_token_length = sink_token_length; - common_enqueue_params.kv_scale_orig_quant = kv_scale_orig_quant; - common_enqueue_params.kv_scale_quant_orig = kv_scale_quant_orig; - common_enqueue_params.attention_output_orig_quant = attention_output_orig_quant; - common_enqueue_params.attention_output_sf_scale = attention_output_sf_scale; - common_enqueue_params.alibi_slopes = alibi_slopes; - common_enqueue_params.context_buf = context_buf_; - common_enqueue_params.context_buf_sf = context_buf_sf_; - common_enqueue_params.key_value_cache = key_value_cache; - common_enqueue_params.block_offsets = block_offsets; - common_enqueue_params.host_primary_pool_pointer = host_primary_pool_pointer; - common_enqueue_params.host_secondary_pool_pointer = host_secondary_pool_pointer; - common_enqueue_params.num_tokens = localNbTokens; - common_enqueue_params.max_blocks_per_sequence = max_blocks_per_sequence; - common_enqueue_params.sequence_lengths = sequence_kv_length; - common_enqueue_params.context_lengths = context_q_lengths; - common_enqueue_params.host_context_lengths = host_context_lengths; - common_enqueue_params.workspace = workspace; - common_enqueue_params.runtime_perf_knobs = runtime_perf_knobs; - - if (isRelativePosition()) - { - common_enqueue_params.relative_attention_bias - = static_cast(inputs[getIdx(IdxEntry::RELATIVE_ATTENTION_BIAS)]); - common_enqueue_params.relative_attention_bias_stride - = inputDesc[getIdx(IdxEntry::RELATIVE_ATTENTION_BIAS)].dims.d[1]; // max_seq_len or num_buckets - } - if (isLognScaling()) - { - common_enqueue_params.logn_scaling_ptr = static_cast(inputs[getIdx(IdxEntry::LOGN_SCALING)]); - } - if (isCrossAttention()) - { - common_enqueue_params.encoder_input_lengths - = reinterpret_cast(inputs[getIdx(IdxEntry::ENCODER_INPUT_LENGTH)]) + seqIdxBeg; - } - - if (is_context) // context stage - { - int const batch_size = localNbSeq; - int const request_batch_size = batch_size; - // num of total tokens (without paddings when remove paddings). - int num_encoder_tokens = 0; - if (isCrossAttention()) - { - if (!mRemovePadding) - { - num_encoder_tokens = request_batch_size * max_encoder_context_len; - } - else - { - num_encoder_tokens = inputDesc[getIdx(IdxEntry::CROSS_KV)].dims.d[0]; - } - } - - common_enqueue_params.input_seq_length = max_context_q_len; - common_enqueue_params.max_past_kv_length = max_context_kv_len; - EnqueueContextParams enqueue_params{common_enqueue_params}; - enqueue_params.attention_packed_mask = attention_packed_mask; - enqueue_params.batch_size = batch_size; - enqueue_params.mrope_rotary_cos_sin = mrope_rotary_cos_sin; - enqueue_params.total_kv_len = enqueue_params.num_tokens; - - if (isCrossAttention()) - { - enqueue_params.cross_kv = static_cast(inputs[getIdx(IdxEntry::CROSS_KV)]); - enqueue_params.cross_kv_length = max_encoder_context_len; - enqueue_params.num_encoder_tokens = num_encoder_tokens; - } - - enqueueContext(enqueue_params, stream); - - { - std::string const afterContexStr = "ctx attention at layer " + std::to_string(mLayerIdx); - TLLM_LOG_TRACE("GPTAttentionPlugin - %s", afterContexStr.c_str()); - - auto progress = static_cast( - inputs[getIdx(IdxEntry::HOST_CONTEXT_PROGRESS)])[0]; - if (progress != nullptr) - { - progress->recordEvent(mLayerIdx, stream); - } - - if (!mFuseFp4Quant) - { - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(localNbTokens, - outputDesc[0].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)], - mFP8ContextFMHA ? nvinfer1::DataType::kFP8 : mType, context_buf_, stream, afterContexStr) - == false, - "Found invalid number (NaN or Inf) in " + afterContexStr); - } - } - } - else // generation stage; max_context_q_len == input_seq_len == 1 - { - TLLM_CHECK_WITH_INFO(useKVCache(), "KV-cache-less is only supported for context"); - int batch_beam = localNbSeq; - TLLM_CHECK(batch_beam % beamWidth == 0); - int32_t const num_requests = batch_beam / beamWidth; - - int const* cache_indir - = beamWidth == 1 ? nullptr : reinterpret_cast(inputs[getIdx(IdxEntry::CACHE_INDIR)]); - - // Medusa: the max input sequence length if variable sequence length is needed. - int const input_seq_length = getGenerationInputSequenceLength(inputDesc, localNbSeq, localNbTokens); - int const max_past_kv_length = isCrossAttention() ? max_encoder_context_len : max_context_kv_len; - auto qkvDims = inputDesc[getIdx(IdxEntry::QKV_TENSOR)].dims; - TLLM_CHECK_WITH_INFO(input_seq_length == 1 || (mIsSpecDecodingEnabled && mUseSpecDecoding), - "Only speculative decoding mode supports input length > 1 in the generation phase, input_seq_length=%d, " - "mIsSpecDecodingEnabled=%s, nDims=%d, (" FMT_DIM ", " FMT_DIM ", " FMT_DIM ")", - input_seq_length, mIsSpecDecodingEnabled ? "true" : "false", qkvDims.nbDims, qkvDims.d[0], qkvDims.d[1], - qkvDims.d[2]); - TLLM_CHECK_WITH_INFO( - input_seq_length == num_decoding_draft_tokens + 1, "The generation input length is not expected."); - common_enqueue_params.input_seq_length = input_seq_length; - common_enqueue_params.max_past_kv_length = max_past_kv_length; - EnqueueGenerationParams enqueue_params{common_enqueue_params}; - enqueue_params.beam_width = beamWidth; - enqueue_params.attention_mask_stride = attention_mask_stride; - enqueue_params.num_requests = num_requests; - enqueue_params.cache_indir = cache_indir; - enqueue_params.semaphores = multiBlockSemaphores(); - enqueue_params.host_past_key_value_lengths = host_past_kv_len_list; - enqueue_params.mrope_position_deltas = mrope_position_deltas; - if (mIsSpecDecodingEnabled && mUseSpecDecoding) - { - enqueue_params.spec_decoding_packed_mask = spec_decoding_packed_mask; - enqueue_params.spec_decoding_position_offsets = spec_decoding_position_offsets; - enqueue_params.spec_decoding_generation_lengths = spec_decoding_generation_lengths; - enqueue_params.spec_decoding_is_generation_length_variable = mSpecDecodingIsGenerationLengthVariable; - enqueue_params.spec_decoding_max_generation_length = mSpecDecodingMaxGenerationLength; - } - if (mFuseFp4Quant) - { - enqueue_params.start_token_idx_sf = tokenIdxBeg; - } - - if (changeSpecDecodingMode) - { - // mUseSpecDecoding is changed, need to re-prepare the DecoderXQARunner - prepareEnqueueGeneration(enqueue_params); - } - - enqueueGeneration(enqueue_params, stream); - - { - std::string const afterGenStr = "gen attention at layer " + std::to_string(mLayerIdx); - { - TLLM_CHECK_DEBUG_WITH_INFO( - tensorrt_llm::runtime::utils::tensorHasInvalid(localNbTokens, - outputDesc[0].dims.d[getPackedTensorHiddenDimIndex(mRemovePadding)], - mFP8ContextFMHA ? nvinfer1::DataType::kFP8 : mType, context_buf_, stream, afterGenStr) - == false, - "Found invalid number (NaN or Inf) in " + afterGenStr); - } - } - } - - return 0; -} - -template -int GPTAttentionPlugin::enqueueDispatchKVCacheType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - if (mPagedKVCache) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - return 0; -} - -int GPTAttentionPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mSkipAttn) - { - bool const* SKIP_ATTN = reinterpret_cast(inputs[getIdx(IdxEntry::SKIP_ATTN)]); - if (SKIP_ATTN[0]) - { - return 0; - } - } - - if (mType == nvinfer1::DataType::kHALF) - { - if (mFuseFp4Quant) - { - return enqueueDispatchKVCacheType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_FP8 - if (mFP8ContextFMHA) - { - return enqueueDispatchKVCacheType( - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return enqueueDispatchKVCacheType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - return enqueueDispatchKVCacheType(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - if (mFuseFp4Quant) - { - return enqueueDispatchKVCacheType<__nv_bfloat16, uint8_t>( - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_FP8 - if (mFP8ContextFMHA) - { - return enqueueDispatchKVCacheType<__nv_bfloat16, __nv_fp8_e4m3>( - inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return enqueueDispatchKVCacheType<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType GPTAttentionPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (mFuseFp4Quant) - { - TLLM_CHECK(index == 0 || index == 1 || (!mPagedKVCache && useKVCache() && index == 2)); - } - else - { - TLLM_CHECK(index == 0 || (!mPagedKVCache && useKVCache() && index == 1)); - } - if (index == 0) - { - if (mFuseFp4Quant) - { - return nvinfer1::DataType::kFP4; - } - return mFP8ContextFMHA && mEnableContextFMHA ? nvinfer1::DataType::kFP8 - : inputTypes[getIdx(IdxEntry::QKV_TENSOR)]; - } - if (mFuseFp4Quant && index == 1) - { - return nvinfer1::DataType::kFP8; - } - return inputTypes[getIdx(IdxEntry::PAST_KEY_VALUE)]; -} - -// IPluginV2 Methods - -char const* GPTAttentionPlugin::getPluginType() const noexcept -{ - return GPT_ATTENTION_PLUGIN_NAME; -} - -char const* GPTAttentionPlugin::getPluginVersion() const noexcept -{ - return GPT_ATTENTION_PLUGIN_VERSION; -} - -int GPTAttentionPlugin::getNbOutputs() const noexcept -{ - int nbOutputs = mFuseFp4Quant ? 2 : 1; - if (!mPagedKVCache && useKVCache()) - { - nbOutputs += 1; - } - return nbOutputs; -} - -size_t GPTAttentionPlugin::getSerializationSize() const noexcept -{ - return GPTAttentionPluginCommon::getCommonSerializationSize(); -} - -void GPTAttentionPlugin::serialize(void* buffer) const noexcept -{ - GPTAttentionPluginCommon::serializeCommon(buffer); -} - -/////////////// - -GPTAttentionPluginCreator::GPTAttentionPluginCreator() - : GPTAttentionPluginCreatorCommon() -{ -} - -char const* GPTAttentionPluginCreator::getPluginName() const noexcept -{ - return GPT_ATTENTION_PLUGIN_NAME; -} - -char const* GPTAttentionPluginCreator::getPluginVersion() const noexcept -{ - return GPT_ATTENTION_PLUGIN_VERSION; -} - -PluginFieldCollection const* GPTAttentionPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* GPTAttentionPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginFieldParser p{fc->nbFields, fc->fields}; - - try - { - auto* obj = new GPTAttentionPlugin(p.getScalar("layer_idx").value(), - p.getScalar("num_heads").value(), p.getScalar("vision_start").value(), - p.getScalar("vision_length").value(), p.getScalar("num_kv_heads").value(), - p.getScalar("num_kv_heads_origin").value(), p.getScalar("head_size").value(), - p.getScalar("unidirectional").value(), p.getScalar("q_scaling").value(), - p.getScalar("attn_logit_softcapping_scale").value(), - static_cast(p.getScalar("position_embedding_type").value()), - p.getScalar("rotary_embedding_dim").value(), p.getScalar("rotary_embedding_base").value(), - static_cast(p.getScalar("rotary_embedding_scale_type").value()), - p.getScalar("rotary_embedding_scale").value(), - p.getScalar("rotary_embedding_short_m_scale").value(), - p.getScalar("rotary_embedding_long_m_scale").value(), - p.getScalar("rotary_embedding_max_positions").value(), - p.getScalar("rotary_embedding_original_max_positions").value(), - static_cast(p.getScalar("tp_size").value()), - static_cast(p.getScalar("tp_rank").value()), - static_cast(p.getScalar("unfuse_qkv_gemm").value()), - static_cast(p.getScalar("use_logn_scaling").value()), - static_cast(p.getScalar("context_fmha_type").value()), - p.getScalar("kv_cache_quant_mode").value(), - static_cast(p.getScalar("remove_input_padding").value()), - static_cast(p.getScalar("mask_type").value()), - BlockSparseParams{p.getScalar("block_sparse_block_size").value(), - static_cast(p.getScalar("block_sparse_homo_head_pattern").value()), - p.getScalar("block_sparse_num_local_blocks").value(), - p.getScalar("block_sparse_vertical_stride").value()}, - static_cast(p.getScalar("paged_kv_cache").value()), - p.getScalar("tokens_per_block").value(), - static_cast(p.getScalar("type_id").value()), - p.getScalar("max_context_length").value(), - static_cast(p.getScalar("qkv_bias_enabled").value()), - static_cast(p.getScalar("do_cross_attention").value()), - static_cast(p.getScalar("max_distance").value()), - static_cast(p.getScalar("pos_shift_enabled").value()), - static_cast(p.getScalar("dense_context_fmha").value()), - static_cast(p.getScalar("use_paged_context_fmha").value()), - static_cast(p.getScalar("use_fp8_context_fmha").value()), - static_cast(p.getScalar("has_full_attention_mask").value()), - static_cast(p.getScalar("use_cache").value()), - static_cast(p.getScalar("is_spec_decoding_enabled").value()), - static_cast(p.getScalar("spec_decoding_is_generation_length_variable").value()), - p.getScalar("spec_decoding_max_generation_length").value(), - static_cast(p.getScalar("is_mla_enabled").value()), - static_cast(p.getScalar("q_lora_rank").value()), - static_cast(p.getScalar("kv_lora_rank").value()), - static_cast(p.getScalar("qk_nope_head_dim").value()), - static_cast(p.getScalar("qk_rope_head_dim").value()), - static_cast(p.getScalar("v_head_dim").value()), - static_cast(p.getScalar("fuse_fp4_quant").value()), - static_cast(p.getScalar("skip_attn").value()), - static_cast(p.getScalar("cp_size").value()), - static_cast(p.getScalar("cp_rank").value()), - static_cast>(p.getSet("cp_group").value())); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* GPTAttentionPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call GPTAttentionPlugin::destroy() - try - { - auto* obj = new GPTAttentionPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h b/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h deleted file mode 100644 index 3e34703c6221..000000000000 --- a/cpp/tensorrt_llm/plugins/gptAttentionPlugin/gptAttentionPlugin.h +++ /dev/null @@ -1,259 +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. - */ -#pragma once - -#include "checkMacrosPlugin.h" -#include "tensorrt_llm/common/cublasMMWrapper.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/common/stringUtils.h" -#include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fmhaRunner.h" -#include "tensorrt_llm/kernels/contextFusedMultiHeadAttention/fused_multihead_attention_common.h" -#include "tensorrt_llm/kernels/gptKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gptAttentionCommon/gptAttentionCommon.h" -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests + num_gen_requests * beam_width -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (beam_width sequences per request). -// Context sequences have to appear first, generation sequences after - -// inputs (see GPTAttentionPlugin::isEntryUsed for when each tensor is actually used) -// 0. input_tensor [batch_size, seq_len, local_hidden_size + 2 * local_num_kv_heads * head_size] or -// [num_tokens, local_hidden_size + 2 * local_num_kv_heads * head_size] when -// enable_remove_input_padding -// 1. sequence_length [batch_size] (optional) -// 2. host_past_key_value_lengths [batch_size] (int32) (optional) -// 3. host_max_attention_window_sizes [num_layers] (int32) -// 4. host_sink_token_length [1] (int32) -// 5. context_lengths [batch_size] -// 6. cache_indir [num_gen_requests, beam_width, memory_max_len] (required in beamsearch) (optional) -// 7. host_request_types [batch_size] int32. 0: context; 1: generation: 2: none. When not in inflight-batching -// mode, -// all elements must be identical. -// 8. past_key_value_pool [batch_size, 2, local_num_kv_heads, max_seq_len, head_size] or -// block_offsets [batch_size, 2, max_blocks_per_seq] if paged kv cache (optional) -// 8.1 host_pool_pointers [2] if paged kv cache (optional) -// 9. kv_cache_quantization_scale [1] (optional) -// 10. kv_cache_dequantization_scale [1] (optional) -// 11. attention_output_quantization_scale [1] (on device, optional) -// 12. attention_mask [num_tokens, kv_seqlen] (on device, bool, optional) -// 13. attention_packed_mask [num_tokens, kv_seqlen / 32] (on device, uint32_t, optional) -// - pack masks by encoding multiple mask positions into a single 32-bit unsigned integer. -// - see kernels/contextMultiHeadAttention/fmhaPackedMask.cpp for more details. -// 14. rotary_inv_freq [head_size / 2] or [head_size] (longrope type) (float) (on device, optional) -// 15. rotary_cos_sin [max_num_embedding_positions, 2] (float) (on device, optional) -// 16. alibi_slopes [num_heads] (optional for ALiBi position embedding) -// 17. relative_attention_bias [num_heads] (optional for ALiBi position embedding) -// 18. host_context_lengths [batch_size] int32. (optional, required when remove_input_padding is true) -// 19. qkv_bias (optional) [local_hidden_size * 3] -// 20. spec_decoding_generation_lengths (optional, required when medusa is enabled) (int32_t) [batch_size] -// 21. spec_decoding_packed_mask (optional, required when medusa is enabled) (int32_t) [num_tokens, packed_mask_dim] -// packed_mask_dim = divUp(max_num_spec_decoding_tokens + 1, 32) -// 22. spec_decoding_position_offsets (optional, required when medusa is enabled) (int32_t) [batch_size, -// max_num_spec_decoding_tokens + 1] -// 23. spec_decoding_use (optional, bool) [1]: If it is set as true, enable speculative decoding -// 24. long_rope_rotary_inv_freq [head / 2] (float) (on device, optional) -// 25. long_rope_rotary_cos_sin [max_num_embedding_positions, 2] (float) (on device, optional) -// 26. host_runtime_perf_knobs (int64) -// 27. host_context_progress (void*) -// 28. position_id_tensor(MLA) [total_tokens], used for rope embedding in MLA -// 29. q_a_proj_tensor(MLA) [hidden_dim, c_q_dim + c_k_dim + ropd_dim], used to proj compacted QKV -// 30. q_a_layernorm_tensor(MLA) [c_q_dim], rmsnorm weight for compacted q -// 31. q_b_proj_tensor(MLA) [c_q_dim, head_num * head_size], weight for companted q to q in context -// 32. kv_a_proj_with_mqa_tensor(MLA) [c_q_dim, head_num * (c_k_dim + rope_dim)], weight for companted q to kdim in -// generation -// 33. kv_a_layernorm_tensor(MLA) [c_k_dim], rmsnorm weight for compacted kv -// 34. kv_b_proj_tensor(MLA) [c_k_dim, head_num * 2 * (head_size - rope_dim)], weight for compacted kv to kv in -// context -// 35. skip_attn (optional, bool) [1]: If it is set as true, skip the atteniton plugin and return -// directly. -// -// outputs -// output_tensor [batch_size, seq_len, local_hidden_size] -// present_key_value_pool (optional if not paged kv cache) [batch_size, 2, local_num_kv_heads, max_seq_len, -// head_size] - -class GPTAttentionPlugin : public GPTAttentionPluginCommon -{ -public: - GPTAttentionPlugin(int layer_idx, int num_heads, int vision_start, int vision_length, int num_kv_heads, - int num_kv_heads_origin, int head_size, int unidirectional, float q_scaling, float attn_logit_softcapping_scale, - tensorrt_llm::kernels::PositionEmbeddingType position_embedding_type, - int rotary_embedding_dim, // for RoPE. 0 for non-RoPE - float rotary_embedding_base, tensorrt_llm::kernels::RotaryScalingType rotary_embedding_scale_type, - float rotary_embedding_scale, float rotary_embedding_short_m_scale, float rotary_embedding_long_m_scale, - int rotary_embedding_max_positions, int rotary_embedding_original_max_positions, int tp_size, - int tp_rank, // for ALiBi - bool unfuse_qkv_gemm, // for AutoPP - bool use_logn_scaling, // for LognScaling - tensorrt_llm::kernels::ContextFMHAType context_fmha_type, int kv_cache_quant_mode, bool remove_input_padding, - tensorrt_llm::kernels::AttentionMaskType mask_type, - tensorrt_llm::kernels::BlockSparseParams block_sparse_params, bool paged_kv_cache, int tokens_per_block, - nvinfer1::DataType type, int32_t max_context_length, bool qkv_bias_enabled, bool cross_attention = false, - int max_distance = 0, bool pos_shift_enabled = false, bool dense_context_fmha = false, - bool use_paged_context_fmha = true, bool use_fp8_context_fmha = true, bool has_full_attention_mask = false, - bool use_cache = true, bool is_spec_decoding_enabled = false, - bool spec_decoding_is_generation_length_variable = false, int spec_decoding_max_generation_length = 1, - bool is_mla_enabled = false, int q_lora_rank = 0, int kv_lora_rank = 0, int qk_nope_head_dim = 0, - int qk_rope_head_dim = 0, int v_head_dim = 0, bool fuse_fp4_quant = false, bool skip_attn = false, - int cp_size = 1, int cp_rank = 0, std::set cp_group = {}); - - GPTAttentionPlugin(void const* data, size_t length); - - ~GPTAttentionPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - template - int enqueueDispatchKVCacheType(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream); - - template - void configurePluginImpl(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept; - template - void configurePluginDispatchKVCacheType(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - - //! This is called on every trt ExecutionContext creation by TRT - //! Note TRT does not call the initialize on cloned plugin, so clone internally should do initialization. - GPTAttentionPlugin* clone() const noexcept override; - - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - -private: - template - int enqueueSome(int32_t seqIdxBeg, int32_t localNbSeq, int32_t tokenIdxBeg, int32_t localNbTokens, - nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - using IndexType = std::int32_t; - - std::vector mEntryIdx; - enum class IdxEntry : size_t - { - QKV_TENSOR, - K_TENSOR, - V_TENSOR, - ATTENTION_MASK, - ATTENTION_PACKED_MASK, - SEQUENCE_LENGTH, - HOST_PAST_KEY_VALUE_LENGTHS, - HOST_MAX_ATTENTION_WINDOW, - HOST_SINK_TOKEN_LENGTH, - CONTEXT_LENGTHS, - CACHE_INDIR, - REQUEST_TYPES, - KV_CACHE_BLOCK_OFFSETS, - HOST_KV_CACHE_BLOCK_OFFSETS, - HOST_KV_CACHE_POOL_POINTERS, - HOST_KV_CACHE_POOL_MAPPING, - PAST_KEY_VALUE, - KV_CACHE_QUANTIZATION_SCALE, - KV_CACHE_DEQUANTIZATION_SCALE, - ATTENTION_OUTPUT_QUANTIZATION_SCALE, - ATTENTION_OUTPUT_SF_SCALE, - ROTARY_INV_FREQ, - ROTARY_COS_SIN, - ALIBI_SLOPES, - RELATIVE_ATTENTION_BIAS, - CROSS_KV, - CROSS_KV_LENGTH, - ENCODER_INPUT_LENGTH, - HOST_CONTEXT_LENGTH, - QKV_BIAS_TENSOR, - SPEC_DECODING_GENERATION_LENGTHS, - SPEC_DECODING_PACKED_MASK, - SPEC_DECODING_POSITION_OFFSETS, - SPEC_DECODING_USE, - LONG_ROPE_ROTARY_INV_FREQ, - LONG_ROPE_ROTARY_COS_SIN, - MROPE_ROTARY_COS_SIN, - MROPE_POSITION_DELTAS, - HOST_RUNTIME_PERF_KNOBS, - HOST_CONTEXT_PROGRESS, - MLA_Q_B_PROJ_TENSOR, - MLA_KV_B_PROJ_TENSOR, - MLA_K_B_PROJ_TRANS_TENSOR, - SKIP_ATTN, - LOGN_SCALING, - ENUM_SIZE, // Used to count the number of IdxEntry, must put in last - }; - - std::string toString(IdxEntry const& entry) const; - bool isEntryUsed(IdxEntry const& entry) const; - void initEntryIdx(); - IndexType getIdx(IdxEntry const& entry) const; - - // Get generation input sequence length (might be larger than 1 in the speculative decoding mode). - int getGenerationInputSequenceLength( - nvinfer1::PluginTensorDesc const* inputDesc, int32_t localNbSeq, int32_t localNbTokens) const; -}; - -class GPTAttentionPluginCreator : public GPTAttentionPluginCreatorCommon -{ -public: - GPTAttentionPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/identityPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/identityPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/identityPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.cpp b/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.cpp deleted file mode 100644 index 109010e7a933..000000000000 --- a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.cpp +++ /dev/null @@ -1,199 +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. - */ -#include "identityPlugin.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/iTensor.h" - -using namespace nvinfer1; -using tensorrt_llm::plugins::IdentityPluginCreator; -using tensorrt_llm::plugins::IdentityPlugin; - -static char const* IDENTITY_PLUGIN_VERSION{"1"}; -static char const* IDENTITY_PLUGIN_NAME{"Identity"}; -PluginFieldCollection IdentityPluginCreator::mFC{}; -std::vector IdentityPluginCreator::mPluginAttributes; - -IdentityPlugin::IdentityPlugin() {} - -// Parameterized constructor -IdentityPlugin::IdentityPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* IdentityPlugin::clone() const noexcept -{ - auto* plugin = new IdentityPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs IdentityPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - return inputs[outputIndex]; -} - -bool IdentityPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - assert(0 <= pos && pos < 2); - PluginTensorDesc const& input = inOut[0]; - PluginTensorDesc const& output = inOut[1]; - switch (pos) - { - case 0: return input.format == nvinfer1::TensorFormat::kLINEAR; - case 1: return output.type == input.type && output.format == nvinfer1::TensorFormat::kLINEAR; - } - return false; -} - -void IdentityPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t IdentityPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int IdentityPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - size_t count = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - count *= inputDesc[0].dims.d[i]; - } - count *= tensorrt_llm::runtime::BufferDataType(inputDesc[0].type).getSize(); - - cudaMemcpyAsync(outputs[0], inputs[0], count, cudaMemcpyDeviceToDevice, stream); - - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType IdentityPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* IdentityPlugin::getPluginType() const noexcept -{ - return IDENTITY_PLUGIN_NAME; -} - -char const* IdentityPlugin::getPluginVersion() const noexcept -{ - return IDENTITY_PLUGIN_VERSION; -} - -int IdentityPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int IdentityPlugin::initialize() noexcept -{ - return 0; -} - -void IdentityPlugin::terminate() noexcept {} - -size_t IdentityPlugin::getSerializationSize() const noexcept -{ - return 0; -} - -void IdentityPlugin::serialize(void* buffer) const noexcept {} - -void IdentityPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -IdentityPluginCreator::IdentityPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* IdentityPluginCreator::getPluginName() const noexcept -{ - return IDENTITY_PLUGIN_NAME; -} - -char const* IdentityPluginCreator::getPluginVersion() const noexcept -{ - return IDENTITY_PLUGIN_VERSION; -} - -PluginFieldCollection const* IdentityPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* IdentityPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - try - { - auto* obj = new IdentityPlugin(); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* IdentityPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call IdentityPlugin::destroy() - try - { - auto* obj = new IdentityPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.h b/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.h deleted file mode 100644 index 9ab10601ae59..000000000000 --- a/cpp/tensorrt_llm/plugins/identityPlugin/identityPlugin.h +++ /dev/null @@ -1,89 +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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class IdentityPlugin : public BasePlugin -{ -public: - IdentityPlugin(); - - IdentityPlugin(void const* data, size_t length); - - ~IdentityPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; -}; - -class IdentityPluginCreator : public BaseCreator -{ -public: - IdentityPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.cpp b/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.cpp deleted file mode 100644 index 02a40a00c919..000000000000 --- a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.cpp +++ /dev/null @@ -1,472 +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. - */ -#include "layernormQuantizationPlugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/layernormKernels.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::LayernormQuantizationPluginCreator; -using tensorrt_llm::plugins::LayernormQuantizationPlugin; - -static char const* LAYERNORM_QUANTIZATION_PLUGIN_VERSION{"1"}; -static char const* LAYERNORM_QUANTIZATION_PLUGIN_NAME{"LayernormQuantization"}; -PluginFieldCollection LayernormQuantizationPluginCreator::mFC{}; -std::vector LayernormQuantizationPluginCreator::mPluginAttributes; - -LayernormQuantizationPlugin::LayernormQuantizationPlugin(float eps, bool useDiffOfSquares, - bool dynamicActivationScaling, bool sumPerToken, bool clampValEnabled, tensorrt_llm::common::QuantMode quantMode, - nvinfer1::DataType type, nvinfer1::DataType outputType) - : mEps(eps) - , mUseDiffOfSquares(useDiffOfSquares) - , mDynActScaling(dynamicActivationScaling) - , mType(type) - , mOutputType(outputType) - , mClampValEnabled(clampValEnabled) - , mQuantMode(quantMode) - , mSumPerToken(sumPerToken) -{ - TLLM_CHECK_WITH_INFO(mOutputType == nvinfer1::DataType::kINT8 || mOutputType == nvinfer1::DataType::kFP8, - "Only int8 or fp8 output type is allowed."); - // Check if the quant mode is valid. - TLLM_CHECK_WITH_INFO(mQuantMode.hasPerTokenScaling(), "The quant mode is not valid."); -} - -// Parameterized constructor -LayernormQuantizationPlugin::LayernormQuantizationPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mEps); - read(d, mUseDiffOfSquares); - read(d, mDynActScaling); - read(d, mSumPerToken); - read(d, mClampValEnabled); - read(d, mQuantMode); - read(d, mType); - read(d, mOutputType); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LayernormQuantizationPlugin::clone() const noexcept -{ - auto* plugin = new LayernormQuantizationPlugin( - mEps, mUseDiffOfSquares, mDynActScaling, mSumPerToken, mClampValEnabled, mQuantMode, mType, mOutputType); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs LayernormQuantizationPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - // Quantized output - return inputs[outputIndex]; - } - - // Dynamic scaling or per-token sum if enabled - try - { - if (outputIndex == 1) - { - TLLM_CHECK(mDynActScaling); - } - else if (outputIndex == 2) - { - TLLM_CHECK(mSumPerToken); - } - else - { - TLLM_CHECK(false); - } - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims - 1; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - ret.d[ret.nbDims - 1] = exprBuilder.constant(1); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LayernormQuantizationPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int const totalPoses - = 6 + static_cast(mClampValEnabled) + static_cast(mDynActScaling) + static_cast(mSumPerToken); - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 4 + static_cast(mClampValEnabled)); - if (pos < nbInputs) - { - if (pos < 3) - { - // activatation, weight, bias - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 3) - { - // scale - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 4 && mClampValEnabled) - { - // clamp_max_v - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - else - { - auto const output_pos = pos - nbInputs; - if (output_pos == 0) - { - // Quantized output - return (inOut[pos].type == mOutputType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (output_pos == 1 && mDynActScaling) - { - // Dynamic scaling if enabled - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (output_pos == 2 && static_cast(mClampValEnabled)) - { - // Clamp value - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - - // We should never reach this point - TLLM_CHECK_WITH_INFO(false, "The input/output is not supported."); - return false; -} - -void LayernormQuantizationPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t LayernormQuantizationPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int LayernormQuantizationPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] - // weight [N, ] - // bias [N, ] - // scale_to_int [1] - // clamp_max_v [2], contains min val, and max val (optional) - // outputs - // output [M(*), N] Normalized activations, potentially with quantization applied. - // dynamic_scaling [M(*), 1] (Optional) Per-token scales if quantization is enabled. - // token_sums [M(*), 1] (Optional) Per-token sums of all the channels (before quantization). - - int64_t m64 = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) - { - m64 *= inputDesc[0].dims.d[i]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - - void const* input = inputs[0]; - void const* weight = inputs[1]; - void const* bias = inputs[2]; - void const* scale = inputs[3]; - void const* clampValPtr = mClampValEnabled ? inputs[4] : nullptr; - void* output = outputs[0]; - void* dynamic_scale = mDynActScaling ? outputs[1] : nullptr; - void* sum_per_token = mSumPerToken ? outputs[2] : nullptr; - - if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kINT8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kFP8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kINT8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, clampValPtr, - scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kFP8) - { - dispatchDataType(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kINT8) - { - dispatchDataType<__nv_bfloat16, int8_t>(nullptr, input, weight, bias, mEps, m, n, stream, mUseDiffOfSquares, - clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kFP8) - { - dispatchDataType<__nv_bfloat16, __nv_fp8_e4m3>(nullptr, input, weight, bias, mEps, m, n, stream, - mUseDiffOfSquares, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#endif // ENABLE_BF16 - sync_check_cuda_error(stream); - return 0; -} - -template -void LayernormQuantizationPlugin::dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, - float const eps, int const tokens, int const hidden_dim, cudaStream_t stream, bool use_diff_of_squares, - void const* clampValPtr, void const* scale, void* dynamic_scale, void* sum_per_token, - void* normed_output_quant) noexcept -{ - // inputs - // activation [dim0(*), dim1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [dim0(*), dim1] - // scale_tokens [dim0(*), 1] - - invokeGeneralLayerNorm(reinterpret_cast(out), reinterpret_cast(input), - reinterpret_cast(gamma), reinterpret_cast(beta), eps, tokens, hidden_dim, mQuantMode, - stream, use_diff_of_squares, reinterpret_cast(clampValPtr), reinterpret_cast(scale), - reinterpret_cast(dynamic_scale), reinterpret_cast(sum_per_token), - reinterpret_cast(normed_output_quant)); -} - -// IPluginV2Ext Methods -nvinfer1::DataType LayernormQuantizationPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index <= 2); - - if (index == 0) - { - // Output 0 quantized output of layer norm - return mOutputType; - } - else if (index == 1) - { - assert(mDynActScaling); - // Output 1 dynamic act scaling - return nvinfer1::DataType::kFLOAT; - } - else if (index == 2) - { - assert(mDynActScaling && mSumPerToken); - // Output 2 per-token sums - return nvinfer1::DataType::kFLOAT; - } - - // We should never reach this point - TLLM_CHECK_WITH_INFO(false, "The output index is not supported."); - return nvinfer1::DataType::kFLOAT; -} - -// IPluginV2 Methods - -char const* LayernormQuantizationPlugin::getPluginType() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* LayernormQuantizationPlugin::getPluginVersion() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_VERSION; -} - -int LayernormQuantizationPlugin::getNbOutputs() const noexcept -{ - return 1 + static_cast(mDynActScaling) + static_cast(mSumPerToken); -} - -int LayernormQuantizationPlugin::initialize() noexcept -{ - return 0; -} - -void LayernormQuantizationPlugin::terminate() noexcept {} - -size_t LayernormQuantizationPlugin::getSerializationSize() const noexcept -{ - return sizeof(mEps) + sizeof(mUseDiffOfSquares) + sizeof(mDynActScaling) + sizeof(mSumPerToken) - + sizeof(mClampValEnabled) + sizeof(mQuantMode) + sizeof(mType) + sizeof(mOutputType); -} - -void LayernormQuantizationPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mEps); - write(d, mUseDiffOfSquares); - write(d, mDynActScaling); - write(d, mSumPerToken); - write(d, mClampValEnabled); - write(d, mQuantMode); - write(d, mType); - write(d, mOutputType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LayernormQuantizationPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -LayernormQuantizationPluginCreator::LayernormQuantizationPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("use_diff_of_squares", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dyn_act_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sum_per_token", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("clamp_val_enabled", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("out_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LayernormQuantizationPluginCreator::getPluginName() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* LayernormQuantizationPluginCreator::getPluginVersion() const noexcept -{ - return LAYERNORM_QUANTIZATION_PLUGIN_VERSION; -} - -PluginFieldCollection const* LayernormQuantizationPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LayernormQuantizationPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - tensorrt_llm::common::QuantMode quantMode{}; - float eps{}; - nvinfer1::DataType type{}; - nvinfer1::DataType outputType{}; - bool useDiffOfSquares{}; - bool dynamicActivationScaling{}; - bool sumPerToken{}; - bool clampValEnabled{}; - - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dyn_act_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dynamicActivationScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "use_diff_of_squares")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - useDiffOfSquares = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sum_per_token")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sumPerToken = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "clamp_val_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - clampValEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "quant_mode")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - quantMode = QuantMode(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "out_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - outputType = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new LayernormQuantizationPlugin( - eps, useDiffOfSquares, dynamicActivationScaling, sumPerToken, clampValEnabled, quantMode, type, outputType); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LayernormQuantizationPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call LayernormQuantizationPlugin::destroy() - try - { - auto* obj = new LayernormQuantizationPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h b/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h deleted file mode 100644 index 5cf3fa7e022f..000000000000 --- a/cpp/tensorrt_llm/plugins/layernormQuantizationPlugin/layernormQuantizationPlugin.h +++ /dev/null @@ -1,110 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class LayernormQuantizationPlugin : public BasePlugin -{ -public: - LayernormQuantizationPlugin(float eps, bool useDiffOfSquares, bool dynamicActivationScaling, bool sumPerToken, - bool clampValEnabled, tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, - nvinfer1::DataType outputType); - - LayernormQuantizationPlugin(void const* data, size_t length); - - ~LayernormQuantizationPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - void dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, float const eps, - int const tokens, int const hidden_dim, cudaStream_t stream, bool use_diff_of_squares, void const* clampValPtr, - void const* scale, void* dynamic_scale, void* sum_per_token, void* normed_output_quant) noexcept; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - float mEps; - bool mUseDiffOfSquares; - bool mDynActScaling; - nvinfer1::DataType mType; - - const std::string mLayerName; - // The quantized output data type - nvinfer1::DataType mOutputType; - // Do we clamp the input tensor? - bool mClampValEnabled; - // The quantization mode - tensorrt_llm::common::QuantMode mQuantMode; - // Should we output the sum of channels per-token? (Used by QServe GEMM) - bool mSumPerToken; -}; - -class LayernormQuantizationPluginCreator : public BaseCreator -{ -public: - LayernormQuantizationPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/lookupPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lookupPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/lookupPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.cpp b/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.cpp deleted file mode 100644 index e4d26f9e5ec6..000000000000 --- a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.cpp +++ /dev/null @@ -1,341 +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. - */ - -#include - -#include "lookupPlugin.h" -#include "tensorrt_llm/kernels/lookupKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::LookupPluginCreator; -using tensorrt_llm::plugins::LookupPlugin; - -static char const* LOOKUP_PLUGIN_VERSION{"1"}; -static char const* LOOKUP_PLUGIN_NAME{"Lookup"}; -PluginFieldCollection LookupPluginCreator::mFC{}; -std::vector LookupPluginCreator::mPluginAttributes; - -LookupPlugin::LookupPlugin(nvinfer1::DataType type, int rank) - : mType(type) - , mRank(rank) -{ - mArch = tensorrt_llm::common::getSMVersion(); -} - -// Parameterized constructor -LookupPlugin::LookupPlugin(void const* data, size_t length) -{ - mArch = tensorrt_llm::common::getSMVersion(); - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mRank); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LookupPlugin::clone() const noexcept -{ - auto* plugin = new LookupPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - plugin->initialize(); - return plugin; -} - -nvinfer1::DimsExprs LookupPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2 || nbInputs == 3); - TLLM_CHECK(outputIndex == 0); - DimsExprs ret; - int const nbDimsInput = inputs[0].nbDims; - int const nbDimsWeight = inputs[1].nbDims; - ret.nbDims = nbDimsInput + 1; - - for (int i = 0; i < nbDimsInput; ++i) - { - ret.d[i] = inputs[0].d[i]; - } - ret.d[nbDimsInput] = inputs[1].d[nbDimsWeight - 1]; - - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LookupPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - bool res = false; - if (nbInputs == 2) - { - switch (pos) - { - case 0: res = ((inOut[0].type == DataType::kINT32) && (inOut[0].format == TensorFormat::kLINEAR)); break; - case 1: res = ((inOut[1].type == mType) && (inOut[1].format == TensorFormat::kLINEAR)); break; - case 2: res = ((inOut[2].type == mType) && (inOut[2].format == TensorFormat::kLINEAR)); break; - default: // should NOT be here! - res = false; - } - } - else - { - TLLM_CHECK_WITH_INFO(mArch == 90, "int8 weight only lookupPlugin is only supported in SM 90 now."); - switch (pos) - { - case 0: res = ((inOut[0].type == DataType::kINT32) && (inOut[0].format == TensorFormat::kLINEAR)); break; - case 1: - res = ((inOut[1].type == DataType::kINT8 || inOut[1].type == mType) - && (inOut[1].format == TensorFormat::kLINEAR)); - break; - case 2: res = ((inOut[2].type == mType) && (inOut[2].format == TensorFormat::kLINEAR)); break; - case 3: res = ((inOut[3].type == mType) && (inOut[3].format == TensorFormat::kLINEAR)); break; - default: // should NOT be here! - res = false; - } - } - return res; -} - -void LookupPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - mNbInputs = nbInputs; -} - -size_t LookupPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int LookupPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - // input [tokenNum] - // weight [localVocabSize, hidden] - // per_token_scales [localVocabSize], optional - // outputs - // embedding [tokenNum, hidden] - - int64_t tokenNum = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - tokenNum *= inputDesc[0].dims.d[i]; - } - - int const localVocabSize = inputDesc[1].dims.d[0]; - int const hidden = inputDesc[1].dims.d[inputDesc[1].dims.nbDims - 1]; - int const* input = reinterpret_cast(inputs[0]); - - int offset = mRank * localVocabSize; - - if (mNbInputs == 3) - { - int8_t const* weight = reinterpret_cast(inputs[1]); - if (mType == DataType::kHALF) - { - half const* per_token_scales = reinterpret_cast(inputs[2]); - half* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, per_token_scales, stream); - } - else if (mType == DataType::kFLOAT) - { - float const* per_token_scales = reinterpret_cast(inputs[2]); - float* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, per_token_scales, stream); - } - else if (mType == DataType::kBF16) - { - __nv_bfloat16 const* per_token_scales = reinterpret_cast<__nv_bfloat16 const*>(inputs[2]); - __nv_bfloat16* output = reinterpret_cast<__nv_bfloat16*>(outputs[0]); - invokeLookUp<__nv_bfloat16, int8_t, int>( - output, input, weight, tokenNum, offset, localVocabSize, hidden, per_token_scales, stream); - } - } - else - { - if (mType == DataType::kHALF) - { - half const* weight = reinterpret_cast(inputs[1]); - half* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, nullptr, stream); - } - else if (mType == DataType::kFLOAT) - { - float const* weight = reinterpret_cast(inputs[1]); - float* output = reinterpret_cast(outputs[0]); - invokeLookUp( - output, input, weight, tokenNum, offset, localVocabSize, hidden, nullptr, stream); - } - else if (mType == DataType::kBF16) - { - __nv_bfloat16 const* weight = reinterpret_cast<__nv_bfloat16 const*>(inputs[1]); - __nv_bfloat16* output = reinterpret_cast<__nv_bfloat16*>(outputs[0]); - invokeLookUp<__nv_bfloat16, __nv_bfloat16, int>( - output, input, weight, tokenNum, offset, localVocabSize, hidden, nullptr, stream); - } - } - sync_check_cuda_error(stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType LookupPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* LookupPlugin::getPluginType() const noexcept -{ - return LOOKUP_PLUGIN_NAME; -} - -char const* LookupPlugin::getPluginVersion() const noexcept -{ - return LOOKUP_PLUGIN_VERSION; -} - -int LookupPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int LookupPlugin::initialize() noexcept -{ - return 0; -} - -void LookupPlugin::destroy() noexcept -{ - delete this; -} - -size_t LookupPlugin::getSerializationSize() const noexcept -{ - return sizeof(mType) + sizeof(mRank); -} - -void LookupPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mRank); - - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LookupPlugin::terminate() noexcept {} - -/////////////// - -LookupPluginCreator::LookupPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("rank", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LookupPluginCreator::getPluginName() const noexcept -{ - return LOOKUP_PLUGIN_NAME; -} - -char const* LookupPluginCreator::getPluginVersion() const noexcept -{ - return LOOKUP_PLUGIN_VERSION; -} - -PluginFieldCollection const* LookupPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LookupPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int rank{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - rank = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new LookupPlugin(type, rank); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LookupPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call LookupPlugin::destroy() - try - { - auto* obj = new LookupPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h b/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h deleted file mode 100644 index 4dddaa1d8bdc..000000000000 --- a/cpp/tensorrt_llm/plugins/lookupPlugin/lookupPlugin.h +++ /dev/null @@ -1,96 +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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class LookupPlugin : public BasePlugin -{ -public: - LookupPlugin() = delete; - - LookupPlugin(nvinfer1::DataType type, int rank); - - LookupPlugin(void const* data, size_t length); - - ~LookupPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - - nvinfer1::DataType mType; - int mRank; - int mNbInputs = 0; - int mArch; -}; - -class LookupPluginCreator : public BaseCreator -{ -public: - LookupPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/loraPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp deleted file mode 100644 index 7a7d925a74f6..000000000000 --- a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.cpp +++ /dev/null @@ -1,525 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "loraPlugin.h" - -#include "pluginUtils.h" -#include "tensorrt_llm/common/assert.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::LoraPluginCreator; -using tensorrt_llm::plugins::LoraPlugin; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* LORA_PLUGIN_VERSION{"1"}; -static char const* LORA_PLUGIN_NAME{"Lora"}; -PluginFieldCollection LoraPluginCreator::mFC{}; -std::vector LoraPluginCreator::mPluginAttributes; - -LoraPlugin::LoraPlugin(int in_hidden_size, std::vector out_hidden_sizes, int transA, int transB, - int num_lora_modules, nvinfer1::DataType type, LoraPlugin::PluginProfilerPtr const& pluginProfiler, - bool remove_input_padding, int max_low_rank, int weight_index) - : mTransA(transA) - , mTransB(transB) - , mType(type) - , mRemoveInputPadding(remove_input_padding) - , mNumLoraModules(num_lora_modules) - , mInHiddenSize(in_hidden_size) - , mMaxLowRank(max_low_rank) - , mWeightIndex(weight_index) - , mPluginProfiler(pluginProfiler) -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - mOutHiddenSizes.resize(mNumLoraModules); - mOutHiddenSizes.assign(out_hidden_sizes.begin(), out_hidden_sizes.end()); - init(); -} - -// Parameterized constructor -LoraPlugin::LoraPlugin(void const* data, size_t length, LoraPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - char const *d = reinterpret_cast(data), *a = d; - read(d, mInHiddenSize); - read(d, mTransA); - read(d, mTransB); - read(d, mNumLoraModules); - read(d, mType); - read(d, mRemoveInputPadding); - read(d, mMaxLowRank); - read(d, mWeightIndex); - mOutHiddenSizes.resize(mNumLoraModules); - for (int i = 0; i < mNumLoraModules; i++) - { - read(d, mOutHiddenSizes[i]); - } - init(); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void LoraPlugin::init() -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - auto cublasWraper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - - mLoraImpl = std::make_shared( - mInHiddenSize, mOutHiddenSizes, mTransA, mTransB, mNumLoraModules, mType, mMaxLowRank, cublasWraper); - - mPluginProfiler->setTranspose(mTransA, mTransB); - mGemmId = GemmIdCublas(mDims.n, mDims.k, mType, mTransA, mTransB, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LoraPlugin::clone() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - auto* plugin = new LoraPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs LoraPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - try - { - TLLM_CHECK(outputIndex < mNumLoraModules); - int const nbDimsA = inputs[getInputTensorIdx()].nbDims; - DimsExprs ret; - ret.nbDims = nbDimsA; - - for (int i = 0; i < ret.nbDims; ++i) - { - ret.d[0] = 0; - } - - if (mTransA) - { - for (int i = 1; i < nbDimsA; ++i) - { - ret.d[i - 1] = inputs[getInputTensorIdx()].d[i]; - } - } - else - { - for (int i = 0; i < nbDimsA - 1; ++i) - { - ret.d[i] = inputs[getInputTensorIdx()].d[i]; - } - } - - auto const* outHiddenSize = exprBuilder.constant(mOutHiddenSizes.at(outputIndex)); - TLLM_CHECK(outHiddenSize != nullptr); - ret.d[ret.nbDims - 1] = outHiddenSize; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LoraPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - if (pos == getHostRequestTypesIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos >= getLoraRanksIdx() && pos < getLoraRanksIdx() + mNumLoraModules) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos >= getLoraWeightsPtrsIdx() && pos < getLoraWeightsPtrsIdx() + mNumLoraModules) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (mRemoveInputPadding && pos == getHostContextLengthsIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void LoraPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - auto const input = in[getInputTensorIdx()]; - - int const nbDimsA = input.max.nbDims; - - auto const minM = utils::computeMDimension(mTransA, input.min); - auto const maxM = utils::computeMDimension(mTransA, input.max); - auto const N = utils::computeNDimension(mTransB, in[getHostRequestTypesIdx()].max); - auto const K = static_cast(mTransA ? input.max.d[0] : input.max.d[nbDimsA - 1]); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - mGemmId.n = N; - mGemmId.k = K; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -size_t LoraPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - int const nbReq = inputs[getLoraRanksIdx()].dims.d[0]; - auto const type = inputs[getInputTensorIdx()].type; - auto const numTokens = getNumTokens(inputs); - return mLoraImpl->getWorkspaceSize(numTokens, nbReq, type); -} - -int64_t LoraPlugin::getNumTokens(nvinfer1::PluginTensorDesc const* input_tensors) const -{ - int ndim = input_tensors[getInputTensorIdx()].dims.nbDims; - TLLM_CHECK_WITH_INFO( - 3 == ndim || 2 == ndim, "hidden_state dimension should be either 2 [numTokens, hidden], or 3 [b, s, hidden]"); - int64_t num_tokens = input_tensors[getInputTensorIdx()].dims.d[0]; - if (ndim == 3) - { - num_tokens *= input_tensors[getInputTensorIdx()].dims.d[1]; - } - return num_tokens; -} - -int LoraPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - if (isBuilding()) - { - return 0; - } - - auto const numReqs = inputDesc[getLoraRanksIdx()].dims.d[0]; - void const* input = inputs[getInputTensorIdx()]; - int const seqLen = mRemoveInputPadding ? 0 : inputDesc[getInputTensorIdx()].dims.d[1]; - int32_t const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - void const* const* loraRanks = &inputs[getLoraRanksIdx()]; - void const* const* loraWeightPtrs = &inputs[getLoraWeightsPtrsIdx()]; - int32_t const* hostContextLengths - = mRemoveInputPadding ? static_cast(inputs[getHostContextLengthsIdx()]) : nullptr; - - int numTokens = getNumTokens(inputDesc); - mExpandLoraWeightPtrs.clear(); - mExpandLoraRanks.clear(); - mExpandLoraWeightPtrs.reserve(mNumLoraModules * numTokens * 2); - mExpandLoraRanks.reserve(mNumLoraModules * numTokens); - - for (int loraModuleIdx = 0; loraModuleIdx < mNumLoraModules; loraModuleIdx++) - { - auto const loraWeightModulePtrs = static_cast(loraWeightPtrs[loraModuleIdx]); - auto const loraRankModule = static_cast(loraRanks[loraModuleIdx]); - - int idx = 0; - for (int reqId = 0; reqId < numReqs; reqId++) - { - // loraWeightModulePtrs has 3 pointers for each module: A,B, and an optional DoRA magnitude - // the current LoRA plugin does not apply DoRA scaling, so the magnitude is ignored - RequestType const reqType = static_cast(reqTypes[reqId]); - if (reqType == RequestType::kGENERATION) - { - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3])); - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3 + 1])); - mExpandLoraRanks.push_back(loraRankModule[reqId]); - idx += 1; - } - else - { - int contextLen = (mRemoveInputPadding ? hostContextLengths[reqId] : seqLen); - - for (int contextId = 0; contextId < contextLen; contextId++) - { - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3])); - mExpandLoraWeightPtrs.push_back(reinterpret_cast(loraWeightModulePtrs[reqId * 3 + 1])); - mExpandLoraRanks.push_back(loraRankModule[reqId]); - idx += 1; - } - } - } - - // In 1st generation phase cross attention qkv lora, cross qkv is skipped by passing an empty encoder_output - // (passing 0 to dim) getNumTokens() will get in cross qkv_lora. Skipping the check for this case. - if (numTokens > 0) - { - TLLM_CHECK_WITH_INFO(idx == numTokens, - fmtstr("LoraParams and input dims don't match, lora tokens %d input tokens %d", idx, numTokens)); - } - } - - // only used for unified gemm - auto bestTactic = mPluginProfiler->getBestConfig(numTokens, mGemmId); - mLoraImpl->setBestTactic(bestTactic); - mLoraImpl->run(numTokens, numReqs, input, mExpandLoraRanks.data(), mExpandLoraWeightPtrs.data(), mWeightIndex, - outputs, workspace, stream); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType LoraPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - TLLM_CHECK(index < mNumLoraModules); - return mType; -} - -// IPluginV2 Methods - -char const* LoraPlugin::getPluginType() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_NAME; -} - -char const* LoraPlugin::getPluginVersion() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_VERSION; -} - -int LoraPlugin::getNbOutputs() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return mNumLoraModules; -} - -int LoraPlugin::initialize() noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - if (!mDims.isInitialized()) - { - return 0; - } - - mLoraImpl->setGemmConfig(); - - mPluginProfiler->profileTactics(mLoraImpl->getCublasWrapper(), mType, mDims, mGemmId); - return 0; -} - -void LoraPlugin::destroy() noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - delete this; -} - -size_t LoraPlugin::getSerializationSize() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return sizeof(mInHiddenSize) + sizeof(mTransA) + sizeof(mTransB) + sizeof(mNumLoraModules) + sizeof(mType) - + mPluginProfiler->getSerializationSize(mGemmId) + sizeof(mRemoveInputPadding) + sizeof(mMaxLowRank) - + sizeof(mWeightIndex) + sizeof(int) * mNumLoraModules; // selected tactics container size -} - -void LoraPlugin::serialize(void* buffer) const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - char *d = static_cast(buffer), *a = d; - write(d, mInHiddenSize); - write(d, mTransA); - write(d, mTransB); - write(d, mNumLoraModules); - write(d, mType); - write(d, mRemoveInputPadding); - write(d, mMaxLowRank); - write(d, mWeightIndex); - for (int i = 0; i < mNumLoraModules; i++) - { - write(d, mOutHiddenSizes.at(i)); - } - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LoraPlugin::terminate() noexcept {} - -/////////////// - -LoraPluginCreator::LoraPluginCreator() -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("transA", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("transB", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("num_lora_modules", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("weight_index", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LoraPluginCreator::getPluginName() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_NAME; -} - -char const* LoraPluginCreator::getPluginVersion() const noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return LORA_PLUGIN_VERSION; -} - -PluginFieldCollection const* LoraPluginCreator::getFieldNames() noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - return &mFC; -} - -IPluginV2* LoraPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int num_lora_modules{}; - int in_hidden_size{}; - int transA{}; - int transB{}; - bool remove_input_padding{}; - int max_low_rank{}; - int weight_index{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "in_hidden_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - in_hidden_size = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "transa")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transA = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "transb")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - transB = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - remove_input_padding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "max_low_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - max_low_rank = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "num_lora_modules")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - num_lora_modules = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "weight_index")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - weight_index = *(static_cast(fields[i].data)); - } - } - std::vector out_hidden_sizes; - out_hidden_sizes.resize(num_lora_modules); - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - for (int j = 0; j < num_lora_modules; j++) - { - if (!strcmp(attrName, fmtstr("out_hidden_size_%d", j).c_str())) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - out_hidden_sizes.at(j) = *(static_cast(fields[i].data)); - } - } - } - try - { - // LoraPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - auto* obj = new LoraPlugin(in_hidden_size, out_hidden_sizes, transA, transB, num_lora_modules, type, - pluginProfiler, remove_input_padding, max_low_rank, weight_index); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LoraPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - TLLM_LOG_DEBUG("%s", __PRETTY_FUNCTION__); - // This object will be deleted when the network is destroyed, which will - // call LoraPlugin::destroy() - try - { - // LoraPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // FIXME enable tactic profiler - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true, /* skip */ true); - auto* obj = new LoraPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h b/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h deleted file mode 100644 index 7795f7b7c76d..000000000000 --- a/cpp/tensorrt_llm/plugins/loraPlugin/loraPlugin.h +++ /dev/null @@ -1,159 +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. - */ -#ifndef TRT_LORA_PLUGIN_H -#define TRT_LORA_PLUGIN_H -#include "tensorrt_llm/kernels/lora/lora.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class LoraPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - using ImplPtr = std::shared_ptr; - using Config = cublasLtMatmulHeuristicResult_t; - - LoraPlugin() = delete; - - LoraPlugin(int in_hidden_size, std::vector out_hidden_sizes, int transA, int transB, int num_lora_modules, - nvinfer1::DataType type, PluginProfilerPtr const& profiler, bool remove_input_padding, int max_low_rank, - int weight_index); - - LoraPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~LoraPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - int64_t getNumTokens(nvinfer1::PluginTensorDesc const* input_tensors) const; - void init(); - - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - } - - IndexType getHostRequestTypesIdx() const - { - return 1; - } - - IndexType getLoraRanksIdx() const - { - return 2; - } - - IndexType getLoraWeightsPtrsIdx() const - { - return 2 + mNumLoraModules; - } - - IndexType getHostContextLengthsIdx() const - { - TLLM_CHECK(mRemoveInputPadding); - return 2 + mNumLoraModules + mNumLoraModules; - } - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - const std::string mLayerName; - - std::vector mOutHiddenSizes; - int mTransA; - int mTransB; - nvinfer1::DataType mType; - bool mRemoveInputPadding; - int mNumLoraModules; - int mInHiddenSize; - int mMaxLowRank; - int mWeightIndex; - - std::vector mExpandLoraWeightPtrs{}; - std::vector mExpandLoraRanks{}; - - GemmDims mDims{}; - GemmIdCublas mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - ImplPtr mLoraImpl; -}; - -class LoraPluginCreator : public BaseCreator -{ -public: - LoraPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_LORA_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/CMakeLists.txt deleted file mode 100644 index b6bd0439cc0c..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.cpp deleted file mode 100644 index 6165d6210f29..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.cpp +++ /dev/null @@ -1,425 +0,0 @@ - -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "lowLatencyGemmPlugin.h" -#include "low_latency_gemm.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaFp8Utils.h" -#include "tensorrt_llm/common/logger.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::internal_cutlass_kernels; -using tensorrt_llm::plugins::LowLatencyGemmPluginCreator; -using tensorrt_llm::plugins::LowLatencyGemmPlugin; -using tensorrt_llm::plugins::LowLatencyGemmPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* LOW_LATENCY_GEMM_PLUGIN_VERSION{"1"}; -static char const* LOW_LATENCY_GEMM_PLUGIN_NAME{"LowLatencyGemm"}; - -PluginFieldCollection LowLatencyGemmPluginCreator::mFC{}; -std::vector LowLatencyGemmPluginCreator::mPluginAttributes; - -using FP8Type = __nv_fp8_e4m3; - -static std::optional getFloatEnv(char const* name) -{ - char const* const env = std::getenv(name); - if (env == nullptr) - { - return std::nullopt; - } - try - { - float value = std::stof(env); - return {value}; - } - catch (std::invalid_argument const& e) - { - return std::nullopt; - } - catch (std::out_of_range const& e) - { - return std::nullopt; - } -}; - -void LowLatencyGemmPluginProfiler::runTactic(int m, int n, int k, LowLatencyGemmPluginProfiler::Config const& tactic, - char* workspace, cudaStream_t const& stream) -{ - - float default_pdl_overlap_ratio = 0.5; - float default_prefetch_ratio = -1.0; - FP8Type* aTmp = reinterpret_cast(workspace); - FP8Type* bTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(aTmp), m * k * sizeof(FP8Type))); - void* cTmp = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(bTmp), n * k * sizeof(FP8Type))); - size_t workspaceSize = mRunner->getWorkspaceSize(m, n, k); - char* workspaceTmp = reinterpret_cast(nextWorkspacePtr( - reinterpret_cast(cTmp), m * n * (mType == nvinfer1::DataType::kFLOAT ? sizeof(float) : sizeof(half)))); - mRunner->gemm(aTmp, bTmp, 1.0f, 0.0f, nullptr, cTmp, m, n, k, default_pdl_overlap_ratio, default_prefetch_ratio, - tactic, workspaceTmp, workspaceSize, stream); -} - -void LowLatencyGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - - std::vector workspaces = {maxM * k * sizeof(FP8Type), n * k * sizeof(FP8Type), - maxM * n * (mType == nvinfer1::DataType::kFLOAT ? sizeof(float) : sizeof(half)), - mRunner->getWorkspaceSize(maxM, n, k)}; - - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector LowLatencyGemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -LowLatencyGemmPlugin::LowLatencyGemmPlugin( - nvinfer1::DataType type, float alpha, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) - , mAplha(alpha) -{ - init(type); -} - -LowLatencyGemmPlugin::LowLatencyGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - read(d, type); - read(d, mAplha); - read(d, mDims); - init(type); - mPluginProfiler->deserialize(d, mDims, mGemmId); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void LowLatencyGemmPlugin::init(nvinfer1::DataType type) -{ - - mType = type; - - if (mType == nvinfer1::DataType::kFLOAT) - { - m_lowLatencyGemmRunner = std::make_shared>(); - } - else if (mType == nvinfer1::DataType::kHALF) - { - m_lowLatencyGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - - else if (mType == nvinfer1::DataType::kBF16) - { - m_lowLatencyGemmRunner = std::make_shared>(); - } -#endif - else - { - TLLM_THROW("Unsupported data type"); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -nvinfer1::DimsExprs LowLatencyGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - // input[1] , weights [n,k] - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LowLatencyGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have fp8 type - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - assert(false); - return false; - } -} - -void LowLatencyGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - m_workspaceMaxSize = m_lowLatencyGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t LowLatencyGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int LowLatencyGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - - // input0 activation [M,K] - // input1 weights [N,K] - // output0 [M,N] - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - int const wsSize = m_lowLatencyGemmRunner->getWorkspaceSize(m, n, k); - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid Low Latency GEMM tactic"); - - auto env_pdl_overlap_ratio = getFloatEnv("TRTLLM_PDL_OVERLAP_RATIO"); - auto env_prefetch_ratio = getFloatEnv("TRTLLM_PREFETCH_RATIO"); - auto valid_ratio = [](std::optional& env_val, float default_val) - { - if (env_val.has_value()) - { - TLLM_CHECK_WITH_INFO(env_val.value() <= 1.0f, "Valid ratio should be less than or equal to 1.0"); - return env_val.value(); - } - return default_val; - }; - float pdl_overlap_ratio = valid_ratio(env_pdl_overlap_ratio, /*default_val=*/0.5); - float prefetch_ratio = valid_ratio(env_prefetch_ratio, /*default_val=*/-1.0); - m_lowLatencyGemmRunner->gemm(const_cast(reinterpret_cast(inputs[0])), - const_cast(reinterpret_cast(inputs[1])), mAplha, 0.0F, nullptr, outputs[0], m, n, k, - pdl_overlap_ratio, prefetch_ratio, *bestTactic, reinterpret_cast(workspace), wsSize, stream); - - return 0; -} - -nvinfer1::DataType LowLatencyGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* LowLatencyGemmPlugin::getPluginType() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_NAME; -} - -char const* LowLatencyGemmPlugin::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_VERSION; -} - -int LowLatencyGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int LowLatencyGemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void LowLatencyGemmPlugin::terminate() noexcept {} - -nvinfer1::IPluginV2DynamicExt* LowLatencyGemmPlugin::clone() const noexcept -{ - auto* plugin = new LowLatencyGemmPlugin(*this); - return plugin; -} - -size_t LowLatencyGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(nvinfer1::DataType) + // dtype - sizeof(float) * 1 + // alpha - sizeof(mDims) + mPluginProfiler->getSerializationSize(mGemmId); -} - -void LowLatencyGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mAplha); - write(d, mDims); - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LowLatencyGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void LowLatencyGemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_lowLatencyGemmRunner, mType, mDims, mGemmId); -} - -LowLatencyGemmPluginCreator::LowLatencyGemmPluginCreator() -{ - - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("alpha", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LowLatencyGemmPluginCreator::getPluginName() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_NAME; -} - -char const* LowLatencyGemmPluginCreator::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* LowLatencyGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LowLatencyGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - float alpha{}; - nvinfer1::DataType type{}; - for (int i = 0; i < fc->nbFields; i++) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "alpha")) - { - - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - alpha = *(static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - - // - // GemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/false); - auto* obj = new LowLatencyGemmPlugin(type, alpha, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LowLatencyGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - try - { - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/true); - auto* obj = new LowLatencyGemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h b/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h deleted file mode 100644 index 98b8f4807174..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmPlugin/lowLatencyGemmPlugin.h +++ /dev/null @@ -1,135 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#pragma once - -#include "low_latency_gemm.h" - -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using LowLatencyGemmRunnerPtr - = std::shared_ptr; - -class LowLatencyGemmPluginProfiler - : public GemmPluginProfiler< - tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmRunnerInterface::ConfigType, - LowLatencyGemmRunnerPtr, GemmIdCore, GemmIdCoreHash> -{ - -public: - using Config = tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmRunnerInterface::ConfigType; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; -}; - -class LowLatencyGemmPlugin : public BasePlugin -{ - -public: - using PluginProfilerPtr = std::shared_ptr; - - LowLatencyGemmPlugin() = delete; - - LowLatencyGemmPlugin(nvinfer1::DataType type, float alpha, PluginProfilerPtr const& pluginProfiler); - - LowLatencyGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - ~LowLatencyGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - void configGemm(); - -private: - std::string const mLayerName; - - LowLatencyGemmRunnerPtr m_lowLatencyGemmRunner; - size_t m_workspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; - float mAplha{1.0F}; -}; - -class LowLatencyGemmPluginCreator : public BaseCreator -{ -public: - LowLatencyGemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/CMakeLists.txt deleted file mode 100644 index b6bd0439cc0c..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.cpp b/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.cpp deleted file mode 100644 index a1aa11c2f165..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.cpp +++ /dev/null @@ -1,468 +0,0 @@ - -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "lowLatencyGemmSwigluPlugin.h" -#include "low_latency_gemm_swiglu.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/cudaFp8Utils.h" -#include "tensorrt_llm/common/logger.h" -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::internal_cutlass_kernels; -using tensorrt_llm::plugins::LowLatencyGemmSwigluPluginCreator; -using tensorrt_llm::plugins::LowLatencyGemmSwigluPlugin; -using tensorrt_llm::plugins::LowLatencyGemmSwigluPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* LOW_LATENCY_GEMM_SWIGLU_PLUGIN_VERSION{"1"}; -static char const* LOW_LATENCY_GEMM_SWIGLU_PLUGIN_NAME{"LowLatencyGemmSwiglu"}; - -PluginFieldCollection LowLatencyGemmSwigluPluginCreator::mFC{}; -std::vector LowLatencyGemmSwigluPluginCreator::mPluginAttributes; - -using FP8Type = __nv_fp8_e4m3; - -static std::optional getFloatEnv(char const* name) -{ - char const* const env = std::getenv(name); - if (env == nullptr) - { - return std::nullopt; - } - try - { - float value = std::stof(env); - return {value}; - } - catch (std::invalid_argument const& e) - { - return std::nullopt; - } - catch (std::out_of_range const& e) - { - return std::nullopt; - } -}; - -static size_t getBytePerElement(nvinfer1::DataType type) -{ - size_t bpe; - if (type == nvinfer1::DataType::kFLOAT) - { - bpe = 4; - } - else if (type == nvinfer1::DataType::kHALF || type == nvinfer1::DataType::kBF16) - { - bpe = 2; - } - else if (type == nvinfer1::DataType::kINT8 || type == nvinfer1::DataType::kFP8) - { - bpe = 1; - } - else - { - TLLM_THROW("Not recognized/implemented"); - } - return bpe; -} - -void LowLatencyGemmSwigluPluginProfiler::runTactic(int m, int n, int k, - LowLatencyGemmSwigluPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - - float default_pdl_overlap_ratio = 0.5; - float default_prefetch_ratio = -1.0; - FP8Type* aTmp = reinterpret_cast(workspace); - FP8Type* bTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(aTmp), m * k * sizeof(FP8Type))); - void* dTmp = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(bTmp), n * k * sizeof(FP8Type))); - size_t workspaceSize = mRunner->getWorkspaceSize(m, n, k); - char* workspaceTmp = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(dTmp), (n / 2 * m * getBytePerElement(mType)))); - mRunner->gemm(aTmp, bTmp, 1.0f, 0.0f, 1.0f, 1.0f, nullptr, dTmp, m, n, k, default_pdl_overlap_ratio, - default_prefetch_ratio, tactic, workspaceTmp, workspaceSize, stream); -} - -int LowLatencyGemmSwigluPluginProfiler::getMaxProfileM() const -{ - return 32768; -} - -void LowLatencyGemmSwigluPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - - std::vector workspaces = {maxM * k * sizeof(FP8Type), // A - n * k * sizeof(FP8Type), // B - maxM * (n / 2) * getBytePerElement(mType), // D - mRunner->getWorkspaceSize(maxM, n, k)}; // workspace - - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector LowLatencyGemmSwigluPluginProfiler::getTactics( - int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -LowLatencyGemmSwigluPlugin::LowLatencyGemmSwigluPlugin(nvinfer1::DataType type, float scale_output, float scale_d0, - float scale_d1, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) - , mScaleOutput(scale_output) - , mScaleD0(scale_d0) - , mScaleD1(scale_d1) -{ - init(type); -} - -LowLatencyGemmSwigluPlugin::LowLatencyGemmSwigluPlugin( - void const* data, size_t length, PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - read(d, type); - read(d, mScaleOutput); - read(d, mScaleD0); - read(d, mScaleD1); - read(d, mDims); - - init(type); - mPluginProfiler->deserialize(d, mDims, mGemmId); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void LowLatencyGemmSwigluPlugin::init(nvinfer1::DataType type) -{ - - mType = type; - - if (mType == nvinfer1::DataType::kFP8) - { - mLowLatencyGemmSwigluRunner = std::make_shared>(); - } - else - { - TLLM_THROW("Unsupported data type"); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* LowLatencyGemmSwigluPlugin::clone() const noexcept -{ - auto* plugin = new LowLatencyGemmSwigluPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs LowLatencyGemmSwigluPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue() / 2); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool LowLatencyGemmSwigluPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have fp8 type - return inOut[pos].type == nvinfer1::DataType::kFP8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - TLLM_CHECK(false); - return false; - } -} - -void LowLatencyGemmSwigluPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[1]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[1]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - mWorkspaceMaxSize = mLowLatencyGemmSwigluRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t LowLatencyGemmSwigluPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return mWorkspaceMaxSize; -} - -int LowLatencyGemmSwigluPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - - // input0 activation [M,K] row-major - // input1 weights [K, N] col-major - // output0 [M,N / 2] row-major - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[1]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - int const wsSize = mLowLatencyGemmSwigluRunner->getWorkspaceSize(m, n, k); - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid Low Latency GEMM SWIGLU tactic"); - - auto env_pdl_overlap_ratio = getFloatEnv("TRTLLM_PDL_OVERLAP_RATIO"); - auto env_prefetch_ratio = getFloatEnv("TRTLLM_PREFETCH_RATIO"); - auto valid_ratio = [](std::optional& env_val, float default_val) - { - if (env_val.has_value()) - { - TLLM_CHECK_WITH_INFO(env_val.value() <= 1.0f, "Valid ratio should be less than or equal to 1.0"); - return env_val.value(); - } - return default_val; - }; - float pdl_overlap_ratio = valid_ratio(env_pdl_overlap_ratio, /*default_val=*/0.5); - float prefetch_ratio = valid_ratio(env_prefetch_ratio, /*default_val=*/-1.0); - mLowLatencyGemmSwigluRunner->gemm(const_cast(reinterpret_cast(inputs[0])), - const_cast(reinterpret_cast(inputs[1])), mScaleOutput, 0.0F, mScaleD0, mScaleD1, - nullptr, outputs[0], m, n, k, pdl_overlap_ratio, prefetch_ratio, *bestTactic, - reinterpret_cast(workspace), wsSize, stream); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType LowLatencyGemmSwigluPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* LowLatencyGemmSwigluPlugin::getPluginType() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* LowLatencyGemmSwigluPlugin::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_VERSION; -} - -int LowLatencyGemmSwigluPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int LowLatencyGemmSwigluPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void LowLatencyGemmSwigluPlugin::terminate() noexcept {} - -size_t LowLatencyGemmSwigluPlugin::getSerializationSize() const noexcept -{ - return sizeof(nvinfer1::DataType) + // dtype - sizeof(float) * 3 + // scales - sizeof(mDims) + mPluginProfiler->getSerializationSize(mGemmId); -} - -void LowLatencyGemmSwigluPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mScaleOutput); - write(d, mScaleD0); - write(d, mScaleD1); - write(d, mDims); - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void LowLatencyGemmSwigluPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void LowLatencyGemmSwigluPlugin::configGemm() -{ - mPluginProfiler->profileTactics(mLowLatencyGemmSwigluRunner, mType, mDims, mGemmId); -} - -////////////////////////////////////////////////////////////////////////// - -LowLatencyGemmSwigluPluginCreator::LowLatencyGemmSwigluPluginCreator() -{ - - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("scale_output", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_d0", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("scale_d1", nullptr, PluginFieldType::kFLOAT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* LowLatencyGemmSwigluPluginCreator::getPluginName() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_NAME; -} - -char const* LowLatencyGemmSwigluPluginCreator::getPluginVersion() const noexcept -{ - return LOW_LATENCY_GEMM_SWIGLU_PLUGIN_VERSION; -} - -PluginFieldCollection const* LowLatencyGemmSwigluPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* LowLatencyGemmSwigluPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - TLLM_CHECK(fc->nbFields == 4); - nvinfer1::DataType type{}; - float scale_output{}; - float scale_d0{}; - float scale_d1{}; - for (int i = 0; i < fc->nbFields; i++) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_output")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_output = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d0")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d0 = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "scale_d1")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - scale_d1 = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - - // - // LowLatencyGemmSwigluPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/false); - auto* obj = new LowLatencyGemmSwigluPlugin(type, scale_output, scale_d0, scale_d1, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* LowLatencyGemmSwigluPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - try - { - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/*inference=*/true); - auto* obj = new LowLatencyGemmSwigluPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h b/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h deleted file mode 100644 index 3f73324e7740..000000000000 --- a/cpp/tensorrt_llm/plugins/lowLatencyGemmSwigluPlugin/lowLatencyGemmSwigluPlugin.h +++ /dev/null @@ -1,140 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#pragma once - -#include "low_latency_gemm_swiglu.h" - -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -using LowLatencyGemmSwigluRunnerPtr - = std::shared_ptr; - -class LowLatencyGemmSwigluPluginProfiler - : public GemmPluginProfiler< - tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmSwigluRunnerInterface::ConfigType, - LowLatencyGemmSwigluRunnerPtr, GemmIdCore, GemmIdCoreHash> -{ - -public: - using Config - = tensorrt_llm::kernels::internal_cutlass_kernels::CutlassLowLatencyFp8GemmSwigluRunnerInterface::ConfigType; - - virtual int getMaxProfileM() const override; - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; -}; - -class LowLatencyGemmSwigluPlugin : public BasePlugin -{ - -public: - using PluginProfilerPtr = std::shared_ptr; - - LowLatencyGemmSwigluPlugin() = delete; - - LowLatencyGemmSwigluPlugin(nvinfer1::DataType type, float scale_output, float scale_d0, float scale_d1, - PluginProfilerPtr const& pluginProfiler); - - LowLatencyGemmSwigluPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - ~LowLatencyGemmSwigluPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - void configGemm(); - -private: - std::string const mLayerName; - - LowLatencyGemmSwigluRunnerPtr mLowLatencyGemmSwigluRunner; - size_t mWorkspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; - float mScaleOutput; - float mScaleD0; - float mScaleD1; -}; - -class LowLatencyGemmSwigluPluginCreator : public BaseCreator -{ -public: - LowLatencyGemmSwigluPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/lruPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/lruPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/lruPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.cpp b/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.cpp deleted file mode 100644 index 9d86b8cb8acd..000000000000 --- a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.cpp +++ /dev/null @@ -1,431 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "lruPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::lruPluginCreator; -using tensorrt_llm::plugins::lruPlugin; - -static char const* LRU_PLUGIN_VERSION{"1"}; -static char const* LRU_PLUGIN_NAME{"LRU"}; -PluginFieldCollection lruPluginCreator::mFC{}; -std::vector lruPluginCreator::mPluginAttributes; - -lruPlugin::lruPlugin(int dim, int block_size, nvinfer1::DataType type, bool removePadding, bool pagedState, - bool yEnabled, bool yBiasEnabled, bool fuseGateEnabled, bool gateBiasEnabled) - : mDim(dim) - , mBlockSize(block_size) - , mType(type) - , mRemovePadding(removePadding) - , mPagedState(pagedState) - , mYEnabled(yEnabled) - , mYBiasEnabled(yBiasEnabled) - , mFuseGateEnabled(fuseGateEnabled) - , mGateBiasEnabled(gateBiasEnabled) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// Parameterized constructor -lruPlugin::lruPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDim); - read(d, mBlockSize); - read(d, mType); - read(d, mRemovePadding); - read(d, mPagedState); - read(d, mYEnabled); - read(d, mYBiasEnabled); - read(d, mFuseGateEnabled); - read(d, mGateBiasEnabled); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* lruPlugin::clone() const noexcept -{ - auto* plugin = new lruPlugin(mDim, mBlockSize, mType, mRemovePadding, mPagedState, mYEnabled, mYBiasEnabled, - mFuseGateEnabled, mGateBiasEnabled); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// state: [batch_size, dim] -nvinfer1::DimsExprs lruPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - return inputs[getXIdx()]; - } - return inputs[getStateIdx()]; -} - -bool lruPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getHostRequestTypesIdx() || pos == getLastTokenIdsIdx() || (mPagedState && pos == getSlotMappingIdx())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (mPagedState && pos == getStateIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (pos == getStateIdx() || pos == (nbInputs + 1)) - { - // Use float for both input and output state - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void lruPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t lruPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -void lruPlugin::setLruParams(lruParams& params, const size_t batch, const size_t dim, const size_t block_size, - const size_t maxSeqLen, void* statePtr, void const* x, void const* gate, void const* gate_bias, void const* gate_x, - void const* gate_x_bias, void const* gate_a, void const* gate_a_bias, void const* y, void const* y_bias, - void const* A, int const* lastTokenIds, int const* slotMapping, void* out, bool removePadding) -{ - // Reset the parameters - memset(¶ms, 0, sizeof(params)); - - params.batch = batch; - params.width = dim; - params.block_size = block_size; - params.max_seqlen = maxSeqLen; - params.remove_padding = removePadding; - - // Set the pointers and strides. - params.A_ptr = const_cast(A); - params.x_ptr = const_cast(x); - params.y_ptr = const_cast(y); - params.y_bias_ptr = const_cast(y_bias); - params.gate_ptr = const_cast(gate); - params.gate_bias_ptr = const_cast(gate_bias); - params.gate_x_ptr = const_cast(gate_x); - params.gate_x_bias_ptr = const_cast(gate_x_bias); - params.gate_a_ptr = const_cast(gate_a); - params.gate_a_bias_ptr = const_cast(gate_a_bias); - params.state_ptr = statePtr; - params.out_ptr = out; - params.last_token_ids_ptr = lastTokenIds; - params.slot_mapping_ptr = slotMapping; -} - -template -int lruPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) -{ - // inputs - // 0. x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. A [dim] - // 2. state [batch_size, dim] or host [1] containing only pointer for paged_state - // 3. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. - // 4. last_token_ids [batch_size] int32 - // 5. state_slot_mapping [batch_size] int32, optional for paged state - // 6. y [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 7. y_bias [dim] - // 8. gate [batch_size, seq_len, 2 * dim] or [num_tokens, 2 * dim] for remove_input_padding - // 9. gate_bias [2 * dim] - // 10. gate_x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 11. gate_a [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 12. gate_x_bias [2 * dim] - // 13. gate_a_bias [2 * dim] - // outputs - // 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. state [batch_size, dim] - auto const batch_size = inputDesc[getHostRequestTypesIdx()].dims.d[0]; - int max_seq_len; - if (mRemovePadding) - { - max_seq_len = -1; - } - else - { - max_seq_len = inputDesc[getXIdx()].dims.d[1]; - } - - // only support context or generation, not for both of them - RequestType const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - - lruParams lru_params; - - int const* slotMapping = mPagedState ? static_cast(inputs[getSlotMappingIdx()]) : nullptr; - void const* y = mYEnabled ? inputs[getYIdx()] : nullptr; - void const* y_bias = mYBiasEnabled ? inputs[getYBiasIdx()] : nullptr; - void const* gate = mFuseGateEnabled ? inputs[getGateIdx()] : nullptr; - void const* gate_bias = (mFuseGateEnabled && mGateBiasEnabled) ? inputs[getGateBiasIdx()] : nullptr; - void const* gate_x = mFuseGateEnabled ? nullptr : inputs[getGateXIdx()]; - void const* gate_a = mFuseGateEnabled ? nullptr : inputs[getGateAIdx()]; - void const* gate_x_bias = (!mFuseGateEnabled && mGateBiasEnabled) ? inputs[getGateXBiasIdx()] : nullptr; - void const* gate_a_bias = (!mFuseGateEnabled && mGateBiasEnabled) ? inputs[getGateABiasIdx()] : nullptr; - - void* statePtr = mPagedState ? *reinterpret_cast(const_cast(inputs[getStateIdx()])) : outputs[1]; - - setLruParams(lru_params, batch_size, mDim, mBlockSize, max_seq_len, statePtr, inputs[getXIdx()], gate, gate_bias, - gate_x, gate_x_bias, gate_a, gate_a_bias, y, y_bias, inputs[getAIdx()], - static_cast(inputs[getLastTokenIdsIdx()]), slotMapping, outputs[0], mRemovePadding); - - if (reqTypes[0] == RequestType::kCONTEXT) - { - invokeRGLRU(lru_params, stream); - } - else if (reqTypes[0] == RequestType::kGENERATION) - { - invokeRGLRUUpdate(lru_params, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -int lruPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType lruPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (index == 0) - { - return inputTypes[getXIdx()]; - } - else - { - return inputTypes[getStateIdx()]; - } -} - -// IPluginV2 Methods - -char const* lruPlugin::getPluginType() const noexcept -{ - return LRU_PLUGIN_NAME; -} - -char const* lruPlugin::getPluginVersion() const noexcept -{ - return LRU_PLUGIN_VERSION; -} - -int lruPlugin::getNbOutputs() const noexcept -{ - return mPagedState ? 1 : 2; -} - -int lruPlugin::initialize() noexcept -{ - return 0; -} - -void lruPlugin::terminate() noexcept {} - -size_t lruPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDim) + sizeof(mBlockSize) + sizeof(mType) + sizeof(mRemovePadding) + sizeof(mPagedState) - + sizeof(mYEnabled) + sizeof(mYBiasEnabled) + sizeof(mFuseGateEnabled) + sizeof(mGateBiasEnabled); -} - -void lruPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDim); - write(d, mBlockSize); - write(d, mType); - write(d, mRemovePadding); - write(d, mPagedState); - write(d, mYEnabled); - write(d, mYBiasEnabled); - write(d, mFuseGateEnabled); - write(d, mGateBiasEnabled); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void lruPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -lruPluginCreator::lruPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("block_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("paged_state", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("y_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("y_bias_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("fuse_gate_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("gate_bias_enabled", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* lruPluginCreator::getPluginName() const noexcept -{ - return LRU_PLUGIN_NAME; -} - -char const* lruPluginCreator::getPluginVersion() const noexcept -{ - return LRU_PLUGIN_VERSION; -} - -PluginFieldCollection const* lruPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* lruPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int dim{}; - int block_size{}; - bool removePadding{}; - bool pagedState{}; - bool yEnabled{}; - bool yBiasEnabled{}; - bool fuseGateEnabled{}; - bool gateBiasEnabled{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "dim")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dim = static_cast(*(static_cast(fields[i].data))); - } - if (!strcmp(attrName, "block_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - block_size = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - removePadding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "paged_state")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - pagedState = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "y_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - yEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "y_bias_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - yBiasEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "fuse_gate_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - fuseGateEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "gate_bias_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - gateBiasEnabled = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new lruPlugin( - dim, block_size, type, removePadding, pagedState, yEnabled, yBiasEnabled, fuseGateEnabled, gateBiasEnabled); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* lruPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call lruPlugin::destroy() - try - { - auto* obj = new lruPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.h b/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.h deleted file mode 100644 index ee4e0b989b34..000000000000 --- a/cpp/tensorrt_llm/plugins/lruPlugin/lruPlugin.h +++ /dev/null @@ -1,239 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_LRU_PLUGIN_H -#define TRT_LRU_PLUGIN_H -#include "tensorrt_llm/kernels/lruKernel.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests or num_gen_requests -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (single sequences per request). -// can not support beam search - -// inputs -// 0. x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. A [dim] -// 2. state [batch_size, dim] or host [1] containing only pointer for paged_state -// 3. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. -// 4. last_token_ids [batch_size] int32 -// 5. state_slot_mapping [batch_size] int32, optional for paged state -// 6. y [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 7. y_bias [dim] -// 8. gate [batch_size, seq_len, 2 * dim] or [num_tokens, 2 * dim] for remove_input_padding -// 9. gate_bias [2 * dim] -// 10. gate_x [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 11. gate_a [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 12. gate_x_bias [2 * dim] -// 13. gate_a_bias [2 * dim] -// outputs -// 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. state [batch_size, dim] - -class lruPlugin : public BasePlugin -{ -public: - lruPlugin(int dim, int block_size, nvinfer1::DataType type, bool removePadding, bool pagedState, bool yEnabled, - bool yBiasEnabled, bool fuseGateEnabled, bool gateBiasEnabled); - - lruPlugin(void const* data, size_t length); - - ~lruPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - using IndexType = std::int32_t; - - IndexType getXIdx() const - { - return 0; - }; - - IndexType getAIdx() const - { - return 1; - }; - - IndexType getStateIdx() const - { - return 2; - }; - - IndexType getHostRequestTypesIdx() const - { - return 3; - }; - - IndexType getLastTokenIdsIdx() const - { - return 4; - }; - - IndexType getSlotMappingIdx() const - { - if (mPagedState) - return 5; - else - return 4; - }; - - IndexType getYIdx() const - { - if (mYEnabled) - return getSlotMappingIdx() + 1; - else - return getSlotMappingIdx(); - }; - - IndexType getYBiasIdx() const - { - if (mYBiasEnabled) - return getYIdx() + 1; - else - return getYIdx(); - }; - - IndexType getGateIdx() const - { - if (mFuseGateEnabled) - return getYBiasIdx() + 1; - else - return getYBiasIdx(); - }; - - IndexType getGateBiasIdx() const - { - if (mFuseGateEnabled && mGateBiasEnabled) - return getGateIdx() + 1; - else - return getGateIdx(); - }; - - IndexType getGateXIdx() const - { - if (mFuseGateEnabled) - return getGateBiasIdx(); - else - return getGateBiasIdx() + 1; - }; - - IndexType getGateAIdx() const - { - if (mFuseGateEnabled) - return getGateXIdx(); - else - return getGateXIdx() + 1; - }; - - IndexType getGateXBiasIdx() const - { - if (!mFuseGateEnabled && mGateBiasEnabled) - return getGateAIdx() + 1; - else - return getGateAIdx(); - }; - - IndexType getGateABiasIdx() const - { - if (!mFuseGateEnabled && mGateBiasEnabled) - return getGateXBiasIdx() + 1; - else - return getGateXBiasIdx(); - }; - - static void setLruParams(tensorrt_llm::kernels::lruParams& params, - // sizes - const size_t batch, const size_t dim, const size_t block_size, const size_t maxSeqLen, - // device pointers - void* statePtr, void const* x, void const* gate, void const* gate_bias, void const* gate_x, - void const* gate_x_bias, void const* gate_a, void const* gate_a_bias, void const* y, void const* y_bias, - void const* A, int const* lastTokenIds, int const* slotMapping, void* out, bool removePadding); - -private: - int mDim; - int mBlockSize; - nvinfer1::DataType mType; - bool mRemovePadding = false; - bool mPagedState = false; - bool mYEnabled = false; - bool mYBiasEnabled = false; - bool mFuseGateEnabled = false; - bool mGateBiasEnabled = false; -}; - -class lruPluginCreator : public BaseCreator -{ -public: - lruPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_LRU_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.cpp b/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.cpp deleted file mode 100644 index 16754248b84d..000000000000 --- a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.cpp +++ /dev/null @@ -1,404 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "mambaConv1dPlugin.h" -#include "tensorrt_llm/common/assert.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::MambaConv1dPluginCreator; -using tensorrt_llm::plugins::MambaConv1dPlugin; - -static char const* MAMBA_CONV1D_PLUGIN_VERSION{"1"}; -static char const* MAMBA_CONV1D_PLUGIN_NAME{"MambaConv1d"}; - -PluginFieldCollection MambaConv1dPluginCreator::mFC{}; -std::vector MambaConv1dPluginCreator::mPluginAttributes; - -MambaConv1dPlugin::MambaConv1dPlugin(int dim, int dconv, int preStride, int postStride, nvinfer1::DataType type, - bool removePadding, bool pagedState, bool applySilu) - : mDim(dim) - , mDConv(dconv) - , mPreStride(preStride) - , mPostStride(postStride) - , mType(type) - , mRemovePadding(removePadding) - , mPagedState(pagedState) - , mApplySilu(applySilu) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// Parameterized constructor -MambaConv1dPlugin::MambaConv1dPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDim); - read(d, mDConv); - read(d, mPreStride); - read(d, mPostStride); - read(d, mType); - read(d, mRemovePadding); - read(d, mPagedState); - read(d, mApplySilu); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* MambaConv1dPlugin::clone() const noexcept -{ - auto* plugin - = new MambaConv1dPlugin(mDim, mDConv, mPreStride, mPostStride, mType, mRemovePadding, mPagedState, mApplySilu); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// state: [batch_size, dconv - 1, dim] -nvinfer1::DimsExprs MambaConv1dPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - auto ret = inputs[getInputTensorIdx()]; - ret.d[mRemovePadding ? 1 : 2] = exprBuilder.constant(mDim); - return ret; - } - return inputs[getConvStateIdx()]; -} - -bool MambaConv1dPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getHostRequestTypesIdx() || pos == getLastTokenIdsIdx() - || (mRemovePadding && pos == getHostContextLengthIdx()) || (mPagedState && pos == getSlotMappingIdx())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (mPagedState && pos == getConvStateIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void MambaConv1dPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t MambaConv1dPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -void MambaConv1dPlugin::setMambaConv1dParams(tensorrt_llm::kernels::MambaConv1dParamsBase& params, const size_t batch, - const size_t dim, const size_t maxSeqLen, const size_t dconv, const size_t preStride, const size_t postStride, - void const* inPtr, void const* stateInPtr, void* stateOutPtr, void const* convWeight, void const* convBias, - void* outPtr, int const* lastTokenIds, int const* stateSlotMapping, bool removePadding, bool applySilu) -{ - // Reset the parameters - memset(¶ms, 0, sizeof(params)); - - params.batch = batch; - params.dim = dim; - params.max_seqlen = maxSeqLen; - params.dconv = dconv; - params.pre_stride = preStride; - params.post_stride = postStride; - - params.remove_padding = removePadding; - params.apply_silu = applySilu; - - // Set the pointers and strides. - params.in_ptr = const_cast(inPtr); - params.state_in_ptr = const_cast(stateInPtr); - params.state_out_ptr = stateOutPtr; - params.weight_ptr = const_cast(convWeight); - params.bias_ptr = const_cast(convBias); - params.out_ptr = outPtr; - params.last_token_ids_ptr = lastTokenIds; - params.state_slot_mapping_ptr = stateSlotMapping; -} - -template -int MambaConv1dPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. conv_state [batch_size, dconv - 1, dim] or host [1] containing only pointer for paged_state - // 2. weight [dim, 1, dconv] - // 3. bias [dim] - // 4. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. - // 5. last_token_ids [batch_size] int32 - // 6. host_context_lengths [batch_size] int32, optional for remove_input_padding - // 7. state_slot_mapping [batch_size] int32, optional - // outputs - // 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // 1. conv_state [batch_size, dconv - 1, dim] - auto const batchSize = inputDesc[getHostRequestTypesIdx()].dims.d[0]; - int maxSeqLen; - if (mRemovePadding) - { - int const* host_context_length = static_cast(inputs[getHostContextLengthIdx()]); - maxSeqLen = *std::max_element(host_context_length, host_context_length + batchSize); - } - else - { - maxSeqLen = inputDesc[getInputTensorIdx()].dims.d[1]; - } - - // only support context or generation, not for both of them - RequestType const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - - MambaConv1dParamsBase mambaConv1dParams; - - int const* slotMapping = mPagedState ? static_cast(inputs[getSlotMappingIdx()]) : nullptr; - void* stateInPtr = mPagedState ? *reinterpret_cast(const_cast(inputs[getConvStateIdx()])) - : const_cast(inputs[getConvStateIdx()]); - void* stateOutPtr - = mPagedState ? *reinterpret_cast(const_cast(inputs[getConvStateIdx()])) : outputs[1]; - - setMambaConv1dParams(mambaConv1dParams, batchSize, mDim, maxSeqLen, mDConv, mPreStride, mPostStride, - inputs[getInputTensorIdx()], stateInPtr, stateOutPtr, inputs[getWeightIdx()], inputs[getBiasIdx()], outputs[0], - static_cast(inputs[getLastTokenIdsIdx()]), slotMapping, mRemovePadding, mApplySilu); - - if (reqTypes[0] == RequestType::kCONTEXT) - { - invokeMambaConv1dContext(mambaConv1dParams, stream); - } - else if (reqTypes[0] == RequestType::kGENERATION) - { - invokeMambaConv1dGeneration(mambaConv1dParams, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -int MambaConv1dPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType MambaConv1dPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - return inputTypes[getInputTensorIdx()]; -} - -// IPluginV2 Methods - -char const* MambaConv1dPlugin::getPluginType() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_NAME; -} - -char const* MambaConv1dPlugin::getPluginVersion() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_VERSION; -} - -int MambaConv1dPlugin::getNbOutputs() const noexcept -{ - return 2; -} - -int MambaConv1dPlugin::initialize() noexcept -{ - return 0; -} - -void MambaConv1dPlugin::terminate() noexcept {} - -size_t MambaConv1dPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDim) + sizeof(mDConv) + sizeof(mPreStride) + sizeof(mPostStride) + sizeof(mType) - + sizeof(mRemovePadding) + sizeof(mPagedState) + sizeof(mApplySilu); -} - -void MambaConv1dPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDim); - write(d, mDConv); - write(d, mPreStride); - write(d, mPostStride); - write(d, mType); - write(d, mRemovePadding); - write(d, mPagedState); - write(d, mApplySilu); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void MambaConv1dPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -MambaConv1dPluginCreator::MambaConv1dPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dconv", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("pre_stride", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("post_stride", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("paged_state", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("apply_silu", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* MambaConv1dPluginCreator::getPluginName() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_NAME; -} - -char const* MambaConv1dPluginCreator::getPluginVersion() const noexcept -{ - return MAMBA_CONV1D_PLUGIN_VERSION; -} - -PluginFieldCollection const* MambaConv1dPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* MambaConv1dPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int dim{}; - int dconv{}; - int pre_stride{}; - int post_stride{}; - bool removePadding{}; - bool pagedState{}; - bool applySilu{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "dim")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dim = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dconv")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dconv = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "pre_stride")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - pre_stride = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "post_stride")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - post_stride = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - removePadding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "paged_state")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - pagedState = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "apply_silu")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - applySilu = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj - = new MambaConv1dPlugin(dim, dconv, pre_stride, post_stride, type, removePadding, pagedState, applySilu); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* MambaConv1dPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call MambaConv1dPlugin::destroy() - try - { - auto* obj = new MambaConv1dPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h b/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h deleted file mode 100644 index d351b1cdc237..000000000000 --- a/cpp/tensorrt_llm/plugins/mambaConv1dPlugin/mambaConv1dPlugin.h +++ /dev/null @@ -1,176 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_MAMBA_CONV1D_PLUGIN_H -#define TRT_MAMBA_CONV1D_PLUGIN_H -#include "tensorrt_llm/kernels/mambaConv1dKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests or num_gen_requests -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (single sequences per request). -// can not support beam search - -// inputs -// 0. input_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. conv_state [batch_size, dconv - 1, dim] or host [1] containing only pointer for paged_state -// 2. weight [1, dconv, dim] -// 3. bias [dim] -// 4. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. -// 5. last_token_ids [batch_size] int32 -// 6. host_context_lengths [batch_size] int32, optional for remove_input_padding -// 7. state_slot_mapping [batch_size] int32, optional -// outputs -// 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. conv_state [batch_size, dconv - 1, dim] - -class MambaConv1dPlugin : public BasePlugin -{ -public: - MambaConv1dPlugin(int dim, int dconv, int preStride, int postStride, nvinfer1::DataType type, bool removePadding, - bool pagedState, bool applySilu); - - MambaConv1dPlugin(void const* data, size_t length); - - ~MambaConv1dPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - - IndexType getConvStateIdx() const - { - return 1; - }; - - IndexType getWeightIdx() const - { - return 2; - }; - - IndexType getBiasIdx() const - { - return 3; - }; - - IndexType getHostRequestTypesIdx() const - { - return 4; - }; - - IndexType getLastTokenIdsIdx() const - { - return 5; - }; - - IndexType getHostContextLengthIdx() const - { - return 6; - }; - - IndexType getSlotMappingIdx() const - { - // if not remove input padding, host_context_length is not used, so the index is 6 - return mRemovePadding ? 7 : 6; - }; - - void setMambaConv1dParams(tensorrt_llm::kernels::MambaConv1dParamsBase& params, - // sizes - const size_t batch, const size_t dim, const size_t maxSeqLen, const size_t dconv, const size_t preStride, - const size_t postStride, - // device pointers - void const* inPtr, void const* stateInPtr, void* stateOutPtr, void const* convWeight, void const* convBias, - void* outPtr, int const* lastTokenIds, int const* stateSlotMapping, bool removePadding, bool applySilu); - -private: - int mDim; - int mDConv; - int mPreStride; - int mPostStride; - nvinfer1::DataType mType; - bool mRemovePadding = false; - bool mPagedState = false; - bool mApplySilu = true; -}; - -class MambaConv1dPluginCreator : public BaseCreator -{ -public: - MambaConv1dPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_MAMBA_CONV1D_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/CMakeLists.txt b/cpp/tensorrt_llm/plugins/mixtureOfExperts/CMakeLists.txt deleted file mode 100644 index 7cc985b60b7a..000000000000 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/CMakeLists.txt +++ /dev/null @@ -1,21 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp deleted file mode 100644 index ccce34850730..000000000000 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.cpp +++ /dev/null @@ -1,1314 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h" -#include "tensorrt_llm/common/cudaBf16Wrapper.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/runtime/iBuffer.h" -#include "tensorrt_llm/runtime/utils/debugUtils.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::plugins; -using tensorrt_llm::common::QuantMode; -using tensorrt_llm::common::nextWorkspacePtr; -using tensorrt_llm::common::calculateTotalWorkspaceSize; -using tensorrt_llm::plugins::MixtureOfExpertsPluginCreator; -using tensorrt_llm::plugins::MixtureOfExpertsPlugin; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -using LoraImpl = tensorrt_llm::kernels::LoraImpl; -using LoraParams = tensorrt_llm::kernels::LoraParams; - -static char const* MIXTURE_OF_EXPERTS_PLUGIN_VERSION{"1"}; -static char const* MIXTURE_OF_EXPERTS_PLUGIN_NAME{"MixtureOfExperts"}; -nvinfer1::PluginFieldCollection MixtureOfExpertsPluginCreator::mFC{}; -std::vector MixtureOfExpertsPluginCreator::mPluginAttributes; - -MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(bool remove_input_padding, int number_of_experts, int experts_per_token, - int expert_hidden_size, int expert_inter_size, int groupwise_quant_algo, int group_size, - ActivationType activation_type, nvinfer1::DataType type, nvinfer1::DataType weight_type, - nvinfer1::DataType output_type, QuantMode quant_mode, bool use_final_scales, bool use_bias, int tp_size, - int tp_rank, int ep_size, int ep_rank, bool force_determinism, int side_stream_id, - MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, nvinfer1::DataType lora_type, - LoraPluginProfilerPtr lora_profiler, int max_low_rank) - : mNumExperts(number_of_experts) - , mExpertsPerToken(experts_per_token) - , mExpertHiddenSize(expert_hidden_size) - , mExpertInterSize(expert_inter_size) - , mGroupwiseQuantAlgo(groupwise_quant_algo) - , mGroupSize(group_size) - , mActivationType(activation_type) - , mType(type) - , mWeightType(weight_type) - , mOutputType(output_type) - , mQuantMode(quant_mode) - , mUseFinalScales(use_final_scales) - , mUseBias(use_bias) - , mParallelismConfig(MOEParallelismConfig{tp_size, tp_rank, ep_size, ep_rank}) - , mUseDeterministicKernels(force_determinism) - , mSideStreamId(side_stream_id) - , mGemmProfiler(std::move(gemm_profiler_ptr)) - , mUseLora(use_lora) - , mLoraType(lora_type) - , mMaxLowRank(max_low_rank) - , mRemoveInputPadding(remove_input_padding) - , mLoraProfiler(std::move(lora_profiler)) -{ - init(); -} - -tensorrt_llm::plugins::MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(MixtureOfExpertsPlugin const& other) - : mMOERunner() - , mNumExperts(other.mNumExperts) - , mExpertsPerToken(other.mExpertsPerToken) - , mExpertHiddenSize(other.mExpertHiddenSize) - , mExpertInterSize(other.mExpertInterSize) - , mGroupwiseQuantAlgo(other.mGroupwiseQuantAlgo) - , mGroupSize(other.mGroupSize) - , mActivationType(other.mActivationType) - , mType(other.mType) - , mWeightType(other.mWeightType) - , mOutputType(other.mOutputType) - , mQuantMode(other.mQuantMode) - , mUseFinalScales(other.mUseFinalScales) - , mUseBias(other.mUseBias) - , mParallelismConfig(other.mParallelismConfig) - , mDims(other.mDims) - , mUseDeterministicKernels(other.mUseDeterministicKernels) - , mSideStreamId(other.mSideStreamId) - , mGemmId1(other.mGemmId1) - , mGemmId2(other.mGemmId2) - , mGemmProfiler(other.mGemmProfiler) - , mUseLora(other.mUseLora) - , mLoraType(other.mLoraType) - , mMaxLowRank(other.mMaxLowRank) - , mRemoveInputPadding(other.mRemoveInputPadding) - , mLoraImpl1(other.mLoraImpl1) - , mLoraImpl2(other.mLoraImpl2) - , mLoraGemmId1(other.mLoraGemmId1) - , mLoraGemmId2(other.mLoraGemmId2) - , mLoraProfiler(other.mLoraProfiler) - , mLayerName(other.mLayerName) - , mNamespace(other.mNamespace) -{ - init(); -} - -size_t MixtureOfExpertsPlugin::getSerializationSize() const noexcept -{ - size_t size = sizeof(mRemoveInputPadding) + sizeof(mNumExperts) + sizeof(mExpertsPerToken) - + sizeof(mExpertHiddenSize) + sizeof(mExpertInterSize) + sizeof(mGroupwiseQuantAlgo) + sizeof(mGroupSize) - + sizeof(mActivationType) + sizeof(mType) + sizeof(mWeightType) + sizeof(mOutputType) - + sizeof(QuantMode::BaseType) + sizeof(mUseFinalScales) + sizeof(mUseBias) + sizeof(mParallelismConfig) - + sizeof(mDims) + sizeof(mUseDeterministicKernels) + sizeof(mSideStreamId) - + mGemmProfiler->getSerializationSize(mGemmId1) + mGemmProfiler->getSerializationSize(mGemmId2) - + sizeof(mUseLora) + sizeof(mLoraType) + sizeof(mMaxLowRank); - - if (hasLora()) - { - size += mLoraProfiler->getSerializationSize(mLoraGemmId1); - size += mLoraProfiler->getSerializationSize(mLoraGemmId2); - } - - return size; -} - -MixtureOfExpertsPlugin::MixtureOfExpertsPlugin(void const* data, size_t length, - MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, LoraPluginProfilerPtr lora_profiler) - : mGemmProfiler(gemm_profiler_ptr) - , mLoraProfiler(lora_profiler) -{ - char const* d = reinterpret_cast(data); - char const* a = d; - read(d, mRemoveInputPadding); - read(d, mNumExperts); - read(d, mExpertsPerToken); - read(d, mExpertHiddenSize); - read(d, mExpertInterSize); - read(d, mGroupwiseQuantAlgo); - read(d, mGroupSize); - read(d, mActivationType); - read(d, mType); - read(d, mWeightType); - read(d, mOutputType); - QuantMode::BaseType quant_mode; - read(d, quant_mode); - mQuantMode = QuantMode{quant_mode}; - read(d, mUseFinalScales); - read(d, mUseBias); - read(d, mParallelismConfig); - read(d, mDims); - read(d, mUseDeterministicKernels); - read(d, mSideStreamId); - read(d, mUseLora); - read(d, mLoraType); - read(d, mMaxLowRank); - - // Call init before deserialising the profiler to initialize mGemmId - init(); - mGemmProfiler->deserialize(d, mDims, mGemmId1); - mGemmProfiler->deserialize(d, mDims, mGemmId2); - - if (hasLora()) - { - mLoraProfiler->deserialize(d, mDims, mLoraGemmId1); - mLoraProfiler->deserialize(d, mDims, mLoraGemmId2); - } - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void MixtureOfExpertsPlugin::serialize(void* buffer) const noexcept -{ - char* d = static_cast(buffer); - char* a = d; - - write(d, mRemoveInputPadding); - write(d, mNumExperts); - write(d, mExpertsPerToken); - write(d, mExpertHiddenSize); - write(d, mExpertInterSize); - write(d, mGroupwiseQuantAlgo); - write(d, mGroupSize); - write(d, mActivationType); - write(d, mType); - write(d, mWeightType); - write(d, mOutputType); - write(d, mQuantMode.value()); - write(d, mUseFinalScales); - write(d, mUseBias); - write(d, mParallelismConfig); - write(d, mDims); - write(d, mUseDeterministicKernels); - write(d, mSideStreamId); - write(d, mUseLora); - write(d, mLoraType); - write(d, mMaxLowRank); - - mGemmProfiler->serialize(d, mGemmId1); - mGemmProfiler->serialize(d, mGemmId2); - - if (hasLora()) - { - mLoraProfiler->serialize(d, mLoraGemmId1); - mLoraProfiler->serialize(d, mLoraGemmId2); - } - - TLLM_CHECK(d == a + getSerializationSize()); -} - -template -std::unique_ptr switch_output_type(nvinfer1::DataType output_type) -{ - switch (output_type) - { - case nvinfer1::DataType::kFP4: - case nvinfer1::DataType::kFP8: - // TODO We need an atomic FP8 reduction for the finalize fusions - TLLM_THROW("Outputting %d directly is not currently supported", static_cast(output_type)); - // return std::make_unique>(); - case nvinfer1::DataType::kHALF: - if constexpr (NeedQuant) - { - return std::make_unique>(); - } - else - { - return std::make_unique>(); - } -#ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: - if constexpr (NeedQuant) - { - return std::make_unique>(); - } - else - { - return std::make_unique>(); - } -#endif - default: TLLM_THROW("Invalid output type %d", static_cast(output_type)); - } -}; - -void MixtureOfExpertsPlugin::init() -{ - TLLM_CHECK_WITH_INFO(mType == DataType::kFP8 || mType == DataType::kFP4 || mOutputType == mType, - "MOE plugin only supports a different output type for FP4/FP8"); - TLLM_CHECK_WITH_INFO(mType != DataType::kFP8 || tensorrt_llm::common::getSMVersion() >= 89, - "MoE FP8 is not supported for architectures less than SM89"); - TLLM_CHECK_WITH_INFO(mType != DataType::kFP4 || (tensorrt_llm::common::getSMVersion() >= 100), - "MoE FP4 is only supported on architecture SM100 or later"); - - TLLM_CHECK_WITH_INFO(!hasLora() || mLoraType == mOutputType, "The LoraType need to keep same with moe OutputType."); - - if (mWeightType == nvinfer1::DataType::kINT8 && mQuantMode.hasInt4Weights()) - { - mWeightType = DataType::kINT4; - } - - if (mType == DataType::kHALF && mWeightType == DataType::kHALF) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kFLOAT && mWeightType == DataType::kFLOAT) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kHALF && mWeightType == DataType::kINT8) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kHALF && mWeightType == DataType::kINT4) - { - mMOERunner = std::make_unique>(); - } -#ifdef ENABLE_FP8 - else if (mType == DataType::kFP8 && mWeightType == DataType::kINT4 && mOutputType == DataType::kHALF) - { - mMOERunner = std::make_unique>(); - } -#endif -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16 && mWeightType == DataType::kBF16) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kBF16 && mWeightType == DataType::kINT8) - { - mMOERunner = std::make_unique>(); - } - else if (mType == DataType::kBF16 && mWeightType == DataType::kINT4) - { - mMOERunner = std::make_unique>(); - } -#ifdef ENABLE_FP8 - else if (mType == DataType::kFP8 && mWeightType == DataType::kINT4 && mOutputType == DataType::kBF16) - { - mMOERunner = std::make_unique< - kernels::CutlassMoeFCRunner<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat16, __nv_bfloat16>>(); - } -#endif -#endif - -#ifdef ENABLE_FP8 - if (mType == DataType::kFP8 && mWeightType == DataType::kFP8) - { - mMOERunner = switch_output_type<__nv_fp8_e4m3>(mOutputType); - } -#endif -#ifdef ENABLE_FP4 - if (mType == DataType::kFP4 && mWeightType == DataType::kFP4) - { - mMOERunner = switch_output_type<__nv_fp4_e2m1, true>(mOutputType); - } -#endif - - if (!mMOERunner) - { - TLLM_THROW( - "Could not construct the mixture of experts plugin with the requested input combination Activation: %d " - "Weight: %d Output: %d", - static_cast(mType), static_cast(mWeightType), static_cast(mOutputType)); - } - - // Finalize fusion should be disabled if Lora is used. - mMOERunner->use_fused_finalize_ - = (mExpertsPerToken < 3 || !mUseDeterministicKernels) && !getEnvMOEDisableFinalizeFusion() && !hasLora(); - - mGemmId1 = GemmIDMoe{1, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - mGemmId2 = GemmIDMoe{2, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - mGemmProfiler->setMaxProfileM(16384 * mNumExperts / mExpertsPerToken); - - if (hasLora()) - { - auto cublasHandle = getCublasHandle(); - auto cublasLtHandle = getCublasLtHandle(); - auto cublasWrapper = std::make_shared(cublasHandle, cublasLtHandle, nullptr, nullptr); - mLoraGemmId1 = GemmIdCublas(mExpertInterSize, mExpertHiddenSize, mLoraType, false, true, mLoraType); - mLoraGemmId2 = GemmIdCublas(mExpertHiddenSize, mExpertInterSize, mLoraType, false, true, mLoraType); - std::vector loraOutSizes1 = {static_cast(mExpertInterSize)}; - mLoraImpl1 = std::make_shared( - mExpertHiddenSize, loraOutSizes1, false, true, 1, mLoraType, mMaxLowRank, cublasWrapper); - std::vector loraOutSizes2 = {static_cast(mExpertHiddenSize)}; - mLoraImpl2 = std::make_shared( - mExpertInterSize, loraOutSizes2, false, true, 1, mLoraType, mMaxLowRank, cublasWrapper); - - TLLM_CUDA_CHECK(cudaEventCreate(&mMemcpyEvent)); - } - mSideStreamPtr = nullptr; - mDebugStallMain = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_MAIN"); - mDebugStallSide = tensorrt_llm::runtime::utils::stallStream("TLLM_DEBUG_MOE_STALL_SIDE"); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* MixtureOfExpertsPlugin::clone() const noexcept -{ - auto* plugin = new MixtureOfExpertsPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs MixtureOfExpertsPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - assert(outputIndex == getOutputTensorIndex() || outputIndex == getOutputDummyTensorIndex()); - return inputs[getInputTensorIndex()]; -} - -bool MixtureOfExpertsPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - TLLM_CHECK(0 <= pos && pos < getNbInputs() + getNbOutputs()); - TLLM_CHECK_WITH_INFO( - nbInputs == getNbInputs(), "Required input to plugin is missing. Expected %d Got %d", getNbInputs(), nbInputs); - TLLM_CHECK_WITH_INFO(nbOutputs == getNbOutputs(), "Required output to plugin is missing. Expected %d Got %d", - getNbOutputs(), nbOutputs); - - if (inOut[pos].format != TensorFormat::kLINEAR) - { - return false; - } - - if (pos == getExpertWeights1Index() || pos == getExpertWeights2Index()) - { - if (mGroupwiseQuantAlgo == 0) - { - auto normalized_weight_type - = mWeightType == nvinfer1::DataType::kINT4 ? nvinfer1::DataType::kINT8 : mWeightType; - return inOut[pos].type == normalized_weight_type; - } - else - { - return inOut[pos].type == mOutputType; - } - } - else if (pos == getTokenSelectedExpertsIndex()) - { - return inOut[pos].type == DataType::kINT32; - } - else if (pos == getTokenFinalScalesIndex()) - { - return inOut[pos].type == DataType::kFLOAT; - } - else if (pos == getExpertBias1Index() || pos == getExpertBias2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (pos == nbInputs + getOutputTensorIndex()) - { - return inOut[pos].type == mOutputType; - } - else if (useSideStream() && pos == nbInputs + getOutputDummyTensorIndex()) - { - return inOut[pos].type == inOut[getInputDummyTensorIndex()].type; - } - else if (useSideStream() && pos == getInputDummyTensorIndex()) - { - return true; - } - else if (hasExpertFp8QuantScales() && getExpertFP8Dequant1Index() <= pos && pos <= getExpertFP8QuantFinalIndex()) - { - return inOut[pos].type == DataType::kFLOAT; - } - else if (hasExpertIntQuantScales() && getExpertIntQuantScale1Index() <= pos - && pos <= getExpertIntQuantScale2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (hasFP4QuantScales() && getFP4GlobalActSF1Index() <= pos && pos <= getFP4GlobalSF2Index()) - { - if (pos == getFP4WeightSF1Index() || pos == getFP4WeightSF2Index()) - return inOut[pos].type == nvinfer1::DataType::kFP8; - else - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (hasLora() && hasExpertFp8QuantScales() && pos == getInputFP8DequantIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kFLOAT; - } - else if (hasExpertWeightQuantZeros() && getExpertIntQuantZeros1Index() <= pos - && pos <= getExpertIntQuantZeros2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (hasExpertPrequantScales() && getExpertPrequantScales1Index() <= pos - && pos <= getExpertPrequantScales2Index()) - { - return inOut[pos].type == mOutputType; - } - else if (hasGroupwiseFp8Alpha() && getExpertFp8Alpha1Index() <= pos && pos <= getExpertFp8Alpha2Index()) - { - return inOut[pos].type == DataType::kFLOAT; - } - else if (hasLora() && pos == getHostRequestTypeIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (hasLora() && (pos == getLoraFC1RanksIndex() || pos == getLoraFC2RanksIndex())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (hasGatedLoraWeightsAndRanks() && pos == getLoraGatedRanksIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (hasLora() && (pos == getLoraFC1WeightPtrsIndex() || pos == getLoraFC2WeightPtrsIndex())) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (hasGatedLoraWeightsAndRanks() && pos == getLoraGatedWeightPtrsIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else if (hasLora() && mRemoveInputPadding && pos == getHostContextLengthIndex()) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if ((hasFP4QuantScales() || hasGroupwiseFp8Alpha()) && pos == getInputTensorIndex()) - { - return inOut[pos].type == mOutputType; - } - else - { - return inOut[pos].type == mType; - } - - return false; -} - -void MixtureOfExpertsPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - TLLM_CHECK_WITH_INFO( - nbInputs == getNbInputs(), "Required input to plugin is missing. Expected %d Got %d", getNbInputs(), nbInputs); - TLLM_CHECK_WITH_INFO(nbOutputs == getNbOutputs(), "Required output to plugin is missing. Expected %d Got %d", - getNbOutputs(), nbOutputs); - - auto in_tensor = in[getInputTensorIndex()]; - - auto const minM - = std::accumulate(in_tensor.min.d, in_tensor.min.d + in_tensor.min.nbDims - 1, 1, std::multiplies()); - auto const maxM - = std::accumulate(in_tensor.max.d, in_tensor.max.d + in_tensor.max.nbDims - 1, 1, std::multiplies()); - - auto weights_1 = in[getExpertWeights1Index()]; - auto weights_2 = in[getExpertWeights2Index()]; - int inner_dim_idx = getGemmShapeInnerDimIndex(); - int const maxK = weights_1.max.d[inner_dim_idx]; - int const maxN = weights_2.max.d[inner_dim_idx]; - int const minK = weights_1.min.d[inner_dim_idx]; - int const minN = weights_2.min.d[inner_dim_idx]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - TLLM_CHECK_WITH_INFO(maxK == mExpertHiddenSize && maxN == mExpertInterSize, - "Configured tensor sizes %dx%d does not match constructor param size %ldx%ld", maxK, maxN, mExpertHiddenSize, - mExpertInterSize); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - - mGemmId1 = GemmIDMoe{1, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - mGemmId2 = GemmIDMoe{2, mNumExperts, mExpertsPerToken, mParallelismConfig, mExpertHiddenSize, mExpertInterSize, - mGroupSize, mActivationType, mType, mWeightType, mQuantMode, !mMOERunner->use_fused_finalize_}; - - if (hasLora()) - { - auto const N = utils::computeNDimension(true, in[getHostRequestTypeIndex()].max); - mLoraGemmId1 = GemmIdCublas(N, mExpertHiddenSize, mLoraType, false, true, mLoraType); - mLoraGemmId2 = GemmIdCublas(N, mExpertInterSize, mLoraType, false, true, mLoraType); - } -} - -auto MixtureOfExpertsPlugin::setupWorkspace(void* base_ptr, int64_t num_tokens, int num_reqs) const -> WorkspaceInfo -{ - size_t moe_workspace_size - = mMOERunner->getWorkspaceSize(num_tokens, mExpertHiddenSize, mExpertInterSize, mNumExperts, mExpertsPerToken, - mActivationType, mParallelismConfig, hasLora(), /*use_deepseek_fp8_block_scale=*/false, - /*min_latency_mode=*/false, hasExpertPrequantScales()); - - // Permutation map - size_t src_to_dest_map_size = mExpertsPerToken * num_tokens * sizeof(int); - - size_t lora_workspace_size = 0; - if (hasLora()) - { - int64_t num_reqs_lora = std::min(num_tokens * mExpertsPerToken, static_cast(num_reqs * mNumExperts)); - lora_workspace_size - = std::max(mLoraImpl1->getWorkspaceSize(num_tokens * mExpertsPerToken, num_reqs_lora, mLoraType), - mLoraImpl2->getWorkspaceSize(num_tokens * mExpertsPerToken, num_reqs_lora, mLoraType)); - } - - std::vector workspaces{ - moe_workspace_size, - src_to_dest_map_size, - lora_workspace_size, - }; - - WorkspaceInfo info{}; - info.size = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - - if (base_ptr) - { - info.workspace = base_ptr; - info.src_to_dest_map = nextWorkspacePtr((int8_t*) info.workspace, moe_workspace_size); - info.lora_workspace = nextWorkspacePtr((int8_t*) info.src_to_dest_map, src_to_dest_map_size); - } - - return info; -} - -int64_t MixtureOfExpertsPlugin::getNumTokens(nvinfer1::PluginTensorDesc const* input_tensors) const -{ - int ndim = input_tensors[getInputTensorIndex()].dims.nbDims; - TLLM_CHECK_WITH_INFO( - 3 == ndim || 2 == ndim, "hidden_state dimension should be either 2 [b*s, hidden], or 3 [b, s, hidden]"); - int64_t num_tokens = input_tensors[getInputTensorIndex()].dims.d[0]; - if (ndim == 3) - { - num_tokens *= input_tensors[getInputTensorIndex()].dims.d[1]; - } - return num_tokens; -} - -size_t MixtureOfExpertsPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - TLLM_CHECK_WITH_INFO( - nbInputs == getNbInputs(), "Required input to plugin is missing. Expected %d Got %d", getNbInputs(), nbInputs); - TLLM_CHECK_WITH_INFO(nbOutputs == getNbOutputs(), "Required output to plugin is missing. Expected %d Got %d", - getNbOutputs(), nbOutputs); - - if (useSideStream()) - { - return 0; - } - int const num_tokens = getNumTokens(inputs); - int const num_lora_reqs = getNumLoraRequests(inputs); - return setupWorkspace(nullptr, num_tokens, num_lora_reqs).size; -} - -MOEParallelismConfig MixtureOfExpertsPlugin::getParallelismConfig() const -{ - return mParallelismConfig; -} - -QuantParams tensorrt_llm::plugins::MixtureOfExpertsPlugin::getQuantParams(nvinfer1::PluginTensorDesc const* inputDesc, - void const* const* inputs, int scale_1_idx, int scale_2_idx, int scale_3_idx, int scale_4_idx, int scale_5_idx, - int scale_6_idx, int scale_7_idx, int scale_8_idx) const -{ - void const* scale_1 = scale_1_idx >= 0 ? inputs[scale_1_idx] : nullptr; - void const* scale_2 = scale_2_idx >= 0 ? inputs[scale_2_idx] : nullptr; - void const* scale_3 = scale_3_idx >= 0 ? inputs[scale_3_idx] : nullptr; - void const* scale_4 = scale_4_idx >= 0 ? inputs[scale_4_idx] : nullptr; - void const* scale_5 = scale_5_idx >= 0 ? inputs[scale_5_idx] : nullptr; - void const* scale_6 = scale_6_idx >= 0 ? inputs[scale_6_idx] : nullptr; - void const* scale_7 = scale_7_idx >= 0 ? inputs[scale_7_idx] : nullptr; - void const* scale_8 = scale_8_idx >= 0 ? inputs[scale_8_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_1 = scale_1_idx >= 0 ? &inputDesc[scale_1_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_2 = scale_2_idx >= 0 ? &inputDesc[scale_2_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_3 = scale_3_idx >= 0 ? &inputDesc[scale_3_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_4 = scale_4_idx >= 0 ? &inputDesc[scale_4_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_5 = scale_5_idx >= 0 ? &inputDesc[scale_5_idx] : nullptr; - nvinfer1::PluginTensorDesc const* desc_6 = scale_6_idx >= 0 ? &inputDesc[scale_6_idx] : nullptr; - auto const gated_inter_size = isGatedActivation(mActivationType) ? mExpertInterSize * 2 : mExpertInterSize; - auto const experts_per_node = mNumExperts / mParallelismConfig.ep_size; - if (hasExpertIntQuantScales()) - { - TLLM_CHECK(scale_1 && scale_2); - if (!hasGroupwiseIntQuantScales()) - { - TLLM_CHECK(!scale_3 && !scale_4 && !scale_5 && !scale_6); - TLLM_CHECK(desc_1->dims.nbDims == 2); - TLLM_CHECK(desc_2->dims.nbDims == 2); - TLLM_CHECK_WITH_INFO( - desc_1->dims.d[0] == experts_per_node, "Incorrect number of experts in int quant scale"); - TLLM_CHECK(desc_1->dims.d[1] == gated_inter_size); - TLLM_CHECK_WITH_INFO( - desc_2->dims.d[0] == experts_per_node, "Incorrect number of experts in int quant scale"); - TLLM_CHECK(desc_2->dims.d[1] == mExpertHiddenSize); - return QuantParams::Int(scale_1, scale_2); - } - else - { - TLLM_CHECK(desc_1->dims.nbDims == 3); - TLLM_CHECK(desc_2->dims.nbDims == 3); - TLLM_CHECK((scale_3 && scale_4) || !hasExpertPrequantScales()); - TLLM_CHECK((scale_5 && scale_6) || !hasExpertWeightQuantZeros()); - TLLM_CHECK((scale_7 && scale_8) || !hasGroupwiseFp8Alpha()); - return QuantParams::GroupWise(mGroupSize, scale_1, scale_2, scale_3, scale_4, scale_5, scale_6, - static_cast(scale_7), static_cast(scale_8)); - } - } - else if (hasExpertFp8QuantScales()) - { - TLLM_CHECK(scale_1 && scale_2 && scale_3); - TLLM_CHECK(scale_4 || !hasExpertFp8FinalQuantScales()); - TLLM_CHECK((scale_5 != nullptr) == hasLora()); - TLLM_CHECK(!scale_6); - TLLM_CHECK(desc_1->dims.nbDims == 2); - TLLM_CHECK(desc_2->dims.nbDims == 1); - TLLM_CHECK(desc_3->dims.nbDims == 2); - TLLM_CHECK_WITH_INFO( - desc_1->dims.d[0] == experts_per_node && desc_1->dims.d[1] == 1, "Incorrect shape for weight FP8 scale"); - TLLM_CHECK(desc_2->dims.d[0] == 1); - TLLM_CHECK_WITH_INFO( - desc_3->dims.d[0] == experts_per_node && desc_3->dims.d[1] == 1, "Incorrect shape for weight FP8 scale"); - return QuantParams::FP8(static_cast(scale_1), static_cast(scale_2), - static_cast(scale_3), static_cast(scale_4), static_cast(scale_5)); - } - else if (hasFP4QuantScales()) - { - TLLM_CHECK(scale_1 && scale_2 && scale_3 && scale_4 && scale_5 && scale_6); - TLLM_CHECK(desc_1->dims.nbDims == 1); - TLLM_CHECK(desc_2->dims.nbDims == 3); - TLLM_CHECK(desc_3->dims.nbDims == 1); - TLLM_CHECK(desc_4->dims.nbDims == 1); - TLLM_CHECK(desc_5->dims.nbDims == 3); - TLLM_CHECK(desc_6->dims.nbDims == 1); - TLLM_CHECK(desc_1->dims.d[0] == 1); - TLLM_CHECK_WITH_INFO(desc_2->dims.d[0] == experts_per_node && desc_2->dims.d[1] == gated_inter_size - && desc_2->dims.d[2] - == mExpertHiddenSize / TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize, - "Incorrect shape for FP4 scale"); - TLLM_CHECK_WITH_INFO(desc_3->dims.d[0] == experts_per_node, "Incorrect shape for FP4 scale"); - TLLM_CHECK(desc_4->dims.d[0] == 1); - TLLM_CHECK_WITH_INFO(desc_5->dims.d[0] == experts_per_node && desc_5->dims.d[1] == mExpertHiddenSize - && desc_5->dims.d[2] - == mExpertInterSize / TmaWarpSpecializedGroupedGemmInput::NVFP4BlockScaleVectorSize, - "Incorrect shape for FP4 scale"); - TLLM_CHECK_WITH_INFO(desc_6->dims.d[0] == experts_per_node, "Incorrect shape for FP4 scale"); - return QuantParams::FP4(static_cast(scale_1), - static_cast(scale_2), - static_cast(scale_3), static_cast(scale_4), - static_cast(scale_5), - static_cast(scale_6)); - } - return {}; -} - -int MixtureOfExpertsPlugin::getNumLoraRequests(nvinfer1::PluginTensorDesc const* input_tensors) const -{ - if (!hasLora()) - return 0; - int num_reqs = input_tensors[getLoraFC1RanksIndex()].dims.d[0]; - return num_reqs; -} - -LoraParams MixtureOfExpertsPlugin::getLoraParams( - nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace) -{ - TLLM_CHECK(hasLora()); - - int const num_reqs = getNumLoraRequests(inputDesc); - int64_t const num_tokens = getNumTokens(inputDesc); - bool is_gated_actiation = isGatedActivation(mActivationType); - - mLoraExpandFC1WeightPtrs.clear(); - mLoraExpandFC2WeightPtrs.clear(); - mLoraExpandFC1Ranks.clear(); - mLoraExpandFC2Ranks.clear(); - - mLoraExpandFC1WeightPtrs.reserve(num_tokens * 2); - mLoraExpandFC2WeightPtrs.reserve(num_tokens * 2); - mLoraExpandFC1Ranks.reserve(num_tokens); - mLoraExpandFC2Ranks.reserve(num_tokens); - - if (is_gated_actiation) - { - mLoraExpandGatedWeightPtrs.clear(); - mLoraExpandGatedRanks.clear(); - mLoraExpandGatedWeightPtrs.reserve(num_tokens * 2); - mLoraExpandGatedRanks.reserve(num_tokens); - } - - int const seq_len = mRemoveInputPadding ? 0 : inputDesc[getInputTensorIndex()].dims.d[1]; - int32_t const* req_types = static_cast(inputs[getHostRequestTypeIndex()]); - int32_t const* host_context_lens - = mRemoveInputPadding ? static_cast(inputs[getHostContextLengthIndex()]) : nullptr; - - auto const fc1_lora_weight_ptrs = static_cast(inputs[getLoraFC1WeightPtrsIndex()]); - auto const fc1_lora_ranks = static_cast(inputs[getLoraFC1RanksIndex()]); - - auto const fc2_lora_weight_ptrs = static_cast(inputs[getLoraFC2WeightPtrsIndex()]); - auto const fc2_lora_ranks = static_cast(inputs[getLoraFC2RanksIndex()]); - - auto const gated_lora_weight_ptrs - = is_gated_actiation ? static_cast(inputs[getLoraGatedWeightPtrsIndex()]) : nullptr; - auto const gated_lora_ranks - = is_gated_actiation ? static_cast(inputs[getLoraGatedRanksIndex()]) : nullptr; - - int idx = 0; - for (int req_id = 0; req_id < num_reqs; req_id++) - { - RequestType const reqType = static_cast(req_types[req_id]); - if (reqType == RequestType::kGENERATION) - { - // lora_weight_ptrs has 3 pointers for each module: A,B, and an optional DoRA magnitude - // the current LoRA implementation does not apply DoRA scaling, so the magnitude is ignored - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC1Ranks.push_back(fc1_lora_ranks[req_id]); - - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC2Ranks.push_back(fc2_lora_ranks[req_id]); - - if (is_gated_actiation) - { - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3]); - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandGatedRanks.push_back(gated_lora_ranks[req_id]); - } - - idx += 1; - } - else - { - int context_len = (mRemoveInputPadding ? host_context_lens[req_id] : seq_len); - - for (int context_id = 0; context_id < context_len; context_id++) - { - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC1WeightPtrs.push_back(fc1_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC1Ranks.push_back(fc1_lora_ranks[req_id]); - - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3]); - mLoraExpandFC2WeightPtrs.push_back(fc2_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandFC2Ranks.push_back(fc2_lora_ranks[req_id]); - - if (is_gated_actiation) - { - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3]); - mLoraExpandGatedWeightPtrs.push_back(gated_lora_weight_ptrs[req_id * 3 + 1]); - mLoraExpandGatedRanks.push_back(gated_lora_ranks[req_id]); - } - } - idx += context_len; - } - } - - TLLM_CHECK_WITH_INFO(idx == num_tokens, fmtstr("idx %d num_tokens %ld", idx, num_tokens)); - - return LoraParams(num_reqs, mLoraExpandFC1Ranks.data(), mLoraExpandFC1WeightPtrs.data(), mLoraExpandFC2Ranks.data(), - mLoraExpandFC2WeightPtrs.data(), mLoraImpl1, mLoraImpl2, workspace, &mMemcpyEvent, mLoraExpandGatedRanks.data(), - mLoraExpandGatedWeightPtrs.data()); -} - -int MixtureOfExpertsPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace_ptr, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - - int64_t const num_tokens = getNumTokens(inputDesc); - int64_t const num_reqs = getNumLoraRequests(inputDesc); - - if (useSideStream()) - { - // Prepare the side stream - if (!mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - nvinfer1::pluginInternal::SideStream side_stream{}; - mSideStreamPtr = reinterpret_cast( - getPluginRegistry()->acquirePluginResource(resource_name.c_str(), &side_stream)); - } - // Debug the code with the main stream stalled (only executed when the environment variable - // TLLM_DEBUG_MOE_STALL_MAIN is set and has a positive value) - mSideStreamPtr->stallMainStream("TLLM_DEBUG_MOE_STALL_MAIN", stream, mDebugStallMain); - // The side stream waits for the inputs managed by the main stream to be ready - mSideStreamPtr->waitMainStreamOnSideStream(stream); - // Provide data dependency for the shared experts running after this plugin by copying inputs on the main stream - size_t count = 1; - for (int i = 0; i < inputDesc[getInputDummyTensorIndex()].dims.nbDims; ++i) - { - count *= inputDesc[getInputDummyTensorIndex()].dims.d[i]; - } - count *= tensorrt_llm::runtime::BufferDataType(inputDesc[getInputDummyTensorIndex()].type).getSize(); - TLLM_CUDA_CHECK(cudaMemcpyAsync(outputs[getOutputDummyTensorIndex()], inputs[getInputDummyTensorIndex()], count, - cudaMemcpyDeviceToDevice, stream)); - // Switch from the main stream to the side stream - stream = mSideStreamPtr->getStream(); - // The workspace is managed by the side stream (otherwise, the lifetime of workspace may be incorrect) - auto const workspace_size = setupWorkspace(nullptr, num_tokens, num_reqs).size; - workspace_ptr = mSideStreamPtr->getWorkspacePtr(workspace_size); - } - auto workspace = setupWorkspace(workspace_ptr, num_tokens, num_reqs); - - auto w1_desc = inputDesc[getExpertWeights1Index()]; - auto w2_desc = inputDesc[getExpertWeights2Index()]; - TLLM_CHECK(w1_desc.dims.nbDims == 3); - auto const experts_per_node = mNumExperts / mParallelismConfig.ep_size; - TLLM_CHECK(w1_desc.dims.d[0] == experts_per_node); - TLLM_CHECK(w2_desc.dims.nbDims == 3); - TLLM_CHECK(w2_desc.dims.d[0] == experts_per_node); - - auto [inner_packed_elements, outer_packed_elements] = getWeightPackedElements(); - int inner_dim_idx = getGemmShapeInnerDimIndex(); - int outer_dim_idx = getGemmShapeOuterDimIndex(); - TLLM_CHECK(w1_desc.dims.d[inner_dim_idx] * inner_packed_elements == mExpertHiddenSize); - if (isGatedActivation(mActivationType)) - { - TLLM_CHECK(w1_desc.dims.d[outer_dim_idx] * outer_packed_elements == mExpertInterSize * 2); - } - else - { - TLLM_CHECK(w1_desc.dims.d[outer_dim_idx] * outer_packed_elements == mExpertInterSize); - } - - TLLM_CHECK(w2_desc.dims.d[inner_dim_idx] * inner_packed_elements == mExpertInterSize); - TLLM_CHECK(w2_desc.dims.d[outer_dim_idx] * outer_packed_elements == mExpertHiddenSize); - - QuantParams quant_params{}; - if (hasExpertIntQuantScales()) - { - if (mGroupSize > 0) - { - quant_params = getQuantParams(inputDesc, inputs, getExpertIntQuantScale1Index(), - getExpertIntQuantScale2Index(), hasExpertPrequantScales() ? getExpertPrequantScales1Index() : -1, - hasExpertPrequantScales() ? getExpertPrequantScales2Index() : -1, - hasExpertWeightQuantZeros() ? getExpertIntQuantZeros1Index() : -1, - hasExpertWeightQuantZeros() ? getExpertIntQuantZeros2Index() : -1, - hasGroupwiseFp8Alpha() ? getExpertFp8Alpha1Index() : -1, - hasGroupwiseFp8Alpha() ? getExpertFp8Alpha2Index() : -1); - } - else - { - quant_params - = getQuantParams(inputDesc, inputs, getExpertIntQuantScale1Index(), getExpertIntQuantScale2Index()); - } - } - else if (hasExpertFp8QuantScales()) - { - quant_params = getQuantParams(inputDesc, inputs, // - getExpertFP8Dequant1Index(), // - getExpertFP8Quant2Index(), // - getExpertFP8Dequant2Index(), // - hasExpertFp8FinalQuantScales() ? getExpertFP8QuantFinalIndex() : -1, - hasLora() ? getInputFP8DequantIndex() : -1); - } - else if (hasFP4QuantScales()) - { - quant_params = getQuantParams(inputDesc, inputs, // - getFP4GlobalActSF1Index(), // - getFP4WeightSF1Index(), // - getFP4GlobalSF1Index(), // - getFP4GlobalActSF2Index(), // - getFP4WeightSF2Index(), // - getFP4GlobalSF2Index() // - ); - } - - LoraParams lora_params{}; - - if (hasLora()) - { - lora_params = getLoraParams(inputDesc, inputs, workspace.lora_workspace); - auto lora_gemm1 = mLoraProfiler->getBestConfig(num_tokens, mLoraGemmId1); - auto lora_gemm2 = mLoraProfiler->getBestConfig(num_tokens, mLoraGemmId2); - - mLoraImpl1->setBestTactic(lora_gemm1); - mLoraImpl2->setBestTactic(lora_gemm2); - } - - std::optional gemm1; - std::optional gemm2; - if (common::getEnvForceDeterministicMOE()) - { - gemm1 = mMOERunner->getTactics(MoeGemmId::GEMM_1)[0]; - gemm2 = mMOERunner->getTactics(MoeGemmId::GEMM_2)[0]; - } - else - { - gemm1 = mGemmProfiler->getBestConfig(num_tokens, mGemmId1); - gemm2 = mGemmProfiler->getBestConfig(num_tokens, mGemmId2); - } - - MoeMinLatencyParams min_latency_params{}; - mMOERunner->setTactic(gemm1, gemm2); -#ifdef USING_OSS_CUTLASS_MOE_GEMM - mMOERunner->runMoe(inputs[getInputTensorIndex()], nullptr, true, - static_cast(inputs[getTokenSelectedExpertsIndex()]), - hasFinalScales() ? static_cast(inputs[getTokenFinalScalesIndex()]) : nullptr, - inputs[getExpertWeights1Index()], hasBias() ? inputs[getExpertBias1Index()] : nullptr, - ActivationParams(mActivationType), inputs[getExpertWeights2Index()], - hasBias() ? inputs[getExpertBias2Index()] : nullptr, quant_params, num_tokens, num_tokens, mExpertHiddenSize, - mExpertHiddenSize /*TRT does not support padding, safe to assume padded/unpadded hidden sizes are the same*/, - mExpertInterSize, mNumExperts, mExpertsPerToken, static_cast(workspace.workspace), - // Outputs - outputs[getOutputTensorIndex()], static_cast(workspace.src_to_dest_map), mParallelismConfig, - /*enable_alltoall=*/false, hasLora(), lora_params, /*use_deepseek_fp8_block_scale=*/false, - /*min_latency_mode=*/false, min_latency_params, stream); -#else - mMOERunner->runMoe(inputs[getInputTensorIndex()], nullptr, true, - static_cast(inputs[getTokenSelectedExpertsIndex()]), - hasFinalScales() ? static_cast(inputs[getTokenFinalScalesIndex()]) : nullptr, - inputs[getExpertWeights1Index()], hasBias() ? inputs[getExpertBias1Index()] : nullptr, - ActivationParams(mActivationType), inputs[getExpertWeights2Index()], - hasBias() ? inputs[getExpertBias2Index()] : nullptr, quant_params, num_tokens, num_tokens, mExpertHiddenSize, - mExpertInterSize, mNumExperts, mExpertsPerToken, static_cast(workspace.workspace), - // Outputs - outputs[getOutputTensorIndex()], static_cast(workspace.src_to_dest_map), mParallelismConfig, hasLora(), - lora_params, /*use_deepseek_fp8_block_scale=*/false, - /*min_latency_mode=*/false, min_latency_params, stream); -#endif - - if (useSideStream()) - { - // Debug the code with the side stream stalled (only executed when the environment variable - // TLLM_DEBUG_MOE_STALL_SIDE is set and has a positive value) - mSideStreamPtr->stallSideStream("TLLM_DEBUG_MOE_STALL_SIDE", mDebugStallSide); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType MixtureOfExpertsPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == getOutputTensorIndex() || index == getOutputDummyTensorIndex()); - if (useSideStream() && index == getOutputDummyTensorIndex()) - { - return inputTypes[getInputDummyTensorIndex()]; - } - return mOutputType; -} - -// IPluginV2 Methods -char const* MixtureOfExpertsPlugin::getPluginType() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_NAME; -} - -char const* MixtureOfExpertsPlugin::getPluginVersion() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_VERSION; -} - -int MixtureOfExpertsPlugin::initialize() noexcept -{ - mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_1); - mGemmProfiler->profileTactics(this, mType, mDims, mGemmId1); - mGemmProfiler->setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile::GEMM_2); - mGemmProfiler->profileTactics(this, mType, mDims, mGemmId2); - - if (hasLora()) - { - mLoraImpl1->setGemmConfig(); - mLoraImpl2->setGemmConfig(); - - mLoraProfiler->profileTactics(mLoraImpl1->getCublasWrapper(), mType, mDims, mLoraGemmId1); - mLoraProfiler->profileTactics(mLoraImpl2->getCublasWrapper(), mType, mDims, mLoraGemmId2); - } - return 0; -} - -void MixtureOfExpertsPlugin::terminate() noexcept -{ - if (mSideStreamPtr) - { - auto const resource_name = nvinfer1::pluginInternal::SideStream::getResourceKey(mSideStreamId); - getPluginRegistry()->releasePluginResource(resource_name.c_str()); - mSideStreamPtr = nullptr; - } -} - -void MixtureOfExpertsPlugin::destroy() noexcept -{ - if (hasLora()) - { - TLLM_CUDA_CHECK(cudaEventDestroy(mMemcpyEvent)); - } - // This gets called when the network containing plugin is destroyed - delete this; -} - -void MixtureOfExpertsPlugin::setPluginNamespace(char const* libNamespace) noexcept -{ - mNamespace = libNamespace; -} - -char const* MixtureOfExpertsPlugin::getPluginNamespace() const noexcept -{ - return mNamespace.c_str(); -} - -/////////////// - -char const* MixtureOfExpertsPluginCreator::getPluginName() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_NAME; -} - -char const* MixtureOfExpertsPluginCreator::getPluginVersion() const noexcept -{ - return MIXTURE_OF_EXPERTS_PLUGIN_VERSION; -} - -nvinfer1::PluginFieldCollection const* MixtureOfExpertsPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -MixtureOfExpertsPluginCreator::MixtureOfExpertsPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(nvinfer1::PluginField("remove_input_padding", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("number_of_experts", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("experts_per_token", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("expert_hidden_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("expert_inter_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("groupwise_quant_algo", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("group_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("activation_type", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("weight_type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("use_final_scales", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("use_bias", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("tp_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("tp_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("ep_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("ep_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("side_stream_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("use_lora", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("lora_type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(nvinfer1::PluginField("max_low_rank", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -IPluginV2* MixtureOfExpertsPluginCreator::createPlugin( - char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept -{ - nvinfer1::PluginField const* fields = fc->fields; - int mRemoveInputPadding{}; - int mNumExperts{}; - int mExpertsPerToken{}; - int mExpertHiddenSize{}; - int mExpertInterSize{}; - int mGroupwiseQuantAlgo{}; - int mGroupSize{}; - int mActivationType{}; - int mType{}; - int mWeightType{}; - int mOutputType{INT_MAX}; - int mQuantMode{}; - int mUseFinalScales{1}; // Default to true - int mUseBias{0}; - int mTPSize{}; - int mTPRank{}; - int mEPSize{}; - int mEPRank{}; - int mRequiresDeterminism{0}; - int mSideStreamId{0}; - int mUseLora{}; - int mLoraType{INT_MAX}; - int mMaxLowRank{0}; - - // Read configurations from each fields - struct MapPair - { - char const* key; - int& field; - bool optional = false; - bool set = false; - }; - - std::array input_map{ - MapPair{"remove_input_padding", std::ref(mRemoveInputPadding)}, - MapPair{"number_of_experts", std::ref(mNumExperts)}, - MapPair{"experts_per_token", std::ref(mExpertsPerToken)}, - MapPair{"expert_hidden_size", std::ref(mExpertHiddenSize)}, - MapPair{"expert_inter_size", std::ref(mExpertInterSize)}, - MapPair{"groupwise_quant_algo", std::ref(mGroupwiseQuantAlgo)}, - MapPair{"group_size", std::ref(mGroupSize)}, - MapPair{"activation_type", std::ref(mActivationType)}, - MapPair{"type_id", std::ref(mType)}, - MapPair{"weight_type_id", std::ref(mWeightType)}, - MapPair{"quant_mode", std::ref(mQuantMode)}, - MapPair{"tp_size", std::ref(mTPSize)}, - MapPair{"tp_rank", std::ref(mTPRank)}, - MapPair{"ep_size", std::ref(mEPSize)}, - MapPair{"ep_rank", std::ref(mEPRank)}, - MapPair{"use_lora", std::ref(mUseLora)}, - MapPair{"use_final_scales", std::ref(mUseFinalScales)}, - - // Optional - MapPair{"use_bias", std::ref(mUseBias), true}, - MapPair{"output_type_id", std::ref(mOutputType), true}, - MapPair{"force_determinism", std::ref(mRequiresDeterminism), true}, - MapPair{"side_stream_id", std::ref(mSideStreamId), true}, - MapPair{"lora_type_id", std::ref(mLoraType), true}, - MapPair{"max_low_rank", std::ref(mMaxLowRank), true}, - }; - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - for (auto& item : input_map) - { - if (!strcmp(item.key, attrName)) - { - TLLM_CHECK(fields[i].type == nvinfer1::PluginFieldType::kINT32); - TLLM_CHECK_WITH_INFO(!item.set, "Parameter %s was set twice", item.key); - item.field = static_cast(*(static_cast(fields[i].data))); - item.set = true; - } - } - } - - for (auto& item : input_map) - { - TLLM_CHECK_WITH_INFO(item.set || item.optional, "Parameter %s is required but not set", item.key); - } - - // Output type is optional, if not set it to the same as mType - if (mOutputType == INT_MAX) - { - mOutputType = mType; - } - - if (mUseLora) - { - TLLM_CHECK_WITH_INFO(mLoraType != INT_MAX && mMaxLowRank != 0, - "MoE fuse lora, lora_type_id and max_low_rank are required but not set"); - } - - try - { - auto gemmProfiler = moePluginProfiler.createGemmPluginProfiler(/* inference */ false); - auto loraProfiler = loraPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - auto* obj = new MixtureOfExpertsPlugin( - // Constructor parameters - mRemoveInputPadding, mNumExperts, mExpertsPerToken, mExpertHiddenSize, mExpertInterSize, - mGroupwiseQuantAlgo, mGroupSize, static_cast(mActivationType), - static_cast(mType), static_cast(mWeightType), - static_cast(mOutputType), QuantMode(mQuantMode), mUseFinalScales != 0, mUseBias != 0, - mTPSize, mTPRank, mEPSize, mEPRank, mRequiresDeterminism != 0, mSideStreamId, gemmProfiler, mUseLora != 0, - static_cast(mLoraType), loraProfiler, mMaxLowRank); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* MixtureOfExpertsPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call MixtureOfExpertsPlugin::destroy() - try - { - auto gemmProfiler = moePluginProfiler.createGemmPluginProfiler(/* inference */ true); - auto loraProfiler = loraPluginProfileManager.createGemmPluginProfiler(/* inference */ false, /* skip */ true); - - auto* obj = new MixtureOfExpertsPlugin( - // Constructor parameters - serialData, serialLength, gemmProfiler, loraProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -void MixtureOfExpertsPluginCreator::setPluginNamespace(char const* libNamespace) noexcept -{ - mNamespace = libNamespace; -} - -char const* MixtureOfExpertsPluginCreator::getPluginNamespace() const noexcept -{ - return mNamespace.c_str(); -} - -void MixtureOfExpertsGemmProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - checkInit(); - size_t bytes = backend.getWorkspaceSize(maxM); - this->setTmpWorkspaceSizeInBytes(bytes); -} - -void MixtureOfExpertsGemmProfiler::runTactic(int m, int n, int k, MixtureOfExpertsGemmProfiler::Config const& tactic, - char* workspace_ptr_char, cudaStream_t const& stream) -{ - checkInit(); - backend.runProfiler(m, tactic, workspace_ptr_char, /*expert_weights*/ nullptr, stream); -} - -auto MixtureOfExpertsGemmProfiler::getTactics(int m, int n, int k) const -> std::vector -{ - assert(mRunner); - return mRunner->mMOERunner->getTactics(backend.mGemmToProfile); -} - -void MixtureOfExpertsGemmProfiler::initTmpData( - int m, int n, int k, char* workspace, size_t ws_size, cudaStream_t stream) -{ - checkInit(); - backend.prepare(m, workspace, /*expert_weights*/ nullptr, stream); -} - -void MixtureOfExpertsGemmProfiler::checkInit() -{ - assert(mRunner); - if (init_backend) - { - return; - } - init_backend = true; - auto& plugin = *mRunner; -#ifdef USING_OSS_CUTLASS_MOE_GEMM - backend.init(*plugin.mMOERunner, backend.mGemmToProfile, plugin.mType, plugin.mWeightType, plugin.mOutputType, - plugin.mNumExperts, plugin.mExpertsPerToken, plugin.mExpertHiddenSize, - plugin.mExpertHiddenSize /*TRT backend does not support unpadded hidden size*/, plugin.mExpertInterSize, - plugin.mGroupSize, plugin.mActivationType, plugin.hasBias(), plugin.hasLora(), /*min_latency_mode=*/false, - /*need_weights=*/true, plugin.getParallelismConfig(), /*enable_alltoall=*/false); -#else - backend.init(*plugin.mMOERunner, backend.mGemmToProfile, plugin.mType, plugin.mWeightType, plugin.mOutputType, - plugin.mNumExperts, plugin.mExpertsPerToken, plugin.mExpertHiddenSize, plugin.mExpertInterSize, - plugin.mGroupSize, plugin.mActivationType, plugin.hasBias(), plugin.hasLora(), /*min_latency_mode=*/false, - /*need_weights=*/true, plugin.getParallelismConfig()); -#endif -} diff --git a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h b/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h deleted file mode 100644 index feb1f10cdc70..000000000000 --- a/cpp/tensorrt_llm/plugins/mixtureOfExperts/mixtureOfExpertsPlugin.h +++ /dev/null @@ -1,637 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-2023 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. - */ -#ifndef TRT_MIXTURE_OF_EXPERTS_PLUGIN_H -#define TRT_MIXTURE_OF_EXPERTS_PLUGIN_H - -#include "NvInferPlugin.h" -#include "tensorrt_llm/kernels/cutlass_kernels/include/cutlass_kernel_selector.h" -#if defined(USING_OSS_CUTLASS_MOE_GEMM) -#include "tensorrt_llm/kernels/cutlass_kernels/include/moe_kernels.h" -#else -#include "moe_kernels.h" -#endif -#include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/lora/lora.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/cudaStreamPlugin/cudaStreamPlugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" -#include "tensorrt_llm/runtime/cudaStream.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -namespace kernels = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE; -using MoeMinLatencyParams = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::MoeMinLatencyParams; -using MOEParallelismConfig = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::MOEParallelismConfig; -using QuantParams = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::QuantParams; -using MoeGemmId = CUTLASS_MOE_GEMM_NAMESPACE::MoeGemmId; -using ActivationType = CUTLASS_MOE_GEMM_NAMESPACE::ActivationType; -using ActivationParams = CUTLASS_MOE_GEMM_KERNELS_NAMESPACE::ActivationParams; -using TmaWarpSpecializedGroupedGemmInput = CUTLASS_MOE_GEMM_NAMESPACE::TmaWarpSpecializedGroupedGemmInput; -using CUTLASS_MOE_GEMM_NAMESPACE::isGatedActivation; - -class MixtureOfExpertsGemmProfiler; -using MixtureOfExpertsPluginProfilerPtr = std::shared_ptr; -using GroupwiseQuantAlgo = tensorrt_llm::common::GroupwiseQuantAlgo; - -struct GemmIDMoe -{ - int gemm_idx; - int num_experts{}; - int experts_per_token{}; - kernels::MOEParallelismConfig parallelism_config{}; - int64_t hidden{}; - int64_t inter{}; - int64_t group_size{}; - ActivationType actfn{}; - nvinfer1::DataType dtype{}; - nvinfer1::DataType wdtype{}; - tensorrt_llm::common::QuantMode quant_mode; - bool determinism_mode = false; - - bool operator==(GemmIDMoe const& id) const - { - return id.gemm_idx == gemm_idx && id.num_experts == num_experts && id.experts_per_token == experts_per_token - && id.parallelism_config == parallelism_config && id.hidden == hidden && id.inter == inter - && id.group_size == group_size && id.actfn == actfn && id.dtype == dtype && id.wdtype == wdtype - && id.quant_mode == quant_mode && id.determinism_mode == determinism_mode; - } - - friend std::ostream& operator<<(std::ostream& out, GemmIDMoe const& id) - { - out << "gemm idx, experts, experts_per_token, parallelism_config, hidden, inter, group_size, actfn, dtype, " - "weight " - "type, parallelism mode, determinism mode=" - - << id.gemm_idx << "," << id.num_experts << "," << id.experts_per_token << "," << id.parallelism_config - << "," << id.hidden << "," << id.inter << "," << id.group_size << "," << static_cast(id.actfn) << "," - << static_cast(id.dtype) << "," << static_cast(id.wdtype) << "," << id.quant_mode.value() << "," - << id.determinism_mode; - return out; - } -}; - -// Hash of GemmIDMoe -struct GemmIDMoeHash -{ - std::size_t operator()(GemmIDMoe const& id) const - { - size_t hash = std::hash{}(id.gemm_idx); - hash ^= std::hash{}(id.num_experts); - hash ^= std::hash{}(id.experts_per_token); - hash ^= std::hash{}(id.parallelism_config.tp_size); - hash ^= std::hash{}(id.parallelism_config.ep_size); - hash ^= std::hash{}(id.parallelism_config.tp_rank); - hash ^= std::hash{}(id.parallelism_config.ep_rank); - hash ^= std::hash{}(id.hidden); - hash ^= std::hash{}(id.inter); - hash ^= std::hash{}(id.group_size); - hash ^= std::hash{}(static_cast(id.actfn)); - hash ^= std::hash{}(static_cast(id.dtype)); - hash ^= std::hash{}(static_cast(id.wdtype)); - hash ^= std::hash{}(static_cast(id.quant_mode.value())); - return hash; - } -}; - -class MixtureOfExpertsPlugin : public nvinfer1::IPluginV2DynamicExt -{ -public: - using LoraPluginProfilerPtr = std::shared_ptr; - using LoraImplPtr = std::shared_ptr; - MixtureOfExpertsPlugin() = delete; - MixtureOfExpertsPlugin(bool remove_input_padding, int number_of_experts, int experts_per_token, - int expert_hidden_size, int expert_inter_size, int groupwise_quant_algo, int group_size, - ActivationType activation_type, nvinfer1::DataType type, nvinfer1::DataType weight_type, - nvinfer1::DataType output_type, tensorrt_llm::common::QuantMode quant_mode, bool use_final_scales, - bool use_bias, int tp_size, int tp_rank, int ep_size, int ep_rank, bool force_determinism, int side_stream_id, - MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, bool use_lora, nvinfer1::DataType lora_type, - LoraPluginProfilerPtr lora_profiler, int max_low_rank); - MixtureOfExpertsPlugin(void const* data, size_t length, MixtureOfExpertsPluginProfilerPtr gemm_profiler_ptr, - LoraPluginProfilerPtr lora_profiler); - MixtureOfExpertsPlugin(MixtureOfExpertsPlugin const&); - - void init(); - - ~MixtureOfExpertsPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - - int getNbOutputs() const noexcept override - { - return 1 + useSideStream(); - } - - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - void setPluginNamespace(char const* pluginNamespace) noexcept override; - char const* getPluginNamespace() const noexcept override; - -private: - friend class MixtureOfExpertsGemmProfiler; - std::unique_ptr mMOERunner{}; - int mNumExperts{}; - int mExpertsPerToken{}; - int64_t mExpertHiddenSize{}; - int64_t mExpertInterSize{}; - int64_t mGroupwiseQuantAlgo{}; - int64_t mGroupSize{}; - ActivationType mActivationType; - nvinfer1::DataType mType{}; - nvinfer1::DataType mWeightType{}; - nvinfer1::DataType mOutputType{}; - tensorrt_llm::common::QuantMode mQuantMode; - bool mUseFinalScales{}; - bool mUseBias{}; - MOEParallelismConfig mParallelismConfig{}; - - GemmDims mDims{}; - bool mUseDeterministicKernels = false; - int mSideStreamId = 0; - - int mDebugStallMain = 0; - int mDebugStallSide = 0; - - GemmIDMoe mGemmId1{}; - GemmIDMoe mGemmId2{}; - - MixtureOfExpertsPluginProfilerPtr mGemmProfiler; - - // lora related - bool mUseLora{}; - nvinfer1::DataType mLoraType{}; - int mMaxLowRank{}; - bool mRemoveInputPadding{}; - - LoraImplPtr mLoraImpl1; - LoraImplPtr mLoraImpl2; - - GemmIdCublas mLoraGemmId1{}; - GemmIdCublas mLoraGemmId2{}; - LoraPluginProfilerPtr mLoraProfiler; - - std::vector mLoraExpandFC1WeightPtrs{}; - std::vector mLoraExpandFC2WeightPtrs{}; - std::vector mLoraExpandGatedWeightPtrs{}; - std::vector mLoraExpandFC1Ranks{}; - std::vector mLoraExpandFC2Ranks{}; - std::vector mLoraExpandGatedRanks{}; - - cudaEvent_t mMemcpyEvent; - nvinfer1::pluginInternal::SideStream* mSideStreamPtr; - - // The below are not serialised - std::string const mLayerName{}; - std::string mNamespace{}; - - struct WorkspaceInfo - { - void* workspace{}; - void* src_to_dest_map{}; - void* lora_workspace{}; - size_t size{}; - }; - - int64_t getNumTokens(nvinfer1::PluginTensorDesc const* input_tensor) const; - WorkspaceInfo setupWorkspace(void* base_ptr, int64_t num_tokens, int num_reqs = 0) const; - - MOEParallelismConfig getParallelismConfig() const; - QuantParams getQuantParams(nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, - int scale_1_idx = -1, int scale_2_idx = -1, int scale_3_idx = -1, int scale_4_idx = -1, int scale_5_idx = -1, - int scale_6_idx = -1, int scale_7_idx = -1, int scale_8_idx = -1) const; - - int getNumLoraRequests(nvinfer1::PluginTensorDesc const* input_tensor) const; - tensorrt_llm::kernels::LoraParams getLoraParams( - nvinfer1::PluginTensorDesc const* inputDesc, void const* const* inputs, void* workspace); - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - - using IndexType = std::int32_t; - - // Inputs - constexpr static IndexType getInputTensorIndex() - { - return 0; - } - - constexpr static IndexType getExpertWeights1Index() - { - return getInputTensorIndex() + 1; - } - - constexpr static IndexType getExpertWeights2Index() - { - return getExpertWeights1Index() + 1; - } - - constexpr static IndexType getTokenSelectedExpertsIndex() - { - return getExpertWeights2Index() + 1; - } - - // Conditional inputs, we only allocate a new index if actually used - bool hasBias() const - { - return mUseBias; - } - - bool hasFinalScales() const - { - return mUseFinalScales; - } - - bool hasExpertIntQuantScales() const - { - return mQuantMode.hasInt4Weights() || mQuantMode.hasInt8Weights(); - } - - bool hasExpertFp8QuantScales() const - { - return mQuantMode.hasFp8Qdq(); - } - - bool hasExpertFp8FinalQuantScales() const - { - return hasExpertFp8QuantScales() && mOutputType == nvinfer1::DataType::kFP8; - } - - bool hasFP4QuantScales() const - { - return mQuantMode.hasNvfp4(); - } - - bool hasGroupwiseIntQuantScales() const - { - return mGroupwiseQuantAlgo > 0; - } - - bool hasExpertWeightQuantZeros() const - { - return mGroupwiseQuantAlgo & GroupwiseQuantAlgo::ZERO; - } - - bool hasExpertPrequantScales() const - { - return mGroupwiseQuantAlgo & GroupwiseQuantAlgo::PRE_QUANT_SCALE; - } - - bool hasGroupwiseFp8Alpha() const - { - return mGroupwiseQuantAlgo & GroupwiseQuantAlgo::FP8_ALPHA; - } - - bool useSideStream() const - { - return mSideStreamId > 0; - } - - bool hasLora() const - { - return mUseLora; - } - - bool hasGatedLoraWeightsAndRanks() const - { - return mUseLora && isGatedActivation(mActivationType); - } - - IndexType getTokenFinalScalesIndex() const - { - return getTokenSelectedExpertsIndex() + hasFinalScales(); - } - - IndexType getExpertBias1Index() const - { - return getTokenFinalScalesIndex() + hasBias(); - } - - IndexType getExpertBias2Index() const - { - return getExpertBias1Index() + hasBias(); - } - - /* - * Weight-Only int quant scales - */ - IndexType getExpertIntQuantScale1Index() const - { - return getExpertBias2Index() + hasExpertIntQuantScales(); - } - - IndexType getExpertIntQuantScale2Index() const - { - return getExpertIntQuantScale1Index() + hasExpertIntQuantScales(); - } - - /* - * FP8 Quant Scales - */ - IndexType getExpertFP8Dequant1Index() const - { - return getExpertIntQuantScale2Index() + hasExpertFp8QuantScales(); - } - - IndexType getExpertFP8Quant2Index() const - { - return getExpertFP8Dequant1Index() + hasExpertFp8QuantScales(); - } - - IndexType getExpertFP8Dequant2Index() const - { - return getExpertFP8Quant2Index() + hasExpertFp8QuantScales(); - } - - IndexType getExpertFP8QuantFinalIndex() const - { - return getExpertFP8Dequant2Index() + hasExpertFp8FinalQuantScales(); - } - - IndexType getInputFP8DequantIndex() const - { - return getExpertFP8QuantFinalIndex() + (hasExpertFp8QuantScales() && hasLora()); - } - - /* - * FP4 Quant Scales - */ - IndexType getFP4GlobalActSF1Index() const - { - return getInputFP8DequantIndex() + hasFP4QuantScales(); - } - - IndexType getFP4WeightSF1Index() const - { - return getFP4GlobalActSF1Index() + hasFP4QuantScales(); - } - - IndexType getFP4GlobalSF1Index() const - { - return getFP4WeightSF1Index() + hasFP4QuantScales(); - } - - IndexType getFP4GlobalActSF2Index() const - { - return getFP4GlobalSF1Index() + hasFP4QuantScales(); - } - - IndexType getFP4WeightSF2Index() const - { - return getFP4GlobalActSF2Index() + hasFP4QuantScales(); - } - - IndexType getFP4GlobalSF2Index() const - { - return getFP4WeightSF2Index() + hasFP4QuantScales(); - } - - /* - * Groupwise Params - */ - IndexType getExpertPrequantScales1Index() const - { - return getFP4GlobalSF2Index() + hasExpertPrequantScales(); - } - - IndexType getExpertPrequantScales2Index() const - { - return getExpertPrequantScales1Index() + hasExpertPrequantScales(); - } - - IndexType getExpertIntQuantZeros1Index() const - { - return getExpertPrequantScales2Index() + hasExpertWeightQuantZeros(); - } - - IndexType getExpertIntQuantZeros2Index() const - { - return getExpertIntQuantZeros1Index() + hasExpertWeightQuantZeros(); - } - - IndexType getExpertFp8Alpha1Index() const - { - return getExpertIntQuantZeros2Index() + hasGroupwiseFp8Alpha(); - } - - IndexType getExpertFp8Alpha2Index() const - { - return getExpertFp8Alpha1Index() + hasGroupwiseFp8Alpha(); - } - - /* - * LoRA params - */ - IndexType getLoraFC1WeightPtrsIndex() const - { - return getExpertFp8Alpha2Index() + hasLora(); - } - - IndexType getLoraFC1RanksIndex() const - { - return getLoraFC1WeightPtrsIndex() + hasLora(); - } - - IndexType getLoraFC2WeightPtrsIndex() const - { - return getLoraFC1RanksIndex() + hasLora(); - } - - IndexType getLoraFC2RanksIndex() const - { - return getLoraFC2WeightPtrsIndex() + hasLora(); - } - - IndexType getLoraGatedWeightPtrsIndex() const - { - return getLoraFC2RanksIndex() + hasGatedLoraWeightsAndRanks(); - } - - IndexType getLoraGatedRanksIndex() const - { - return getLoraGatedWeightPtrsIndex() + hasGatedLoraWeightsAndRanks(); - } - - IndexType getHostRequestTypeIndex() const - { - return getLoraGatedRanksIndex() + hasLora(); - } - - IndexType getHostContextLengthIndex() const - { - return getHostRequestTypeIndex() + (mRemoveInputPadding && hasLora()); - } - - IndexType getInputDummyTensorIndex() const - { - return getHostContextLengthIndex() + useSideStream(); - } - - IndexType getNbInputs() const - { - return getInputDummyTensorIndex() + 1; - } - - // Outputs - constexpr static IndexType getOutputTensorIndex() - { - return 0; - } - - IndexType getOutputDummyTensorIndex() const - { - return getOutputTensorIndex() + useSideStream(); - } - - /** - * Get the index of the expert shape tuple that represents the inner dimension - */ - int getGemmShapeInnerDimIndex() const - { - // In weight only mode the shape is transposed - return hasExpertIntQuantScales() ? 1 : 2; - } - - /** - * Get the index of the expert shape tuple that represents the outer dimension - */ - int getGemmShapeOuterDimIndex() const - { - // In weight only mode the shape is transposed - return hasExpertIntQuantScales() ? 2 : 1; - } - - /** - * Get quantization dimension scaling factor - */ - std::pair getWeightPackedElements() const - { - if (mGroupwiseQuantAlgo == 0) - { - return {1, mQuantMode.hasInt4Weights() ? 2 : 1}; - } - else - { - return {1, 4}; - } - } -}; - -class MixtureOfExpertsGemmProfiler - : public tensorrt_llm::plugins::GemmPluginProfiler -{ -public: - MixtureOfExpertsGemmProfiler() - { - // NOTE: Do not access mPlugin here, since we are called from the constructor before all fields are init - } - - void setGemmToProfile(kernels::GemmProfilerBackend::GemmToProfile gemm_to_profile) - { - // Just set the backend directly. This will just be reused in checkInit(). - backend.mGemmToProfile = gemm_to_profile; - // We need to set the backend to reinitialise itself with the new GEMM - init_backend = false; - } - - void setMaxProfileM(int maxProfileM) - { - mMaxProfileM = maxProfileM; - } - - virtual int getMaxProfileM() const override - { - return mMaxProfileM; - } - -protected: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - std::vector getTactics(int m, int n, int k) const override; - void initTmpData(int maxM, int n, int k, char* workspace, size_t size, cudaStream_t stream) override; - - void checkInit(); - - bool init_backend = false; - kernels::GemmProfilerBackend backend{}; - -private: - int mMaxProfileM = 0; -}; - -class MixtureOfExpertsPluginCreator : public nvinfer1::IPluginCreator -{ -public: - MixtureOfExpertsPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - - void setPluginNamespace(char const* pluginNamespace) noexcept override; - - char const* getPluginNamespace() const noexcept override; - -private: - GemmPluginProfilerManager moePluginProfiler; - GemmPluginProfilerManager loraPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; - std::string mNamespace; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_MIXTURE_OF_EXPERTS_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/ncclPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.cpp deleted file mode 100644 index 4825dd51bbab..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.cpp +++ /dev/null @@ -1,253 +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. - */ -#include "allgatherPlugin.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::AllgatherPluginCreator; -using tensorrt_llm::plugins::AllgatherPlugin; - -static char const* ALLGATHER_PLUGIN_VERSION{"1"}; -static char const* ALLGATHER_PLUGIN_NAME{"AllGather"}; -PluginFieldCollection AllgatherPluginCreator::mFC{}; -std::vector AllgatherPluginCreator::mPluginAttributes; - -AllgatherPlugin::AllgatherPlugin(std::set group, nvinfer1::DataType type) - : mGroup(std::move(group)) - , mType(type) -{ -} - -// Parameterized constructor -AllgatherPlugin::AllgatherPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - mGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* AllgatherPlugin::clone() const noexcept -{ - auto* plugin = new AllgatherPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs AllgatherPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - auto ret = inputs[0]; - auto groupSize = exprBuilder.constant(mGroup.size()); - ret.d[0] = exprBuilder.operation(DimensionOperation::kPROD, *ret.d[0], *groupSize); - return ret; -} - -bool AllgatherPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void AllgatherPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t AllgatherPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int AllgatherPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclAllGather(inputs[0], outputs[0], size, (*getDtypeMap())[inputDesc[0].type], *mNcclComm, stream)); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType AllgatherPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* AllgatherPlugin::getPluginType() const noexcept -{ - return ALLGATHER_PLUGIN_NAME; -} - -char const* AllgatherPlugin::getPluginVersion() const noexcept -{ - return ALLGATHER_PLUGIN_VERSION; -} - -int AllgatherPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int AllgatherPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - mNcclComm = getComm(mGroup); - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - return 0; -} - -void AllgatherPlugin::terminate() noexcept {} - -size_t AllgatherPlugin::getSerializationSize() const noexcept -{ - return sizeof(int) * mGroup.size() + sizeof(mType); -} - -void AllgatherPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - for (auto it = mGroup.begin(); it != mGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void AllgatherPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -AllgatherPluginCreator::AllgatherPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("group", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* AllgatherPluginCreator::getPluginName() const noexcept -{ - return ALLGATHER_PLUGIN_NAME; -} - -char const* AllgatherPluginCreator::getPluginVersion() const noexcept -{ - return ALLGATHER_PLUGIN_VERSION; -} - -PluginFieldCollection const* AllgatherPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* AllgatherPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - std::set group; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - group.insert(*r); - ++r; - } - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new AllgatherPlugin(group, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* AllgatherPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call AllgatherPlugin::destroy() - try - { - auto* obj = new AllgatherPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h deleted file mode 100644 index 3d7810e6bd49..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allgatherPlugin.h +++ /dev/null @@ -1,92 +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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class AllgatherPlugin : public BasePlugin -{ -public: - AllgatherPlugin(std::set group, nvinfer1::DataType type); - - AllgatherPlugin(void const* data, size_t length); - - ~AllgatherPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - std::set mGroup; - nvinfer1::DataType mType; - std::shared_ptr mNcclComm; -}; - -class AllgatherPluginCreator : public BaseCreator -{ -public: - AllgatherPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp deleted file mode 100644 index 24d9aff418f7..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.cpp +++ /dev/null @@ -1,986 +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. - */ -#include "allreducePlugin.h" - -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/customAllReduceUtils.h" -#include "tensorrt_llm/common/dataType.h" -#include "tensorrt_llm/common/nvmlWrapper.h" -#include "tensorrt_llm/kernels/customAllReduceKernels.h" -#include "tensorrt_llm/kernels/userbuffers/ub_interface.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::AllreducePluginCreator; -using tensorrt_llm::plugins::AllreducePlugin; -using tensorrt_llm::kernels::AllReduceFusionOp; -using tensorrt_llm::kernels::AllReduceStrategyType; -using tensorrt_llm::kernels::AllReduceStrategyConfig; -using tensorrt_llm::mpi::MpiTag; - -static char const* ALLREDUCE_PLUGIN_VERSION{"1"}; -static char const* ALLREDUCE_PLUGIN_NAME{"AllReduce"}; -PluginFieldCollection AllreducePluginCreator::mFC{}; -std::vector AllreducePluginCreator::mPluginAttributes; - -AllreducePlugin::AllreducePlugin(std::set group, nvinfer1::DataType type, AllReduceStrategyType strategy, - AllReduceStrategyConfig config, AllReduceFusionOp op, int32_t counter, float eps, int8_t affine, int8_t bias, - int8_t scale) - : mGroup(std::move(group)) - , mType(type) - , mStrategy(strategy) - , mConfig(config) - , mOp(op) - , mEps(eps) - , mAffine(affine) - , mBias(bias) - , mScale(scale) -{ - check(); -} - -// Parameterized constructor -AllreducePlugin::AllreducePlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mStrategy); - read(d, mConfig); - read(d, mOp); - read(d, mEps); - read(d, mAffine); - read(d, mBias); - read(d, mScale); - mGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); - check(); -} - -void AllreducePlugin::check() noexcept -{ - if (mStrategy != AllReduceStrategyType::UB) - { - TLLM_CHECK(mOp != AllReduceFusionOp::LAST_PROCESS_FOR_UB); - } -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* AllreducePlugin::clone() const noexcept -{ - auto* plugin = new AllreducePlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs AllreducePlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4 && mStrategy == AllReduceStrategyType::UB && mScale) - { - if (outputIndex == 0) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - return ret; - } - else if (outputIndex == 2) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - auto dimM = exprBuilder.operation( - DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 2], *exprBuilder.constant(128)); - ret.d[ret.nbDims - 2] = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - ret.d[ret.nbDims - 1] = exprBuilder.operation( - DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16)); - return ret; - } - } - return inputs[0]; -} - -bool AllreducePlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int base_inputs = 0; - switch (mStrategy) - { - case AllReduceStrategyType::NCCL: - case AllReduceStrategyType::UB: - case AllReduceStrategyType::NCCL_SYMMETRIC: base_inputs = 1; break; - default: base_inputs = 2; break; - } - int fusion_op_extra_inputs = 0; - int scale_idx = 0; - if (mOp != AllReduceFusionOp::NONE) - { - ++fusion_op_extra_inputs; - if (mAffine) - { - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - ++fusion_op_extra_inputs; - ++fusion_op_extra_inputs; - } - if (mBias) - { - ++fusion_op_extra_inputs; - } - if (mScale) - { - scale_idx = base_inputs + fusion_op_extra_inputs; - ++fusion_op_extra_inputs; - } - } - - TLLM_CHECK(nbInputs == (base_inputs + fusion_op_extra_inputs)); - - if (pos == 1) - { - switch (mStrategy) - { - case AllReduceStrategyType::NCCL: - case AllReduceStrategyType::UB: - case AllReduceStrategyType::NCCL_SYMMETRIC: break; - default: return (inOut[pos].type == nvinfer1::DataType::kINT64) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - if (mStrategy == AllReduceStrategyType::UB) - { - if (mScale && pos == scale_idx) - { - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - if (pos == nbInputs) - { - return (inOut[pos].type == nvinfer1::DataType::kFP4) && (inOut[pos].format == TensorFormat::kLINEAR); - } - if (pos == (nbInputs + 2)) - { - return (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8) - { - if (pos == nbInputs) - { - return (inOut[pos].type == nvinfer1::DataType::kFP8) && (inOut[pos].format == TensorFormat::kLINEAR); - } - } - } - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void AllreducePlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t AllreducePlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -AllReduceStrategyType AllreducePlugin::selectImplementation( - size_t messageSize, int worldSize, nvinfer1::DataType type) noexcept -{ - bool const isAuto = (mStrategy == AllReduceStrategyType::AUTO); - - bool const forceDeterministic = common::getEnvForceDeterministicAllReduce(); - if (!mIsP2PSupported) - { - if (!isAuto) - { - TLLM_LOG_INFO("Since Peer to Peer not supported, fallback to AllReduceStrategy: NCCL_SYMMETRIC"); - } - else if (forceDeterministic) - { - TLLM_LOG_WARNING( - "Since Peer to Peer not supported, fallback to AllReduceStrategy: NCCL_SYMMETRIC. NCCL_SYMMETRIC might " - "produce " - "non-deterministic results."); - } - return AllReduceStrategyType::NCCL_SYMMETRIC; - } - - if (isAuto && !mIsNVLINKSupported && !forceDeterministic) - { - return AllReduceStrategyType::NCCL_SYMMETRIC; - } - - auto const maxWorkspaceSize = utils::customAllReduceUtils::getMaxRequiredWorkspaceSize(worldSize); - - AllReduceStrategyType strat = AllReduceStrategyType::NCCL_SYMMETRIC; - auto const messageSizeBytes = messageSize * common::getDTypeSize(type); - - if (messageSizeBytes <= maxWorkspaceSize) - { - // In some instances, the two-shot strategy has exhibited significant performance issues. - // As a temporary measure, we have disabled the two-shot strategy. - // TODO: remove this WAR after https://nvbugspro.nvidia.com/bug/4718747 is fixed. - if (!isAuto) - { - strat = mStrategy; - } - else if (forceDeterministic) - { - strat = AllReduceStrategyType::ONESHOT; - } - else if (worldSize <= 2) - { - strat = AllReduceStrategyType::ONESHOT; - } - else if (worldSize <= 4) - { - if (messageSizeBytes < 1 * 1000 * 1000) - { - strat = AllReduceStrategyType::ONESHOT; - } - else - { - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - } - else - { - if (messageSizeBytes < 500 * 1000) - { - strat = AllReduceStrategyType::ONESHOT; - } - else - { - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - } - - if (!kernels::configurationSupported(strat, messageSize, worldSize, type)) - { - if (!isAuto) - { - TLLM_LOG_WARNING("Since not aligned, fallback to AllReduceStrategy: NCCL_SYMMETRIC"); - } - else if (forceDeterministic) - { - TLLM_LOG_WARNING( - "Since not aligned, fallback to AllReduceStrategy: NCCL_SYMMETRIC. NCCL_SYMMETRIC might produce " - "non-deterministic results."); - } - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - } - else - { - if (!isAuto) - { - TLLM_LOG_WARNING("Since messageSize > maxWorkspace, fallback to AllReduceStrategy: NCCL_SYMMETRIC"); - } - else if (forceDeterministic) - { - TLLM_LOG_WARNING( - "Since messageSize > maxWorkspace, fallback to AllReduceStrategy: NCCL_SYMMETRIC. NCCL_SYMMETRIC might " - "produce " - "non-deterministic results."); - } - strat = AllReduceStrategyType::NCCL_SYMMETRIC; - } - - return strat; -} - -int AllreducePlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - - kernels::AllReduceStrategyType runtimeStrategy; - - static char* forceNcclAllReduceStrategyChar = std::getenv("FORCE_NCCL_ALL_REDUCE_STRATEGY"); - bool forceNcclAllReduceStrategy = (forceNcclAllReduceStrategyChar != nullptr); - if (forceNcclAllReduceStrategy || mStrategy == AllReduceStrategyType::NCCL) - { - runtimeStrategy = AllReduceStrategyType::NCCL; - } - else if (mStrategy == AllReduceStrategyType::NCCL_SYMMETRIC) - { - runtimeStrategy = AllReduceStrategyType::NCCL_SYMMETRIC; - } - else if (mStrategy == AllReduceStrategyType::UB) - { - runtimeStrategy = AllReduceStrategyType::UB; - } - else - { - runtimeStrategy = selectImplementation(size, mGroup.size(), mType); - } - - // Log runtime strategy - auto const rank = COMM_SESSION.getRank(); - switch (runtimeStrategy) - { - case AllReduceStrategyType::NCCL: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: NCCL", rank); - break; - } - case AllReduceStrategyType::NCCL_SYMMETRIC: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: NCCL_SYMMETRIC", rank); - break; - } - case AllReduceStrategyType::ONESHOT: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: ONESHOT", rank); - break; - } - case AllReduceStrategyType::TWOSHOT: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: TWOSHOT", rank); - break; - } - case AllReduceStrategyType::UB: - { - TLLM_LOG_DEBUG("AllReducePlugin strategy for rank %d: UB", rank); - break; - } - default: break; - } - - if (runtimeStrategy == AllReduceStrategyType::NCCL || runtimeStrategy == AllReduceStrategyType::NCCL_SYMMETRIC) - { - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM || mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - { - NCCLCHECK(ncclAllReduce(inputs[0], outputs[1], size, (*getDtypeMap())[mType], ncclSum, *mNcclComm, stream)); - tensorrt_llm::kernels::AllReduceParams params; - int fusion_ptr_idx = 0; - if (mStrategy == AllReduceStrategyType::NCCL || mStrategy == AllReduceStrategyType::NCCL_SYMMETRIC) - { - fusion_ptr_idx = 1; - } - else - { - fusion_ptr_idx = 2; - } - params.fusion_params.bias_buffer = mBias ? inputs[fusion_ptr_idx++] : nullptr; - params.fusion_params.residual_buffer = inputs[fusion_ptr_idx++]; - params.fusion_params.weight_buffer = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - { - params.fusion_params.weight_buffer_pre_residual_norm = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - } - params.local_output_buffer_ptr = outputs[0]; - params.elts_total = size; - params.fusion_params.hidden_size = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - params.fusion_params.eps = mEps; - params.fusion_params.intermediate_buffer = outputs[1]; - TLLM_LOG_DEBUG("residualRmsNorm called"); - tensorrt_llm::kernels::residualRmsNorm(params, mType, stream, mOp); - } - else - { - NCCLCHECK(ncclAllReduce(inputs[0], outputs[0], size, (*getDtypeMap())[mType], ncclSum, *mNcclComm, stream)); - } - } - else if (runtimeStrategy == AllReduceStrategyType::UB) - { - TLLM_CHECK(!mBias); - - size_t dtype_size = tensorrt_llm::common::getDTypeSize(mType); - int hidden_size = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - - TLLM_CHECK_WITH_INFO(tensorrt_llm::runtime::ub::ub_is_initialized(), "UserBuffer has not been initialized!"); - auto ub_buffer0 = tensorrt_llm::runtime::ub::ub_get(0); - auto ub_buffer1 = tensorrt_llm::runtime::ub::ub_get(1); - TLLM_CHECK(inputs[0] == ub_buffer0.addr); - auto ub_comm = tensorrt_llm::runtime::ub::ub_comm(); - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8) - { - TLLM_CHECK(mAffine); - TLLM_CHECK(mScale); - TLLM_CHECK(outputs[0] == ub_buffer1.addr); - void* residual = const_cast(inputs[1]); - void* gamma = const_cast(inputs[2]); - float* scale = const_cast(reinterpret_cast(inputs[3])); - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_rmsnorm_quant_launcher(ub_buffer0.handle, 0, - ub_buffer1.handle, 0, size, hidden_size, nullptr, gamma, mEps, scale, residual, outputs[1], mType, - ub_comm, stream); - } - else if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - auto ub_buffer2 = tensorrt_llm::runtime::ub::ub_get(2); - TLLM_CHECK(mAffine); - TLLM_CHECK(mScale); - TLLM_CHECK(outputs[0] == ub_buffer1.addr); - TLLM_CHECK(outputs[2] == ub_buffer2.addr); - void* residual = const_cast(inputs[1]); - void* gamma = const_cast(inputs[2]); - float* scale = const_cast(reinterpret_cast(inputs[3])); - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_rmsnorm_quant_fp4_launcher(ub_buffer0.handle, 0, - ub_buffer1.handle, 0, ub_buffer2.handle, 0, size, hidden_size, nullptr, gamma, mEps, scale, residual, - outputs[1], mType, ub_comm, stream); - } - else if (mOp == AllReduceFusionOp::LAST_PROCESS_FOR_UB) - { - TLLM_CHECK(outputs[1] == ub_buffer1.addr); - void* residual = const_cast(inputs[1]); - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_launcher( - ub_buffer0.handle, 0, size, mType, ub_comm, stream); - tensorrt_llm::kernels::ub::allgather2_userbuff_residual_launcher( - ub_buffer1.handle, 0, size, hidden_size, residual, mType, ub_comm, stream); - TLLM_CUDA_CHECK( - cudaMemcpyAsync(outputs[0], ub_buffer0.addr, size * dtype_size, cudaMemcpyDeviceToDevice, stream)); - } - else if (mOp == AllReduceFusionOp::NONE) - { - tensorrt_llm::kernels::ub::allreduce2_userbuff_inplace_launcher( - ub_buffer0.handle, 0, size, mType, ub_comm, stream); - TLLM_CUDA_CHECK( - cudaMemcpyAsync(outputs[0], ub_buffer0.addr, size * dtype_size, cudaMemcpyDeviceToDevice, stream)); - } - else - { - TLLM_CHECK_WITH_INFO(false, "Unsupported UB allreduce fusion op"); - } - } - else - { - auto const tpSize = mGroup.size(); - int tpRank = 0; - for (auto const& currentRank : mGroup) - { - if (rank == currentRank) - break; - ++tpRank; - } - - int token_num = size / inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - int hidden_size = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - auto params = tensorrt_llm::kernels::AllReduceParams::deserialize( - reinterpret_cast(const_cast(inputs[1])), tpSize, tpRank, mType, token_num, hidden_size, - mOp); - - params.local_output_buffer_ptr = outputs[0]; - params.local_input_buffer_ptr = inputs[0]; - params.elts_total = size; - - int fusion_ptr_idx = 2; - params.fusion_params.bias_buffer = mBias ? inputs[fusion_ptr_idx++] : nullptr; - params.fusion_params.residual_buffer = inputs[fusion_ptr_idx++]; - params.fusion_params.weight_buffer = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_PREPOST_NORM) - params.fusion_params.weight_buffer_pre_residual_norm = mAffine ? inputs[fusion_ptr_idx++] : nullptr; - params.fusion_params.hidden_size = hidden_size; - params.fusion_params.eps = mEps; - params.fusion_params.intermediate_buffer = outputs[1]; - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM) - { - for (size_t i = 0; i < tpSize; ++i) - { - params.fusion_params.lamport_peer_comm_buffer_ptrs[i] - = reinterpret_cast(const_cast(inputs[1]))[tpSize * 4 + i]; - params.fusion_params.lamport_peer_comm_buffer_ptrs[i + tensorrt_llm::kernels::MAX_RANKS_PER_NODE] - = reinterpret_cast(const_cast(inputs[1]))[tpSize * 5 + i]; - params.fusion_params.lamport_peer_comm_buffer_ptrs[i + tensorrt_llm::kernels::MAX_RANKS_PER_NODE * 2] - = reinterpret_cast(const_cast(inputs[1]))[tpSize * 6 + i]; - } - } - TLLM_LOG_DEBUG("customAllReduce called"); - tensorrt_llm::kernels::customAllReduce(params, mType, runtimeStrategy, mConfig, mOp, stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType AllreducePlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index < getNbOutputs()); - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - if (index == 0) - { - return nvinfer1::DataType::kFP4; - } - else if (index == 2) - { - return nvinfer1::DataType::kFP8; - } - } - if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_FP8) - { - if (index == 0) - { - return nvinfer1::DataType::kFP8; - } - } - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* AllreducePlugin::getPluginType() const noexcept -{ - return ALLREDUCE_PLUGIN_NAME; -} - -char const* AllreducePlugin::getPluginVersion() const noexcept -{ - return ALLREDUCE_PLUGIN_VERSION; -} - -int AllreducePlugin::getNbOutputs() const noexcept -{ - if (mOp == AllReduceFusionOp::NONE) - { - return 1; - } - else if (mOp == AllReduceFusionOp::RESIDUAL_RMS_NORM_QUANT_NVFP4) - { - return 3; - } - else - { - return 2; - } -} - -bool AllreducePlugin::isCustomAllReduceSupported(int ranks_per_node) const noexcept -{ - constexpr bool isCudaVersionSupported = -#if defined(CUDART_VERSION) && CUDART_VERSION >= 11020 - true; -#else - false; -#endif - - return isCudaVersionSupported && (ranks_per_node % 2 == 0) - && (static_cast(ranks_per_node) <= kernels::MAX_RANKS_PER_NODE) && (ranks_per_node > 0); -} - -using tensorrt_llm::common::NvmlManager; -using tensorrt_llm::common::NVMLWrapper; - -std::set getLocalGroup(std::set const& group) -{ - auto const myRank = COMM_SESSION.getRank(); - auto const myLocalRank = LOCAL_COMM_SESSION.getRank(); - auto const localSize = LOCAL_COMM_SESSION.getSize(); - - std::vector ranks(localSize, 0); - std::vector localRanks(localSize, 0); - if (group.size() >= static_cast(localSize)) - { - LOCAL_COMM_SESSION.allgather(&myRank, ranks.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); - LOCAL_COMM_SESSION.allgather(&myLocalRank, localRanks.data(), 1, tensorrt_llm::mpi::MpiType::kINT32); - } - else - { - if (myRank == *group.begin()) - { - ranks.clear(); - int rank; - ranks.push_back(myRank); - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.recvValue(rank, *it, MpiTag::kDefault); - ranks.push_back(rank); - } - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.send(ranks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *it, MpiTag::kDefault); - } - - localRanks.clear(); - localRanks.push_back(myLocalRank); - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.recvValue(rank, *it, MpiTag::kDefault); - localRanks.push_back(rank); - } - for (auto it = std::next(std::begin(group), 1); it != group.end(); ++it) - { - COMM_SESSION.send( - localRanks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *it, MpiTag::kDefault); - } - } - else - { - COMM_SESSION.sendValue(myRank, *group.begin(), MpiTag::kDefault); - COMM_SESSION.recv( - ranks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *group.begin(), MpiTag::kDefault); - - COMM_SESSION.sendValue(myLocalRank, *group.begin(), MpiTag::kDefault); - COMM_SESSION.recv( - localRanks.data(), localSize, tensorrt_llm::mpi::MpiType::kINT32, *group.begin(), MpiTag::kDefault); - } - } - - std::set localGroup; - for (size_t i = 0; i < ranks.size(); ++i) - { - auto rank = ranks[i]; - if (group.find(rank) != group.end()) - { - localGroup.insert(localRanks[i]); - } - } - return localGroup; -} - -void AllreducePlugin::initGroupTopology() noexcept -{ - static std::map, std::tuple> cache; - if (cache.find(mGroup) != cache.end()) - { - auto [isNVLINKSupported, isP2PSupported] = cache[mGroup]; - mIsNVLINKSupported = isNVLINKSupported; - mIsP2PSupported = isP2PSupported; - return; - } - setGroupTopology(); - cache[mGroup] = {mIsNVLINKSupported, mIsP2PSupported}; -} - -void AllreducePlugin::setGroupTopology() noexcept -{ - auto const rank = COMM_SESSION.getRank(); - TLLM_LOG_INFO("Detecting local TP group for rank %d", rank); - std::set localGroup = getLocalGroup(mGroup); - if (mGroup.size() != localGroup.size()) - { - mIsP2PSupported = false; - mIsNVLINKSupported = false; - TLLM_LOG_INFO("Found inter-node TP group for rank %d", rank); - return; - } - TLLM_LOG_INFO("TP group is intra-node for rank %d", rank); - - NvmlManager nvmlManager; - auto const& nvml = nvmlManager.sharedWrapper(); - std::unordered_set visitedDevice; - mIsP2PSupported = true; - mIsNVLINKSupported = true; - - // Use cudaDeviceCanAccessPeer to determine whether p2p is supported, - // and use nvml to determine whether there are nvlink links between ranks. - for (int firstDeviceId : localGroup) - { - for (int secondDeviceId : localGroup) - { - if (firstDeviceId == secondDeviceId || visitedDevice.find(secondDeviceId) != visitedDevice.end()) - { - continue; - } - - int canAccessPeer = 0; - TLLM_CUDA_CHECK(cudaDeviceCanAccessPeer(&canAccessPeer, firstDeviceId, secondDeviceId)); - - if (!canAccessPeer) - { - mIsP2PSupported = false; - mIsNVLINKSupported = false; - - return; - } - - nvmlDevice_t firstDevice; - NVML_CHECK(nvml->nvmlDeviceGetHandleByIndex(firstDeviceId, &firstDevice)); - - bool isNVLINK = false; - - for (unsigned int link = 0; link < NVML_NVLINK_MAX_LINKS; link++) - { - nvmlPciInfo_t remotePciInfo; - if (nvml->nvmlDeviceGetNvLinkRemotePciInfo(firstDevice, link, &remotePciInfo) != NVML_SUCCESS) - { - continue; - } - - nvmlDevice_t remoteDevice; - auto const result = nvml->nvmlDeviceGetHandleByPciBusId(remotePciInfo.busId, &remoteDevice); - - if (result == NVML_SUCCESS) - { - // Two GPUs are connected directly through nvlink - unsigned int remoteDeviceId; - NVML_CHECK(nvml->nvmlDeviceGetIndex(remoteDevice, &remoteDeviceId)); - - if (remoteDeviceId == static_cast(secondDeviceId)) - { - isNVLINK = true; - } - } - else if (result == NVML_ERROR_NOT_FOUND) - { - // Maybe Two GPUs are connected via nvswitch, - // now remotePciInfo represents the pci information of nvswitch, - // determine whether nvlink is supported by whether two GPUs are connected to the same nvswitch. - nvmlDevice_t secondDevice; - NVML_CHECK(nvml->nvmlDeviceGetHandleByIndex(secondDeviceId, &secondDevice)); - - for (unsigned int secondLink = 0; secondLink < NVML_NVLINK_MAX_LINKS; secondLink++) - { - nvmlPciInfo_t secondRemotePciInfo; - if (nvml->nvmlDeviceGetNvLinkRemotePciInfo(secondDevice, secondLink, &secondRemotePciInfo) - != NVML_SUCCESS) - { - continue; - } - - if (strcmp(remotePciInfo.busId, secondRemotePciInfo.busId) == 0) - { - isNVLINK = true; - break; - } - } - } - else - { - NVML_CHECK(result); - } - - if (isNVLINK) - { - break; - } - } - - mIsNVLINKSupported &= isNVLINK; - } - visitedDevice.insert(firstDeviceId); - } -} - -int AllreducePlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - - TLLM_LOG_TRACE("%s start for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - mNcclComm = getComm(mGroup); - if (mStrategy != AllReduceStrategyType::NCCL) - { - initGroupTopology(); - } - - TLLM_LOG_TRACE("%s stop for rank %d", __PRETTY_FUNCTION__, COMM_SESSION.getRank()); - return 0; -} - -void AllreducePlugin::terminate() noexcept {} - -size_t AllreducePlugin::getSerializationSize() const noexcept -{ - return sizeof(int) * mGroup.size() + sizeof(mType) + sizeof(mStrategy) + sizeof(mConfig) + sizeof(mOp) - + sizeof(mEps) + sizeof(mAffine) + sizeof(mBias) + sizeof(mScale); -} - -void AllreducePlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mStrategy); - write(d, mConfig); - write(d, mOp); - write(d, mEps); - write(d, mAffine); - write(d, mBias); - write(d, mScale); - for (auto it = mGroup.begin(); it != mGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void AllreducePlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -AllreducePluginCreator::AllreducePluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("group", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("strategy", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("config", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("fusion_op", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("counter", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("affine", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("scale", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* AllreducePluginCreator::getPluginName() const noexcept -{ - return ALLREDUCE_PLUGIN_NAME; -} - -char const* AllreducePluginCreator::getPluginVersion() const noexcept -{ - return ALLREDUCE_PLUGIN_VERSION; -} - -PluginFieldCollection const* AllreducePluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* AllreducePluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - std::set group; - nvinfer1::DataType type{}; - AllReduceStrategyType strategy{}; - AllReduceStrategyConfig config{}; - AllReduceFusionOp fusion_op{}; - int32_t counter{}; - float eps{}; - int8_t affine{}; - int8_t bias{}; - int8_t scale{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - group.insert(*r); - ++r; - } - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "strategy")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - strategy = static_cast(*static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "config")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - config = static_cast(*static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "fusion_op")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - fusion_op = static_cast(*static_cast(fields[i].data)); - } - else if (!strcmp(attrName, "counter")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - counter = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "affine")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - affine = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "bias")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - bias = *static_cast(fields[i].data); - } - else if (!strcmp(attrName, "scale")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - scale = *static_cast(fields[i].data); - } - } - try - { - auto* obj = new AllreducePlugin(group, type, strategy, config, fusion_op, counter, eps, affine, bias, scale); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* AllreducePluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call AllreducePlugin::destroy() - try - { - auto* obj = new AllreducePlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h deleted file mode 100644 index 881fbf3b89a5..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/allreducePlugin.h +++ /dev/null @@ -1,114 +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. - */ -#pragma once - -#include "tensorrt_llm/kernels/customAllReduceKernels.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ -namespace tk = ::tensorrt_llm::kernels; - -class AllreducePlugin : public BasePlugin -{ -public: - AllreducePlugin(std::set group, nvinfer1::DataType type, tk::AllReduceStrategyType strategy, - tk::AllReduceStrategyConfig config, tk::AllReduceFusionOp op, int32_t counter, float eps, int8_t affine, - int8_t bias, int8_t scale); - - AllreducePlugin(void const* data, size_t length); - - ~AllreducePlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - bool isCustomAllReduceSupported(int ranks_per_node) const noexcept; - void initGroupTopology() noexcept; - void setGroupTopology() noexcept; - tk::AllReduceStrategyType selectImplementation(size_t messageSize, int worldSize, nvinfer1::DataType type) noexcept; - void check() noexcept; - -private: - std::string const mLayerName; - std::set mGroup; - bool mIsNVLINKSupported; - bool mIsP2PSupported; - nvinfer1::DataType mType; - tk::AllReduceStrategyType mStrategy; - tk::AllReduceStrategyConfig mConfig; - tk::AllReduceFusionOp mOp; - float mEps; - std::shared_ptr mNcclComm; - int8_t mAffine; - int8_t mBias; - int8_t mScale; -}; - -class AllreducePluginCreator : public BaseCreator -{ -public: - AllreducePluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.cpp deleted file mode 100644 index 089ed31175b2..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.cpp +++ /dev/null @@ -1,252 +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. - */ -#include "recvPlugin.h" - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::RecvPluginCreator; -using tensorrt_llm::plugins::RecvPlugin; -using tensorrt_llm::mpi::MpiTag; - -static char const* RECV_PLUGIN_VERSION{"1"}; -static char const* RECV_PLUGIN_NAME{"Recv"}; -PluginFieldCollection RecvPluginCreator::mFC{}; -std::vector RecvPluginCreator::mPluginAttributes; - -RecvPlugin::RecvPlugin(int srcRank, nvinfer1::DataType type) - : mSrcRank(srcRank) - , mType(type) -{ -} - -// Parameterized constructor -RecvPlugin::RecvPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mSrcRank); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* RecvPlugin::clone() const noexcept -{ - auto* plugin = new RecvPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs RecvPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - return inputs[0]; -} - -bool RecvPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void RecvPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t RecvPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int RecvPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - TLLM_LOG_DEBUG("start ncclRecv with size %d", size); - NCCLCHECK(ncclRecv(outputs[0], size, (*getDtypeMap())[inputDesc[0].type], 0, mComm, stream)); - TLLM_LOG_DEBUG("end ncclRecv with size %d", size); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType RecvPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* RecvPlugin::getPluginType() const noexcept -{ - return RECV_PLUGIN_NAME; -} - -char const* RecvPlugin::getPluginVersion() const noexcept -{ - return RECV_PLUGIN_VERSION; -} - -int RecvPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int RecvPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - ncclUniqueId id; - COMM_SESSION.recvValue(id, mSrcRank, MpiTag::kDefault); -// Need static connection initialization for accurate KV cache size estimation -#if defined(_WIN32) - if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) - _putenv_s("NCCL_RUNTIME_CONNECT", "0"); -#else - setenv("NCCL_RUNTIME_CONNECT", "0", 0); -#endif // _WIN32 - NCCLCHECK(ncclCommInitRank(&mComm, 2, id, 1)); - return 0; -} - -void RecvPlugin::terminate() noexcept -{ - if (isBuilding()) - { - return; - } - NCCLCHECK(ncclCommDestroy(mComm)); -} - -size_t RecvPlugin::getSerializationSize() const noexcept -{ - return sizeof(mSrcRank) + sizeof(mType); -} - -void RecvPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mSrcRank); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void RecvPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -RecvPluginCreator::RecvPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("src_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* RecvPluginCreator::getPluginName() const noexcept -{ - return RECV_PLUGIN_NAME; -} - -char const* RecvPluginCreator::getPluginVersion() const noexcept -{ - return RECV_PLUGIN_VERSION; -} - -PluginFieldCollection const* RecvPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* RecvPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int srcRank{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "src_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - srcRank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new RecvPlugin(srcRank, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* RecvPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call RecvPlugin::destroy() - try - { - auto* obj = new RecvPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.h deleted file mode 100644 index 5c8eedfb5218..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/recvPlugin.h +++ /dev/null @@ -1,90 +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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class RecvPlugin : public BasePlugin -{ -public: - RecvPlugin(int srcRank, nvinfer1::DataType type); - - RecvPlugin(void const* data, size_t length); - - ~RecvPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - ncclComm_t mComm; // TODO: Remove this - int mSrcRank; - nvinfer1::DataType mType; -}; - -class RecvPluginCreator : public BaseCreator -{ -public: - RecvPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.cpp deleted file mode 100644 index fe17c44fc418..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.cpp +++ /dev/null @@ -1,252 +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. - */ -#include "reduceScatterPlugin.h" - -#include -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::ReduceScatterPluginCreator; -using tensorrt_llm::plugins::ReduceScatterPlugin; - -static char const* REDUCE_SCATTER_PLUGIN_VERSION{"1"}; -static char const* REDUCE_SCATTER_PLUGIN_NAME{"ReduceScatter"}; -PluginFieldCollection ReduceScatterPluginCreator::mFC{}; -std::vector ReduceScatterPluginCreator::mPluginAttributes; - -ReduceScatterPlugin::ReduceScatterPlugin(std::set group, nvinfer1::DataType type) - : mGroup(std::move(group)) - , mType(type) -{ -} - -// Parameterized constructor -ReduceScatterPlugin::ReduceScatterPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - mGroup.clear(); - int groupItem = 0; - while (d != a + length) - { - read(d, groupItem); - mGroup.insert(groupItem); - } - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* ReduceScatterPlugin::clone() const noexcept -{ - auto* plugin = new ReduceScatterPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs ReduceScatterPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - auto output = inputs[0]; - output.d[0] - = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *output.d[0], *exprBuilder.constant(mGroup.size())); - return output; -} - -bool ReduceScatterPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void ReduceScatterPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t ReduceScatterPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int ReduceScatterPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < outputDesc[0].dims.nbDims; ++i) - { - size *= outputDesc[0].dims.d[i]; - } - - TLLM_CHECK_WITH_INFO(mNcclComm.get() != nullptr, "mNcclComm should be initialized before used"); - NCCLCHECK(ncclReduceScatter( - inputs[0], outputs[0], size, (*getDtypeMap())[inputDesc[0].type], ncclSum, *mNcclComm, stream)); - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType ReduceScatterPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* ReduceScatterPlugin::getPluginType() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_NAME; -} - -char const* ReduceScatterPlugin::getPluginVersion() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_VERSION; -} - -int ReduceScatterPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int ReduceScatterPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - mNcclComm = getComm(mGroup); - return 0; -} - -void ReduceScatterPlugin::terminate() noexcept {} - -size_t ReduceScatterPlugin::getSerializationSize() const noexcept -{ - return sizeof(int) * mGroup.size() + sizeof(mType); -} - -void ReduceScatterPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - for (auto it = mGroup.begin(); it != mGroup.end(); ++it) - { - write(d, *it); - } - TLLM_CHECK(d == a + getSerializationSize()); -} - -void ReduceScatterPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -ReduceScatterPluginCreator::ReduceScatterPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("group", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* ReduceScatterPluginCreator::getPluginName() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_NAME; -} - -char const* ReduceScatterPluginCreator::getPluginVersion() const noexcept -{ - return REDUCE_SCATTER_PLUGIN_VERSION; -} - -PluginFieldCollection const* ReduceScatterPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* ReduceScatterPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - std::set group; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "group")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - auto const* r = static_cast(fields[i].data); - for (int j = 0; j < fields[i].length; ++j) - { - group.insert(*r); - ++r; - } - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new ReduceScatterPlugin(group, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* ReduceScatterPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call ReduceScatterPlugin::destroy() - try - { - auto* obj = new ReduceScatterPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h deleted file mode 100644 index c630b57a2b98..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/reduceScatterPlugin.h +++ /dev/null @@ -1,91 +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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class ReduceScatterPlugin : public BasePlugin -{ -public: - ReduceScatterPlugin(std::set group, nvinfer1::DataType type); - - ReduceScatterPlugin(void const* data, size_t length); - - ~ReduceScatterPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - std::set mGroup; - nvinfer1::DataType mType; - std::shared_ptr mNcclComm; -}; - -class ReduceScatterPluginCreator : public BaseCreator -{ -public: - ReduceScatterPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.cpp b/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.cpp deleted file mode 100644 index 81d66aa8211e..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.cpp +++ /dev/null @@ -1,255 +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. - */ -#include "sendPlugin.h" - -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" - -#include -#include - -using namespace nvinfer1; -using tensorrt_llm::plugins::SendPluginCreator; -using tensorrt_llm::plugins::SendPlugin; -using tensorrt_llm::mpi::MpiTag; - -static char const* SEND_PLUGIN_VERSION{"1"}; -static char const* SEND_PLUGIN_NAME{"Send"}; -PluginFieldCollection SendPluginCreator::mFC{}; -std::vector SendPluginCreator::mPluginAttributes; - -SendPlugin::SendPlugin(int tgtRank, nvinfer1::DataType type) - : mTgtRank(tgtRank) - , mType(type) -{ -} - -// Parameterized constructor -SendPlugin::SendPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mTgtRank); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* SendPlugin::clone() const noexcept -{ - auto* plugin = new SendPlugin(*this); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs SendPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - return inputs[0]; -} - -bool SendPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); -} - -void SendPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t SendPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int SendPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - size_t size = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims; ++i) - { - size *= inputDesc[0].dims.d[i]; - } - - TLLM_LOG_DEBUG("start ncclSend with size %d", size); - NCCLCHECK(ncclSend(inputs[0], size, (*getDtypeMap())[inputDesc[0].type], 1, mComm, stream)); - TLLM_LOG_DEBUG("end ncclSend with size %d", size); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType SendPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index == 0); - return inputTypes[0]; -} - -// IPluginV2 Methods - -char const* SendPlugin::getPluginType() const noexcept -{ - return SEND_PLUGIN_NAME; -} - -char const* SendPlugin::getPluginVersion() const noexcept -{ - return SEND_PLUGIN_VERSION; -} - -int SendPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int SendPlugin::initialize() noexcept -{ - if (isBuilding()) - { - return 0; - } - - ncclUniqueId id; - ncclGetUniqueId(&id); - COMM_SESSION.sendValue(id, mTgtRank, MpiTag::kDefault); -// Need static connection initialization for accurate KV cache size estimation -#if defined(_WIN32) - if (getenv("NCCL_RUNTIME_CONNECT") == nullptr) - _putenv_s("NCCL_RUNTIME_CONNECT", "0"); -#else - setenv("NCCL_RUNTIME_CONNECT", "0", 0); -#endif // _WIN32 - NCCLCHECK(ncclCommInitRank(&mComm, 2, id, 0)); - return 0; -} - -void SendPlugin::terminate() noexcept -{ - if (isBuilding()) - { - return; - } - NCCLCHECK(ncclCommDestroy(mComm)); -} - -size_t SendPlugin::getSerializationSize() const noexcept -{ - return sizeof(mTgtRank) + sizeof(mType); -} - -void SendPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mTgtRank); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void SendPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -SendPluginCreator::SendPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("tgt_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* SendPluginCreator::getPluginName() const noexcept -{ - return SEND_PLUGIN_NAME; -} - -char const* SendPluginCreator::getPluginVersion() const noexcept -{ - return SEND_PLUGIN_VERSION; -} - -PluginFieldCollection const* SendPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* SendPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int tgtRank{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "tgt_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - tgtRank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - - try - { - auto* obj = new SendPlugin(tgtRank, type); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* SendPluginCreator::deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call SendPlugin::destroy() - try - { - auto* obj = new SendPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.h b/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.h deleted file mode 100644 index 0d36b0ebff28..000000000000 --- a/cpp/tensorrt_llm/plugins/ncclPlugin/sendPlugin.h +++ /dev/null @@ -1,89 +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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include - -namespace tensorrt_llm::plugins -{ - -class SendPlugin : public BasePlugin -{ -public: - SendPlugin(int tgtRank, nvinfer1::DataType type); - - SendPlugin(void const* data, size_t length); - - ~SendPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - ncclComm_t mComm; // TODO: Remove this - int mTgtRank; - nvinfer1::DataType mType; -}; - -class SendPluginCreator : public BaseCreator -{ -public: - SendPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/qserveGemmPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.cpp deleted file mode 100644 index 166f1cc32cbe..000000000000 --- a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.cpp +++ /dev/null @@ -1,416 +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. - */ - -#include "qserveGemmPlugin.h" -#include "tensorrt_llm/kernels/qserveGemm.h" -#include -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::QServeGemmPluginCreator; -using tensorrt_llm::plugins::QServeGemmPlugin; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; -using namespace tensorrt_llm::kernels::qserve; - -static char const* QSERVE_GEMM_PLUGIN_VERSION{"1"}; -static char const* QSERVE_GEMM_PLUGIN_NAME{"QServeGemm"}; - -PluginFieldCollection QServeGemmPluginCreator::mFC{}; -std::vector QServeGemmPluginCreator::mPluginAttributes; - -namespace tensorrt_llm::plugins -{ - -QServeGemmPlugin::QServeGemmPlugin( - // QuantMode quantMode, - nvinfer1::DataType dtype, int groupSize) -{ - init(dtype, groupSize); -} - -QServeGemmPlugin::QServeGemmPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - - nvinfer1::DataType type; - unsigned int quantMode; - int groupSize; - - read(d, quantMode); - read(d, type); - read(d, groupSize); - - read(d, mDims); - - // mQuantMode = QuantMode(quantMode); - - init(type, groupSize); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void QServeGemmPlugin::init(nvinfer1::DataType dtype, int groupSize) -{ - if (groupSize <= 0) - groupSize = -1; // Per-channel - mGroupSize = groupSize; - mType = dtype; - mRunner = std::make_shared(); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QServeGemmPlugin::clone() const noexcept -{ - auto* plugin = new QServeGemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs QServeGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 6); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool QServeGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (mGroupSize != -1) - { // Per-group - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // uint4 weights packed in int8 - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // int8 weight s2_zeros - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 3: - // int8 weight s2_scales - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // fp16 weight s1_scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 5: - // fp16 activation scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 6: - // fp16 output activation - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - default: return false; - } - } - - else - { // Per-channel - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // uint4 weights packed in int8 - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // fp16 s1_scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 3: - // fp16 s1_szeros - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // fp16 act_sums - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 5: - // fp16 act_scales - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - case 6: - // fp16 output activation - return inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format == TensorFormat::kLINEAR; - default: return false; - } - } -} - -void QServeGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - m_workspaceMaxSize = mRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t QServeGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int QServeGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept -{ - // inputs - - // Per group: - // activation [M, K] int8_t Quantized sint8 activations - // weights [N, K/2] int8_t Quantized uint4 weights (packed as int8_t) - // s2_zeros [K/group_size, N] int8_t Level-2 sint8 scaled zeros of weights - // s2_scales [K/group_size, N] int8_t Level-2 sint8 scales of weights - // s1_scales [N] half Level-1 fp16 scales of weights - // act_scales [M] half Scales of activations - - // Per channel: - // activation [M, K] int8_t Quantized sint8 activations - // weights [N, K/2] int8_t Quantized uint4 weights (packed as int8_t) - // s1_scales [N] half Level-1 scales of weights - // s1_szeros [N] half Level-1 scaled zeros of weights - // act_sums [M] half Per-token sums of activations - // act_scales [M] half Scales of activations - - // outputs - // mat [M(*), N] half - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - // TODO: Implement optimized kernels if (m <= 4) - - if (mGroupSize != -1) - { - ParamsPerGroup params = {reinterpret_cast(inputs[0]), // A - reinterpret_cast(inputs[1]), // B - reinterpret_cast(inputs[2]), // s2_zeros - reinterpret_cast(inputs[3]), // s2_scales - reinterpret_cast(inputs[4]), // s1_scales - reinterpret_cast(inputs[5]), // act_scales - reinterpret_cast(outputs[0]), // C - m, n, k}; - mRunner->gemmPerGroup(params, stream); - } - else - { - ParamsPerChannel params = {reinterpret_cast(inputs[0]), // A - reinterpret_cast(inputs[1]), // B - reinterpret_cast(inputs[2]), // s1_scales - reinterpret_cast(inputs[3]), // s1_szeros - reinterpret_cast(inputs[4]), // act_sums - reinterpret_cast(inputs[5]), // act_scales - reinterpret_cast(outputs[0]), // C - m, n, k}; - mRunner->gemmPerChannel(params, stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QServeGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* QServeGemmPlugin::getPluginType() const noexcept -{ - return QSERVE_GEMM_PLUGIN_NAME; -} - -char const* QServeGemmPlugin::getPluginVersion() const noexcept -{ - return QSERVE_GEMM_PLUGIN_VERSION; -} - -int QServeGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int QServeGemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void QServeGemmPlugin::terminate() noexcept {} - -size_t QServeGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(mQuantMode) + // QuantMode - sizeof(mType) + // dtype - sizeof(mGroupSize) + // GroupSize - sizeof(mDims); // Dimensions -} - -void QServeGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mGroupSize); - write(d, mDims); - - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QServeGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void QServeGemmPlugin::configGemm() {} - -/////////////// - -QServeGemmPluginCreator::QServeGemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.push_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.push_back(PluginField("group_size", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QServeGemmPluginCreator::getPluginName() const noexcept -{ - return QSERVE_GEMM_PLUGIN_NAME; -} - -char const* QServeGemmPluginCreator::getPluginVersion() const noexcept -{ - return QSERVE_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* QServeGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QServeGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - // We do not use any fields for now. - - PluginField const* fields = fc->fields; - - // bool perTokenScaling, perChannelScaling; - DataType dtype{}; - int group_size = -1; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dtype = static_cast(*(static_cast(fields[i].data))); - // Only supports fp16 for now. - assert(dtype == nvinfer1::DataType::kHALF); - } - else if (!strcmp(attrName, "group_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - group_size = *static_cast(fields[i].data); - // Currently only support per-channel or g128. - assert(group_size == -1 || group_size == 128); - } - } - try - { - // QServeGemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - // auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - // QuantMode quantMode = QuantMode::fromQuantAlgo("W4A8_QSERVE"); - auto* obj = new QServeGemmPlugin(dtype, group_size); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QServeGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QServeGemmPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - // auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new QServeGemmPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h b/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h deleted file mode 100644 index 086460863c4f..000000000000 --- a/cpp/tensorrt_llm/plugins/qserveGemmPlugin/qserveGemmPlugin.h +++ /dev/null @@ -1,112 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using QServeGemmRunnerPtr = std::shared_ptr; - -class QServeGemmPlugin : public BasePlugin -{ -public: - // using PluginProfilerPtr = std::shared_ptr; - - QServeGemmPlugin(void const* data, size_t length); - - QServeGemmPlugin(nvinfer1::DataType dtype, int groupSize); - - ~QServeGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType dtype, int groupSize); - - void configGemm(); - - std::string const mLayerName; - - QServeGemmRunnerPtr mRunner; - - tensorrt_llm::common::QuantMode mQuantMode; // Not used for now - GemmDims mDims{}; - - size_t m_workspaceMaxSize; - - // Only supports fp16 output for now. - nvinfer1::DataType mType; - - int mGroupSize; -}; - -class QServeGemmPluginCreator : public BaseCreator -{ -public: - QServeGemmPluginCreator(); - - QServeGemmPluginCreator(void const* data, size_t length); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.cpp b/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.cpp deleted file mode 100644 index 23d0b80390e3..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.cpp +++ /dev/null @@ -1,353 +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. - */ -#include "quantizePerTokenPlugin.h" -#include "tensorrt_llm/kernels/quantization.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels; -using tensorrt_llm::plugins::QuantizePerTokenPluginCreator; -using tensorrt_llm::plugins::QuantizePerTokenPlugin; - -static char const* QUANTIZE_PER_TOKEN_PLUGIN_VERSION{"1"}; -static char const* QUANTIZE_PER_TOKEN_PLUGIN_NAME{"QuantizePerToken"}; -PluginFieldCollection QuantizePerTokenPluginCreator::mFC{}; -std::vector QuantizePerTokenPluginCreator::mPluginAttributes; - -QuantizePerTokenPlugin::QuantizePerTokenPlugin( - nvinfer1::DataType outputType, QuantMode quantMode, bool clampValEnabled, bool sumPerToken) - : mOutputType{outputType} - , mQuantMode{quantMode} - , mClampValEnabled{clampValEnabled} - , mSumPerToken{sumPerToken} -{ - TLLM_CHECK_WITH_INFO(mOutputType == nvinfer1::DataType::kINT8 || mOutputType == nvinfer1::DataType::kFP8, - "Only int8 or fp8 output type is allowed."); - // Check if the quant mode is valid. - TLLM_CHECK_WITH_INFO(mQuantMode.hasPerTokenScaling(), "The quant mode is not valid."); -} - -// Parameterized constructor -QuantizePerTokenPlugin::QuantizePerTokenPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mOutputType); - read(d, mQuantMode); - read(d, mClampValEnabled); - read(d, mSumPerToken); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QuantizePerTokenPlugin::clone() const noexcept -{ - auto* plugin = new QuantizePerTokenPlugin(mOutputType, mQuantMode, mClampValEnabled, mSumPerToken); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs QuantizePerTokenPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs <= 2); - TLLM_CHECK(outputIndex <= 2); - if (outputIndex == 2) - { - // Per token sums. - TLLM_CHECK(mSumPerToken); - } - - if (outputIndex == 0) - { - // Quantized input - return inputs[0]; - } - - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int ii = 0; ii < ret.nbDims - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[ret.nbDims - 1] = exprBuilder.constant(1); - // [M(*), 1] dynamic per token scales or sums - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool QuantizePerTokenPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == 0) - { - // activation - return (inOut[pos].type == nvinfer1::DataType::kFLOAT || inOut[pos].type == nvinfer1::DataType::kHALF -#ifdef ENABLE_BF16 - || inOut[pos].type == nvinfer1::DataType::kBF16 -#endif - ) - && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 1 && mClampValEnabled) - { - // clamp_max_v - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 1 + int(mClampValEnabled)) - { - // quantized activation - return inOut[pos].type == mOutputType && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 2 + int(mClampValEnabled)) - { - // scales - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - else if (pos == 3 + int(mClampValEnabled)) - { - TLLM_CHECK(mSumPerToken); - // per-token sums - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - - // Never should be here - assert(false); - return false; -} - -void QuantizePerTokenPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t QuantizePerTokenPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -template -void QuantizePerTokenPlugin::dispatchDataType(void* output, void const* input, void const* clampValPtr, void* scalePtr, - void* sumPtr, int dim0, int dim1, cudaStream_t stream) noexcept -{ - // inputs - // activation [dim0(*), dim1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [dim0(*), dim1] - // scale_tokens [dim0(*), 1] - - invokePerTokenQuantization(reinterpret_cast(output), reinterpret_cast(input), dim0, dim1, - reinterpret_cast(clampValPtr), reinterpret_cast(scalePtr), - reinterpret_cast(sumPtr), mQuantMode, stream); -} - -int QuantizePerTokenPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // activation [M(*), K] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [M(*), K] Quantized activations. - // scale_tokens [M(*), 1] Per-token scales. - // token_sums [M(*), 1] (Optional) Per-token sums of all the channels (before quantization). - - int64_t m = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m *= inputDesc[0].dims.d[ii]; - } - int64_t const k = inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]; - - void const* clampValPtr = mClampValEnabled ? inputs[1] : nullptr; - void* sumPtr = mSumPerToken ? outputs[2] : nullptr; - - if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kINT8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kFP8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#endif // ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kINT8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kFP8) - { - dispatchDataType(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#endif // ENABLE_FP8 -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kINT8) - { - dispatchDataType<__nv_bfloat16, int8_t>(outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kFP8) - { - dispatchDataType<__nv_bfloat16, __nv_fp8_e4m3>( - outputs[0], inputs[0], clampValPtr, outputs[1], sumPtr, m, k, stream); - } -#endif // ENABLE_FP8 -#endif // ENABLE_BF16 - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QuantizePerTokenPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(nbInputs >= 1); - TLLM_CHECK(index <= 2); - if (index == 2) - { - // Per token sums. - TLLM_CHECK(mSumPerToken); - } - return index == 0 ? mOutputType : nvinfer1::DataType::kFLOAT; -} - -// IPluginV2 Methods - -char const* QuantizePerTokenPlugin::getPluginType() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_NAME; -} - -char const* QuantizePerTokenPlugin::getPluginVersion() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_VERSION; -} - -int QuantizePerTokenPlugin::getNbOutputs() const noexcept -{ - return 2 + static_cast(mSumPerToken); -} - -int QuantizePerTokenPlugin::initialize() noexcept -{ - return 0; -} - -void QuantizePerTokenPlugin::terminate() noexcept {} - -size_t QuantizePerTokenPlugin::getSerializationSize() const noexcept -{ - return sizeof(mOutputType) + sizeof(mQuantMode) + sizeof(mClampValEnabled) + sizeof(mSumPerToken); -} - -void QuantizePerTokenPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mOutputType); - write(d, mQuantMode); - write(d, mClampValEnabled); - write(d, mSumPerToken); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QuantizePerTokenPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -QuantizePerTokenPluginCreator::QuantizePerTokenPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("clamp_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("sum_per_token", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QuantizePerTokenPluginCreator::getPluginName() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_NAME; -} - -char const* QuantizePerTokenPluginCreator::getPluginVersion() const noexcept -{ - return QUANTIZE_PER_TOKEN_PLUGIN_VERSION; -} - -PluginFieldCollection const* QuantizePerTokenPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QuantizePerTokenPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginFieldParser p{fc->nbFields, fc->fields}; - try - { - auto* obj = new QuantizePerTokenPlugin(static_cast(p.getScalar("type_id").value()), - QuantMode(p.getScalar("quant_mode").value()), - static_cast(p.getScalar("clamp_enabled").value()), - static_cast(p.getScalar("sum_per_token").value())); - - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QuantizePerTokenPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QuantizePerTokenPlugin::destroy() - try - { - auto* obj = new QuantizePerTokenPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h b/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h deleted file mode 100644 index 47b218acfd28..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizePerTokenPlugin/quantizePerTokenPlugin.h +++ /dev/null @@ -1,104 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class QuantizePerTokenPlugin : public BasePlugin -{ -public: - QuantizePerTokenPlugin(nvinfer1::DataType outputType, tensorrt_llm::common::QuantMode quantMode, - bool clampValEnabled, bool sumPerToken); - - QuantizePerTokenPlugin(void const* data, size_t length); - - ~QuantizePerTokenPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - void dispatchDataType(void* output, void const* input, void const* clampValPtr, void* scalePtr, void* sumPtr, - int dim0, int dim1, cudaStream_t stream) noexcept; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - std::string const mLayerName; - // The quantized output data type. - nvinfer1::DataType mOutputType; - // The quantization mode. - tensorrt_llm::common::QuantMode mQuantMode; - // Do we clamp the input tensor ? - bool mClampValEnabled; - // Do we output the per-token sum? - bool mSumPerToken; -}; - -class QuantizePerTokenPluginCreator : public BaseCreator -{ -public: - QuantizePerTokenPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp b/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp deleted file mode 100644 index cacb32b809bf..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp +++ /dev/null @@ -1,250 +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. - */ -#include "quantizeTensorPlugin.h" -#include "tensorrt_llm/kernels/quantization.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using tensorrt_llm::plugins::QuantizeTensorPluginCreator; -using tensorrt_llm::plugins::QuantizeTensorPlugin; - -static char const* QUANTIZE_TENSOR_PLUGIN_VERSION{"1"}; -static char const* QUANTIZE_TENSOR_PLUGIN_NAME{"QuantizeTensor"}; -PluginFieldCollection QuantizeTensorPluginCreator::mFC{}; -std::vector QuantizeTensorPluginCreator::mPluginAttributes; - -QuantizeTensorPlugin::QuantizeTensorPlugin() {} - -// Parameterized constructor -QuantizeTensorPlugin::QuantizeTensorPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QuantizeTensorPlugin::clone() const noexcept -{ - return new QuantizeTensorPlugin(*this); -} - -nvinfer1::DimsExprs QuantizeTensorPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(outputIndex < 1); - // Quantized input - return inputs[0]; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool QuantizeTensorPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return (inOut[pos].type == nvinfer1::DataType::kFLOAT || inOut[pos].type == nvinfer1::DataType::kHALF -#ifdef ENABLE_BF16 - || inOut[pos].type == nvinfer1::DataType::kBF16 -#endif - ) - && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // scales - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // quantized activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - TLLM_CHECK(false); - return false; - } -} - -void QuantizeTensorPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t QuantizeTensorPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int QuantizeTensorPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // activation [M(*), K] - // scale [1, 1] - // outputs - // quant [M(*), K] - - int64_t numElts = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims; ++ii) - { - numElts *= inputDesc[0].dims.d[ii]; - } - - if (inputDesc[0].type == DataType::kFLOAT) - { - invokeQuantization(reinterpret_cast(outputs[0]), reinterpret_cast(inputs[0]), - numElts, reinterpret_cast(inputs[1]), stream, mProp.maxGridSize[0]); - } - else if (inputDesc[0].type == DataType::kHALF) - { - invokeQuantization(reinterpret_cast(outputs[0]), reinterpret_cast(inputs[0]), - numElts, reinterpret_cast(inputs[1]), stream, mProp.maxGridSize[0]); - } -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16) - { - invokeQuantization<__nv_bfloat16>(reinterpret_cast(outputs[0]), - reinterpret_cast<__nv_bfloat16 const*>(inputs[0]), numElts, reinterpret_cast(inputs[1]), - stream, mProp.maxGridSize[0]); - } -#endif - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QuantizeTensorPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(nbInputs == 2); - TLLM_CHECK(index == 0); - return nvinfer1::DataType::kINT8; -} - -// IPluginV2 Methods - -char const* QuantizeTensorPlugin::getPluginType() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_NAME; -} - -char const* QuantizeTensorPlugin::getPluginVersion() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_VERSION; -} - -int QuantizeTensorPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int QuantizeTensorPlugin::initialize() noexcept -{ - int deviceId = 0; - tensorrt_llm::common::check_cuda_error(cudaGetDevice(&deviceId)); - tensorrt_llm::common::check_cuda_error(cudaGetDeviceProperties(&mProp, deviceId)); - return 0; -} - -void QuantizeTensorPlugin::terminate() noexcept {} - -size_t QuantizeTensorPlugin::getSerializationSize() const noexcept -{ - return 0; -} - -void QuantizeTensorPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QuantizeTensorPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -QuantizeTensorPluginCreator::QuantizeTensorPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QuantizeTensorPluginCreator::getPluginName() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_NAME; -} - -char const* QuantizeTensorPluginCreator::getPluginVersion() const noexcept -{ - return QUANTIZE_TENSOR_PLUGIN_VERSION; -} - -PluginFieldCollection const* QuantizeTensorPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QuantizeTensorPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - try - { - auto* obj = new QuantizeTensorPlugin(); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QuantizeTensorPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QuantizeTensorPlugin::destroy() - try - { - auto* obj = new QuantizeTensorPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h b/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h deleted file mode 100644 index 6f1ce864ec35..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.h +++ /dev/null @@ -1,92 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class QuantizeTensorPlugin : public BasePlugin -{ -public: - QuantizeTensorPlugin(); - - QuantizeTensorPlugin(void const* data, size_t length); - - ~QuantizeTensorPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - cudaDeviceProp mProp; -}; - -class QuantizeTensorPluginCreator : public BaseCreator -{ -public: - QuantizeTensorPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.cpp b/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.cpp deleted file mode 100644 index b5eaffeeda2a..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.cpp +++ /dev/null @@ -1,301 +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. - */ -#include "quantizeToFP4Plugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/quantization.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::QuantizeToFP4PluginCreator; -using tensorrt_llm::plugins::QuantizeToFP4Plugin; - -constexpr nvinfer1::DataType FP4_DTYPE = nvinfer1::DataType::kFP4; -constexpr nvinfer1::DataType FP8_DTYPE = nvinfer1::DataType::kFP8; - -static char const* QUANT_FP4_PLUGIN_VERSION{"1"}; -static char const* QUANT_FP4_PLUGIN_NAME{"QuantizeToFP4"}; -PluginFieldCollection QuantizeToFP4PluginCreator::mFC{}; -std::vector QuantizeToFP4PluginCreator::mPluginAttributes; - -QuantizeToFP4Plugin::QuantizeToFP4Plugin(){}; - -// Parameterized constructor -QuantizeToFP4Plugin::QuantizeToFP4Plugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* QuantizeToFP4Plugin::clone() const noexcept -{ - auto* plugin = new QuantizeToFP4Plugin(); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs QuantizeToFP4Plugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - // Quantized output in FP4 datatype. - if (outputIndex == 0) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - // // Div up by 16 as the storage type has 16 FP4 values per element. - // ret.d[ret.nbDims - 1] - // = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], - // *exprBuilder.constant(16)); - return ret; - } - // Scaling Factors in FP8. - else if (outputIndex == 1) - { - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - // Sequence dimension or token dimension. - // Pad to multiple of 128. - auto dimM - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 2], *exprBuilder.constant(128)); - ret.d[ret.nbDims - 2] = exprBuilder.operation(DimensionOperation::kPROD, *dimM, *exprBuilder.constant(128)); - // Hidden size dimension. - // Div (rounding up) by 16 since 16 elements share one SF and SF padded to k%4==0. - ret.d[ret.nbDims - 1] - = exprBuilder.operation(DimensionOperation::kCEIL_DIV, *ret.d[ret.nbDims - 1], *exprBuilder.constant(16)); - return ret; - } - return DimsExprs{}; -} - -bool QuantizeToFP4Plugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - // half input + float global_sf + fp4 output (e2m1) + fp8 SF output. - int const totalPoses = 2 + 2; - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 2); - switch (pos) - { - case 0: - return (inOut[pos].type == nvinfer1::DataType::kHALF || inOut[pos].type == nvinfer1::DataType::kBF16 - || inOut[pos].type == nvinfer1::DataType::kFP8) - && (inOut[pos].format == TensorFormat::kLINEAR); - case 1: return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - case 2: return (inOut[pos].type == FP4_DTYPE) && (inOut[pos].format == TensorFormat::kLINEAR); - case 3: return (inOut[pos].type == FP8_DTYPE) && (inOut[pos].format == TensorFormat::kLINEAR); - default: break; - } - return false; -} - -void QuantizeToFP4Plugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t QuantizeToFP4Plugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -int QuantizeToFP4Plugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] half data type - // SF scale [1] float data type - // used to scale SF from input range to fp8 range (448.f / (MaxVal of input / 6.f)) - // outputs - // output [M(*), N] fp4 storage (E2M1) - // SF output [M, N / 16] fp8 storage (UE4M3) - - int64_t m64 = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) - { - m64 *= inputDesc[0].dims.d[i]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - TLLM_CHECK_WITH_INFO(n % 16 == 0, "the N dimension must be multiple of 16."); - - float const* SFScale = static_cast(inputs[1]); - int64_t* output = reinterpret_cast(outputs[0]); - int32_t* SFoutput = reinterpret_cast(outputs[1]); - - DataType inputDtype = inputDesc[0].type; - - switch (inputDtype) - { - case DataType::kHALF: - { - auto input = reinterpret_cast(inputs[0]); - invokeFP4Quantization(1, m, n, input, SFScale, output, SFoutput, false, QuantizationSFLayout::SWIZZLED, - mMultiProcessorCount, stream); - break; - } - - case DataType::kBF16: - { - auto input = reinterpret_cast<__nv_bfloat16 const*>(inputs[0]); - invokeFP4Quantization(1, m, n, input, SFScale, output, SFoutput, false, QuantizationSFLayout::SWIZZLED, - mMultiProcessorCount, stream); - break; - } - - case DataType::kFP8: - { - auto input = reinterpret_cast<__nv_fp8_e4m3 const*>(inputs[0]); - invokeFP4Quantization(1, m, n, input, SFScale, output, SFoutput, false, QuantizationSFLayout::SWIZZLED, - mMultiProcessorCount, stream); - break; - } - - default: TLLM_LOG_ERROR("only half, bfloat16 and fp8 data type are supported."); break; - } - - // Use UE4M3 scales by default. - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType QuantizeToFP4Plugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (index == 0) - { - // Output 0 quantized output. - return FP4_DTYPE; - } - // Output 1 SF (scaling factors). - return FP8_DTYPE; -} - -// IPluginV2 Methods - -char const* QuantizeToFP4Plugin::getPluginType() const noexcept -{ - return QUANT_FP4_PLUGIN_NAME; -} - -char const* QuantizeToFP4Plugin::getPluginVersion() const noexcept -{ - return QUANT_FP4_PLUGIN_VERSION; -} - -int QuantizeToFP4Plugin::getNbOutputs() const noexcept -{ - return 2; -} - -int QuantizeToFP4Plugin::initialize() noexcept -{ - return 0; -} - -void QuantizeToFP4Plugin::terminate() noexcept {} - -size_t QuantizeToFP4Plugin::getSerializationSize() const noexcept -{ - return 0; -} - -void QuantizeToFP4Plugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - TLLM_CHECK(d == a + getSerializationSize()); -} - -void QuantizeToFP4Plugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -QuantizeToFP4PluginCreator::QuantizeToFP4PluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* QuantizeToFP4PluginCreator::getPluginName() const noexcept -{ - return QUANT_FP4_PLUGIN_NAME; -} - -char const* QuantizeToFP4PluginCreator::getPluginVersion() const noexcept -{ - return QUANT_FP4_PLUGIN_VERSION; -} - -PluginFieldCollection const* QuantizeToFP4PluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* QuantizeToFP4PluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - try - { - auto* obj = new QuantizeToFP4Plugin(); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* QuantizeToFP4PluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call QuantizeToFP4Plugin::destroy() - try - { - auto* obj = new QuantizeToFP4Plugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h b/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h deleted file mode 100644 index b584837a447a..000000000000 --- a/cpp/tensorrt_llm/plugins/quantizeToFP4Plugin/quantizeToFP4Plugin.h +++ /dev/null @@ -1,90 +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. - */ -#pragma once - -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class QuantizeToFP4Plugin : public BasePlugin -{ -public: - QuantizeToFP4Plugin(); - - QuantizeToFP4Plugin(void const* data, size_t length); - - ~QuantizeToFP4Plugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - const std::string mLayerName; - int const mMultiProcessorCount = tensorrt_llm::common::getMultiProcessorCount(); -}; - -class QuantizeToFP4PluginCreator : public BaseCreator -{ -public: - QuantizeToFP4PluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.cpp b/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.cpp deleted file mode 100644 index 16d0bf2dc356..000000000000 --- a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.cpp +++ /dev/null @@ -1,452 +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. - */ -#include "rmsnormQuantizationPlugin.h" -#include "pluginUtils.h" -#include "tensorrt_llm/kernels/rmsnormKernels.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::RmsnormQuantizationPluginCreator; -using tensorrt_llm::plugins::RmsnormQuantizationPlugin; - -static char const* RMSNORM_QUANTIZATION_PLUGIN_VERSION{"1"}; -static char const* RMSNORM_QUANTIZATION_PLUGIN_NAME{"RmsnormQuantization"}; -PluginFieldCollection RmsnormQuantizationPluginCreator::mFC{}; -std::vector RmsnormQuantizationPluginCreator::mPluginAttributes; - -RmsnormQuantizationPlugin::RmsnormQuantizationPlugin(float eps, bool dynamicActivationScaling, bool sumPerToken, - bool clampValEnabled, QuantMode quantMode, nvinfer1::DataType type, nvinfer1::DataType outputType) - : mEps(eps) - , mDynActScaling(dynamicActivationScaling) - , mType(type) - , mOutputType{outputType} - , mClampValEnabled{clampValEnabled} - , mQuantMode{quantMode} - , mSumPerToken(sumPerToken) -{ - TLLM_CHECK_WITH_INFO(mOutputType == nvinfer1::DataType::kINT8 || mOutputType == nvinfer1::DataType::kFP8, - "Only int8 or fp8 output type is allowed."); - // Check if the quant mode is valid. - TLLM_CHECK_WITH_INFO(mQuantMode.hasPerTokenScaling(), "The quant mode is not valid."); -} - -// Parameterized constructor -RmsnormQuantizationPlugin::RmsnormQuantizationPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mEps); - read(d, mDynActScaling); - read(d, mSumPerToken); - read(d, mClampValEnabled); - read(d, mQuantMode); - read(d, mType); - read(d, mOutputType); - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* RmsnormQuantizationPlugin::clone() const noexcept -{ - auto* plugin = new RmsnormQuantizationPlugin( - mEps, mDynActScaling, mSumPerToken, mClampValEnabled, mQuantMode, mType, mOutputType); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -nvinfer1::DimsExprs RmsnormQuantizationPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - // Quantized output - return inputs[outputIndex]; - } - - // Dynamic scaling or per-token sum if enabled. - try - { - if (outputIndex == 1) - { - TLLM_CHECK(mDynActScaling); - } - else if (outputIndex == 2) - { - TLLM_CHECK(mSumPerToken); - } - else - { - TLLM_CHECK(false); - } - - DimsExprs ret; - ret.nbDims = inputs[0].nbDims; - for (int di = 0; di < ret.nbDims - 1; ++di) - { - ret.d[di] = inputs[0].d[di]; - } - ret.d[ret.nbDims - 1] = exprBuilder.constant(1); - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool RmsnormQuantizationPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - int const totalPoses - = 6 + static_cast(mClampValEnabled) + static_cast(mDynActScaling) + static_cast(mSumPerToken); - TLLM_CHECK(0 <= pos && pos < totalPoses); - TLLM_CHECK(nbInputs == 4 + static_cast(mClampValEnabled)); - if (pos < nbInputs) - { - if (pos < 3) - { - // activation, weight, bias - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 3) - { - // scale - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 4 && mClampValEnabled) - { - // clamp_max_v - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - } - } - else if (pos == 4 + int(mClampValEnabled)) - { - // Quantized output - return (inOut[pos].type == mOutputType) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 5 + int(mClampValEnabled)) - { - // Dynamic scaling if enabled - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (pos == 6 + int(mClampValEnabled)) - { - // Per-token activation sum if enabled - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - - // Never should be here - TLLM_CHECK_WITH_INFO(false, "The input/output is not supported."); - return false; -} - -void RmsnormQuantizationPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t RmsnormQuantizationPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return 0; -} - -template -void RmsnormQuantizationPlugin::dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, - float const eps, int const tokens, int const hidden_dim, cudaStream_t stream, void const* clampValPtr, - void const* scale, void* dynamic_scale, void* sum_per_token, void* normed_output_quant) noexcept -{ - // inputs - // activation [dim0(*), dim1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // quant [dim0(*), dim1] - // scale_tokens [dim0(*), 1] - - invokeGeneralRmsNorm(reinterpret_cast(out), reinterpret_cast(input), - reinterpret_cast(gamma), reinterpret_cast(beta), eps, tokens, hidden_dim, mQuantMode, - stream, reinterpret_cast(clampValPtr), reinterpret_cast(scale), - reinterpret_cast(dynamic_scale), reinterpret_cast(sum_per_token), - reinterpret_cast(normed_output_quant)); -} - -int RmsnormQuantizationPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // input [M(*), N] - // weight [N, ] - // bias [N, ] - // scale_to_int [1] - // clamp_value [2], contains min val, and max val (optional) - // outputs - // output [M(*), N] Normalized activations, potentially with quantization applied. - // dynamic_scaling [M(*), 1] (Optional) Per-token scales if quantization is enabled. - // token_sums [M(*), 1] (Optional) Per-token sums of all the channels (before quantization). - - int64_t m64 = 1; - for (int i = 0; i < inputDesc[0].dims.nbDims - 1; ++i) - { - m64 *= inputDesc[0].dims.d[i]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - - void const* input = inputs[0]; - void const* weight = inputs[1]; - void const* bias = inputs[2]; - void const* scale = inputs[3]; - void const* clampValPtr = mClampValEnabled ? inputs[4] : nullptr; - void* output = outputs[0]; - void* dynamic_scale = mDynActScaling ? outputs[1] : nullptr; - void* sum_per_token = mSumPerToken ? outputs[2] : nullptr; - - if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kINT8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kFLOAT && mOutputType == DataType::kFP8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kINT8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kHALF && mOutputType == DataType::kFP8) - { - dispatchDataType( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#ifdef ENABLE_BF16 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kINT8) - { - dispatchDataType<__nv_bfloat16, int8_t>( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#ifdef ENABLE_FP8 - else if (inputDesc[0].type == DataType::kBF16 && mOutputType == DataType::kFP8) - { - dispatchDataType<__nv_bfloat16, __nv_fp8_e4m3>( - nullptr, input, weight, bias, mEps, m, n, stream, clampValPtr, scale, dynamic_scale, sum_per_token, output); - } -#endif // ENABLE_FP8 -#endif // ENABLE_BF16 - sync_check_cuda_error(stream); - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType RmsnormQuantizationPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - assert(index <= 2); - - if (index == 0) - { - // Output 0 quantized output of layer norm - return mOutputType; - } - if (index == 1) - { - assert(mDynActScaling); - // Output 1 dynamic act scaling - return nvinfer1::DataType::kFLOAT; - } - // index == 2 - { - assert(mDynActScaling && mSumPerToken); - // Output 2 per token sum - return nvinfer1::DataType::kFLOAT; - } -} - -// IPluginV2 Methods - -char const* RmsnormQuantizationPlugin::getPluginType() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* RmsnormQuantizationPlugin::getPluginVersion() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_VERSION; -} - -int RmsnormQuantizationPlugin::getNbOutputs() const noexcept -{ - return 1 + static_cast(mDynActScaling) + static_cast(mSumPerToken); -} - -int RmsnormQuantizationPlugin::initialize() noexcept -{ - return 0; -} - -void RmsnormQuantizationPlugin::terminate() noexcept {} - -size_t RmsnormQuantizationPlugin::getSerializationSize() const noexcept -{ - return sizeof(mOutputType) + sizeof(mClampValEnabled) + sizeof(mEps) + sizeof(mDynActScaling) + sizeof(mSumPerToken) - + sizeof(mType) + sizeof(mQuantMode); -} - -void RmsnormQuantizationPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mEps); - write(d, mDynActScaling); - write(d, mSumPerToken); - write(d, mClampValEnabled); - write(d, mQuantMode); - write(d, mType); - write(d, mOutputType); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void RmsnormQuantizationPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -RmsnormQuantizationPluginCreator::RmsnormQuantizationPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32)); - mPluginAttributes.emplace_back(PluginField("dyn_act_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("sum_per_token", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("clamp_enabled", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_mode", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("out_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* RmsnormQuantizationPluginCreator::getPluginName() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_NAME; -} - -char const* RmsnormQuantizationPluginCreator::getPluginVersion() const noexcept -{ - return RMSNORM_QUANTIZATION_PLUGIN_VERSION; -} - -PluginFieldCollection const* RmsnormQuantizationPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* RmsnormQuantizationPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType outputType{}; - QuantMode quantMode; - bool clampValEnabled = false; - float eps{}; - nvinfer1::DataType type{}; - bool dynamicActivationScaling{}; - bool sumPerToken{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "quant_mode")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - quantMode = QuantMode(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "out_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - outputType = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "clamp_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - clampValEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "eps")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - eps = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dyn_act_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dynamicActivationScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "sum_per_token")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - sumPerToken = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new RmsnormQuantizationPlugin( - eps, dynamicActivationScaling, sumPerToken, clampValEnabled, quantMode, type, outputType); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* RmsnormQuantizationPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call RmsnormQuantizationPlugin::destroy() - try - { - auto* obj = new RmsnormQuantizationPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h b/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h deleted file mode 100644 index 762a9bb8de1b..000000000000 --- a/cpp/tensorrt_llm/plugins/rmsnormQuantizationPlugin/rmsnormQuantizationPlugin.h +++ /dev/null @@ -1,108 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -class RmsnormQuantizationPlugin : public BasePlugin -{ -public: - RmsnormQuantizationPlugin(float eps, bool dynamicActivationScaling, bool sumPerToken, bool clampValEnabled, - tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, nvinfer1::DataType outputType); - - RmsnormQuantizationPlugin(void const* data, size_t length); - - ~RmsnormQuantizationPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - template - void dispatchDataType(void* out, void const* input, void const* gamma, void const* beta, float const eps, - int const tokens, int const hidden_dim, cudaStream_t stream, void const* clampValPtr, void const* scale, - void* dynamic_scale, void* normed_output_quant, void* act_sum) noexcept; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - float mEps; - bool mDynActScaling; - nvinfer1::DataType mType; - - std::string const mLayerName; - // The quantized output data type. - nvinfer1::DataType mOutputType; - // Do we clamp the input tensor ? - bool mClampValEnabled; - // The quantization mode. - tensorrt_llm::common::QuantMode mQuantMode; - // Should we output the sum of channels per-token? (Used by QServe GEMM) - bool mSumPerToken; -}; - -class RmsnormQuantizationPluginCreator : public BaseCreator -{ -public: - RmsnormQuantizationPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/selectiveScanPlugin/CMakeLists.txt deleted file mode 100644 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.cpp b/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.cpp deleted file mode 100644 index 3e60182f28c2..000000000000 --- a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.cpp +++ /dev/null @@ -1,594 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "selectiveScanPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::SelectiveScanPluginCreator; -using tensorrt_llm::plugins::SelectiveScanPlugin; - -static char const* SELECTIVE_SCAN_PLUGIN_VERSION{"1"}; -static char const* SELECTIVE_SCAN_PLUGIN_NAME{"SelectiveScan"}; -PluginFieldCollection SelectiveScanPluginCreator::mFC{}; -std::vector SelectiveScanPluginCreator::mPluginAttributes; - -SelectiveScanPlugin::SelectiveScanPlugin(int dim, int dstate, int dtRank, int nHeads, int nGroups, int chunkSize, - bool deltaSoftplus, nvinfer1::DataType type, bool removePadding, bool pagedState, bool zEnabled, bool isMamba2) - : mDim(dim) - , mDState(dstate) - , mDtRank(dtRank) - , mNHeads(nHeads) - , mNGroups(nGroups) - , mChunkSize(chunkSize) - , mDeltaSoftplus(deltaSoftplus) - , mType(type) - , mRemovePadding(removePadding) - , mPagedState(pagedState) - , mZEnabled(zEnabled) - , mIsMamba2(isMamba2) - , mDriver(tensorrt_llm::common::CUDADriverWrapper::getInstance()) -{ - TLLM_CHECK_WITH_INFO( - (mChunkSize == 256 || mChunkSize == 128) || (!mIsMamba2), "Only support CHUNK_SIZE 256 or 128"); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// Parameterized constructor -SelectiveScanPlugin::SelectiveScanPlugin(void const* data, size_t length) - : mDriver(tensorrt_llm::common::CUDADriverWrapper::getInstance()) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mDim); - read(d, mDState); - read(d, mDtRank); - read(d, mNHeads); - read(d, mNGroups); - read(d, mChunkSize); - read(d, mDeltaSoftplus); - read(d, mType); - read(d, mRemovePadding); - read(d, mPagedState); - read(d, mZEnabled); - read(d, mIsMamba2); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO( - (mChunkSize == 256 || mChunkSize == 128) || (!mIsMamba2), "Only support CHUNK_SIZE 256 or 128"); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF), - "Only support float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* SelectiveScanPlugin::clone() const noexcept -{ - auto* plugin = new SelectiveScanPlugin(mDim, mDState, mDtRank, mNHeads, mNGroups, mChunkSize, mDeltaSoftplus, mType, - mRemovePadding, mPagedState, mZEnabled, mIsMamba2); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// output_tensor: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// state: [batch_size, dstate, dim] -nvinfer1::DimsExprs SelectiveScanPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - if (outputIndex == 0) - { - if (mIsMamba2) - { - auto ret = inputs[getInputTensorIdx()]; - ret.d[mRemovePadding ? 1 : 2] = exprBuilder.constant(mDim); - return ret; - } - else - { - return inputs[getInputTensorIdx()]; - } - } - return inputs[getStateIdx()]; -} - -bool SelectiveScanPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos == getHostRequestTypesIdx() || pos == getLastTokenIdsIdx() - || (mRemovePadding && pos == getHostContextLengthIdx()) || (mPagedState && pos == getSlotMappingIdx())) - { - return inOut[pos].type == nvinfer1::DataType::kINT32; - } - else if (pos == getAIdx() || pos == getDeltaBiasIdx() || pos == getDIdx()) - { - return (inOut[pos].type == nvinfer1::DataType::kFLOAT) && (inOut[pos].format == TensorFormat::kLINEAR); - } - else if (mPagedState && pos == getStateIdx()) - { - return inOut[pos].type == nvinfer1::DataType::kINT64; - } - else - { - return (inOut[pos].type == mType) && (inOut[pos].format == TensorFormat::kLINEAR); - } -} - -void SelectiveScanPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t SelectiveScanPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - if (!mIsMamba2) - return 0; - - int const NUM_BUFFERS = 6; - size_t workspaces[NUM_BUFFERS]; - - if (mRemovePadding) - { - int B = inputs[getLastTokenIdsIdx()].dims.d[0]; - int BxL = inputs[getInputTensorIdx()].dims.d[0]; // num_tokens - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int BxC = (BxL + Q - 1) / Q + B; - - workspaces[0] = long(BxC) * H * N * P * 2; // g_mxOs_ - workspaces[1] = long(BxC) * H * N * P * 4; // g_mxSt_ in float - workspaces[2] = long(BxC) * H * Q * 4; // g_mxdc_ in float - workspaces[3] = long(BxC) * H * Q * 4; // g_mxdA_ in float - workspaces[4] = long(BxC) * G * Q * Q * 2; // g_mxCB_ - workspaces[5] = 1024; // TMA descs - } - else - { - int B = inputs[getInputTensorIdx()].dims.d[0]; - int L = inputs[getInputTensorIdx()].dims.d[1]; - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int C = (L + Q - 1) / Q; - - workspaces[0] = long(B * C) * H * N * P * 2; // g_mxOs_ - workspaces[1] = long(B * C) * H * N * P * 4; // g_mxSt_ in float - workspaces[2] = long(B * C) * H * Q * 4; // g_mxdc_ in float - workspaces[3] = long(B * C) * H * Q * 4; // g_mxdA_ in float - workspaces[4] = long(B * C) * G * Q * Q * 2; // g_mxCB_ - workspaces[5] = 1024; // TMA descs - } - - return calculateTotalWorkspaceSize(workspaces, NUM_BUFFERS); -} - -void SelectiveScanPlugin::setSSMParams(SSMParamsBase& params, const size_t batch, const size_t dim, - const size_t maxSeqLen, const size_t numTokens, const size_t dstate, const size_t dtRank, const size_t nHeads, - const size_t nGroups, const size_t chunkSize, void* statePtr, void const* x, void const* delta, - void const* deltaBias, void const* A, void const* BC, void const* D, void const* z, void* osPtr, void* stPtr, - void* dcPtr, void* dAPtr, void* cbPtr, void* descPtr, int const* lastTokenIds, int const* slotMapping, void* out, - bool deltaSoftplus, bool removePadding) -{ - // Reset the parameters - memset(¶ms, 0, sizeof(params)); - - params.batch = batch; - params.dim = dim; - params.max_seqlen = maxSeqLen; - params.num_tokens = numTokens; - params.dstate = dstate; - params.dt_rank = dtRank; - params.nheads = nHeads; - params.ngroups = nGroups; - params.chunk_size = chunkSize; - - params.delta_softplus = deltaSoftplus; - params.remove_padding = removePadding; - params.is_mamba2 = mIsMamba2; - - // Set the pointers and strides. - params.u_ptr = const_cast(x); - params.delta_ptr = const_cast(delta); - params.A_ptr = const_cast(A); - params.BC_ptr = const_cast(BC); - params.D_ptr = const_cast(D); - params.delta_bias_ptr = const_cast(deltaBias); - params.out_ptr = out; - params.x_ptr = statePtr; - params.z_ptr = const_cast(z); - params.Os_ptr = osPtr; - params.St_ptr = stPtr; - params.dc_ptr = dcPtr; - params.dA_ptr = dAPtr; - params.CB_ptr = cbPtr; - params.desc_ptr = descPtr; - params.last_token_ids_ptr = lastTokenIds; - params.slot_mapping_ptr = slotMapping; -} - -template -int SelectiveScanPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, max_seq_len, dim] or [num_tokens, dim] - // 1. state mamba: [batch_size, dstate, dim] or host [1] containing only pointer for paged_state - // mamba2: [batch_size, nheads, dstate, dim] or host [1] containing only pointer for paged_state - // 2. delta, mamba: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding - // mamba2: [batch_size, seq_len, nheads] or [num_tokens, nheads] for remove_input_padding - // 3. delta_bias, [dim] for mamba, [nheads] for mamba2 - // 4. A, [dstate, dim] for mamba, [nheads] for mamba2 - // 5. BC, mamba: [batch_size, seq_len, dstate * 2] or [num_tokens, dstate * 2] for remove_input_padding - // mamba2: [batch_size, seq_len, ngroups * dstate * 2] or [num_tokens, ngroups * dstate * 2] for - // remove_input_padding - // 6. D, [dim] for mamba, [nheads] for mamba2 - // 7. host_request_types [batch_size] int32. 0: context; 1: generation. - // 8. last_token_ids [batch_size] int32 - // 9. host_context_lengths [batch_size] int32, optional for remove_input_padding - // 10. state_slot_mapping [batch_size] int32, optional for paged state - // 11. z [batch_size, max_seq_len, dim] or [num_tokens, dim] - // outputs - // 0. output_tensor [batch_size, max_seq_len, dim] or [num_tokens, dim] - // 1. state, [batch_size, dstate, dim] for mamba, [batch_size, nheads, dstate, dim] for mamba2 - auto const batch_size = inputDesc[getHostRequestTypesIdx()].dims.d[0]; - int max_seq_len; - if (mRemovePadding) - { - int const* host_context_length = static_cast(inputs[getHostContextLengthIdx()]); - max_seq_len = *std::max_element(host_context_length, host_context_length + batch_size); - } - else - { - max_seq_len = inputDesc[getInputTensorIdx()].dims.d[1]; - } - - // only support context or generation, not for both of them - RequestType const* reqTypes = static_cast(inputs[getHostRequestTypesIdx()]); - - SSMParamsBase ssm_params; - - int const* slotMapping = mPagedState ? static_cast(inputs[getSlotMappingIdx()]) : nullptr; - void const* z = mZEnabled ? inputs[getZIdx()] : nullptr; - - void* statePtr = mPagedState ? *reinterpret_cast(const_cast(inputs[getStateIdx()])) : outputs[1]; - - // Workspace pointer shift - int8_t* workspace_byte_ptr = reinterpret_cast(workspace); - size_t offset = 0; - - T* mxOs = nullptr; - float* mxSt = nullptr; - float* mxdc = nullptr; - float* mxdA = nullptr; - T* mxCB = nullptr; - void* descs = nullptr; - - if (!mIsMamba2 || reqTypes[0] == RequestType::kGENERATION) /* no workspace needed */ - ; - else if (mRemovePadding) - { - int B = inputDesc[getLastTokenIdsIdx()].dims.d[0]; - int BxL = inputDesc[getInputTensorIdx()].dims.d[0]; // num_tokens - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int BxC = (BxL + Q - 1) / Q + B; - - mxOs = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * N * P * 2)); - mxSt = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * N * P * 4)); - mxdc = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * Q * 4)); - mxdA = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * H * Q * 4)); - mxCB = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(BxC) * G * Q * Q * 2)); - descs = nextWorkspacePtr(workspace_byte_ptr, offset, 1024); - } - else - { - int B = inputDesc[getInputTensorIdx()].dims.d[0]; - int L = inputDesc[getInputTensorIdx()].dims.d[1]; - int H = mNHeads; - int P = mDim / H; - int G = mNGroups; - int N = mDState; - int Q = mChunkSize; - int C = (L + Q - 1) / Q; - - mxOs = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * N * P * 2)); - mxSt = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * N * P * 4)); - mxdc = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * Q * 4)); - mxdA = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * H * Q * 4)); - mxCB = reinterpret_cast(nextWorkspacePtr(workspace_byte_ptr, offset, long(B * C) * G * Q * Q * 2)); - descs = nextWorkspacePtr(workspace_byte_ptr, offset, 1024); - } - - int numTokens = inputDesc[getInputTensorIdx()].dims.d[0]; - if (!mRemovePadding) - numTokens *= inputDesc[getInputTensorIdx()].dims.d[1]; - - setSSMParams(ssm_params, batch_size, mDim, max_seq_len, numTokens, mDState, mDtRank, mNHeads, mNGroups, mChunkSize, - statePtr, inputs[getInputTensorIdx()], inputs[getDeltaIdx()], inputs[getDeltaBiasIdx()], inputs[getAIdx()], - inputs[getBCIdx()], inputs[getDIdx()], z, mxOs, mxSt, mxdc, mxdA, mxCB, descs, - static_cast(inputs[getLastTokenIdsIdx()]), slotMapping, outputs[0], mDeltaSoftplus, mRemovePadding); - - if (reqTypes[0] == RequestType::kCONTEXT) - { - if (mIsMamba2) - { - invokeChunkScan(ssm_params, stream, mDriver.get()); - } - else - { - invokeSelectiveScan(ssm_params, stream); - } - } - else if (reqTypes[0] == RequestType::kGENERATION) - { - invokeSelectiveScanUpdate(ssm_params, stream); - } - sync_check_cuda_error(stream); - return 0; -} - -int SelectiveScanPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (isBuilding()) - { - return 0; - } - if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType SelectiveScanPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - if (index == 0) - { - return inputTypes[getInputTensorIdx()]; - } - else - { - return inputTypes[getStateIdx()]; - } -} - -// IPluginV2 Methods - -char const* SelectiveScanPlugin::getPluginType() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_NAME; -} - -char const* SelectiveScanPlugin::getPluginVersion() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_VERSION; -} - -int SelectiveScanPlugin::getNbOutputs() const noexcept -{ - return mPagedState ? 1 : 2; -} - -int SelectiveScanPlugin::initialize() noexcept -{ - return 0; -} - -void SelectiveScanPlugin::terminate() noexcept {} - -size_t SelectiveScanPlugin::getSerializationSize() const noexcept -{ - return sizeof(mDim) + sizeof(mDState) + sizeof(mDtRank) + sizeof(mNHeads) + sizeof(mNGroups) + sizeof(mChunkSize) - + sizeof(mDeltaSoftplus) + sizeof(mType) + sizeof(mRemovePadding) + sizeof(mPagedState) + sizeof(mZEnabled) - + sizeof(mIsMamba2); -} - -void SelectiveScanPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mDim); - write(d, mDState); - write(d, mDtRank); - write(d, mNHeads); - write(d, mNGroups); - write(d, mChunkSize); - write(d, mDeltaSoftplus); - write(d, mType); - write(d, mRemovePadding); - write(d, mPagedState); - write(d, mZEnabled); - write(d, mIsMamba2); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void SelectiveScanPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -SelectiveScanPluginCreator::SelectiveScanPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("dim", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dstate", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("dt_rank", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("nheads", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("ngroups", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("chunk_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("delta_softplus", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("remove_input_padding", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("paged_state", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("z_enabled", nullptr, PluginFieldType::kINT8)); - mPluginAttributes.emplace_back(PluginField("is_mamba2", nullptr, PluginFieldType::kINT8)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* SelectiveScanPluginCreator::getPluginName() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_NAME; -} - -char const* SelectiveScanPluginCreator::getPluginVersion() const noexcept -{ - return SELECTIVE_SCAN_PLUGIN_VERSION; -} - -PluginFieldCollection const* SelectiveScanPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* SelectiveScanPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - int dim{}; - int dstate{}; - int dtRank{}; - int nHeads{}; - int nGroups{}; - int chunkSize{}; - bool deltaSoftplus{}; - bool removePadding{}; - bool pagedState{}; - bool zEnabled{}; - bool isMamab2{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "dim")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dim = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dstate")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dstate = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "dt_rank")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - dtRank = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "nheads")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - nHeads = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "ngroups")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - nGroups = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "chunk_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - chunkSize = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "delta_softplus")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - deltaSoftplus = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "remove_input_padding")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - removePadding = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "paged_state")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - pagedState = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "z_enabled")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - zEnabled = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "is_mamba2")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT8); - isMamab2 = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - auto* obj = new SelectiveScanPlugin(dim, dstate, dtRank, nHeads, nGroups, chunkSize, deltaSoftplus, type, - removePadding, pagedState, zEnabled, isMamab2); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* SelectiveScanPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call SelectiveScanPlugin::destroy() - try - { - auto* obj = new SelectiveScanPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h b/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h deleted file mode 100644 index 96cb86fc4cbb..000000000000 --- a/cpp/tensorrt_llm/plugins/selectiveScanPlugin/selectiveScanPlugin.h +++ /dev/null @@ -1,218 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_SELECTIVE_SCAN_PLUGIN_H -#define TRT_SELECTIVE_SCAN_PLUGIN_H -#include "tensorrt_llm/kernels/selectiveScan/selectiveScan.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -// batch_size = num_ctx_requests or num_gen_requests -// num_ctx_requests = number of context requests (single sequence per request). -// num_gen_requests = number of generation requests (single sequences per request). -// can not support beam search - -// inputs -// 0. input_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. state, mamba: [batch_size, dstate, dim] or host [1] containing only pointer for paged_state -// mamba2: [batch_size, nheads, dstate, dim] or host [1] containing only pointer for paged_state -// 2. delta, mamba: [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// mamba2: [batch_size, seq_len, nheads] or [num_tokens, nheads] for remove_input_padding -// 3. delta_bias, [dim] for mamba, [nheads] for mamba2 -// 4. A, [dstate, dim] for mamba, [nheads] for mamba2 -// 5. BC, mamba: [batch_size, seq_len, dstate * 2] or [num_tokens, dstate * 2] for remove_input_padding -// mamba2: [batch_size, seq_len, ngroups * dstate * 2] or [num_tokens, ngroups * dstate * 2] for -// remove_input_padding -// 6. D, [dim] for mamba, [nheads] for mamba2 -// 7. host_request_types [batch_size] int32. 0: context; 1: generation; 2: none. -// 8. last_token_ids [batch_size] int32 -// 9. host_context_lengths [batch_size] int32, optional for remove_input_padding -// 10. state_slot_mapping [batch_size] int32, optional for paged state -// 11. z [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// outputs -// 0. output_tensor [batch_size, seq_len, dim] or [num_tokens, dim] for remove_input_padding -// 1. state, [batch_size, dstate, dim] for mamba, [batch_size, nheads, dstate, dim] for mamba2 - -class SelectiveScanPlugin : public BasePlugin -{ -public: - SelectiveScanPlugin(int dim, int dstate, int dtRank, int nHeads, int nGroups, int chunkSize, bool deltaSoftplus, - nvinfer1::DataType type, bool removePadding, bool pagedState, bool zEnabled, bool isMamba2); - - SelectiveScanPlugin(void const* data, size_t length); - - ~SelectiveScanPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - - enum class RequestType : int32_t - { - kCONTEXT = 0, - kGENERATION = 1 - }; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - - IndexType getStateIdx() const - { - return 1; - }; - - IndexType getDeltaIdx() const - { - return 2; - }; - - IndexType getDeltaBiasIdx() const - { - return 3; - }; - - IndexType getAIdx() const - { - return 4; - }; - - IndexType getBCIdx() const - { - return 5; - }; - - IndexType getDIdx() const - { - return 6; - }; - - IndexType getHostRequestTypesIdx() const - { - return 7; - }; - - IndexType getLastTokenIdsIdx() const - { - return 8; - }; - - IndexType getHostContextLengthIdx() const - { - if (mRemovePadding) - return 9; - else - return 8; - }; - - IndexType getSlotMappingIdx() const - { - if (mPagedState) - return getHostContextLengthIdx() + 1; - else - return getHostContextLengthIdx(); - }; - - IndexType getZIdx() const - { - if (mZEnabled) - return getSlotMappingIdx() + 1; - else - return getSlotMappingIdx(); - }; - - void setSSMParams(tensorrt_llm::kernels::SSMParamsBase& params, - // sizes - const size_t batch, const size_t dim, const size_t maxSeqLen, const size_t numTokens, const size_t dstate, - const size_t dtRank, const size_t nHeads, const size_t nGroups, const size_t chunkSize, - // device pointers - void* statePtr, void const* x, void const* delta, void const* deltaBias, void const* A, void const* BC, - void const* D, void const* z, void* osPtr, void* stPtr, void* dcPtr, void* dAPtr, void* cbPtr, void* descs, - int const* lastTokenIds, int const* slotMapping, void* out, bool deltaSoftplus, bool removePadding); - -private: - int mDim; - int mDState; - int mDtRank; - int mNHeads; - int mNGroups; - int mChunkSize; - bool mDeltaSoftplus; - nvinfer1::DataType mType; - bool mRemovePadding = false; - bool mPagedState = false; - bool mZEnabled = true; - bool mIsMamba2 = false; - std::shared_ptr mDriver; -}; - -class SelectiveScanPluginCreator : public BaseCreator -{ -public: - SelectiveScanPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif // TRT_SELECTIVE_SCAN_PLUGIN_H diff --git a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.cpp b/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.cpp deleted file mode 100644 index 718d8b7e830d..000000000000 --- a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.cpp +++ /dev/null @@ -1,431 +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. - */ -#include "smoothQuantGemmPlugin.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/int8SQ.h" -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::SmoothQuantGemmPluginCreator; -using tensorrt_llm::plugins::SmoothQuantGemmPlugin; -using tensorrt_llm::plugins::SmoothQuantGemmPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* SQ_GEMM_PLUGIN_VERSION{"1"}; -static char const* SQ_GEMM_PLUGIN_NAME{"SmoothQuantGemm"}; -PluginFieldCollection SmoothQuantGemmPluginCreator::mFC{}; -std::vector SmoothQuantGemmPluginCreator::mPluginAttributes; - -void SmoothQuantGemmPluginProfiler::runTactic(int m, int n, int k, SmoothQuantGemmPluginProfiler::Config const& tactic, - char* workspace, cudaStream_t const& stream) -{ - int8_t* aTmp = reinterpret_cast(workspace); - int8_t* bTmp = nextWorkspacePtr(aTmp, m * k * sizeof(int8_t)); - void* cTmp = reinterpret_cast(nextWorkspacePtr(bTmp, n * k * sizeof(int8_t))); - float* alphaRowTmp = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(cTmp), m * n * (mType == nvinfer1::DataType::kFLOAT ? 4 : 2))); - float* alphaColTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(alphaRowTmp), m * sizeof(float))); - char* workspaceTmp - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(alphaColTmp), n * sizeof(float))); - - int const wsSize = mRunner->getWorkspaceSize(m, n, k); - - mRunner->gemm( - aTmp, bTmp, mQuantMode, alphaColTmp, alphaRowTmp, cTmp, m, n, k, tactic, workspaceTmp, wsSize, stream); -} - -void SmoothQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - std::vector workspaces = { - maxM * k * sizeof(int8_t), // A - n * k * sizeof(int8_t), // B - maxM * n * (mType == nvinfer1::DataType::kFLOAT ? 4u : 2u), // C - maxM * sizeof(float), // alphaRow - n * sizeof(float), // alphaCol - mRunner->getWorkspaceSize(maxM, n, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector SmoothQuantGemmPluginProfiler::getTactics(int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -SmoothQuantGemmPlugin::SmoothQuantGemmPlugin( - QuantMode quantMode, nvinfer1::DataType type, SmoothQuantGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mQuantMode(quantMode) - , mPluginProfiler(pluginProfiler) -{ - init(type); -} - -// Parameterized constructor -SmoothQuantGemmPlugin::SmoothQuantGemmPlugin( - void const* data, size_t length, SmoothQuantGemmPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - unsigned int quantMode; - read(d, quantMode); - read(d, type); - read(d, mDims); - - mQuantMode = QuantMode(quantMode); - - init(type); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void SmoothQuantGemmPlugin::init(nvinfer1::DataType type) -{ - mType = type; - if (mType == nvinfer1::DataType::kHALF) - { - m_sqGemmRunner = std::make_shared>(); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - m_sqGemmRunner = std::make_shared>(); - } - else if (mType == nvinfer1::DataType::kINT32) - { - m_sqGemmRunner = std::make_shared>(); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - m_sqGemmRunner = std::make_shared>(); - } -#endif - - mPluginProfiler->setQuantMode(mQuantMode); - - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* SmoothQuantGemmPlugin::clone() const noexcept -{ - auto* plugin = new SmoothQuantGemmPlugin(*this); - return plugin; -} - -nvinfer1::DimsExprs SmoothQuantGemmPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - try - { - TLLM_CHECK(nbInputs == 4); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - TLLM_CHECK(nbDimsA >= 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - ret.d[nbDimsA - 1] = inputs[1].d[0]; - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool SmoothQuantGemmPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights stored in checkpoint must have int8 type - return inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format == TensorFormat::kLINEAR; - case 2: - // scales channels - case 3: - // scales tokens - return inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format == TensorFormat::kLINEAR; - case 4: - // out - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - default: - // Never should be here - assert(false); - return false; - } -} - -void SmoothQuantGemmPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[0]; - int const minK = in[0].min.d[in[0].min.nbDims - 1]; - int const minN = in[1].min.d[0]; - - TLLM_CHECK_WITH_INFO(minN == maxN, "Variable out channels is not allowed"); - TLLM_CHECK_WITH_INFO(minK == maxK, "Variable in channels is not allowed"); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, maxN, maxK}; - } - mGemmId = {maxN, maxK, mType}; - - m_workspaceMaxSize = m_sqGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t SmoothQuantGemmPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int SmoothQuantGemmPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M(*), K] - // mat2 [N, K] - // scale_tokens [M, 1] if has_per_token_scaling else [1, 1] - // scale_channels [1, N] if has_per_channel_scaling else [1, 1] - // outputs - // mat [M(*), N] - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[0]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - int const wsSize = m_sqGemmRunner->getWorkspaceSize(m, n, k); - if (m <= 4) - { - tensorrt_llm::kernels::smooth_quant::Params params(reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[1]), reinterpret_cast(inputs[2]), - reinterpret_cast(inputs[3]), reinterpret_cast(outputs[0]), m, n, k, mQuantMode); - if (mType == nvinfer1::DataType::kHALF) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher(params, stream); - } - else if (mType == nvinfer1::DataType::kFLOAT) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher(params, stream); - } -#ifdef ENABLE_BF16 - else if (mType == nvinfer1::DataType::kBF16) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher<__nv_bfloat16>(params, stream); - } -#endif - else if (mType == nvinfer1::DataType::kINT32) - { - tensorrt_llm::kernels::smooth_quant::int8_sq_launcher(params, stream); - } - } - else - { - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, "No valid SQ GEMM tactic"); - m_sqGemmRunner->gemm(reinterpret_cast(inputs[0]), reinterpret_cast(inputs[1]), - mQuantMode, reinterpret_cast(inputs[3]), reinterpret_cast(inputs[2]), - reinterpret_cast(outputs[0]), m, n, k, *bestTactic, reinterpret_cast(workspace), wsSize, - stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType SmoothQuantGemmPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* SmoothQuantGemmPlugin::getPluginType() const noexcept -{ - return SQ_GEMM_PLUGIN_NAME; -} - -char const* SmoothQuantGemmPlugin::getPluginVersion() const noexcept -{ - return SQ_GEMM_PLUGIN_VERSION; -} - -int SmoothQuantGemmPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int SmoothQuantGemmPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void SmoothQuantGemmPlugin::terminate() noexcept {} - -size_t SmoothQuantGemmPlugin::getSerializationSize() const noexcept -{ - return sizeof(unsigned int) + // QuantMode - sizeof(nvinfer1::DataType) + // dtype - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void SmoothQuantGemmPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mQuantMode.value()); - write(d, mType); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void SmoothQuantGemmPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -void SmoothQuantGemmPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_sqGemmRunner, mType, mDims, mGemmId); -} - -/////////////// - -SmoothQuantGemmPluginCreator::SmoothQuantGemmPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("has_per_channel_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("has_per_token_scaling", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* SmoothQuantGemmPluginCreator::getPluginName() const noexcept -{ - return SQ_GEMM_PLUGIN_NAME; -} - -char const* SmoothQuantGemmPluginCreator::getPluginVersion() const noexcept -{ - return SQ_GEMM_PLUGIN_VERSION; -} - -PluginFieldCollection const* SmoothQuantGemmPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* SmoothQuantGemmPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - bool perTokenScaling{}; - bool perChannelScaling{}; - nvinfer1::DataType type{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "has_per_channel_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - perChannelScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "has_per_token_scaling")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - perTokenScaling = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // SmoothQuantGemmPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - QuantMode quantMode = QuantMode::fromDescription(true, true, perTokenScaling, perChannelScaling, false, false, - false, false, false, false, false, false, false, false, false, false); - auto* obj = new SmoothQuantGemmPlugin(quantMode, type, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* SmoothQuantGemmPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call SmoothQuantGemmPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new SmoothQuantGemmPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h b/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h deleted file mode 100644 index 3cabf558076b..000000000000 --- a/cpp/tensorrt_llm/plugins/smoothQuantGemmPlugin/smoothQuantGemmPlugin.h +++ /dev/null @@ -1,140 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/cutlass_kernels/int8_gemm/int8_gemm.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include -#include -#include -#include -#include - -namespace tensorrt_llm::plugins -{ - -using perfMapType = std::unordered_map; -using SqGemmRunnerPtr = std::shared_ptr; - -class SmoothQuantGemmPluginProfiler : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantMode(tensorrt_llm::common::QuantMode const& quantMode) - { - mQuantMode = quantMode; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - -private: - tensorrt_llm::common::QuantMode mQuantMode; -}; - -class SmoothQuantGemmPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - SmoothQuantGemmPlugin() = delete; - - SmoothQuantGemmPlugin( - tensorrt_llm::common::QuantMode quantMode, nvinfer1::DataType type, PluginProfilerPtr const& pluginProfiler); - - SmoothQuantGemmPlugin(void const* data, size_t length, PluginProfilerPtr const& pluginProfiler); - - ~SmoothQuantGemmPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type); - - void configGemm(); - -private: - const std::string mLayerName; - - SqGemmRunnerPtr m_sqGemmRunner; - tensorrt_llm::common::QuantMode mQuantMode; - size_t m_workspaceMaxSize; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; - - nvinfer1::DataType mType; -}; - -class SmoothQuantGemmPluginCreator : public BaseCreator -{ -public: - SmoothQuantGemmPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/topkLastDimPlugin/CMakeLists.txt deleted file mode 100644 index 6b4e3d8d9e0f..000000000000 --- a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/CMakeLists.txt +++ /dev/null @@ -1,22 +0,0 @@ -# -# SPDX-FileCopyrightText: Copyright (c) 1993-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. -# - -file(GLOB SRCS *.cpp) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.cpp b/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.cpp deleted file mode 100644 index 072bfc9c8fc4..000000000000 --- a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.cpp +++ /dev/null @@ -1,316 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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 "topkLastDimPlugin.h" -#include "tensorrt_llm/common/assert.h" - -using namespace nvinfer1; -using namespace tensorrt_llm::kernels; -using namespace tensorrt_llm::common; -using tensorrt_llm::plugins::TopkLastDimPluginCreator; -using tensorrt_llm::plugins::TopkLastDimPlugin; - -static char const* TOPK_LAST_DIM_PLUGIN_VERSION{"1"}; -static char const* TOPK_LAST_DIM_PLUGIN_NAME{"TopkLastDim"}; -PluginFieldCollection TopkLastDimPluginCreator::mFC{}; -std::vector TopkLastDimPluginCreator::mPluginAttributes; - -TopkLastDimPlugin::TopkLastDimPlugin(nvinfer1::DataType type, int32_t k, bool is_largest) - : mType(type) - , mK(k) // To avoid data-dependent shape, enforce K to be non-dynamic - , mIsLargest(is_largest) -{ - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); -} - -// Parameterized constructor -TopkLastDimPlugin::TopkLastDimPlugin(void const* data, size_t length) -{ - char const *d = reinterpret_cast(data), *a = d; - read(d, mType); - read(d, mK); - read(d, mIsLargest); - TLLM_CHECK(d == a + length); - TLLM_CHECK_WITH_INFO((mType == DataType::kBF16) || (mType == DataType::kFLOAT) || (mType == DataType::kHALF) - || (mType == DataType::kINT32), - "Only support int, float, half, and bfloat16."); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* TopkLastDimPlugin::clone() const noexcept -{ - auto* plugin = new TopkLastDimPlugin(mType, mK, mIsLargest); - plugin->setPluginNamespace(mNamespace.c_str()); - return plugin; -} - -// Outputs -// out_val or out_idx: [batch_size, K] -nvinfer1::DimsExprs TopkLastDimPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - TLLM_CHECK_WITH_INFO(outputIndex < 2, "Only 2 outputs."); - nvinfer1::DimsExprs output(inputs[0]); - int numDim = output.nbDims; - output.d[numDim - 1] = exprBuilder.constant(mK); - return output; -} - -bool TopkLastDimPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - bool res = inOut[pos].format == TensorFormat::kLINEAR; - if (pos < 2) // input and out_val tensor must be the same type as the plugin - { - res = res && inOut[pos].type == mType; - } - else if (pos == 2) // out_idx must be int32 - { - res = res && inOut[pos].type == DataType::kINT32; - } - return res; -} - -void TopkLastDimPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ -} - -size_t TopkLastDimPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - // extract shape info and then call helper - auto const batchSize = inputs[getInputTensorIdx()].dims.d[0]; - auto const inputLength = inputs[getInputTensorIdx()].dims.d[1]; - size_t tempStorageBytes{}; - if (mType == DataType::kINT32) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize(batchSize, inputLength, mK, mIsLargest); - } - else if (mType == DataType::kHALF) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize(batchSize, inputLength, mK, mIsLargest); - } - else if (mType == DataType::kFLOAT) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize(batchSize, inputLength, mK, mIsLargest); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - tempStorageBytes = invokeComputeTopkLastDimWorkspaceSize<__nv_bfloat16>(batchSize, inputLength, mK, mIsLargest); - } -#endif - return tempStorageBytes; -} - -template -int TopkLastDimPlugin::enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) -{ - // inputs - // 0. input_tensor [batch_size, inputLength] - // outputs - // 0. output_values [batch_size, k] - // 1. output_indices [batch_size, k] - auto const batchSize = inputDesc[getInputTensorIdx()].dims.d[0]; - auto const inputLength = inputDesc[getInputTensorIdx()].dims.d[1]; - if (batchSize == 0) - { - // nothing to do for empty tensor - return 0; - } - - invokeTopkLastDim( - batchSize, inputLength, mK, mIsLargest, inputs[getInputTensorIdx()], outputs[0], outputs[1], workspace, stream); - - sync_check_cuda_error(stream); - return 0; -} - -int TopkLastDimPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - if (mType == DataType::kINT32) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kHALF) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } - else if (mType == DataType::kFLOAT) - { - return enqueueImpl(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#ifdef ENABLE_BF16 - else if (mType == DataType::kBF16) - { - return enqueueImpl<__nv_bfloat16>(inputDesc, outputDesc, inputs, outputs, workspace, stream); - } -#endif - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType TopkLastDimPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK_WITH_INFO(index < 2, "Only 2 outputs."); - nvinfer1::DataType data_type; - if (index == 1) - { - data_type = DataType::kINT32; - } - else - { - data_type = inputTypes[getInputTensorIdx()]; - } - return data_type; -} - -// IPluginV2 Methods - -char const* TopkLastDimPlugin::getPluginType() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_NAME; -} - -char const* TopkLastDimPlugin::getPluginVersion() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_VERSION; -} - -int TopkLastDimPlugin::getNbOutputs() const noexcept -{ - return 2; -} - -int TopkLastDimPlugin::initialize() noexcept -{ - return 0; -} - -void TopkLastDimPlugin::terminate() noexcept {} - -size_t TopkLastDimPlugin::getSerializationSize() const noexcept -{ - return sizeof(mType) + sizeof(mK) + sizeof(mIsLargest); -} - -void TopkLastDimPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mK); - write(d, mIsLargest); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void TopkLastDimPlugin::destroy() noexcept -{ - delete this; -} - -/////////////// - -TopkLastDimPluginCreator::TopkLastDimPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("k", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("is_largest", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* TopkLastDimPluginCreator::getPluginName() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_NAME; -} - -char const* TopkLastDimPluginCreator::getPluginVersion() const noexcept -{ - return TOPK_LAST_DIM_PLUGIN_VERSION; -} - -PluginFieldCollection const* TopkLastDimPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* TopkLastDimPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int32_t k{}; - bool is_largest{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "k")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - k = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "is_largest")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - is_largest = static_cast(*(static_cast(fields[i].data))) != 0; - } - } - try - { - auto* obj = new TopkLastDimPlugin(type, k, is_largest); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* TopkLastDimPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call TopkLastDimPlugin::destroy() - try - { - auto* obj = new TopkLastDimPlugin(serialData, serialLength); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h b/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h deleted file mode 100644 index 0ca38ccfe105..000000000000 --- a/cpp/tensorrt_llm/plugins/topkLastDimPlugin/topkLastDimPlugin.h +++ /dev/null @@ -1,99 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 1993-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. - */ - -#ifndef TRT_TOPK_LAST_DIM_PLUGIN_H -#define TRT_TOPK_LAST_DIM_PLUGIN_H - -#include "tensorrt_llm/kernels/topkLastDim.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include - -namespace tensorrt_llm::plugins -{ -class TopkLastDimPlugin : public BasePlugin -{ -public: - TopkLastDimPlugin(nvinfer1::DataType type, int32_t k, bool largest); - TopkLastDimPlugin(void const* data, size_t length); - ~TopkLastDimPlugin() override = default; - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - template - int enqueueImpl(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream); - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - using IndexType = std::int32_t; - - IndexType getInputTensorIdx() const - { - return 0; - }; - -private: - nvinfer1::DataType mType; - int32_t mK; - bool mIsLargest; -}; - -class TopkLastDimPluginCreator : public BaseCreator -{ -public: - TopkLastDimPluginCreator(); - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins - -#endif diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp deleted file mode 100644 index 85f0cf011293..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.cpp +++ /dev/null @@ -1,657 +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. - */ -#include "weightOnlyGroupwiseQuantMatmulPlugin.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::WeightOnlyGroupwiseQuantMatmulPluginCreator; -using tensorrt_llm::plugins::WeightOnlyGroupwiseQuantMatmulPlugin; -using tensorrt_llm::plugins::WeightOnlyGroupwiseQuantGemmPluginProfiler; -using tensorrt_llm::plugins::WeightOnlyGemmRunnerPtr; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* WOQ_GROUPWISE_MATMUL_PLUGIN_VERSION{"1"}; -static char const* WOQ_GROUPWISE_MATMUL_PLUGIN_NAME{"WeightOnlyGroupwiseQuantMatmul"}; -PluginFieldCollection WeightOnlyGroupwiseQuantMatmulPluginCreator::mFC{}; -std::vector WeightOnlyGroupwiseQuantMatmulPluginCreator::mPluginAttributes; - -void WeightOnlyGroupwiseQuantGemmPluginProfiler::runTactic(int m, int n, int k, - WeightOnlyGroupwiseQuantGemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const originalN = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; - half* actPtr = reinterpret_cast(workspace); - void* weightPtr = nextWorkspacePtr(reinterpret_cast(actPtr), m * k * sizeof(half)); - half* inputScalesPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(weightPtr), n * k * sizeof(float))); - half* zerosPtr = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(inputScalesPtr), k * originalN * sizeof(half) / mGroupSize)); - half* biasesPtr = reinterpret_cast( - nextWorkspacePtr(reinterpret_cast(zerosPtr), k * originalN * sizeof(half) / mGroupSize)); - half* outputPtr = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(biasesPtr), n * sizeof(half))); - char* workspacePtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(outputPtr), m * originalN * sizeof(half))); - if ((mQuantAlgo & GroupwiseQuantAlgo::ZERO) == 0) - { - zerosPtr = nullptr; - } - if ((mQuantAlgo & GroupwiseQuantAlgo::BIAS) == 0) - { - biasesPtr = nullptr; - } - - if (tactic.enableCudaKernel) - { - // run CUDA kernel - void const* pre_quant_scale_ptr = nullptr; - bool apply_alpha_in_advance = false; - float alpha = 1.0; - tensorrt_llm::kernels::weight_only::Params params{actPtr, pre_quant_scale_ptr, weightPtr, inputScalesPtr, - zerosPtr, biasesPtr, outputPtr, alpha, m, originalN, k, mGroupSize, mCudaKernelType, - apply_alpha_in_advance}; - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - // run CUTLASS kernel - int const wsSize = mRunner->getWorkspaceSize(m, originalN, k); - if (mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - mRunner->gemm(actPtr, reinterpret_cast(weightPtr), inputScalesPtr, zerosPtr, biasesPtr, outputPtr, - m, originalN, k, mGroupSize, tactic, workspacePtr, wsSize, stream); - } - else - { - mRunner->gemm(actPtr, reinterpret_cast(weightPtr), inputScalesPtr, zerosPtr, biasesPtr, - outputPtr, m, originalN, k, mGroupSize, tactic, workspacePtr, wsSize, stream); - } - } -} - -void WeightOnlyGroupwiseQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const originalN = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; - std::vector workspaces = { - maxM * k * sizeof(half), // A - k * n * sizeof(float), // B - k * originalN * sizeof(half) / mGroupSize, // scales - k * originalN * sizeof(half) / mGroupSize, // zeros - originalN * sizeof(half), // biases - maxM * originalN * sizeof(half), // C - mRunner->getWorkspaceSize(maxM, originalN, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector WeightOnlyGroupwiseQuantGemmPluginProfiler::getTactics( - int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -bool WeightOnlyGroupwiseQuantGemmPluginProfiler::checkTactic(int m, int n, int k, Config const& tactic) const -{ - // stop to profile Cuda kernel for m >= 16 - if (tactic.enableCudaKernel) - { - return m < 16; - } - return true; -} - -WeightOnlyGroupwiseQuantMatmulPlugin::WeightOnlyGroupwiseQuantMatmulPlugin(nvinfer1::DataType type, int quant_algo, - int group_size, float alpha, WeightOnlyGroupwiseQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - init(type, quant_algo, group_size, alpha); -} - -// Parameterized constructor -WeightOnlyGroupwiseQuantMatmulPlugin::WeightOnlyGroupwiseQuantMatmulPlugin( - void const* data, size_t length, WeightOnlyGroupwiseQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - int quant_algo = 0; - int group_size = 0; - float alpha = 1.0f; - read(d, type); - read(d, quant_algo); - read(d, group_size); - read(d, alpha); - read(d, mDims); - - init(type, quant_algo, group_size, alpha); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -template -using GemmRunner = tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunner; - -template -WeightOnlyGemmRunnerPtr selectGemmRunnerForZERO(int quant_algo) -{ - if (quant_algo & GroupwiseQuantAlgo::ZERO) - { - return std::make_shared>(); - } - else - { - return std::make_shared>(); - } -} - -template -WeightOnlyGemmRunnerPtr selectGemmRunnerForWeightType(int quant_algo) -{ - if (quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - return selectGemmRunnerForZERO(quant_algo); - } - else - { - return selectGemmRunnerForZERO(quant_algo); - } -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::init(nvinfer1::DataType type, int quant_algo, int group_size, float alpha) -{ - mArch = tensorrt_llm::common::getSMVersion(); - mType = type; - mQuantAlgo = quant_algo; - mGroupSize = group_size; - - // quant_algo = int8_weight * 16 + fp8_alpha * 8 + pre_quant_scale * 4 + zero * 2 + bias - mPreQuantScaleInputIdx = (quant_algo & GroupwiseQuantAlgo::PRE_QUANT_SCALE) ? 1 : 0; - mWeightInputIdx = mPreQuantScaleInputIdx + 1; - mScalesInputIdx = mWeightInputIdx + 1; - mZerosInputIdx = (quant_algo & GroupwiseQuantAlgo::ZERO) ? mScalesInputIdx + 1 : mScalesInputIdx; - mBiasesInputIdx = (quant_algo & GroupwiseQuantAlgo::BIAS) ? mZerosInputIdx + 1 : mZerosInputIdx; - - if (mType == nvinfer1::DataType::kHALF) - { - // CUTLASS kernel selection - if (quant_algo & GroupwiseQuantAlgo::FP8_ALPHA) - { - mAlpha = alpha; - - // Ada & Hopper style kernels - if (mArch < 89) - { - TLLM_THROW("W4A(fp)8 kernel is unsupported on pre-Ada (sm<89) architectures!"); - } - assert(!(quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) && "W4A(fp)8 kernel requires INT4 weight!"); - m_weightOnlyGroupwiseGemmRunner - = selectGemmRunnerForZERO<__nv_fp8_e4m3, cutlass::uint4b_t, half>(quant_algo); - } - else - { - m_weightOnlyGroupwiseGemmRunner = selectGemmRunnerForWeightType(quant_algo); - } - // CUDA kernel selection - if (quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - // INT8 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int8Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int8Groupwise; - } - else - { - // INT4 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int4Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int4Groupwise; - } - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - // CUTLASS kernel selection - if (quant_algo & GroupwiseQuantAlgo::FP8_ALPHA) - { - mAlpha = alpha; - - // FP8 requires at least sm89 devices - if (mArch < 89) - { - TLLM_THROW("W4A(fp)8 kernel is unsupported on pre-Ada (sm<89) architectures!"); - } - assert(!(quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) && "W4A(fp)8 kernel requires INT4 weight!"); - m_weightOnlyGroupwiseGemmRunner - = selectGemmRunnerForZERO<__nv_fp8_e4m3, cutlass::uint4b_t, __nv_bfloat16, half>(quant_algo); - } - else - { - m_weightOnlyGroupwiseGemmRunner = selectGemmRunnerForWeightType<__nv_bfloat16>(quant_algo); - } - // CUDA kernel selection - if (quant_algo & GroupwiseQuantAlgo::INT8_WEIGHT) - { - // INT8 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int8Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int8Groupwise; - } - else - { - // INT4 weight - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int4Groupwise); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int4Groupwise; - } - } -#endif - else - { - TLLM_THROW("Unsupported data type"); - } - mPluginProfiler->setQuantAlgo(mQuantAlgo); - mPluginProfiler->setGroupSize(mGroupSize); - if (mCudaKernelEnabled) - { - mPluginProfiler->setCudaKernelType(mCudaKernelType, mArch); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* WeightOnlyGroupwiseQuantMatmulPlugin::clone() const noexcept -{ - auto* plugin = new WeightOnlyGroupwiseQuantMatmulPlugin(*this); - return plugin; -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_weightOnlyGroupwiseGemmRunner, mType, mDims, mGemmId, mCudaKernelEnabled); -} - -nvinfer1::DimsExprs WeightOnlyGroupwiseQuantMatmulPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - - // inputs - // 0 activations [M, K] - // 1 pre-quant scales [K] (optional) - // 2 weights [K, N/2] - // 3 scales [K // group_size, N] - // 4 zeros [K // group_size, N] (optional) - // 5 biases [N] (optional) - - try - { - TLLM_CHECK(nbInputs == mBiasesInputIdx + 1); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - int const nbDimsB = inputs[mWeightInputIdx].nbDims; - TLLM_CHECK(nbDimsA >= 2); - TLLM_CHECK(nbDimsB == 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - - // int4/int8 weight only quant (INT4*4 -> FP16, INT8*2 -> FP16) - int const weight_multiplier = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? FP16_INT8_RATIO : FP16_INT4_RATIO; - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[mWeightInputIdx].d[1]->getConstantValue() * weight_multiplier); - - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool WeightOnlyGroupwiseQuantMatmulPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - if (pos < nbInputs + 1) - { - return inOut[pos].type == mType && inOut[pos].format == TensorFormat::kLINEAR; - } - else - { - // Never should be here - assert(false); - return false; - } -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int const weight_multiplier = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? FP16_INT8_RATIO : FP16_INT4_RATIO; - int const maxN = in[mWeightInputIdx].max.d[1] * weight_multiplier; - - auto const K = maxK; - auto const N = maxN / weight_multiplier; - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - mGemmId = {N, K, mType}; - - size_t smoothedActSize = static_cast(maxM) * static_cast(maxK) - * (in[0].desc.type == nvinfer1::DataType::kFLOAT ? sizeof(float) : sizeof(half)); - m_workspaceMaxSize = smoothedActSize + m_weightOnlyGroupwiseGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t WeightOnlyGroupwiseQuantMatmulPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -template -void pre_quant_scale_for_act(int const m, int const k, int const mQuantAlgo, int const mPreQuantScaleInputIdx, - void const* const* inputs, void* workspace, cudaStream_t stream) -{ - // Apply pre-quant per channel scale on activations - if (mQuantAlgo & GroupwiseQuantAlgo::FP8_ALPHA) - { - tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher( - reinterpret_cast<__nv_fp8_e4m3*>(workspace), reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream); - } - else - { - tensorrt_llm::kernels::apply_per_channel_scale_kernel_launcher( - reinterpret_cast(workspace), reinterpret_cast(inputs[0]), - reinterpret_cast(inputs[mPreQuantScaleInputIdx]), m, k, nullptr, stream); - } -} - -int WeightOnlyGroupwiseQuantMatmulPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // 0 activations [M, K] - // 1 pre-quant scales [K] - // 2 weights [K, N/2] - // 3 scales [K // group_size, N] - // 4 zeros [K // group_size, N] - // 5 biases [N] - // outputs - // mat [M, N] - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[mWeightInputIdx].dims.d[1]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - // get best tactic and check if CUDA kernel should be used - bool use_cuda_kernel = false; - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, - "No valid weight only groupwise GEMM tactic(It is usually caused by the failure to execute all " - "candidate configurations of the CUTLASS kernel, please pay attention to the warning information " - "when building the engine.)"); - use_cuda_kernel = bestTactic->enableCudaKernel; - - bool use_pre_quant_scale = mQuantAlgo & GroupwiseQuantAlgo::PRE_QUANT_SCALE; - half const* zeros_ptr - = (mQuantAlgo & GroupwiseQuantAlgo::ZERO) ? reinterpret_cast(inputs[mZerosInputIdx]) : nullptr; - half const* biases_ptr - = (mQuantAlgo & GroupwiseQuantAlgo::BIAS) ? reinterpret_cast(inputs[mBiasesInputIdx]) : nullptr; - half const* act_ptr = reinterpret_cast(inputs[0]); - - if (use_pre_quant_scale && !use_cuda_kernel) - { - // Apply pre-quant per channel scale on activations - act_ptr = reinterpret_cast(workspace); - if (mType == nvinfer1::DataType::kHALF) - { - pre_quant_scale_for_act(m, k, mQuantAlgo, mPreQuantScaleInputIdx, inputs, workspace, stream); - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - pre_quant_scale_for_act<__nv_bfloat16>(m, k, mQuantAlgo, mPreQuantScaleInputIdx, inputs, workspace, stream); - } -#endif - } - -#if defined(ENABLE_BF16) - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16, - "No valid weightOnlyGropwiseQuantMatmul configuration"); -#else - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF, "No valid weightOnlyGropwiseQuantMatmul configuration"); -#endif - - // Quantized weights are packed in FP16 format (INT4*4 -> FP16, INT8*2 -> FP16) - int real_n = mQuantAlgo & GroupwiseQuantAlgo::INT8_WEIGHT ? n * FP16_INT8_RATIO : n * FP16_INT4_RATIO; - - if (use_cuda_kernel) - { - // Apply CUDA kernel - void const* pre_quant_scale_ptr = nullptr; - if (use_pre_quant_scale) - pre_quant_scale_ptr = inputs[mPreQuantScaleInputIdx]; - void const* cuda_kernel_act_ptr = inputs[0]; - void const* cuda_kernel_weight_ptr = inputs[mWeightInputIdx]; - void const* cuda_kernel_scales_ptr = inputs[mScalesInputIdx]; - void* cuda_kernel_out_ptr = outputs[0]; - tensorrt_llm::kernels::weight_only::Params params{cuda_kernel_act_ptr, pre_quant_scale_ptr, - cuda_kernel_weight_ptr, cuda_kernel_scales_ptr, zeros_ptr, biases_ptr, cuda_kernel_out_ptr, mAlpha, m, - real_n, k, mGroupSize, mCudaKernelType, static_cast(mQuantAlgo & GroupwiseQuantAlgo::FP8_ALPHA)}; - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - // Apply CUTLASS kernel - int const ws_bytes = m_weightOnlyGroupwiseGemmRunner->getWorkspaceSize(m, real_n, k); - int32_t* weight_ptr = const_cast(reinterpret_cast(inputs[mWeightInputIdx])); - m_weightOnlyGroupwiseGemmRunner->gemm(act_ptr, weight_ptr, inputs[mScalesInputIdx], zeros_ptr, biases_ptr, - mAlpha, outputs[0], m, real_n, k, mGroupSize, *bestTactic, - reinterpret_cast(workspace) + m * k * sizeof(half), ws_bytes, stream); - } - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType WeightOnlyGroupwiseQuantMatmulPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* WeightOnlyGroupwiseQuantMatmulPlugin::getPluginType() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyGroupwiseQuantMatmulPlugin::getPluginVersion() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_VERSION; -} - -int WeightOnlyGroupwiseQuantMatmulPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int WeightOnlyGroupwiseQuantMatmulPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::terminate() noexcept {} - -size_t WeightOnlyGroupwiseQuantMatmulPlugin::getSerializationSize() const noexcept -{ - return sizeof(nvinfer1::DataType) + // mType - sizeof(int) + // mQuantAlgo - sizeof(int) + // mGroupSize - sizeof(float) + // mAlpha - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mQuantAlgo); - write(d, mGroupSize); - write(d, mAlpha); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void WeightOnlyGroupwiseQuantMatmulPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -WeightOnlyGroupwiseQuantMatmulPluginCreator::WeightOnlyGroupwiseQuantMatmulPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("quant_algo", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("group_size", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("alpha", nullptr, PluginFieldType::kFLOAT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* WeightOnlyGroupwiseQuantMatmulPluginCreator::getPluginName() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyGroupwiseQuantMatmulPluginCreator::getPluginVersion() const noexcept -{ - return WOQ_GROUPWISE_MATMUL_PLUGIN_VERSION; -} - -PluginFieldCollection const* WeightOnlyGroupwiseQuantMatmulPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* WeightOnlyGroupwiseQuantMatmulPluginCreator::createPlugin( - char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - int QuantAlgo{}; - int GroupSize{}; - float Alpha{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "quant_algo")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - QuantAlgo = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "group_size")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - GroupSize = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "alpha")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kFLOAT32); - Alpha = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // WeightOnlyGroupwiseQuantMatmulPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - auto* obj = new WeightOnlyGroupwiseQuantMatmulPlugin(type, QuantAlgo, GroupSize, Alpha, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* WeightOnlyGroupwiseQuantMatmulPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call weightOnlyGroupwiseQuantMatmulPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new WeightOnlyGroupwiseQuantMatmulPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h b/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h deleted file mode 100644 index 94e98ce0f5c0..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin/weightOnlyGroupwiseQuantMatmulPlugin.h +++ /dev/null @@ -1,186 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h" -#include "tensorrt_llm/kernels/preQuantScaleKernel.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv//kernelLauncher.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h" - -#include - -#include -#include -#include -#include -#include -#include - -// The blank line here is to avoid clang-format -sort-includes option reordering these two cutlass header files and -// breaking dependencies -#include "cutlass/integer_subbyte.h" - -namespace tensorrt_llm::plugins -{ - -using WeightOnlyGemmRunner = tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunnerInterface; -using WeightOnlyGemmRunnerPtr = std::shared_ptr; -using KernelType = tensorrt_llm::kernels::weight_only::KernelType; - -class WeightOnlyGroupwiseQuantGemmPluginProfiler - : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setQuantAlgo(int quantAlgo) - { - mQuantAlgo = quantAlgo; - } - - void setGroupSize(int groupSize) - { - mGroupSize = groupSize; - } - - void setCudaKernelType(KernelType cudaKernelType, int arch) - { - mCudaKernelType = cudaKernelType; - mArch = arch; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - - bool checkTactic(int m, int n, int k, Config const& tactic) const override; - -private: - int mQuantAlgo; - int mGroupSize; - KernelType mCudaKernelType; - int mArch; -}; - -class WeightOnlyGroupwiseQuantMatmulPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - - WeightOnlyGroupwiseQuantMatmulPlugin() = delete; - - WeightOnlyGroupwiseQuantMatmulPlugin( - nvinfer1::DataType type, int quant_algo, int group_size, float alpha, PluginProfilerPtr const& profiler); - - WeightOnlyGroupwiseQuantMatmulPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~WeightOnlyGroupwiseQuantMatmulPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - // group_size: 64, 128 - void init(nvinfer1::DataType type, int quant_algo, int group_size, float alpha); - - void configGemm(); - -private: - const std::string mLayerName; - - WeightOnlyGemmRunnerPtr m_weightOnlyGroupwiseGemmRunner; - size_t m_workspaceMaxSize; - nvinfer1::DataType mType; - bool mCudaKernelEnabled; - tensorrt_llm::kernels::weight_only::KernelType mCudaKernelType; - int mArch; - - // When M is smaller than this value, we trigger a fast path - // I.e. a tailored kernel instead of cutlass. - - int mQuantAlgo; - - int mGroupSize; - - float mAlpha = 1.0f; - - int mPreQuantScaleInputIdx; - int mWeightInputIdx; - int mScalesInputIdx; - int mZerosInputIdx; - int mBiasesInputIdx; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; -}; - -class WeightOnlyGroupwiseQuantMatmulPluginCreator : public BaseCreator -{ -public: - WeightOnlyGroupwiseQuantMatmulPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/CMakeLists.txt b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/CMakeLists.txt deleted file mode 100755 index 86876224fccd..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/CMakeLists.txt +++ /dev/null @@ -1,21 +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) -set(PLUGIN_SOURCES ${PLUGIN_SOURCES} ${SRCS}) -set(PLUGIN_SOURCES - ${PLUGIN_SOURCES} - PARENT_SCOPE) diff --git a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp deleted file mode 100644 index f3ed07fafaff..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.cpp +++ /dev/null @@ -1,507 +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. - */ -#include "weightOnlyQuantMatmulPlugin.h" - -#include - -using namespace nvinfer1; -using namespace tensorrt_llm::common; -using namespace tensorrt_llm::kernels::cutlass_kernels; -using tensorrt_llm::plugins::WeightOnlyQuantMatmulPluginCreator; -using tensorrt_llm::plugins::WeightOnlyQuantMatmulPlugin; -using tensorrt_llm::plugins::WeightOnlyQuantGemmPluginProfiler; -using tensorrt_llm::plugins::read; -using tensorrt_llm::plugins::write; - -static char const* WOQ_MATMUL_PLUGIN_VERSION{"1"}; -static char const* WOQ_MATMUL_PLUGIN_NAME{"WeightOnlyQuantMatmul"}; -PluginFieldCollection WeightOnlyQuantMatmulPluginCreator::mFC{}; -std::vector WeightOnlyQuantMatmulPluginCreator::mPluginAttributes; - -void WeightOnlyQuantGemmPluginProfiler::runTactic(int m, int n, int k, - WeightOnlyQuantGemmPluginProfiler::Config const& tactic, char* workspace, cudaStream_t const& stream) -{ - int const originalN = n * getWeightTypeMultiplier(mWeightTypeId); - half* actPtr = reinterpret_cast(workspace); - int8_t* weightPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(actPtr), m * k * sizeof(half))); - half* scalesPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(weightPtr), n * k * sizeof(int8_t))); - half* outputPtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(scalesPtr), originalN * sizeof(half))); - char* workspacePtr - = reinterpret_cast(nextWorkspacePtr(reinterpret_cast(outputPtr), m * originalN * sizeof(half))); - - int const wsSize = mRunner->getWorkspaceSize(m, originalN, k); - - if (tactic.enableCudaKernel) - { - // run CUDA kernel - tensorrt_llm::kernels::weight_only::Params params{actPtr, nullptr, weightPtr, scalesPtr, nullptr, nullptr, - outputPtr, 1.f, m, originalN, k, 0, mCudaKernelType}; - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - // run CUTLASS kernel - if (mWeightTypeId == WeightTypeId::INT8) - { - mRunner->gemm( - actPtr, weightPtr, scalesPtr, outputPtr, m, originalN, k, tactic, workspacePtr, wsSize, stream); - } - else - { - mRunner->gemm(actPtr, reinterpret_cast(weightPtr), scalesPtr, outputPtr, m, originalN, - k, tactic, workspacePtr, wsSize, stream); - } - } -} - -void WeightOnlyQuantGemmPluginProfiler::computeTmpSize(size_t maxM, size_t n, size_t k) -{ - int const originalN = n * getWeightTypeMultiplier(mWeightTypeId); - std::vector workspaces = { - maxM * k * sizeof(half), // A - n * k * sizeof(int8_t), // B - originalN * sizeof(half), // scales - maxM * originalN * sizeof(half), // C - mRunner->getWorkspaceSize(maxM, originalN, k) // workspace - }; - size_t bytes = calculateTotalWorkspaceSize(workspaces.data(), workspaces.size()); - setTmpWorkspaceSizeInBytes(bytes); -} - -std::vector WeightOnlyQuantGemmPluginProfiler::getTactics( - int m, int n, int k) const -{ - return mRunner->getConfigs(); -} - -bool WeightOnlyQuantGemmPluginProfiler::checkTactic(int m, int n, int k, Config const& tactic) const -{ - // stop to profile Cuda kernel for m >= 16 - if (tactic.enableCudaKernel) - { - return m < 16; - } - return true; -} - -WeightOnlyQuantMatmulPlugin::WeightOnlyQuantMatmulPlugin(nvinfer1::DataType type, WeightTypeId weightTypeId, - WeightOnlyQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - init(type, weightTypeId); -} - -// Parameterized constructor -WeightOnlyQuantMatmulPlugin::WeightOnlyQuantMatmulPlugin( - void const* data, size_t length, WeightOnlyQuantMatmulPlugin::PluginProfilerPtr const& pluginProfiler) - : mPluginProfiler(pluginProfiler) -{ - char const *d = reinterpret_cast(data), *a = d; - nvinfer1::DataType type; - WeightTypeId weightTypeId; - read(d, type); - read(d, weightTypeId); - read(d, mDims); - - init(type, weightTypeId); - - mPluginProfiler->deserialize(d, mDims, mGemmId); - - TLLM_CHECK_WITH_INFO(d == a + length, - "Expected length (%d) != real length (%d). This is often " - "caused by using different TensorRT LLM version to build " - "engine and run engine.", - (int) length, (int) (d - a)); -} - -void WeightOnlyQuantMatmulPlugin::init(nvinfer1::DataType type, WeightTypeId weightTypeId) -{ - mArch = tensorrt_llm::common::getSMVersion(); - mType = type; - mWeightTypeId = weightTypeId; - - if (mWeightTypeId == WeightTypeId::INT8) - { - if (mType == nvinfer1::DataType::kHALF) - { - m_weightOnlyGemmRunner = std::make_shared< - CutlassFpAIntBGemmRunner>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int8PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int8PerChannel; - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - m_weightOnlyGemmRunner = std::make_shared< - CutlassFpAIntBGemmRunner<__nv_bfloat16, uint8_t, cutlass::WeightOnlyQuantOp::PER_COLUMN_SCALE_ONLY>>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int8PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int8PerChannel; - } -#endif - else - { - TLLM_CHECK(false); - } - } - else if (mWeightTypeId == WeightTypeId::INT4) - { - if (mType == nvinfer1::DataType::kHALF) - { - m_weightOnlyGemmRunner = std::make_shared< - CutlassFpAIntBGemmRunner>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::FP16Int4PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::FP16Int4PerChannel; - } -#if defined(ENABLE_BF16) - else if (mType == nvinfer1::DataType::kBF16) - { - m_weightOnlyGemmRunner = std::make_shared>(); - mCudaKernelEnabled = tensorrt_llm::kernels::weight_only::is_supported( - mArch, tensorrt_llm::kernels::weight_only::KernelType::BF16Int4PerChannel); - mCudaKernelType = tensorrt_llm::kernels::weight_only::KernelType::BF16Int4PerChannel; - } -#endif - else - { - TLLM_CHECK(false); - } - } - else - { - TLLM_CHECK(false); - } - - mPluginProfiler->setWeightTypeId(mWeightTypeId); - if (mCudaKernelEnabled) - { - mPluginProfiler->setCudaKernelType(mCudaKernelType, mArch); - } - mGemmId = GemmIdCore(mDims.n, mDims.k, mType); -} - -// IPluginV2DynamicExt Methods -nvinfer1::IPluginV2DynamicExt* WeightOnlyQuantMatmulPlugin::clone() const noexcept -{ - auto* plugin = new WeightOnlyQuantMatmulPlugin(*this); - return plugin; -} - -void WeightOnlyQuantMatmulPlugin::configGemm() -{ - mPluginProfiler->profileTactics(m_weightOnlyGemmRunner, mType, mDims, mGemmId, mCudaKernelEnabled); -} - -nvinfer1::DimsExprs WeightOnlyQuantMatmulPlugin::getOutputDimensions( - int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept -{ - // input [m1, m2, m3, ... , k] - // weight [k, n] for int8, [k, n/2] for int4 - - try - { - TLLM_CHECK(nbInputs == 3); - TLLM_CHECK(outputIndex == 0); - int const nbDimsA = inputs[0].nbDims; - int const nbDimsB = inputs[1].nbDims; - TLLM_CHECK(nbDimsA >= 2); - TLLM_CHECK(nbDimsB == 2); - DimsExprs ret; - ret.nbDims = nbDimsA; - for (int ii = 0; ii < nbDimsA - 1; ++ii) - { - ret.d[ii] = inputs[0].d[ii]; - } - if (mWeightTypeId == WeightTypeId::INT8) - { - // int8 weight only quant - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue()); - } - else - { - // int4 weight only quant - ret.d[nbDimsA - 1] = exprBuilder.constant(inputs[1].d[1]->getConstantValue() * INT8_INT4_RATIO); - } - return ret; - } - catch (std::exception const& e) - { - caughtError(e); - } - return DimsExprs{}; -} - -bool WeightOnlyQuantMatmulPlugin::supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept -{ - switch (pos) - { - case 0: - // activation - return inOut[0].type == mType && inOut[0].format == TensorFormat::kLINEAR; - case 1: - // weights - // Weights are required to be int8, but will be reinterpreted as int4 in enqueue if required - // Weights stored in checkpoint should have int8/int4 type - return inOut[1].type == nvinfer1::DataType::kINT8 && inOut[1].format == TensorFormat::kLINEAR; - case 2: - // scales channels - return inOut[2].type == mType && inOut[2].format == TensorFormat::kLINEAR; - case 3: - // out - return inOut[3].type == mType && inOut[3].format == TensorFormat::kLINEAR; - default: - // Never should be here - assert(false); - return false; - } -} - -void WeightOnlyQuantMatmulPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept -{ - auto const minM = std::accumulate(in[0].min.d, in[0].min.d + in[0].min.nbDims - 1, 1, std::multiplies()); - auto const maxM = std::accumulate(in[0].max.d, in[0].max.d + in[0].max.nbDims - 1, 1, std::multiplies()); - - int const maxK = in[0].max.d[in[0].max.nbDims - 1]; - int const maxN = in[1].max.d[1] * getWeightTypeMultiplier(mWeightTypeId); - - auto const K = maxK; - auto const N = maxN / getWeightTypeMultiplier(mWeightTypeId); - - if (!mDims.isInitialized()) - { - mDims = {minM, maxM, N, K}; - } - - mGemmId = {N, K, mType}; - - m_workspaceMaxSize = m_weightOnlyGemmRunner->getWorkspaceSize(maxM, maxN, maxK); -} - -size_t WeightOnlyQuantMatmulPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept -{ - return m_workspaceMaxSize; -} - -int WeightOnlyQuantMatmulPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, - nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace, - cudaStream_t stream) noexcept -{ - // inputs - // mat1 [M1, M2,..., K] - // mat2 [K, N] for int8, [K, N/2] for int4 - // scale_channels [N] - // outputs - // mat [M, N] - - int64_t m64 = 1; - for (int ii = 0; ii < inputDesc[0].dims.nbDims - 1; ++ii) - { - m64 *= inputDesc[0].dims.d[ii]; - } - int const m = TLLM_INT32_CAST(m64); - int const n = TLLM_INT32_CAST(inputDesc[1].dims.d[1]); - int const k = TLLM_INT32_CAST(inputDesc[0].dims.d[inputDesc[0].dims.nbDims - 1]); - - if (m == 0) - return 0; - -#if defined(ENABLE_BF16) - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16, - "No valid weightOnlyQuantMatmul configuration"); -#else - TLLM_CHECK_WITH_INFO(mType == nvinfer1::DataType::kHALF, "No valid weightOnlyQuantMatmul configuration"); -#endif - int real_n = mWeightTypeId == WeightTypeId::INT4 ? n * INT8_INT4_RATIO : n; - - // get best tactic and check if CUDA kernel should be used - bool use_cuda_kernel = false; - auto const& bestTactic = mPluginProfiler->getBestConfig(m, mGemmId); - TLLM_CHECK_WITH_INFO(bestTactic, - "No valid weight only per-channel GEMM tactic(It is usually caused by the failure to execute all candidate " - "configurations of the CUTLASS kernel, please pay attention to the warning information when building the " - "engine.)"); - use_cuda_kernel = bestTactic->enableCudaKernel; - if (use_cuda_kernel) - { - void const* cuda_kernel_act_ptr = inputs[0]; - void const* cuda_kernel_weight_ptr = inputs[1]; - void const* cuda_kernel_scales_ptr = inputs[2]; - void* cuda_kernel_out_ptr = outputs[0]; - tensorrt_llm::kernels::weight_only::Params params(cuda_kernel_act_ptr, nullptr, cuda_kernel_weight_ptr, - cuda_kernel_scales_ptr, nullptr, nullptr, cuda_kernel_out_ptr, 1.f, m, real_n, k, 0, mCudaKernelType); - tensorrt_llm::kernels::weight_only::kernel_launcher(mArch, params, stream); - } - else - { - int const ws_size = m_weightOnlyGemmRunner->getWorkspaceSize(m, real_n, k); - - m_weightOnlyGemmRunner->gemm(inputs[0], inputs[1], inputs[2], outputs[0], m, real_n, k, *bestTactic, - reinterpret_cast(workspace), ws_size, stream); - } - - return 0; -} - -// IPluginV2Ext Methods -nvinfer1::DataType WeightOnlyQuantMatmulPlugin::getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept -{ - TLLM_CHECK(index == 0); - return mType; -} - -// IPluginV2 Methods - -char const* WeightOnlyQuantMatmulPlugin::getPluginType() const noexcept -{ - return WOQ_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyQuantMatmulPlugin::getPluginVersion() const noexcept -{ - return WOQ_MATMUL_PLUGIN_VERSION; -} - -int WeightOnlyQuantMatmulPlugin::getNbOutputs() const noexcept -{ - return 1; -} - -int WeightOnlyQuantMatmulPlugin::initialize() noexcept -{ - configGemm(); - return 0; -} - -void WeightOnlyQuantMatmulPlugin::terminate() noexcept {} - -size_t WeightOnlyQuantMatmulPlugin::getSerializationSize() const noexcept -{ - return sizeof(mWeightTypeId) + // mWeightTypeId - sizeof(nvinfer1::DataType) + // mType - sizeof(mDims) + // Dimensions - mPluginProfiler->getSerializationSize(mGemmId); // selected tactics container size -} - -void WeightOnlyQuantMatmulPlugin::serialize(void* buffer) const noexcept -{ - char *d = static_cast(buffer), *a = d; - write(d, mType); - write(d, mWeightTypeId); - write(d, mDims); - - mPluginProfiler->serialize(d, mGemmId); - TLLM_CHECK(d == a + getSerializationSize()); -} - -void WeightOnlyQuantMatmulPlugin::destroy() noexcept -{ - // This gets called when the network containing plugin is destroyed - delete this; -} - -/////////////// - -WeightOnlyQuantMatmulPluginCreator::WeightOnlyQuantMatmulPluginCreator() -{ - // Fill PluginFieldCollection with PluginField arguments metadata - mPluginAttributes.clear(); - mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32)); - mPluginAttributes.emplace_back(PluginField("weight_type_id", nullptr, PluginFieldType::kINT32)); - mFC.nbFields = mPluginAttributes.size(); - mFC.fields = mPluginAttributes.data(); -} - -char const* WeightOnlyQuantMatmulPluginCreator::getPluginName() const noexcept -{ - return WOQ_MATMUL_PLUGIN_NAME; -} - -char const* WeightOnlyQuantMatmulPluginCreator::getPluginVersion() const noexcept -{ - return WOQ_MATMUL_PLUGIN_VERSION; -} - -PluginFieldCollection const* WeightOnlyQuantMatmulPluginCreator::getFieldNames() noexcept -{ - return &mFC; -} - -IPluginV2* WeightOnlyQuantMatmulPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept -{ - PluginField const* fields = fc->fields; - nvinfer1::DataType type{}; - WeightTypeId weightTypeId{}; - // Read configurations from each fields - for (int i = 0; i < fc->nbFields; ++i) - { - char const* attrName = fields[i].name; - if (!strcmp(attrName, "weight_type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - weightTypeId = static_cast(*(static_cast(fields[i].data))); - } - else if (!strcmp(attrName, "type_id")) - { - TLLM_CHECK(fields[i].type == PluginFieldType::kINT32); - type = static_cast(*(static_cast(fields[i].data))); - } - } - try - { - // WeightOnlyGroupwiseQuantMatmulPluginCreator is unique and shared for an engine generation - // Create plugin profiler with shared tactics map - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ false); - auto* obj = new WeightOnlyQuantMatmulPlugin(type, weightTypeId, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} - -IPluginV2* WeightOnlyQuantMatmulPluginCreator::deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept -{ - // This object will be deleted when the network is destroyed, which will - // call WeightOnlyQuantMatmulPlugin::destroy() - try - { - // Create plugin profiler with private tactics map which is read from the serialized engine - auto pluginProfiler = gemmPluginProfileManager.createGemmPluginProfiler(/* inference */ true); - auto* obj = new WeightOnlyQuantMatmulPlugin(serialData, serialLength, pluginProfiler); - obj->setPluginNamespace(mNamespace.c_str()); - return obj; - } - catch (std::exception const& e) - { - caughtError(e); - } - return nullptr; -} diff --git a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h b/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h deleted file mode 100644 index 3177d8297d2d..000000000000 --- a/cpp/tensorrt_llm/plugins/weightOnlyQuantMatmulPlugin/weightOnlyQuantMatmulPlugin.h +++ /dev/null @@ -1,175 +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. - */ -#pragma once - -#include "tensorrt_llm/common/quantization.h" -#include "tensorrt_llm/kernels/cutlass_kernels/fpA_intB_gemm/fpA_intB_gemm.h" -#include "tensorrt_llm/kernels/weightOnlyBatchedGemv/kernelLauncher.h" -#include "tensorrt_llm/plugins/common/gemmPluginProfiler.h" -#include "tensorrt_llm/plugins/common/plugin.h" - -#include -#include -#include -#include -#include -#include - -// The blank line here is to avoid clang-format -sort-includes option reordering these two cutlass header files and -// breaking dependencies -#include "cutlass/integer_subbyte.h" - -namespace tensorrt_llm::plugins -{ -enum class WeightTypeId -{ - INT8 = 1, - INT4 = 2, -}; - -constexpr int32_t FP16_BITS = 16; -constexpr int32_t INT8_BITS = 8; -constexpr int32_t INT4_BITS = 4; -constexpr int32_t INT8_INT4_RATIO = INT8_BITS / INT4_BITS; -constexpr int32_t FP16_INT4_RATIO = FP16_BITS / INT4_BITS; -constexpr int32_t FP16_INT8_RATIO = FP16_BITS / INT8_BITS; - -inline int32_t getWeightTypeMultiplier(WeightTypeId weightTypeId) -{ - return weightTypeId == WeightTypeId::INT8 ? 1 : INT8_INT4_RATIO; -} - -using WeightOnlyGemmRunner = tensorrt_llm::kernels::cutlass_kernels::CutlassFpAIntBGemmRunnerInterface; -using WeightOnlyGemmRunnerPtr = std::shared_ptr; -using KernelType = tensorrt_llm::kernels::weight_only::KernelType; - -class WeightOnlyQuantGemmPluginProfiler : public GemmPluginProfiler -{ -public: - using Config = tensorrt_llm::cutlass_extensions::CutlassGemmConfig; - - void setWeightTypeId(WeightTypeId weightId) - { - mWeightTypeId = weightId; - } - - void setCudaKernelType(KernelType cudaKernelType, int arch) - { - mCudaKernelType = cudaKernelType; - mArch = arch; - } - -protected: - void runTactic(int m, int n, int k, Config const& tactic, char* workspace, cudaStream_t const& stream) override; - - void computeTmpSize(size_t maxM, size_t n, size_t k) override; - - std::vector getTactics(int m, int n, int k) const override; - - bool checkTactic(int m, int n, int k, Config const& tactic) const override; - -private: - WeightTypeId mWeightTypeId; - KernelType mCudaKernelType; - int mArch; -}; - -class WeightOnlyQuantMatmulPlugin : public BasePlugin -{ -public: - using PluginProfilerPtr = std::shared_ptr; - WeightOnlyQuantMatmulPlugin() = delete; - - WeightOnlyQuantMatmulPlugin(nvinfer1::DataType type, WeightTypeId weightTypeId, PluginProfilerPtr const& profiler); - - WeightOnlyQuantMatmulPlugin(void const* data, size_t length, PluginProfilerPtr const& profiler); - - ~WeightOnlyQuantMatmulPlugin() override = default; - - // IPluginV2DynamicExt Methods - nvinfer1::IPluginV2DynamicExt* clone() const noexcept override; - nvinfer1::DimsExprs getOutputDimensions(int outputIndex, nvinfer1::DimsExprs const* inputs, int nbInputs, - nvinfer1::IExprBuilder& exprBuilder) noexcept override; - bool supportsFormatCombination( - int pos, nvinfer1::PluginTensorDesc const* inOut, int nbInputs, int nbOutputs) noexcept override; - void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int nbInputs, - nvinfer1::DynamicPluginTensorDesc const* out, int nbOutputs) noexcept override; - size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int nbInputs, - nvinfer1::PluginTensorDesc const* outputs, int nbOutputs) const noexcept override; - int enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc, - void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; - - // IPluginV2Ext Methods - nvinfer1::DataType getOutputDataType( - int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept override; - - // IPluginV2 Methods - char const* getPluginType() const noexcept override; - char const* getPluginVersion() const noexcept override; - int getNbOutputs() const noexcept override; - int initialize() noexcept override; - void terminate() noexcept override; - size_t getSerializationSize() const noexcept override; - void serialize(void* buffer) const noexcept override; - void destroy() noexcept override; - -private: - void init(nvinfer1::DataType type, WeightTypeId weightTypeId); - - void configGemm(); - -private: - const std::string mLayerName; - - WeightOnlyGemmRunnerPtr m_weightOnlyGemmRunner; - size_t m_workspaceMaxSize; - nvinfer1::DataType mType; - WeightTypeId mWeightTypeId; - bool mCudaKernelEnabled; - tensorrt_llm::kernels::weight_only::KernelType mCudaKernelType; - int mArch; - - GemmDims mDims{}; - GemmIdCore mGemmId{}; - - PluginProfilerPtr mPluginProfiler; -}; - -class WeightOnlyQuantMatmulPluginCreator : public BaseCreator -{ -public: - WeightOnlyQuantMatmulPluginCreator(); - - char const* getPluginName() const noexcept override; - - char const* getPluginVersion() const noexcept override; - - nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override; - - nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override; - - nvinfer1::IPluginV2* deserializePlugin( - char const* name, void const* serialData, size_t serialLength) noexcept override; - -private: - GemmPluginProfilerManager gemmPluginProfileManager; - static nvinfer1::PluginFieldCollection mFC; - static std::vector mPluginAttributes; -}; - -} // namespace tensorrt_llm::plugins diff --git a/cpp/tensorrt_llm/runtime/CMakeLists.txt b/cpp/tensorrt_llm/runtime/CMakeLists.txt index ca81fbb0f6cd..11a9391c0e69 100644 --- a/cpp/tensorrt_llm/runtime/CMakeLists.txt +++ b/cpp/tensorrt_llm/runtime/CMakeLists.txt @@ -26,7 +26,6 @@ set(SRCS eagleBuffers.cpp explicitDraftTokensBuffers.cpp lookaheadBuffers.cpp - layerProfiler.cpp loraManager.cpp loraUtils.cpp loraModule.cpp @@ -51,9 +50,6 @@ set(SRCS promptTuningParams.cpp runtimeKernels.cu tllmBuffers.cpp - tllmRuntime.cpp - tllmStreamReaders.cpp - tllmLogger.cpp workerPool.cpp worldConfig.cpp virtualMemory.cpp) diff --git a/cpp/tensorrt_llm/runtime/bufferManager.cpp b/cpp/tensorrt_llm/runtime/bufferManager.cpp index 3de42a253158..5ca6ccdaa9e2 100644 --- a/cpp/tensorrt_llm/runtime/bufferManager.cpp +++ b/cpp/tensorrt_llm/runtime/bufferManager.cpp @@ -37,7 +37,7 @@ BufferManager::BufferManager(CudaStreamPtr stream, bool trimPool) mPool = CudaMemPool::getPrimaryPoolForDevice(mStream->getDevice()); } -BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, nvinfer1::DataType type) const +BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, tensorrt_llm::DataType type) const { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -51,7 +51,7 @@ BufferManager::IBufferPtr BufferManager::gpu(std::size_t size, nvinfer1::DataTyp return gpuSync(size, type); } -BufferManager::ITensorPtr BufferManager::gpu(nvinfer1::Dims dims, nvinfer1::DataType type) const +BufferManager::ITensorPtr BufferManager::gpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) const { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -65,7 +65,7 @@ BufferManager::ITensorPtr BufferManager::gpu(nvinfer1::Dims dims, nvinfer1::Data return gpuSync(dims, type); } -BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, tensorrt_llm::DataType type) { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -74,7 +74,7 @@ BufferManager::IBufferPtr BufferManager::gpuSync(std::size_t size, nvinfer1::Dat return std::make_unique(size, type, CudaAllocator{}); } -BufferManager::ITensorPtr BufferManager::gpuSync(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::gpuSync(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { if (auto vmAllocator = getVirtualMemoryAllocator()) { @@ -83,47 +83,47 @@ BufferManager::ITensorPtr BufferManager::gpuSync(nvinfer1::Dims dims, nvinfer1:: return std::make_unique(dims, type, CudaAllocator{}); } -BufferManager::IBufferPtr BufferManager::cpu(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::cpu(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::cpu(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::cpu(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::IBufferPtr BufferManager::pinned(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::pinned(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::pinned(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::pinned(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::IBufferPtr BufferManager::pinnedPool(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::pinnedPool(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::pinnedPool(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::pinnedPool(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::IBufferPtr BufferManager::managed(std::size_t size, nvinfer1::DataType type) +BufferManager::IBufferPtr BufferManager::managed(std::size_t size, tensorrt_llm::DataType type) { return std::make_unique(size, type); } -BufferManager::ITensorPtr BufferManager::managed(nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::managed(tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type); } -BufferManager::ITensorPtr BufferManager::ipcNvls(std::set ranks, nvinfer1::Dims dims, nvinfer1::DataType type) +BufferManager::ITensorPtr BufferManager::ipcNvls(std::set ranks, tensorrt_llm::Dims dims, tensorrt_llm::DataType type) { return std::make_unique(dims, type, ranks); } @@ -187,7 +187,7 @@ void BufferManager::copy(IBuffer const& src, IBuffer& dst) const } BufferManager::IBufferPtr BufferManager::allocate( - MemoryType memoryType, std::size_t size, nvinfer1::DataType type) const + MemoryType memoryType, std::size_t size, tensorrt_llm::DataType type) const { switch (memoryType) { @@ -202,7 +202,7 @@ BufferManager::IBufferPtr BufferManager::allocate( } BufferManager::ITensorPtr BufferManager::allocate( - MemoryType memoryType, nvinfer1::Dims dims, nvinfer1::DataType type) const + MemoryType memoryType, tensorrt_llm::Dims dims, tensorrt_llm::DataType type) const { switch (memoryType) { diff --git a/cpp/tensorrt_llm/runtime/bufferView.h b/cpp/tensorrt_llm/runtime/bufferView.h index 236b89d7d455..aabf9c7ff793 100644 --- a/cpp/tensorrt_llm/runtime/bufferView.h +++ b/cpp/tensorrt_llm/runtime/bufferView.h @@ -70,7 +70,7 @@ class BufferView : virtual public IBuffer return mBuffer->getCapacity() - mOffset; } - [[nodiscard]] nvinfer1::DataType getDataType() const override + [[nodiscard]] tensorrt_llm::DataType getDataType() const override { return mBuffer->getDataType(); } diff --git a/cpp/tensorrt_llm/runtime/decoderState.cpp b/cpp/tensorrt_llm/runtime/decoderState.cpp index b5851dc1c2d2..8e317cc721b7 100644 --- a/cpp/tensorrt_llm/runtime/decoderState.cpp +++ b/cpp/tensorrt_llm/runtime/decoderState.cpp @@ -27,10 +27,10 @@ using TensorPtr = DecoderState::TensorPtr; BeamSearchBuffers::BeamSearchBuffers(BufferManager const& bufferManager) : mOutputBeamHypotheses{} - , mCumLogProbsTmp(bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT)) + , mCumLogProbsTmp(bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT)) { mOutputBeamHypotheses.empty(bufferManager); - mCumLogProbsTmp = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kFLOAT); + mCumLogProbsTmp = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kFLOAT); int device; cudaGetDevice(&device); @@ -54,7 +54,7 @@ DecoderState::DecoderState() } void DecoderState::setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, SizeType32 maxAttentionWindow, - SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, nvinfer1::DataType dtype, ModelConfig const& modelConfig, + SizeType32 sinkTokenLength, SizeType32 maxSequenceLength, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -64,7 +64,7 @@ void DecoderState::setup(SizeType32 maxNumSequences, SizeType32 maxBeamWidth, Si TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); } -void DecoderState::setupBuffers(nvinfer1::DataType dtype, BufferManager const& bufferManager) +void DecoderState::setupBuffers(tensorrt_llm::DataType dtype, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); auto constexpr nvTokenIdType = TRTDataType::value; @@ -114,7 +114,7 @@ void DecoderState::setupBuffers(nvinfer1::DataType dtype, BufferManager const& b } void DecoderState::setupSpeculativeDecoding(SpeculativeDecodingMode const& speculativeDecodingMode, - SizeType32 maxTokensPerEngineStep, nvinfer1::DataType dtype, ModelConfig const& modelConfig, + SizeType32 maxTokensPerEngineStep, tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -125,7 +125,7 @@ void DecoderState::setupSpeculativeDecoding(SpeculativeDecodingMode const& specu } void DecoderState::setupSpeculativeDecodingBuffers( - SpeculativeDecodingMode const speculativeDecodingMode, nvinfer1::DataType dtype, BufferManager const& bufferManager) + SpeculativeDecodingMode const speculativeDecodingMode, tensorrt_llm::DataType dtype, BufferManager const& bufferManager) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); @@ -151,13 +151,13 @@ void DecoderState::setupSpeculativeDecodingBuffers( if (speculativeDecodingMode.predictsDraftTokens()) { speculativeDecodingOutputs.nextDraftTokens - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); if (speculativeDecodingMode.variableDraftLength()) { speculativeDecodingOutputs.nextDraftTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); speculativeDecodingOutputs.prevDraftTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); } } if (speculativeDecodingMode.isLookaheadDecoding()) @@ -167,11 +167,11 @@ void DecoderState::setupSpeculativeDecodingBuffers( if (speculativeDecodingMode.needsKVCacheRewind()) { speculativeDecodingOutputs.acceptedTokensLen - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); speculativeDecodingOutputs.acceptedLengthsCumSum - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); speculativeDecodingOutputs.pathsOffsets - = bufferManager.emptyTensor(MemoryType::kGPU, nvinfer1::DataType::kINT32); + = bufferManager.emptyTensor(MemoryType::kGPU, tensorrt_llm::DataType::kINT32); } dOutput->speculativeDecodingOutputs = speculativeDecodingOutputs; diff --git a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp index f8a4fa4e7467..c3f8b76bb827 100644 --- a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp +++ b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.cpp @@ -19,7 +19,7 @@ #include tensorrt_llm::runtime::DecodingLayerWorkspace::DecodingLayerWorkspace(std::shared_ptr bufferManager, - tensorrt_llm::layers::DecoderDomain const& decoderDomain, nvinfer1::DataType logitsType, + tensorrt_llm::layers::DecoderDomain const& decoderDomain, tensorrt_llm::DataType logitsType, size_t workspaceBufferSizeInBytes) : mBufferManager(std::move(bufferManager)) , mBatchSlotsDevice( @@ -82,7 +82,7 @@ void tensorrt_llm::runtime::DecodingLayerWorkspace::resize(size_t minSize) } tensorrt_llm::runtime::DecodingLayerWorkspace::TensorPtr -tensorrt_llm::runtime::DecodingLayerWorkspace::getWorkspaceAsDeviceTensor(ITensor::Shape shape, nvinfer1::DataType type) +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>>( diff --git a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h index c2688b51139f..8f82bda01b72 100644 --- a/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h +++ b/cpp/tensorrt_llm/runtime/decodingLayerWorkspace.h @@ -39,7 +39,7 @@ class DecodingLayerWorkspace using BufferPtr = IBuffer::SharedPtr; DecodingLayerWorkspace(std::shared_ptr bufferManager, layers::DecoderDomain const& decoderDomain, - nvinfer1::DataType logitsType, size_t workspaceBufferSizeInBytes); + tensorrt_llm::DataType logitsType, size_t workspaceBufferSizeInBytes); DecodingLayerWorkspace() = delete; @@ -71,7 +71,7 @@ class DecodingLayerWorkspace [[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, nvinfer1::DataType type); + 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 @@ -112,7 +112,7 @@ class DecodingLayerWorkspace { size_t lastTensorOffset = 0; auto alignedSizeCalculator - = [&lastTensorOffset](std::pair const& tensorDescriptor) + = [&lastTensorOffset](std::pair const& tensorDescriptor) { auto const& [shape, type] = tensorDescriptor; auto const sizeInBytes = ITensor::volume(shape) * tensorrt_llm::common::getDTypeSize(type); diff --git a/cpp/tensorrt_llm/runtime/eagleBuffers.cpp b/cpp/tensorrt_llm/runtime/eagleBuffers.cpp index 097fd95f49aa..0750339cfcff 100644 --- a/cpp/tensorrt_llm/runtime/eagleBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/eagleBuffers.cpp @@ -41,50 +41,50 @@ void EagleBuffers::Inputs::create(SizeType32 maxNumSequences, BufferManager cons auto const numEagleLayers = speculativeDecodingModule.getMaxDraftPathLen(); auto constexpr TRTTokenIdType = runtime::TRTDataType::value; - temperatures = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kFLOAT); - randomDataSample = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kFLOAT); + 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}), nvinfer1::DataType::kFLOAT); + = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens}), tensorrt_llm::DataType::kFLOAT); draftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), TRTTokenIdType); - draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); draftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), tensorrt_llm::DataType::kINT32); draftPathsHost = BufferManager::pinnedPool( - ITensor::makeShape({maxNumSequences, maxNumPaths, maxPathLen}), nvinfer1::DataType::kINT32); - specDecodingGenerationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + 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}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); specDecodingPackedMasks = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingTokens, common::ceilDiv(maxDecodingTokens, 32)}), - nvinfer1::DataType::kINT32); + tensorrt_llm::DataType::kINT32); specDecodingPositionOffsets - = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), tensorrt_llm::DataType::kINT32); eagleNetCtxRequestTypesHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetCtxContextLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetCtxPastKeyValueLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetGenRequestTypesHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetGenContextLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); eagleNetGenPastKeyValueLengthsHost - = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); inputGenTokensHost = BufferManager::pinnedPool( - ITensor::makeShape({maxNumSequences * maxDecodingTokens}), nvinfer1::DataType::kINT32); - chunkedContextNextTokens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + 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}), nvinfer1::DataType::kINT32); - dynamicTreeMaxTopKHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - prevScores = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), nvinfer1::DataType::kFLOAT); + 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}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); allLayersDraftTokenIds = manager.gpu( ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), TRTTokenIdType); @@ -114,58 +114,58 @@ EagleBuffers::EagleBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, run auto constexpr TRTTokenIdType = runtime::TRTDataType::value; // input tensors - engineInputs.temperatures = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - engineInputs.posteriorAlpha = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - engineInputs.posteriorThreshold = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); - posteriorAlphaHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kFLOAT); - posteriorThresholdHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kFLOAT); - greedySamplingHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + 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}), nvinfer1::DataType::kINT32); + engineInputs.draftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); engineInputs.draftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), tensorrt_llm::DataType::kINT32); engineInputs.specDecodingGenerationLengths - = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); engineInputs.specDecodingPositionOffsets - = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.specDecodingPackedMasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + = 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, nvinfer1::DataType::kFLOAT); - engineInputs.randomDataValidation = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kFLOAT); + 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, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetCtxContextLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetCtxPastKeyValueLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetGenRequestTypesHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetGenContextLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); + = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); engineInputs.eagleNetGenPastKeyValueLengthsHost - = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); - engineInputs.inputGenTokensHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, nvinfer1::DataType::kINT32); - engineInputs.chunkedContextNextTokens = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.useSpecDecoding = BufferManager::cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + = 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, nvinfer1::DataType::kINT32); + chunkedContextNextTokensHost = manager.emptyTensor(runtime::MemoryType::kPINNEDPOOL, tensorrt_llm::DataType::kINT32); // Eagle-2 - engineInputs.useDynamicTreeHost = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + engineInputs.useDynamicTreeHost = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); engineInputs.dynamicTreeMaxTopKHost - = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); engineInputs.prevScores - = manager.gpu(ITensor::makeShape({maxNumSequences, maxDecodingDraftTokens}), nvinfer1::DataType::kFLOAT); + = 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}), - nvinfer1::DataType::kFLOAT); + tensorrt_llm::DataType::kFLOAT); engineInputs.allLayersDraftTokenIds = manager.gpu( ITensor::makeShape({maxNumSequences, numEagleLayers, maxDecodingDraftTokens * maxDecodingDraftTokens}), TRTTokenIdType); @@ -176,24 +176,24 @@ EagleBuffers::EagleBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, run // output tensors engineOutputs.nextDraftTokens = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), TRTTokenIdType); - engineOutputs.nextDraftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + engineOutputs.nextDraftLens = manager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32); engineOutputs.nextDraftPaths - = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), nvinfer1::DataType::kINT32); + = manager.gpu(ITensor::makeShape({maxNumSequences, numPaths, pathLen}), tensorrt_llm::DataType::kINT32); engineOutputs.acceptedTokens - = manager.gpu(ITensor::makeShape({maxNumSequences, pathLen}), nvinfer1::DataType::kINT32); - engineOutputs.acceptedLens = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - engineOutputs.acceptedPaths = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + = 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}), nvinfer1::DataType::kINT32); + = 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, nvinfer1::DataType::kINT32); - maxGenerationLength = manager.gpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + 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); @@ -520,15 +520,15 @@ void EagleBuffers::setFromInputs(RequestVector const& contextRequests, RequestVe switch (dtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: setFromInputs( contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: setFromInputs( contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); break; - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: setFromInputs<__nv_bfloat16>( contextRequests, genRequests, vocabSizePadded, seqSlots, draftBuffers, *eagleModule, manager); break; diff --git a/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp b/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp index ed205ca0e117..036cb4aa50a1 100644 --- a/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/explicitDraftTokensBuffers.cpp @@ -40,23 +40,23 @@ void ExplicitDraftTokensBuffers::Inputs::create(SizeType32 maxNumSequences, Buff auto constexpr TRTTokenIdType = runtime::TRTDataType::value; auto const dtype = modelConfig.getDataType(); - maxGenLengthHost = manager.pinned(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + maxGenLengthHost = manager.pinned(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); temperatures = manager.gpu(ITensor::makeShape({maxNumSequences}), dtype); - positionIdsBase = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - generationLengths = manager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); - generationLengthsHost = manager.pinned(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32); + 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}), nvinfer1::DataType::kINT32); + = 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)}), - nvinfer1::DataType::kINT32); - positionIds = manager.gpu(ITensor::makeShape({maxNumSequences * maxDecodingTokens}), nvinfer1::DataType::kINT32); - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + 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, @@ -81,52 +81,52 @@ ExplicitDraftTokensBuffers::ExplicitDraftTokensBuffers(SizeType32 maxBatchSize, auto const dtype = modelConfig.getDataType(); // input tensors - engineInputs.requestTypesDevice = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + 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}), nvinfer1::DataType::kINT32); + = 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, nvinfer1::DataType::kINT32); - engineInputs.positionIds = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.positionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineInputs.packedMasks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + 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, nvinfer1::DataType::kINT32); - engineInputs.useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + 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}), nvinfer1::DataType::kINT32); + = 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}), nvinfer1::DataType::kINT32); - engineOutputs.totalGenToken = manager.gpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + 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, nvinfer1::DataType::kINT32); - engineOutputs.nextPositionOffsets = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineOutputs.masks = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kBOOL); + 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, nvinfer1::DataType::kINT32); - engineOutputs.bestPathIndices = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); - engineOutputs.packedPositionIds = manager.emptyTensor(runtime::MemoryType::kGPU, nvinfer1::DataType::kINT32); + 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, nvinfer1::DataType::kINT32); + cumSumGenerationLengths = manager.emptyTensor(runtime::MemoryType::kGPU, tensorrt_llm::DataType::kINT32); // pre-allocate empty tensors reshape(0, maxNumSequences, modelConfig); @@ -295,15 +295,15 @@ void ExplicitDraftTokensBuffers::setFromInputs(SizeType32 numCtxSequences, SizeT switch (dtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: setFromInputs(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, contextPositionIds, *explicitDraftTokensModule, stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: setFromInputs(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, contextPositionIds, *explicitDraftTokensModule, stream); break; - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: setFromInputs<__nv_bfloat16>(numCtxSequences, numGenSequences, vocabSizePadded, seqSlots, draftBuffers, contextPositionIds, *explicitDraftTokensModule, stream); break; diff --git a/cpp/tensorrt_llm/runtime/gptDecoder.cpp b/cpp/tensorrt_llm/runtime/gptDecoder.cpp index 930877206462..e1ac1717af45 100644 --- a/cpp/tensorrt_llm/runtime/gptDecoder.cpp +++ b/cpp/tensorrt_llm/runtime/gptDecoder.cpp @@ -21,7 +21,7 @@ #include "tensorrt_llm/layers/dynamicDecodeLayer.h" #include "tensorrt_llm/runtime/decodingLayerWorkspace.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include @@ -121,7 +121,7 @@ void GptDecoder::disableLookahead( template void GptDecoder::setup(SamplingConfig const& samplingConfig, size_t batchSize, TensorConstPtr const& batchSlots, - std::optional const& output, std::optional explicitDraftTokensDType, + std::optional const& output, std::optional explicitDraftTokensDType, std::optional> const& lookaheadPrompt, std::optional> const& lookaheadAlgoConfigs) { diff --git a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp b/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp index c55d02093afc..1bddb324f89b 100644 --- a/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp +++ b/cpp/tensorrt_llm/runtime/gptDecoderBatched.cpp @@ -74,7 +74,7 @@ void GptDecoderBatched::disableLookahead(RequestVector const& genRequests, Tenso } void GptDecoderBatched::setup(executor::DecodingMode const& mode, SizeType32 maxNumSequences, SizeType32 maxBeamWidth, - nvinfer1::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) + tensorrt_llm::DataType dtype, ModelConfig const& modelConfig, WorldConfig const& worldConfig) { TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); TLLM_CHECK(maxNumSequences > 0); diff --git a/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp b/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp index 311f63eaf1e7..6110fa13e993 100644 --- a/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp +++ b/cpp/tensorrt_llm/runtime/gptJsonConfig.cpp @@ -80,14 +80,14 @@ std::optional parseJsonFieldOptional(Json const& json, std::string_vi return value; } -nvinfer1::DataType strToDType(std::string type) +tensorrt_llm::DataType strToDType(std::string type) { - static std::map const typeMap = {{"int64", nvinfer1::DataType::kINT64}, - {"int32", nvinfer1::DataType::kINT32}, {"int", nvinfer1::DataType::kINT32}, - {"float32", nvinfer1::DataType::kFLOAT}, {"bfloat16", nvinfer1::DataType::kBF16}, - {"float16", nvinfer1::DataType::kHALF}, {"bool", nvinfer1::DataType::kBOOL}, - {"uint8", nvinfer1::DataType::kUINT8}, {"int8", nvinfer1::DataType::kINT8}, {"fp8", nvinfer1::DataType::kFP8}, - {"int4", nvinfer1::DataType::kINT4}}; + static std::map const typeMap = {{"int64", tensorrt_llm::DataType::kINT64}, + {"int32", tensorrt_llm::DataType::kINT32}, {"int", tensorrt_llm::DataType::kINT32}, + {"float32", tensorrt_llm::DataType::kFLOAT}, {"bfloat16", tensorrt_llm::DataType::kBF16}, + {"float16", tensorrt_llm::DataType::kHALF}, {"bool", tensorrt_llm::DataType::kBOOL}, + {"uint8", tensorrt_llm::DataType::kUINT8}, {"int8", tensorrt_llm::DataType::kINT8}, {"fp8", tensorrt_llm::DataType::kFP8}, + {"int4", tensorrt_llm::DataType::kINT4}}; TLLM_CHECK_WITH_INFO(typeMap.count(type) > 0, type + " not found in strToDtype."); return typeMap.at(type); @@ -140,14 +140,14 @@ std::vector buildLayerTypes( return result; } -ModelConfig parseMultimodalConfig(Json const& json, nvinfer1::DataType dataType) +ModelConfig parseMultimodalConfig(Json const& json, tensorrt_llm::DataType dataType) { return ModelConfig{128, 10, 10, 0, 1, 128, dataType}; // use dummy values because vision engines of multimodal models does not record this info in config } ModelConfig createModelConfig(Json const& json, bool engineVersionNone, SizeType32 tensorParallelism, - SizeType32 contextParallelism, nvinfer1::DataType dataType) + SizeType32 contextParallelism, tensorrt_llm::DataType dataType) { auto const& config = engineVersionNone ? json.at("builder_config") : json.at("pretrained_config"); auto const multiModalName = parseJsonFieldOptional(config, "model_name"); @@ -248,14 +248,14 @@ ModelConfig createModelConfig(Json const& json, bool engineVersionNone, SizeType modelConfig.setLayerTypes(layerTypes); // Set logits datatype - auto logitsDtype = nvinfer1::DataType::kFLOAT; + auto logitsDtype = tensorrt_llm::DataType::kFLOAT; if (logitsDtypeStr == "float32") { - logitsDtype = nvinfer1::DataType::kFLOAT; + logitsDtype = tensorrt_llm::DataType::kFLOAT; } else if (logitsDtypeStr == "float16") { - logitsDtype = nvinfer1::DataType::kHALF; + logitsDtype = tensorrt_llm::DataType::kHALF; } else { @@ -490,15 +490,15 @@ GptJsonConfig parseJson(InputType&& input) { if (precision == "float32") { - return nvinfer1::DataType::kFLOAT; + return tensorrt_llm::DataType::kFLOAT; } if (precision == "float16") { - return nvinfer1::DataType::kHALF; + return tensorrt_llm::DataType::kHALF; } if (precision == "bfloat16") { - return nvinfer1::DataType::kBF16; + return tensorrt_llm::DataType::kBF16; } TLLM_THROW("Model data type '%s' not supported", precision.c_str()); }(); diff --git a/cpp/tensorrt_llm/runtime/iBuffer.cpp b/cpp/tensorrt_llm/runtime/iBuffer.cpp index 77707a0e4cf8..6f12d817cdab 100644 --- a/cpp/tensorrt_llm/runtime/iBuffer.cpp +++ b/cpp/tensorrt_llm/runtime/iBuffer.cpp @@ -48,7 +48,7 @@ IBuffer::UniquePtr IBuffer::slice(IBuffer::SharedPtr buffer, std::size_t offset, return std::make_unique(std::move(buffer), offset, size); } -IBuffer::UniquePtr IBuffer::wrap(void* data, nvinfer1::DataType type, std::size_t size, std::size_t capacity) +IBuffer::UniquePtr IBuffer::wrap(void* data, tensorrt_llm::DataType type, std::size_t size, std::size_t capacity) { TLLM_CHECK_WITH_INFO(size <= capacity, "Requested size is larger than capacity"); auto memoryType = IBuffer::memoryType(data); @@ -91,17 +91,17 @@ char const* IBuffer::getDataTypeName(DataType dataType) { switch (dataType) { - case nvinfer1::DataType::kINT64: return DataTypeTraits::name; - case nvinfer1::DataType::kINT32: return DataTypeTraits::name; - case nvinfer1::DataType::kFLOAT: return DataTypeTraits::name; - case nvinfer1::DataType::kBF16: return DataTypeTraits::name; - case nvinfer1::DataType::kHALF: return DataTypeTraits::name; - case nvinfer1::DataType::kBOOL: return DataTypeTraits::name; - case nvinfer1::DataType::kUINT8: return DataTypeTraits::name; - case nvinfer1::DataType::kINT8: return DataTypeTraits::name; - case nvinfer1::DataType::kFP8: return DataTypeTraits::name; - case nvinfer1::DataType::kINT4: [[fallthrough]] /* do nothing */; - case nvinfer1::DataType::kFP4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kINT64: return DataTypeTraits::name; + case tensorrt_llm::DataType::kINT32: return DataTypeTraits::name; + case tensorrt_llm::DataType::kFLOAT: return DataTypeTraits::name; + case tensorrt_llm::DataType::kBF16: return DataTypeTraits::name; + case tensorrt_llm::DataType::kHALF: return DataTypeTraits::name; + case tensorrt_llm::DataType::kBOOL: return DataTypeTraits::name; + case tensorrt_llm::DataType::kUINT8: return DataTypeTraits::name; + case tensorrt_llm::DataType::kINT8: return DataTypeTraits::name; + case tensorrt_llm::DataType::kFP8: return DataTypeTraits::name; + case tensorrt_llm::DataType::kINT4: [[fallthrough]] /* do nothing */; + case tensorrt_llm::DataType::kFP4: [[fallthrough]] /* do nothing */; default: TLLM_THROW("Unknown data type"); } } diff --git a/cpp/tensorrt_llm/runtime/iTensor.cpp b/cpp/tensorrt_llm/runtime/iTensor.cpp index f78b25fdb19a..13d13c9e9a13 100644 --- a/cpp/tensorrt_llm/runtime/iTensor.cpp +++ b/cpp/tensorrt_llm/runtime/iTensor.cpp @@ -72,22 +72,22 @@ ITensor::UniquePtr ITensor::slice(SharedPtr tensor, Shape const& offsetDims, ITe return std::make_unique(std::move(tensor), offset, volume(dims), dims); } -ITensor::UniquePtr ITensor::view(IBuffer::SharedPtr buffer, nvinfer1::Dims const& dims) +ITensor::UniquePtr ITensor::view(IBuffer::SharedPtr buffer, tensorrt_llm::Dims const& dims) { auto const size = buffer->getSize(); return std::make_unique(std::move(buffer), 0, size, dims); } -nvinfer1::Dims ITensor::makeShape(std::initializer_list const& dims) +tensorrt_llm::Dims ITensor::makeShape(std::initializer_list const& dims) { - TLLM_CHECK_WITH_INFO(dims.size() <= nvinfer1::Dims::MAX_DIMS, "Number of dimensions is too large"); - nvinfer1::Dims shape{}; + TLLM_CHECK_WITH_INFO(dims.size() <= tensorrt_llm::Dims::MAX_DIMS, "Number of dimensions is too large"); + tensorrt_llm::Dims shape{}; shape.nbDims = static_cast(dims.size()); std::copy(dims.begin(), dims.end(), shape.d); return shape; } -std::string ITensor::toString(nvinfer1::Dims const& dims) +std::string ITensor::toString(tensorrt_llm::Dims const& dims) { if (dims.nbDims < 0) { @@ -103,7 +103,7 @@ std::string ITensor::toString(nvinfer1::Dims const& dims) } } -ITensor::UniquePtr ITensor::wrap(void* data, nvinfer1::DataType type, nvinfer1::Dims const& shape, std::size_t capacity) +ITensor::UniquePtr ITensor::wrap(void* data, tensorrt_llm::DataType type, tensorrt_llm::Dims const& shape, std::size_t capacity) { auto const size = volumeNonNegative(shape); TLLM_CHECK_WITH_INFO(size <= capacity, "Requested size is larger than capacity"); @@ -230,18 +230,18 @@ std::ostream& tensorrt_llm::runtime::operator<<(std::ostream& out, ITensor const { switch (tensor.getDataType()) { - case nvinfer1::DataType::kFLOAT: printTensor(tensor, out); break; - case nvinfer1::DataType::kHALF: printTensor(tensor, out); break; - case nvinfer1::DataType::kBOOL: printTensor(tensor, out); break; - case nvinfer1::DataType::kINT8: printTensor(tensor, out); break; - case nvinfer1::DataType::kINT32: printTensor(tensor, out); break; - case nvinfer1::DataType::kINT64: printTensor(tensor, out); break; - case nvinfer1::DataType::kUINT8: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kFLOAT: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kHALF: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kBOOL: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kINT8: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kINT32: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kINT64: printTensor(tensor, out); break; + case tensorrt_llm::DataType::kUINT8: printTensor(tensor, out); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: printTensor<__nv_bfloat16, float>(tensor, out); break; + case tensorrt_llm::DataType::kBF16: printTensor<__nv_bfloat16, float>(tensor, out); break; #endif #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: printTensor<__nv_fp8_e4m3, float>(tensor, out); break; + case tensorrt_llm::DataType::kFP8: printTensor<__nv_fp8_e4m3, float>(tensor, out); break; #endif default: TLLM_THROW("Unsupported data type"); } diff --git a/cpp/tensorrt_llm/runtime/ipcUtils.cpp b/cpp/tensorrt_llm/runtime/ipcUtils.cpp index 23a7e28a4f27..48368844f850 100644 --- a/cpp/tensorrt_llm/runtime/ipcUtils.cpp +++ b/cpp/tensorrt_llm/runtime/ipcUtils.cpp @@ -20,7 +20,7 @@ #include "tensorrt_llm/common/workspace.h" #include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include namespace tensorrt_llm::runtime @@ -83,7 +83,7 @@ void IpcMemory::allocateIpcMemory(std::size_t bufferSize, BufferManager const& m // IPC handles. If we want to support stream-ordered allocations here, we need to create another pool with the // correct handle type. auto const ipcAlignedBufferSize = common::alignSize(bufferSize, 1LU << 21); - mBuffer = BufferManager::gpuSync(ipcAlignedBufferSize, nvinfer1::DataType::kUINT8); + mBuffer = BufferManager::gpuSync(ipcAlignedBufferSize, tensorrt_llm::DataType::kUINT8); manager.setZero(*mBuffer); auto* bufferPtr = mBuffer->data(); @@ -149,7 +149,7 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi { auto const tpSize = worldConfig.getTensorParallelism(); mAllReduceCommPtrs = BufferManager::cpu( - ITensor::makeShape({static_cast(7) * tpSize + 3}), nvinfer1::DataType::kINT64); + ITensor::makeShape({static_cast(7) * tpSize + 3}), tensorrt_llm::DataType::kINT64); } else { @@ -178,7 +178,7 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi mAllReduceCommPtrs = BufferManager::cpu(ITensor::makeShape({static_cast(mIpcMemoryHandles.size()) * tpSize + 3}), - nvinfer1::DataType::kINT64); + tensorrt_llm::DataType::kINT64); auto commPtrs = BufferRange(*mAllReduceCommPtrs); // Start from 1 since 0 represents released state for barrier at the beginning of the all_reduce. // The last element is the barrier flag counter. @@ -211,9 +211,9 @@ AllReduceBuffers::AllReduceBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWi void lamportInitializeAll(void* buffer_0, void* buffer_1, void* buffer_2, size_t size) { #if ENABLE_MULTI_DEVICE - tensorrt_llm::kernels::lamportInitialize(buffer_0, size / sizeof(half), nvinfer1::DataType::kHALF, 0); - tensorrt_llm::kernels::lamportInitialize(buffer_1, size / sizeof(half), nvinfer1::DataType::kHALF, 0); - tensorrt_llm::kernels::lamportInitialize(buffer_2, size / sizeof(half), nvinfer1::DataType::kHALF, 0); + tensorrt_llm::kernels::lamportInitialize(buffer_0, size / sizeof(half), tensorrt_llm::DataType::kHALF, 0); + tensorrt_llm::kernels::lamportInitialize(buffer_1, size / sizeof(half), tensorrt_llm::DataType::kHALF, 0); + tensorrt_llm::kernels::lamportInitialize(buffer_2, size / sizeof(half), tensorrt_llm::DataType::kHALF, 0); cudaDeviceSynchronize(); #endif } diff --git a/cpp/tensorrt_llm/runtime/layerProfiler.cpp b/cpp/tensorrt_llm/runtime/layerProfiler.cpp deleted file mode 100644 index 4c3c9779cedb..000000000000 --- a/cpp/tensorrt_llm/runtime/layerProfiler.cpp +++ /dev/null @@ -1,97 +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/layerProfiler.h" -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; - -void LayerProfiler::reportLayerTime(char const* layerName, float timeMs) noexcept -{ - if (mIterator == mLayers.end()) - { - bool const first = !mLayers.empty() && mLayers.begin()->name == layerName; - mUpdatesCount += mLayers.empty() || first; - if (first) - { - mIterator = mLayers.begin(); - } - else - { - mLayers.emplace_back(); - mLayers.back().name = layerName; - mIterator = mLayers.end() - 1; - } - } - - mIterator->timeMs.push_back(timeMs); - ++mIterator; -} - -float LayerProfiler::getTotalTime() const noexcept -{ - auto const plusLayerTime = [](float accumulator, LayerProfile const& lp) - { return accumulator + std::accumulate(lp.timeMs.begin(), lp.timeMs.end(), 0.F, std::plus()); }; - return std::accumulate(mLayers.begin(), mLayers.end(), 0.0F, plusLayerTime); -} - -std::string LayerProfiler::getLayerProfile() noexcept -{ - std::string const nameHdr(" Layer"); - std::string const timeHdr(" Time(ms)"); - - float const totalTimeMs = getTotalTime(); - - auto const timeLength = timeHdr.size(); - - std::unordered_map layer2times; - std::vector layer_order; - for (auto const& p : mLayers) - { - if (!layer2times.count(p.name)) - { - layer2times[p.name] = 0; - layer_order.push_back(p.name); - } - for (auto const& t : p.timeMs) - { - layer2times[p.name] += t; - } - } - - std::stringstream ss; - ss << "\n=== Per-layer Profile ===\n" << timeHdr << nameHdr << "\n"; - - for (auto const& name : layer_order) - { - if (layer2times[name] == 0.0f) - { - continue; - } - ss << std::setw(timeLength) << std::fixed << std::setprecision(2) << layer2times[name] << " " << name << "\n"; - } - - ss << std::setw(timeLength) << std::fixed << std::setprecision(2) << totalTimeMs << " Total\n"; - ss << "\n"; - - // clear data - mLayers.clear(); - - return ss.str(); -} diff --git a/cpp/tensorrt_llm/runtime/layerProfiler.h b/cpp/tensorrt_llm/runtime/layerProfiler.h deleted file mode 100644 index bcae1546de3d..000000000000 --- a/cpp/tensorrt_llm/runtime/layerProfiler.h +++ /dev/null @@ -1,47 +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/runtime/common.h" -#include - -#include - -namespace tensorrt_llm::runtime -{ -struct LayerProfile -{ - std::string name; - std::vector timeMs; -}; - -class LayerProfiler : public nvinfer1::IProfiler -{ - -public: - void reportLayerTime(char const* layerName, float timeMs) noexcept override; - - std::string getLayerProfile() noexcept; - -private: - [[nodiscard]] float getTotalTime() const noexcept; - - std::vector mLayers; - std::vector::iterator mIterator{mLayers.begin()}; - int32_t mUpdatesCount{0}; -}; -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp b/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp index ef800ef218e4..c41244dfa7df 100644 --- a/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/lookaheadBuffers.cpp @@ -16,208 +16,22 @@ */ #include "tensorrt_llm/runtime/lookaheadBuffers.h" -#include "tensorrt_llm/layers/lookaheadDecodingUtils.h" +#include "tensorrt_llm/common/cudaUtils.h" namespace tensorrt_llm::runtime { LookaheadDecodingBuffers::LookaheadDecodingBuffers( SizeType32 maxNumSequences, SizeType32 maxTokensPerStep, BufferManager const& bufferManager) - : generationLengths(bufferManager.gpu(ITensor::makeShape({maxNumSequences}), nvinfer1::DataType::kINT32)) + : generationLengths(bufferManager.gpu(ITensor::makeShape({maxNumSequences}), tensorrt_llm::DataType::kINT32)) , positionOffsets( - bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), nvinfer1::DataType::kINT32)) + bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), tensorrt_llm::DataType::kINT32)) , packedMasks(bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep, static_cast(common::divUp(maxTokensPerStep, 32))}), - nvinfer1::DataType::kINT32)) + tensorrt_llm::DataType::kINT32)) , positionIds( - bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), nvinfer1::DataType::kINT32)) + bufferManager.gpu(ITensor::makeShape({maxNumSequences, maxTokensPerStep}), tensorrt_llm::DataType::kINT32)) { } -LookaheadRuntimeBuffers::LookaheadRuntimeBuffers(SizeType32 maxBatchSize, SizeType32 maxBeamWidth, - BufferManager const& manager, ModelConfig const& modelConfig, WorldConfig const& worldConfig, - executor::DecodingConfig const& /* decodingConfig */, TllmRuntime const& runtime) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - TLLM_CHECK_WITH_INFO(maxBeamWidth == 1, "Lookahead decoding does not support beam search"); - - auto const tokensPerStep = modelConfig.getMaxDecodingTokens(); - auto const numPackedMasks = static_cast(tensorrt_llm::common::divUp(tokensPerStep, 32)); - - cumSumLength = manager.pinned(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - - packedMasksDevice - = manager.gpu(ITensor::makeShape({maxBatchSize * tokensPerStep, numPackedMasks}), nvinfer1::DataType::kINT32); - positionOffsetsDevice = manager.gpu(ITensor::makeShape({maxBatchSize, tokensPerStep}), nvinfer1::DataType::kINT32); - generationLengthsDevice = manager.gpu(ITensor::makeShape({maxBatchSize}), nvinfer1::DataType::kINT32); - positionIdsDevice = manager.gpu(ITensor::makeShape({maxBatchSize, tokensPerStep}), nvinfer1::DataType::kINT32); - - packedMaskHost = manager.cpu(packedMasksDevice->getShape(), nvinfer1::DataType::kINT32); - positionOffsetsHost = manager.cpu(positionOffsetsDevice->getShape(), nvinfer1::DataType::kINT32); - generationLengthsHost = manager.cpu(generationLengthsDevice->getShape(), nvinfer1::DataType::kINT32); - positionIdsHost = manager.cpu(positionIdsDevice->getShape(), nvinfer1::DataType::kINT32); - - packedMaskHostCopy = manager.cpu(packedMasksDevice->getShape(), nvinfer1::DataType::kINT32); - positionOffsetsHostCopy = manager.cpu(positionOffsetsDevice->getShape(), nvinfer1::DataType::kINT32); - generationLengthsHostCopy = manager.cpu(generationLengthsDevice->getShape(), nvinfer1::DataType::kINT32); - positionIdsHostCopy = manager.cpu(positionIdsDevice->getShape(), nvinfer1::DataType::kINT32); - - batchSlotsHostCopy = manager.cpu(generationLengthsDevice->getShape(), nvinfer1::DataType::kINT32); - - useSpecDecoding = manager.cpu(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); - bufferCast(*useSpecDecoding)[0] = 1; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::setFromInputs(SizeType32 numCtxSequences, SizeType32 numGenSequences, - ITensor const& requestTypes, ITensor const& seqSlots, LookaheadDecodingBuffers const& decoderLookaheadBuffers, - TllmRuntime const& runtime, ModelConfig const& modelConfig, WorldConfig const& worldConfig) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const& manager = runtime.getBufferManager(); - - auto const tokensPerStep = modelConfig.getMaxDecodingTokens(); - - manager.copy(seqSlots, *batchSlotsHostCopy); - manager.copy(*decoderLookaheadBuffers.generationLengths, *generationLengthsHostCopy); - manager.copy(*decoderLookaheadBuffers.positionOffsets, *positionOffsetsHostCopy); - manager.copy(*decoderLookaheadBuffers.packedMasks, *packedMaskHostCopy); - manager.copy(*decoderLookaheadBuffers.positionIds, *positionIdsHostCopy); - - manager.getStream().synchronize(); - - BufferRange batchSlotsRange(*batchSlotsHostCopy); - BufferRange cumSumLengthRange(*cumSumLength); - - SizeType32 maxGenerationLength = 0; - for (SizeType32 bi = 0; bi < numGenSequences; bi++) - { - SizeType32 gbi = batchSlotsRange[bi + numCtxSequences]; - SizeType32 theLength = BufferRange(*generationLengthsHostCopy)[gbi]; - maxGenerationLength = std::max(maxGenerationLength, theLength); - } - - auto positionOffsetShape = positionOffsetsHost->getShape(); - positionOffsetShape.d[1] = maxGenerationLength; - positionOffsetsHost->reshape(positionOffsetShape); - positionOffsetsDevice->reshape(positionOffsetShape); - - auto positionIdsShape = positionIdsHostCopy->getShape(); - auto positionIdsShape1D = ITensor::makeShape({ITensor::volume(positionIdsShape)}); - positionIdsHostCopy->reshape(positionIdsShape1D); - positionIdsHost->reshape(positionIdsShape1D); - - cumSumLengthRange[0] = 0; - for (SizeType32 bi = 0; bi < numGenSequences; bi++) - { - SizeType32 gbi = batchSlotsRange[bi + numCtxSequences]; - SizeType32 theLength = BufferRange(*generationLengthsHostCopy)[gbi]; - - manager.copy(*ITensor::at(generationLengthsHostCopy, {gbi}), *ITensor::at(generationLengthsHost, {bi})); - - manager.copy(*ITensor::slice(positionOffsetsHostCopy, {gbi, 0}, theLength), - *ITensor::slice(positionOffsetsHost, {bi, 0}, theLength)); - - manager.copy(*ITensor::slice(packedMaskHostCopy, gbi * tokensPerStep, theLength), - *ITensor::slice(packedMaskHost, cumSumLengthRange[0], theLength)); - - manager.copy(*ITensor::slice(positionIdsHostCopy, gbi * tokensPerStep, theLength), - *ITensor::slice(positionIdsHost, cumSumLengthRange[0], theLength)); - - cumSumLengthRange[0] += theLength; - } - - positionIdsHostCopy->reshape(positionIdsShape); - positionIdsHost->reshape(positionIdsShape); - positionIdsDevice->reshape(positionIdsShape); - - manager.copy(*ITensor::slice(generationLengthsHost, 0, numGenSequences), - *ITensor::slice(generationLengthsDevice, 0, numGenSequences)); - manager.copy(*ITensor::slice(positionOffsetsHost, 0, numGenSequences), - *ITensor::slice(positionOffsetsDevice, 0, numGenSequences)); - manager.copy(*ITensor::slice(packedMaskHost, 0, numGenSequences * tokensPerStep), - *ITensor::slice(packedMasksDevice, 0, numGenSequences * tokensPerStep)); - manager.copy( - *ITensor::slice(positionIdsHost, 0, numGenSequences), *ITensor::slice(positionIdsDevice, 0, numGenSequences)); - positionIdsDevice->reshape(ITensor::makeShape({cumSumLengthRange[0]})); - - manager.getStream().synchronize(); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::reshape(SizeType32 numCtxSequences, SizeType32 numGenSequences, SizeType32 tokensPerStep) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const numSequences = numGenSequences; - - auto packedMaskShape = packedMasksDevice->getShape(); - packedMaskShape.d[0] = numSequences * tokensPerStep; - packedMasksDevice->reshape(packedMaskShape); - packedMaskHost->reshape(packedMaskShape); - - auto generationLengthsShape = generationLengthsDevice->getShape(); - generationLengthsShape.d[0] = numSequences; - generationLengthsDevice->reshape(generationLengthsShape); - generationLengthsHost->reshape(generationLengthsShape); - - auto positionOffsetsShape = positionOffsetsDevice->getShape(); - positionOffsetsShape.d[0] = numSequences; - positionOffsetsDevice->reshape(positionOffsetsShape); - positionOffsetsHost->reshape(positionOffsetsShape); - - auto positionIdsShape = positionIdsDevice->getShape(); - positionIdsShape.d[0] = numSequences; - positionIdsDevice->reshape(positionIdsShape); - positionIdsHost->reshape(positionIdsShape); - - auto batchSlotsShape = batchSlotsHostCopy->getShape(); - batchSlotsShape.d[0] = numCtxSequences + numGenSequences; - batchSlotsHostCopy->reshape(batchSlotsShape); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::enableLookaheadDecoding(SizeType32 maxBatchSize, SizeType32 tokensPerStep) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - auto const numPackedMasks = static_cast(tensorrt_llm::common::divUp(tokensPerStep, 32)); - packedMasksDevice->reshape(ITensor::makeShape({maxBatchSize * tokensPerStep, numPackedMasks})); - generationLengthsDevice->reshape(ITensor::makeShape({maxBatchSize})); - positionOffsetsDevice->reshape(ITensor::makeShape({maxBatchSize, tokensPerStep})); - bufferCast(*useSpecDecoding)[0] = 1; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::disableLookaheadDecoding() -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - packedMasksDevice->reshape(ITensor::makeShape({1, 1})); - generationLengthsDevice->reshape(ITensor::makeShape({1})); - positionOffsetsDevice->reshape(ITensor::makeShape({1, 1})); - bufferCast(*useSpecDecoding)[0] = 0; - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void LookaheadRuntimeBuffers::insertInputTensors( - TensorMap& inputBuffers, TensorMap& /* outputBuffers */, WorldConfig const& /* worldConfig */) const -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - - inputBuffers.insert_or_assign("spec_decoding_packed_mask", packedMasksDevice); - inputBuffers.insert_or_assign("spec_decoding_generation_lengths", generationLengthsDevice); - inputBuffers.insert_or_assign("spec_decoding_position_offsets", positionOffsetsDevice); - inputBuffers.insert_or_assign("spec_decoding_use", useSpecDecoding); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - } // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/loraCache.cpp b/cpp/tensorrt_llm/runtime/loraCache.cpp index 3dbb814f058b..cf1558bb1de2 100644 --- a/cpp/tensorrt_llm/runtime/loraCache.cpp +++ b/cpp/tensorrt_llm/runtime/loraCache.cpp @@ -537,15 +537,15 @@ void LoraCache::splitTransposeCpu(ITensor& output, ITensor const& input, SizeTyp switch (input.getDataType()) { - case nvinfer1::DataType::kINT32: splitTransposeCpuInner(output, input, tpSize, tpRank); break; - case nvinfer1::DataType::kFLOAT: splitTransposeCpuInner(output, input, tpSize, tpRank); break; - case nvinfer1::DataType::kHALF: splitTransposeCpuInner(output, input, tpSize, tpRank); break; - case nvinfer1::DataType::kINT8: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kINT32: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kFLOAT: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kHALF: splitTransposeCpuInner(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kINT8: splitTransposeCpuInner(output, input, tpSize, tpRank); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: splitTransposeCpuInner<__nv_fp8_e4m3>(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kFP8: splitTransposeCpuInner<__nv_fp8_e4m3>(output, input, tpSize, tpRank); break; #endif // ENABLE_FP8 #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: splitTransposeCpuInner<__nv_bfloat16>(output, input, tpSize, tpRank); break; + case tensorrt_llm::DataType::kBF16: splitTransposeCpuInner<__nv_bfloat16>(output, input, tpSize, tpRank); break; #endif // ENABLE_BF16 default: TLLM_CHECK_WITH_INFO(false, "data type not supported"); } diff --git a/cpp/tensorrt_llm/runtime/loraManager.cpp b/cpp/tensorrt_llm/runtime/loraManager.cpp index 8d7ebe389853..05eecb1aab9a 100644 --- a/cpp/tensorrt_llm/runtime/loraManager.cpp +++ b/cpp/tensorrt_llm/runtime/loraManager.cpp @@ -26,7 +26,7 @@ #include "tensorrt_llm/runtime/utils/runtimeUtils.h" #include "tensorrt_llm/runtime/worldConfig.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" namespace tensorrt_llm::runtime { diff --git a/cpp/tensorrt_llm/runtime/loraUtils.cpp b/cpp/tensorrt_llm/runtime/loraUtils.cpp index 3c5e95162474..3104fcb4a783 100644 --- a/cpp/tensorrt_llm/runtime/loraUtils.cpp +++ b/cpp/tensorrt_llm/runtime/loraUtils.cpp @@ -57,7 +57,7 @@ void loraValidateRequestTensorDims(std::optional const& optR keys->getShape().d[0] == expectedBatchSize, "Expected batch dimension to be 1 for each lora request"); TLLM_CHECK_WITH_INFO(weights->getMemoryType() != MemoryType::kGPU, "Expected lora weights to be in CPU memory"); TLLM_CHECK_WITH_INFO(keys->getMemoryType() != MemoryType::kGPU, "Expected lora weights to be in CPU memory"); - TLLM_CHECK_WITH_INFO(keys->getDataType() == nvinfer1::DataType::kINT32, + TLLM_CHECK_WITH_INFO(keys->getDataType() == tensorrt_llm::DataType::kINT32, "Expected lora keys to have TYPE_INT32 but was " + std::string(keys->getDataTypeName())); TLLM_CHECK_WITH_INFO(keys->getShape().d[1] == weights->getShape().d[1], diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp index dc79fd5f48e2..f1c263485826 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.cpp @@ -29,19 +29,19 @@ namespace { #if ENABLE_MULTI_DEVICE -ncclDataType_t toNcclType(nvinfer1::DataType dataType) +ncclDataType_t toNcclType(tensorrt_llm::DataType dataType) { switch (dataType) { - case nvinfer1::DataType::kFLOAT: return ncclFloat32; - case nvinfer1::DataType::kHALF: return ncclHalf; - case nvinfer1::DataType::kINT8: return ncclInt8; - case nvinfer1::DataType::kINT32: return ncclInt32; - case nvinfer1::DataType::kUINT8: return ncclUint8; - case nvinfer1::DataType::kINT64: return ncclInt64; - case nvinfer1::DataType::kFP8: return ncclUint8; + case tensorrt_llm::DataType::kFLOAT: return ncclFloat32; + case tensorrt_llm::DataType::kHALF: return ncclHalf; + case tensorrt_llm::DataType::kINT8: return ncclInt8; + case tensorrt_llm::DataType::kINT32: return ncclInt32; + case tensorrt_llm::DataType::kUINT8: return ncclUint8; + case tensorrt_llm::DataType::kINT64: return ncclInt64; + case tensorrt_llm::DataType::kFP8: return ncclUint8; #if ENABLE_BF16 - case nvinfer1::DataType::kBF16: return ncclBfloat16; + case tensorrt_llm::DataType::kBF16: return ncclBfloat16; #endif // ENABLE_BF16 default: TLLM_THROW("Unsupported data type: %d", static_cast(dataType)); } @@ -50,7 +50,7 @@ ncclDataType_t toNcclType(nvinfer1::DataType dataType) } // namespace void NcclCommunicator::send( - void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const + void const* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE TLLM_NCCL_CHECK(ncclSend(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); @@ -60,7 +60,7 @@ void NcclCommunicator::send( } void NcclCommunicator::receive( - void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const + void* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const { #if ENABLE_MULTI_DEVICE TLLM_NCCL_CHECK(ncclRecv(sendbuff, count, toNcclType(dataType), peer, mComm, stream.get())); diff --git a/cpp/tensorrt_llm/runtime/ncclCommunicator.h b/cpp/tensorrt_llm/runtime/ncclCommunicator.h index 76cce4beab8a..7414e162db0c 100644 --- a/cpp/tensorrt_llm/runtime/ncclCommunicator.h +++ b/cpp/tensorrt_llm/runtime/ncclCommunicator.h @@ -57,9 +57,9 @@ class NcclCommunicator private: void send( - void const* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; + void const* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const; - void receive(void* sendbuff, size_t count, nvinfer1::DataType dataType, int peer, CudaStream const& stream) const; + void receive(void* sendbuff, size_t count, tensorrt_llm::DataType dataType, int peer, CudaStream const& stream) const; static ncclComm_t createComm(int worldSize, int rank, mpi::MpiComm const& mpiComm); diff --git a/cpp/tensorrt_llm/runtime/runtimeKernels.cu b/cpp/tensorrt_llm/runtime/runtimeKernels.cu index 3b3dbcac894a..4404558be0d4 100644 --- a/cpp/tensorrt_llm/runtime/runtimeKernels.cu +++ b/cpp/tensorrt_llm/runtime/runtimeKernels.cu @@ -21,7 +21,7 @@ #include "tensorrt_llm/kernels/speculativeDecoding/kvCacheUpdateKernels.h" #include "tensorrt_llm/runtime/runtimeKernels.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include using namespace tensorrt_llm::runtime; @@ -333,13 +333,13 @@ void invokeFillBatch(IBuffer& buffer, IBuffer const& slotIndices, std::size_t sl { switch (buffer.getDataType()) { - case nvinfer1::DataType::kINT32: + case tensorrt_llm::DataType::kINT32: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; - case nvinfer1::DataType::kINT8: + case tensorrt_llm::DataType::kINT8: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; - case nvinfer1::DataType::kFLOAT: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; + case tensorrt_llm::DataType::kFLOAT: invokeFillBatch(buffer, slotIndices, slotStride, values, stream); break; default: TLLM_THROW("data type not supported"); } } @@ -349,13 +349,13 @@ void invokeGatherBatch(IBuffer& buffer, IBuffer const& values, IBuffer const& sl { switch (buffer.getDataType()) { - case nvinfer1::DataType::kINT32: + case tensorrt_llm::DataType::kINT32: invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); break; - case nvinfer1::DataType::kINT8: + case tensorrt_llm::DataType::kINT8: invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); break; - case nvinfer1::DataType::kFLOAT: invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); break; + case tensorrt_llm::DataType::kFLOAT: invokeGatherBatch(buffer, values, slotIndices, slotStride, stream); break; default: TLLM_THROW("data type not supported"); } } @@ -408,12 +408,12 @@ void scatterTensor(ITensor& output, ITensor const& input, SizeType32 beamWidth, { switch (input.getDataType()) { - case nvinfer1::DataType::kINT32: invokeScatterTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kFLOAT: invokeScatterTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kHALF: invokeScatterTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kINT8: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT32: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFLOAT: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kHALF: invokeScatterTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT8: invokeScatterTensor(output, input, beamWidth, stream); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: invokeScatterTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFP8: invokeScatterTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; #endif // ENABLE_FP8 default: TLLM_THROW("data type not supported"); } @@ -423,15 +423,15 @@ void tileTensor(ITensor& output, ITensor const& input, SizeType32 beamWidth, Cud { switch (input.getDataType()) { - case nvinfer1::DataType::kINT32: invokeTileTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kFLOAT: invokeTileTensor(output, input, beamWidth, stream); break; - case nvinfer1::DataType::kHALF: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT32: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFLOAT: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kHALF: invokeTileTensor(output, input, beamWidth, stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: invokeTileTensor<__nv_bfloat16>(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kBF16: invokeTileTensor<__nv_bfloat16>(output, input, beamWidth, stream); break; #endif // ENABLE_BF16 - case nvinfer1::DataType::kINT8: invokeTileTensor(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kINT8: invokeTileTensor(output, input, beamWidth, stream); break; #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: invokeTileTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; + case tensorrt_llm::DataType::kFP8: invokeTileTensor<__nv_fp8_e4m3>(output, input, beamWidth, stream); break; #endif // ENABLE_FP8 default: TLLM_THROW("data type not supported"); } @@ -444,22 +444,22 @@ void mergeLogitsFragments(BufferManager const& bufferManager, ITensor& output, { switch (output.getDataType()) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: invokeMergeLogitsFragments(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: invokeMergeLogitsFragments(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: invokeMergeLogitsFragments<__nv_bfloat16>(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; #endif // ENABLE_BF16 #ifdef ENABLE_FP8 - case nvinfer1::DataType::kFP8: + case tensorrt_llm::DataType::kFP8: invokeMergeLogitsFragments<__nv_fp8_e4m3>(bufferManager, output, fragmentsVector, cachePointerDevice, cachePointerHost, firstBatchSlotIdx, microBatchSize, beamWidth, stream, stepOffset); break; diff --git a/cpp/tensorrt_llm/runtime/tensorView.h b/cpp/tensorrt_llm/runtime/tensorView.h index 17e7fb719415..4cf51adf3d6d 100644 --- a/cpp/tensorrt_llm/runtime/tensorView.h +++ b/cpp/tensorrt_llm/runtime/tensorView.h @@ -45,19 +45,19 @@ class TensorView : virtual public ITensor, public BufferView mDims.d[0] = size; } - TensorView(IBuffer::SharedPtr const& buffer, size_t offset, size_t size, nvinfer1::Dims const& dims) + TensorView(IBuffer::SharedPtr const& buffer, size_t offset, size_t size, tensorrt_llm::Dims const& dims) : BufferView{buffer, offset, size} , mDims{dims} { Base::resize(ITensor::volumeNonNegative(dims)); } - [[nodiscard]] nvinfer1::Dims const& getShape() const override + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override { return mDims; } - void reshape(nvinfer1::Dims const& dims) override + void reshape(tensorrt_llm::Dims const& dims) override { Base::resize(ITensor::volumeNonNegative(dims)); mDims = dims; @@ -81,6 +81,6 @@ class TensorView : virtual public ITensor, public BufferView return shape.nbDims > 0 && shape.d[0] > 0 ? ITensor::volume(shape) / shape.d[0] : 0; } - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; } // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/tllmBuffers.cpp b/cpp/tensorrt_llm/runtime/tllmBuffers.cpp index ff7ed04001d3..32152da10c29 100644 --- a/cpp/tensorrt_llm/runtime/tllmBuffers.cpp +++ b/cpp/tensorrt_llm/runtime/tllmBuffers.cpp @@ -62,12 +62,12 @@ std::shared_ptr MulticastTensorView::lock() const /////////////////////////////////////// // MulticastTensorView ITensor methods /////////////////////////////////////// -nvinfer1::Dims const& MulticastTensorView::getShape() const +tensorrt_llm::Dims const& MulticastTensorView::getShape() const { return mDims; } -void MulticastTensorView::reshape(nvinfer1::Dims const& dims) +void MulticastTensorView::reshape(tensorrt_llm::Dims const& dims) { auto new_size = nonNegative(volume(dims)); if (new_size > getCapacity()) @@ -102,7 +102,7 @@ std::size_t MulticastTensorView::getCapacity() const return lock()->getCapacity(); } -nvinfer1::DataType MulticastTensorView::getDataType() const +tensorrt_llm::DataType MulticastTensorView::getDataType() const { return lock()->getDataType(); } diff --git a/cpp/tensorrt_llm/runtime/tllmBuffers.h b/cpp/tensorrt_llm/runtime/tllmBuffers.h index faed36537e5c..e1f34a1a11af 100644 --- a/cpp/tensorrt_llm/runtime/tllmBuffers.h +++ b/cpp/tensorrt_llm/runtime/tllmBuffers.h @@ -27,7 +27,7 @@ #include "tensorrt_llm/runtime/memoryCounters.h" #include "tensorrt_llm/runtime/virtualMemory.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -550,7 +550,7 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca //! //! \brief Construct an empty buffer. //! - explicit GenericBuffer(nvinfer1::DataType type, TAllocator allocator = {}) // NOLINT(*-pro-type-member-init) + explicit GenericBuffer(tensorrt_llm::DataType type, TAllocator allocator = {}) // NOLINT(*-pro-type-member-init) : GenericBuffer{0, type, std::move(allocator)} { } @@ -559,7 +559,7 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca //! \brief Construct a buffer with the specified allocation size in number of elements. //! explicit GenericBuffer( // NOLINT(*-pro-type-member-init) - std::size_t size, nvinfer1::DataType type, TAllocator allocator = {}) + std::size_t size, tensorrt_llm::DataType type, TAllocator allocator = {}) : GenericBuffer{size, size, type, std::move(allocator)} { } @@ -636,7 +636,7 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca //! //! \brief Returns the type of the buffer. //! - [[nodiscard]] nvinfer1::DataType getDataType() const override + [[nodiscard]] tensorrt_llm::DataType getDataType() const override { return mType; } @@ -687,7 +687,7 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca } protected: - explicit GenericBuffer(std::size_t size, std::size_t capacity, nvinfer1::DataType type, TAllocator allocator = {}) + explicit GenericBuffer(std::size_t size, std::size_t capacity, tensorrt_llm::DataType type, TAllocator allocator = {}) : TAllocator{std::move(allocator)} , mSize{size} , mCapacity{capacity} @@ -700,14 +700,14 @@ class GenericBuffer : virtual public IBuffer, TAllocator // Inherit from TAlloca private: std::size_t mSize{0}, mCapacity{0}; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; void* mBuffer; }; class MulticastBuffer : virtual public IBuffer { public: - explicit MulticastBuffer(nvinfer1::DataType type, std::set const& ranks) + explicit MulticastBuffer(tensorrt_llm::DataType type, std::set const& ranks) : mSize(0) , mCapacity(0) , mType(type) @@ -716,7 +716,7 @@ class MulticastBuffer : virtual public IBuffer TLLM_CHECK(ranks.size() > 1); } - explicit MulticastBuffer(size_t size, nvinfer1::DataType type, std::set const& ranks) + explicit MulticastBuffer(size_t size, tensorrt_llm::DataType type, std::set const& ranks) : mSize(0) , mCapacity(0) , mType(type) @@ -817,7 +817,7 @@ class MulticastBuffer : virtual public IBuffer return mCapacity; } - [[nodiscard]] nvinfer1::DataType getDataType() const override + [[nodiscard]] tensorrt_llm::DataType getDataType() const override { return mType; } @@ -853,7 +853,7 @@ class MulticastBuffer : virtual public IBuffer private: std::size_t mSize = 0; std::size_t mCapacity = 0; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; std::set mRanks; IpcNvlsHandle* mHandle; }; @@ -882,7 +882,7 @@ class GenericTensor : virtual public ITensor, public GenericBuffer //! //! \brief Construct an empty tensor. //! - explicit GenericTensor(nvinfer1::DataType type, TAllocator allocator = {}) + explicit GenericTensor(tensorrt_llm::DataType type, TAllocator allocator = {}) : Base{type, std::move(allocator)} { mDims.nbDims = 0; @@ -891,14 +891,14 @@ class GenericTensor : virtual public ITensor, public GenericBuffer //! //! \brief Construct a tensor with the specified allocation dimensions. //! - explicit GenericTensor(nvinfer1::Dims const& dims, nvinfer1::DataType type, TAllocator allocator = {}) + explicit GenericTensor(tensorrt_llm::Dims const& dims, tensorrt_llm::DataType type, TAllocator allocator = {}) : Base{nonNegative(volume(dims)), type, std::move(allocator)} , mDims{dims} { } explicit GenericTensor( - nvinfer1::Dims const& dims, std::size_t capacity, nvinfer1::DataType type, TAllocator allocator = {}) + tensorrt_llm::Dims const& dims, std::size_t capacity, tensorrt_llm::DataType type, TAllocator allocator = {}) : Base{nonNegative(volume(dims)), capacity, type, std::move(allocator)} , mDims{dims} { @@ -923,12 +923,12 @@ class GenericTensor : virtual public ITensor, public GenericBuffer return *this; } - [[nodiscard]] nvinfer1::Dims const& getShape() const override + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override { return mDims; } - void reshape(nvinfer1::Dims const& dims) override + void reshape(tensorrt_llm::Dims const& dims) override { Base::resize(nonNegative(volume(dims))); mDims = dims; @@ -946,7 +946,7 @@ class GenericTensor : virtual public ITensor, public GenericBuffer } private: - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; // Forward declaration @@ -971,9 +971,9 @@ class MulticastTensorView : virtual public ITensor ///////////////////// // ITensor methods ///////////////////// - [[nodiscard]] nvinfer1::Dims const& getShape() const override; + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override; - void reshape(nvinfer1::Dims const& dims) override; + void reshape(tensorrt_llm::Dims const& dims) override; ///////////////////// // IBuffer methods @@ -983,7 +983,7 @@ class MulticastTensorView : virtual public ITensor [[nodiscard]] std::size_t getCapacity() const override; - [[nodiscard]] nvinfer1::DataType getDataType() const override; + [[nodiscard]] tensorrt_llm::DataType getDataType() const override; [[nodiscard]] MemoryType getMemoryType() const override; @@ -1016,7 +1016,7 @@ class MulticastTensorView : virtual public ITensor std::weak_ptr mTensor; ViewType mViewType; - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; class MulticastTensor : virtual public ITensor, @@ -1026,13 +1026,13 @@ class MulticastTensor : virtual public ITensor, public: using Base = MulticastBuffer; - explicit MulticastTensor(nvinfer1::DataType type, std::set const& ranks) + explicit MulticastTensor(tensorrt_llm::DataType type, std::set const& ranks) : Base(type, ranks) { mDims.nbDims = 0; } - explicit MulticastTensor(nvinfer1::Dims const& dims, nvinfer1::DataType type, std::set const& ranks) + explicit MulticastTensor(tensorrt_llm::Dims const& dims, tensorrt_llm::DataType type, std::set const& ranks) : Base(nonNegative(volume(dims)), type, ranks) , mDims(dims) { @@ -1068,12 +1068,12 @@ class MulticastTensor : virtual public ITensor, ///////////////////// // ITensor methods ///////////////////// - [[nodiscard]] nvinfer1::Dims const& getShape() const override + [[nodiscard]] tensorrt_llm::Dims const& getShape() const override { return mDims; } - void reshape(nvinfer1::Dims const& dims) override + void reshape(tensorrt_llm::Dims const& dims) override { Base::resize(nonNegative(volume(dims))); mDims = dims; @@ -1091,7 +1091,7 @@ class MulticastTensor : virtual public ITensor, } private: - nvinfer1::Dims mDims{}; + tensorrt_llm::Dims mDims{}; }; using DeviceTensor = GenericTensor; diff --git a/cpp/tensorrt_llm/runtime/tllmLogger.cpp b/cpp/tensorrt_llm/runtime/tllmLogger.cpp deleted file mode 100644 index 586ab2f4ae95..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmLogger.cpp +++ /dev/null @@ -1,62 +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 "tllmLogger.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" - -using namespace tensorrt_llm::runtime; -namespace tc = tensorrt_llm::common; - -void TllmLogger::log(nvinfer1::ILogger::Severity severity, nvinfer1::AsciiChar const* msg) noexcept -{ - switch (severity) - { - case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: - case nvinfer1::ILogger::Severity::kERROR: TLLM_LOG_ERROR(msg); break; - case nvinfer1::ILogger::Severity::kWARNING: TLLM_LOG_WARNING(msg); break; - case nvinfer1::ILogger::Severity::kINFO: TLLM_LOG_INFO(msg); break; - case nvinfer1::ILogger::Severity::kVERBOSE: TLLM_LOG_DEBUG(msg); break; - default: TLLM_LOG_TRACE(msg); break; - } -} - -nvinfer1::ILogger::Severity TllmLogger::getLevel() -{ - auto* const logger = tc::Logger::getLogger(); - switch (logger->getLevel()) - { - case tc::Logger::Level::ERROR: return nvinfer1::ILogger::Severity::kERROR; - case tc::Logger::Level::WARNING: return nvinfer1::ILogger::Severity::kWARNING; - case tc::Logger::Level::INFO: return nvinfer1::ILogger::Severity::kINFO; - case tc::Logger::Level::DEBUG: - case tc::Logger::Level::TRACE: return nvinfer1::ILogger::Severity::kVERBOSE; - default: return nvinfer1::ILogger::Severity::kINTERNAL_ERROR; - } -} - -void TllmLogger::setLevel(nvinfer1::ILogger::Severity level) -{ - auto* const logger = tc::Logger::getLogger(); - switch (level) - { - case nvinfer1::ILogger::Severity::kINTERNAL_ERROR: - case nvinfer1::ILogger::Severity::kERROR: logger->setLevel(tc::Logger::Level::ERROR); break; - case nvinfer1::ILogger::Severity::kWARNING: logger->setLevel(tc::Logger::Level::WARNING); break; - case nvinfer1::ILogger::Severity::kINFO: logger->setLevel(tc::Logger::Level::INFO); break; - case nvinfer1::ILogger::Severity::kVERBOSE: logger->setLevel(tc::Logger::Level::TRACE); break; - default: TLLM_THROW("Unsupported severity"); - } -} diff --git a/cpp/tensorrt_llm/runtime/tllmRuntime.cpp b/cpp/tensorrt_llm/runtime/tllmRuntime.cpp deleted file mode 100644 index 7c2ca4747213..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmRuntime.cpp +++ /dev/null @@ -1,831 +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 "tllmRuntime.h" -#include "common.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" -#include "tensorrt_llm/common/nvtxUtils.h" -#include "tensorrt_llm/common/safetensors.h" -#include "tensorrt_llm/executor/tensor.h" -#include "tensorrt_llm/kernels/userbuffers/ub_interface.h" -#include "tensorrt_llm/runtime/utils/mpiUtils.h" -#include "tllmLogger.h" -#include "tllmStreamReaders.h" - -#include "nlohmann/json.hpp" -#include - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -using namespace tensorrt_llm::runtime; -using TensorMap = StringPtrMap; - -namespace -{ -static_assert(std::is_signed::value, "SizeType32 must be signed"); - -nvinfer1::Dims shapeToDims(std::vector const& shape) -{ - TLLM_CHECK(shape.size() <= nvinfer1::Dims::MAX_DIMS); - nvinfer1::Dims dims; - auto constexpr dim_max = std::numeric_limits::max(); - dims.nbDims = static_cast(shape.size()); - for (std::size_t i = 0; i < shape.size(); ++i) - { - // shape[i] >= 0 because it has unsigned type. Check upper bound: - TLLM_CHECK(shape[i] <= static_cast(dim_max)); - dims.d[i] = static_cast(shape[i]); - } - return dims; -} - -std::vector dimsToShape(nvinfer1::Dims const& dims) -{ - TLLM_CHECK(dims.nbDims >= 0); - std::vector shape(dims.nbDims); - for (std::int32_t i = 0; i < dims.nbDims; ++i) - { - TLLM_CHECK(dims.d[i] >= 0); - shape[i] = static_cast(dims.d[i]); - } - return shape; -} - -tensorrt_llm::runtime::TllmLogger defaultLogger{}; - -void setWeightStreaming(nvinfer1::ICudaEngine& engine, float const gpuWeightsPercent) -{ - if (gpuWeightsPercent < 1) - { - int64_t streamableSize = engine.getStreamableWeightsSize(); - int64_t budget = gpuWeightsPercent * streamableSize; - TLLM_LOG_INFO("Set gpu weights percent to %f, which is %lld bytes. Valid range: %lld bytes - %lld bytes.", - gpuWeightsPercent, budget, 0, streamableSize); - engine.setWeightStreamingBudgetV2(budget); - } -} - -class LayerInfo -{ -public: - LayerInfo(std::optional name, std::string type) - : name(std::move(name)) - , type(std::move(type)){}; - std::optional name; - std::string type; -}; - -void assessLikelihoodOfRuntimeAllocation( - nvinfer1::ICudaEngine const& engine, nvinfer1::IEngineInspector const& engineInspector) - -{ - TLLM_LOG_INFO("Inspecting the engine to identify potential runtime issues..."); - auto const profilingVerbosity = engine.getProfilingVerbosity(); - if (profilingVerbosity != nvinfer1::ProfilingVerbosity::kDETAILED) - { - TLLM_LOG_INFO( - "The profiling verbosity of the engine does not allow this analysis to proceed. Re-build the engine with " - "'detailed' profiling verbosity to get more diagnostics."); - return; - } - auto const* const layerTypeKey = "LayerType"; - auto const* const nameKey = "Name"; - auto const numLayers = engine.getNbLayers(); - TLLM_LOG_INFO("Model has %i layers.", numLayers); - std::vector indexes(numLayers); - std::iota(indexes.begin(), indexes.end(), 0); - std::vector> layerInfos(numLayers); - std::transform(indexes.cbegin(), indexes.cend(), layerInfos.begin(), - [&](SizeType32 const idx) - { - auto const* const layerInfo - = engineInspector.getLayerInformation(idx, nvinfer1::LayerInformationFormat::kJSON); - - // Needs to be copied explicitly, see documentation of `getLayerInformation`. - auto const layerInfoCopy = std::string(layerInfo); - auto const jsonLayerInfo = nlohmann::json::parse(layerInfoCopy); - auto const layerJsonType = jsonLayerInfo.type(); - if (layerJsonType != nlohmann::detail::value_t::object) - { - return std::optional{}; - } - if (!jsonLayerInfo.contains(layerTypeKey)) - { - return std::optional{}; - } - auto const& typeJson = jsonLayerInfo.at(layerTypeKey); - if (typeJson.type() != nlohmann::detail::value_t::string) - { - return std::optional{}; - } - std::optional name{}; - if (jsonLayerInfo.contains(nameKey)) - { - auto const& nameJson = jsonLayerInfo.at(nameKey); - auto const nameJsonType = nameJson.type(); - if (nameJsonType == nlohmann::detail::value_t::string) - { - name = nameJson.get(); - } - } - return std::make_optional(LayerInfo{name, typeJson.get()}); - }); - auto const layersWithInfoEnd = std::partition( - layerInfos.begin(), layerInfos.end(), [](std::optional const& info) { return info.has_value(); }); - if (layersWithInfoEnd == layerInfos.begin()) - { - TLLM_LOG_INFO("Engine layer infos could not be parsed into useful information."); - return; - } - auto const allocateLayersEnd = std::partition(layerInfos.begin(), layersWithInfoEnd, - [](std::optional const& info) { return info.value().type == "allocate"; }); - auto numWarnings = 0; - for (auto layerInfo = layerInfos.begin(); layerInfo != allocateLayersEnd; layerInfo++) - { - auto constexpr maxNumWarnings = 25; - if (numWarnings < maxNumWarnings) - { - auto const layerName = layerInfo->value().name.value_or(""); - TLLM_LOG_WARNING( - "Layer '%s' has type '%s', which could lead to large runtime memory allocations. Performance " - "might be degraded and / or you might run out of memory.", - layerName.c_str(), layerInfo->value().type.c_str()); - } - numWarnings++; - } - if (numWarnings > 0) - { - TLLM_LOG_WARNING( - "There were a total of %i layers with type 'allocate'. Some warnings might have been silenced to keep the " - "output concise.", - numWarnings); - } -} - -} // namespace - -TllmRuntime::TllmRuntime(RawEngine const& rawEngine, nvinfer1::ILogger* logger, bool useGpuDirectStorage, - float gpuWeightsPercent, bool useShapeInference) - : mStream(std::make_shared()) - , mBufferManager{mStream, true} // Ensure to trim the memory pool on destruction. - , mRuntime{nvinfer1::createInferRuntime(static_cast(logger) ? *logger : defaultLogger)} - , mUseShapeInference{useShapeInference} - , mUserBufferEnabled{false} -{ - auto const startTime = std::chrono::high_resolution_clock::now(); - - switch (rawEngine.getType()) - { - case RawEngine::Type::FilePath: - { - if (useGpuDirectStorage) - { - TLLM_LOG_INFO("GDS is used to load the engine!"); - auto reader = GDSStreamReader(rawEngine.getPath()); - mEngine.reset(mRuntime->deserializeCudaEngine(reader)); - } - else - { - auto reader = StreamReader(rawEngine.getPath()); - mEngine.reset(mRuntime->deserializeCudaEngine(reader)); - } - break; - } - case RawEngine::Type::AddressWithSize: - mEngine.reset(mRuntime->deserializeCudaEngine(rawEngine.getAddress(), rawEngine.getSize())); - break; - case RawEngine::Type::HostMemory: - mEngine.reset( - mRuntime->deserializeCudaEngine(rawEngine.getHostMemory()->data(), rawEngine.getHostMemory()->size())); - break; - default: TLLM_THROW("Unsupported raw engine type."); - } - - auto const elapsedMs - = std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - startTime); - - TLLM_LOG_INFO("Engine load time %lld ms", elapsedMs); - - TLLM_CHECK_WITH_INFO(mEngine != nullptr, "Failed to deserialize cuda engine."); - mEngineInspector.reset(mEngine->createEngineInspector()); - assessLikelihoodOfRuntimeAllocation(*mEngine, *mEngineInspector); - setWeightStreaming(getEngine(), gpuWeightsPercent); - auto const devMemorySize = mEngine->getDeviceMemorySizeV2(); - mEngineBuffer = mBufferManager.gpu(devMemorySize); - // Print context memory size for CI/CD to track. - TLLM_LOG_INFO("[MemUsageChange] Allocated %.2f MiB for execution context memory.", - static_cast(devMemorySize) / 1048576.0); - - cacheTensorNames(); -} - -void TllmRuntime::cacheTensorNames() -{ - for (std::int32_t i = 0; i < mEngine->getNbIOTensors(); ++i) - { - auto const* const name = mEngine->getIOTensorName(i); - if (mEngine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kINPUT) - { - mInputTensorNames.emplace_back(name); - } - else if (mEngine->getTensorIOMode(name) == nvinfer1::TensorIOMode::kOUTPUT) - { - mOutputTensorNames.emplace_back(name); - } - } -} - -nvinfer1::IExecutionContext& TllmRuntime::addContext(std::int32_t profileIndex) -{ - TLLM_CHECK(0 <= profileIndex && profileIndex < mEngine->getNbOptimizationProfiles()); - mContexts.emplace_back(mEngine->createExecutionContextWithoutDeviceMemory()); - if (!mContexts.back()) - { - if (mEngine->getStreamableWeightsSize() > 0) - { - TLLM_THROW("Failed to allocate memory for weights. Please try reducing --gpu_weights_percent."); - } - else - { - TLLM_THROW("Internal Error: Failed to create an execution context."); - } - } - auto& context = *mContexts.back(); - context.setDeviceMemoryV2(mEngineBuffer->data(), static_cast(mEngineBuffer->getCapacity())); - - if (tensorrt_llm::common::Logger::getLogger()->isEnabled(tensorrt_llm::common::Logger::TRACE) - && mContexts.size() == 1) - { - // Print engine information only once - printEngineInfo(); - } - - context.setOptimizationProfileAsync(profileIndex, mStream->get()); - // If nvtx verbosity is DETAILED, print an info about potential perf overhead. - if (context.getNvtxVerbosity() == nvinfer1::ProfilingVerbosity::kDETAILED) - { - TLLM_LOG_INFO( - "The engine was built with kDETAILED profiling verbosity, which may result in small overheads at runtime."); - } - return context; -} - -void TllmRuntime::printEngineInfo() -{ - auto& context = *(mContexts[0]); - int const nIO = mEngine->getNbIOTensors(); // Count of input / output tensor - int const nOP = mEngine->getNbOptimizationProfiles(); // Count of Optimization Profile - std::size_t maxNameWidth = 0; - std::size_t maxShapeWidth = 0; - - // Get information of engine input / output - std::vector tensorNameList{}; - tensorNameList.reserve(nIO); - for (int i = 0; i < nIO; ++i) - { - tensorNameList.emplace_back(mEngine->getIOTensorName(i)); - } - std::vector> tensorInfo(nIO); // Tensor Information Vector - std::vector>> profileInfo(nIO); // Tensor Optimization Profile Vector - for (int i = 0; i < nIO; ++i) - { - auto const& name = tensorNameList[i]; - char const* nameC{name.c_str()}; // name of C-style - maxNameWidth = std::max(maxNameWidth, name.size()); - tensorInfo[i]["mode"] = mEngine->getTensorIOMode(nameC) == nvinfer1::TensorIOMode::kINPUT ? "I" : "O"; - tensorInfo[i]["location"] - = mEngine->getTensorLocation(nameC) == nvinfer1::TensorLocation::kDEVICE ? "GPU" : "CPU"; - tensorInfo[i]["data_type"] = dataTypeToString(mEngine->getTensorDataType(nameC)); - tensorInfo[i]["build_shape"] = shapeToString(mEngine->getTensorShape(nameC)); - maxShapeWidth = std::max(maxShapeWidth, tensorInfo[i]["build_shape"].size()); - if (tensorInfo[i]["mode"] == "I") - { - std::vector> topPerTensor(nOP); - for (int k = 0; k < nOP; ++k) - { - if (tensorInfo[i]["location"] == std::string("GPU")) - { - std::vector top(3); - top[0] = mEngine->getProfileShape(nameC, k, nvinfer1::OptProfileSelector::kMIN); - top[1] = mEngine->getProfileShape(nameC, k, nvinfer1::OptProfileSelector::kOPT); - top[2] = mEngine->getProfileShape(nameC, k, nvinfer1::OptProfileSelector::kMAX); - topPerTensor[k] = top; - maxShapeWidth = std::max(maxShapeWidth, shapeToString(top[2]).size()); - } - else - { - // Shape input tensor, not used in TRT-LLM support yet - std::vector top(3); - int const nDim = mEngine->getTensorShape(nameC).nbDims; - nvinfer1::Dims64 tensorShape{nDim, {-1}}; - int const* pos = nullptr; - pos = mEngine->getProfileTensorValues(nameC, k, nvinfer1::OptProfileSelector::kMIN); - std::copy(pos, pos + nDim, tensorShape.d); - top[0] = tensorShape; - pos = mEngine->getProfileTensorValues(nameC, k, nvinfer1::OptProfileSelector::kOPT); - std::copy(pos, pos + nDim, tensorShape.d); - top[1] = tensorShape; - pos = mEngine->getProfileTensorValues(nameC, k, nvinfer1::OptProfileSelector::kMAX); - std::copy(pos, pos + nDim, tensorShape.d); - top[2] = tensorShape; - topPerTensor[k] = top; - } - } - profileInfo[i] = topPerTensor; - } - else - { - profileInfo[i] = std::vector>(nOP); - } - } - // Set input shape to get output shape - for (int k = 0; k < nOP; ++k) - { - for (int j = 0; j < 3; ++j) // Min, Opt, Max - { - for (int i = 0; i < nIO; ++i) - { - auto const& name = tensorNameList[i]; - char const* nameC = name.c_str(); - if (tensorInfo[i]["mode"] == "I") - { - if (tensorInfo[i]["location"] == std::string("GPU")) - { - context.setInputShape(nameC, profileInfo[i][k][j]); - } - else - { - // Shape input tensor, not used in TRT-LLM support yet - context.setInputTensorAddress(nameC, profileInfo[i][k][j].d); - } - } - else - { - TLLM_CHECK_WITH_INFO(context.allInputDimensionsSpecified(), "Input dimensions not specified"); - TLLM_CHECK_WITH_INFO(context.allInputShapesSpecified(), "Input shapes not specified"); - if (tensorInfo[i]["location"] == std::string("GPU")) - { - profileInfo[i][k].push_back(context.getTensorShape(nameC)); - } - else - { - // Shape input tensor, not used in TRT-LLM support yet - int const nDim = mEngine->getTensorShape(nameC).nbDims; - nvinfer1::Dims64 tensorShape{nDim, {}}; - int const* pos = reinterpret_cast(context.getTensorAddress(nameC)); - std::copy(pos, pos + nDim, tensorShape.d); - profileInfo[i][k].push_back(tensorShape); - } - } - } - } - } - - // Print information of engine input / output - std::string info; - TLLM_LOG_TRACE("Information of engine input / output."); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 24, '=')); - info = alignText("Name", maxNameWidth) + "|I/O|Location|DataType|" + alignText("Shape", maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 24, '-')); - for (int i = 0; i < nIO; ++i) - { - info = alignText(tensorNameList[i], maxNameWidth, false) + "|"; - info += alignText(tensorInfo[i]["mode"], 3) + "|"; - info += alignText(tensorInfo[i]["location"], 8) + "|"; - info += alignText(tensorInfo[i]["data_type"], 8) + "|"; - info += alignText(tensorInfo[i]["build_shape"], maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - } - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 24, '=')); - // Print information of optimization profile - TLLM_LOG_TRACE("Information of optimization profile."); - for (int k = 0; k < nOP; ++k) - { - TLLM_LOG_TRACE("Optimization Profile %d:", k); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth * 3 + 4, '=')); - info = alignText("Name", maxNameWidth) + "|"; - info += alignText("Min", maxShapeWidth) + "|"; - info += alignText("Opt", maxShapeWidth) + "|"; - info += alignText("Max", maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth * 3 + 4, '-')); - for (int i = 0; i < nIO; ++i) - { - auto const& top = profileInfo[i][k]; - info = alignText(tensorNameList[i], maxNameWidth, false) + "|"; - info += alignText(shapeToString(top[0]), maxShapeWidth) + "|"; - info += alignText(shapeToString(top[1]), maxShapeWidth) + "|"; - info += alignText(shapeToString(top[2]), maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - } - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth * 3 + 4, '=')); - } -} - -void TllmRuntime::printContextInfo(SizeType32 contextIndex) -{ - auto const& context = *(mContexts[contextIndex]); - int const nIO = mEngine->getNbIOTensors(); // Count of input / output tensor - std::size_t maxNameWidth = 0; - std::size_t maxShapeWidth = 0; - std::vector> tensorInfo(nIO); - for (int i = 0; i < nIO; ++i) - { - auto const name = std::string(mEngine->getIOTensorName(i)); - bool const isInput = mEngine->getTensorIOMode(name.c_str()) == nvinfer1::TensorIOMode::kINPUT; - auto const shape = shapeToString(context.getTensorShape(name.c_str())); - tensorInfo[i] = std::make_tuple(name, isInput, shape); - maxNameWidth = std::max(maxNameWidth, name.size()); - maxShapeWidth = std::max(maxShapeWidth, shape.size()); - // Shape input tensor is not considered in TRT-LLM yet - } - - TLLM_LOG_TRACE("Information of context input / output."); - TLLM_LOG_TRACE("Using Optimization Profile: %d", contextIndex); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 6, '=')); - std::string info = alignText("Name", maxNameWidth) + "|I/O|" + alignText("Shape", maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 6, '-')); - for (int i = 0; i < nIO; ++i) - { - auto const& [name, isInput, shape] = tensorInfo[i]; - info = alignText(name, maxNameWidth, false) + "|"; - info += alignText(isInput ? "I" : "O", 3) + "|"; - info += alignText(shape, maxShapeWidth) + "|"; - TLLM_LOG_TRACE(info.c_str()); - } - TLLM_LOG_TRACE(std::string(maxNameWidth + maxShapeWidth + 6, '=')); -} - -void TllmRuntime::clearContexts() -{ - for (auto& context : mContexts) - { - context.reset(); - } - mContexts.clear(); -} - -bool TllmRuntime::executeContext(SizeType32 contextIndex) const -{ - NVTX3_FUNC_RANGE(); - auto& context = getContext(contextIndex); - auto res = context.enqueueV3(mStream->get()); - sync_check_cuda_error(mStream->get()); - return res; -} - -void TllmRuntime::setInputTensorsImpl(SizeType32 contextIndex, TensorMap const& tensorMap, bool throwOnMiss) -{ - NVTX3_FUNC_RANGE(); - auto& context = getContext(contextIndex); - for (auto const& name : mInputTensorNames) - { - auto const pos = tensorMap.find(name); - if (pos == tensorMap.end()) - { - if (throwOnMiss) - { - auto expectedShape = mEngine->getTensorShape(name.c_str()); - TLLM_THROW("Input tensor '%s' not found; expected shape: %s", name.c_str(), - ITensor::toString(expectedShape).c_str()); - } - else - { - continue; - } - } - - auto const& tensor = pos->second; - auto const tensorDtype = tensor->getDataType(); - auto const engineDtype = mEngine->getTensorDataType(name.c_str()); - // WAR: TRT does not support mixed FP8 and FP16 input, so engine expects FP16 tensors. - TLLM_CHECK_WITH_INFO(tensorDtype == engineDtype - || (tensorDtype == nvinfer1::DataType::kFP8 && engineDtype == nvinfer1::DataType::kHALF), - "%s: expected type %d, provided type %d", name.c_str(), static_cast(engineDtype), - static_cast(tensorDtype)); - - auto tensorShape = tensor->getShape(); - - // Change shape of `cache_indirection` for Variable-Beam-Width-Search - // TODO: remove this hack if beamWidth of each request are passed into GptAttentionPlugin by input tensor - if (name == "cache_indirection" && mCurrentBeamWidths.size() > 0) - { - SizeType32 const beamWidth = getCurrentBeamWidth(); - if (tensorShape.d[1] != beamWidth) - { - tensorShape.d[1] = beamWidth; - TLLM_LOG_TRACE("Change shape of cache_indirection to %s", ITensor::toString(tensorShape).c_str()); - } - } - - auto const setInputShapeSuccess = context.setInputShape(name.c_str(), tensorShape); - if (!setInputShapeSuccess) - { - auto const minShape - = mEngine->getProfileShape(name.c_str(), contextIndex, nvinfer1::OptProfileSelector::kMIN); - auto const maxShape - = mEngine->getProfileShape(name.c_str(), contextIndex, nvinfer1::OptProfileSelector::kMAX); - - TLLM_THROW("Tensor '%s' has invalid shape %s, expected in range min %s, max %s", name.c_str(), - ITensor::toString(tensorShape).c_str(), ITensor::toString(minShape).c_str(), - ITensor::toString(maxShape).c_str()); - } - auto* const data = tensor->data(); - if (static_cast(data)) - { - context.setInputTensorAddress(name.c_str(), data); - } - else - { - TLLM_CHECK_WITH_INFO(tensor->getSize() == 0, std::string("Invalid data for tensor: ") + name); - // TensorRT runtime does not support nullptr. - if (!mDummyTensor) - { - mDummyTensor = mBufferManager.gpu(ITensor::makeShape({1})); - } - context.setInputTensorAddress(name.c_str(), mDummyTensor->data()); - } - } -} - -void TllmRuntime::setStaticInputTensors(TensorMap const& tensorMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - - TLLM_CHECK_WITH_INFO(getNbContexts() > 0, "Contexts should be created before calling setStaticInputTensors"); - for (auto contextIndex = 0; contextIndex < getNbContexts(); ++contextIndex) - { - setInputTensorsImpl(contextIndex, tensorMap, false); - } - - // move static input tensor names to separate vector - auto const begin = mInputTensorNames.begin(); - auto end = mInputTensorNames.end(); - for (auto const& [name, tensor] : tensorMap) - { - end = std::remove(begin, end, name); - } - mInputTensorNames.erase(end, mInputTensorNames.end()); - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TllmRuntime::setInputTensors(SizeType32 contextIndex, TensorMap const& tensorMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - setInputTensorsImpl(contextIndex, tensorMap, true); - - auto& context = getContext(contextIndex); - if (mUseShapeInference) - { - NVTX3_SCOPED_RANGE(infer_shapes); - char const* missing = nullptr; - auto const nbMissing = context.inferShapes(1, &missing); - if (nbMissing > 0) - { - TLLM_THROW("Input shape not specified: %s", missing); - } - else if (nbMissing < 0) - { - TLLM_THROW("Invalid input shape"); - } - } - - { - NVTX3_SCOPED_RANGE(final_checks); - TLLM_CHECK_WITH_INFO(context.allInputDimensionsSpecified(), "Input dimensions not specified"); - TLLM_CHECK_WITH_INFO(context.allInputShapesSpecified(), "Input shapes not specified"); - } - - // Print shape of input / output tensors for the TRT engine - if (tensorrt_llm::common::Logger::getLogger()->isEnabled(tensorrt_llm::common::Logger::TRACE)) - { - printContextInfo(contextIndex); - } - - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TllmRuntime::setOutputTensors(SizeType32 contextIndex, TensorMap& tensorMap) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - if (isUserBufferEnabled()) - { - // This function will identify the output tensors in the network that need to be bound as UB buffers - // and bind the corresponding buffers to them based on their names. - setUserBufferTensors(contextIndex, tensorMap); - } - - auto& context = getContext(contextIndex); - for (auto const& name : mOutputTensorNames) - { - auto const engineDtype = mEngine->getTensorDataType(name.c_str()); - auto const pos = tensorMap.find(name); - if (pos != tensorMap.end()) - { - auto const& tensor = pos->second; - auto const tensorDtype = tensor->getDataType(); - // WAR: TRT does not support mixed FP8 and FP16 input, so engine expects FP16 tensors. - TLLM_CHECK_WITH_INFO(tensorDtype == engineDtype - || (tensorDtype == nvinfer1::DataType::kFP8 && engineDtype == nvinfer1::DataType::kHALF), - "%s: expected type %d, provided type %d", name.c_str(), static_cast(engineDtype), - static_cast(tensorDtype)); - - if (mUseShapeInference) - { - auto const dims = context.getTensorShape(name.c_str()); - tensor->reshape(dims); - } - context.setTensorAddress(name.c_str(), tensor->data()); - } - else if (mUseShapeInference) - { - auto const dims = context.getTensorShape(name.c_str()); - auto tensor = ITensor::SharedPtr(mBufferManager.gpu(dims, engineDtype)); - tensorMap.insert(pos, std::make_pair(name, tensor)); - context.setTensorAddress(name.c_str(), tensor->data()); - } - else - { - TLLM_THROW("Tensor %s is not found in tensorMap and shape inference is not allowed", name.c_str()); - } - } - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} - -void TllmRuntime::setUserBufferTensors(SizeType32 contextIndex, TensorMap& tensorMap) -{ - auto startsWith = [](std::string const& str, std::string const& prefix) -> bool - { return str.size() > prefix.size() && str.compare(0, prefix.size(), prefix) == 0; }; - std::string const prefix(tensorrt_llm::runtime::ub::tensor_prefix); - auto& context = getContext(contextIndex); - for (auto const& name : mOutputTensorNames) - { - auto const pos = tensorMap.find(name); - if (pos != tensorMap.end() || !startsWith(name, prefix)) - { - continue; - } - auto const engineDtype = mEngine->getTensorDataType(name.c_str()); - auto const dims = context.getTensorShape(name.c_str()); - void* ubBuffer = nullptr; - if (name[prefix.size()] == '0') - { - ubBuffer = tensorrt_llm::runtime::ub::ub_get(0).addr; - } - else if (name[prefix.size()] == '1') - { - ubBuffer = tensorrt_llm::runtime::ub::ub_get(1).addr; - } - else if (name[prefix.size()] == '2') - { - ubBuffer = tensorrt_llm::runtime::ub::ub_get(2).addr; - } - else - { - TLLM_CHECK(false); - } - auto tensor = ITensor::SharedPtr(ITensor::wrap(ubBuffer, engineDtype, dims)); - tensorMap.insert(pos, std::make_pair(name, tensor)); - context.setTensorAddress(name.c_str(), ubBuffer); - } -} - -void TllmRuntime::initializeUserBuffer(tensorrt_llm::runtime::WorldConfig const& world_config, SizeType32 maxBatchSize, - SizeType32 maxBeamWidth, SizeType32 maxSequenceLength, SizeType32 hiddenSize, - std::optional maxNumTokens) -{ - auto startsWith = [](std::string const& str, std::string const& prefix) -> bool - { return str.size() > prefix.size() && str.compare(0, prefix.size(), prefix) == 0; }; - std::string const prefix(tensorrt_llm::runtime::ub::tensor_prefix); - bool useNVFP4Model = false; - for (auto const& name : mOutputTensorNames) - { - if (startsWith(name, prefix)) - { - mUserBufferEnabled = true; - if (name[prefix.size()] == '2') - { - useNVFP4Model = true; - break; - } - } - } - if (!mUserBufferEnabled) - { - return; - } - // The hidden size returned by ModelConfig is the real hidden size divided by the TP size. - auto const tpSize = world_config.getTensorParallelism(); - size_t const realHiddenSize = hiddenSize * tpSize; - size_t const tokensNum = maxNumTokens.value_or(maxBatchSize * maxBeamWidth * maxSequenceLength); - TLLM_CHECK(tokensNum > 0); - size_t const elemNum = tokensNum * realHiddenSize; - TLLM_LOG_INFO("[UserBuffer] MaxBatchSize %d, maxBeamWidth %d, maxSequenceLength %d, maxNumTokens %d, select %lu", - maxBatchSize, maxBeamWidth, maxSequenceLength, maxNumTokens.has_value() ? maxNumTokens.value() : 0, tokensNum); - tensorrt_llm::runtime::ub::ub_initialize(world_config); - tensorrt_llm::runtime::ub::ub_allocate(elemNum * sizeof(half)); - tensorrt_llm::runtime::ub::ub_allocate(elemNum * sizeof(half)); - if (useNVFP4Model) - { - tensorrt_llm::runtime::ub::ub_allocate(elemNum * sizeof(uint8_t) / 16); - } -} - -CudaStream const& TllmRuntime::getStream() const -{ - return *mStream; -} - -bool TllmRuntime::hasLayerProfiler(SizeType32 contextId) const -{ - return mContexts[contextId]->getProfiler() != nullptr; -} - -void TllmRuntime::setLayerProfiler() -{ - mLayerProfiler = std::make_unique(); - for (auto& context : mContexts) - { - context->setProfiler(mLayerProfiler.get()); - context->setEnqueueEmitsProfile(false); - } -} - -std::string TllmRuntime::getLayerProfileInfo() const -{ - TLLM_CHECK(mLayerProfiler); - return mLayerProfiler->getLayerProfile(); -} - -void TllmRuntime::reportToProfiler(SizeType32 contextId) -{ - mContexts[contextId]->reportToProfiler(); -} - -void TllmRuntime::loadManagedWeights(RawEngine const& rawEngine, int localRank) -{ - TLLM_LOG_TRACE("%s start", __PRETTY_FUNCTION__); - NVTX3_FUNC_RANGE(); - auto& engine = getEngine(); - auto& manager = getBufferManager(); - if (rawEngine.getManagedWeightsMapOpt().has_value()) - { - TLLM_LOG_DEBUG("Loading managed weights from raw engine"); - auto executorMap = rawEngine.getManagedWeightsMapOpt().value(); - for (auto const& [name, weight] : executorMap) - { - TLLM_LOG_DEBUG("Loading managed weight: %s", name.c_str()); - auto iTensor = tensorrt_llm::executor::detail::toITensor(weight); - auto weightsDevice = std::shared_ptr{manager.copyFrom(*iTensor, MemoryType::kGPU)}; - mManagedWeightsMap.insert(std::make_pair(name, weightsDevice)); - } - } - else - { - TLLM_LOG_DEBUG("Loading managed weights from file"); - auto const enginePath = rawEngine.getPathOpt(); - TLLM_CHECK_WITH_INFO(enginePath.has_value(), "Engine path is not set."); - auto weightPath - = enginePath->parent_path() / ("rank" + std::to_string(localRank) + "_managed_weights.safetensors"); - auto managed_weights = common::safetensors::ISafeTensor::open(weightPath.string().c_str()); - for (auto const& name : managed_weights->keys()) - { - TLLM_LOG_DEBUG("Loading managed weight: %s", name.c_str()); - auto const weight = managed_weights->getTensor(name.c_str()); - TLLM_CHECK(weight->dtype() == engine.getTensorDataType(name.c_str())); - auto weightsDevice - = std::shared_ptr{manager.allocate(MemoryType::kGPU, weight->trtDims(), weight->dtype())}; - manager.copy(weight->data(), *weightsDevice, MemoryType::kCPU); - mManagedWeightsMap.insert(std::make_pair(name, weightsDevice)); - } - } - setStaticInputTensors(mManagedWeightsMap); - TLLM_LOG_TRACE("%s stop", __PRETTY_FUNCTION__); -} diff --git a/cpp/tensorrt_llm/runtime/tllmRuntime.h b/cpp/tensorrt_llm/runtime/tllmRuntime.h deleted file mode 100644 index dfef06d8b45f..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmRuntime.h +++ /dev/null @@ -1,243 +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/batch_manager/llmRequest.h" -#include "tensorrt_llm/runtime/bufferManager.h" -#include "tensorrt_llm/runtime/common.h" -#include "tensorrt_llm/runtime/iTensor.h" -#include "tensorrt_llm/runtime/layerProfiler.h" -#include "tensorrt_llm/runtime/rawEngine.h" -#include "tensorrt_llm/runtime/worldConfig.h" -#include - -#include -#include -#include -#include -#include - -namespace tensorrt_llm::runtime -{ -class TllmRuntime -{ -public: - using TensorMap = StringPtrMap; - - explicit TllmRuntime(RawEngine const& rawEngine, nvinfer1::ILogger* logger, bool useGpuDirectStorage = false, - float gpuWeightsPercent = 1.0f, bool useShapeInference = true); - - SizeType32 getNbContexts() const - { - return static_cast(mContexts.size()); - } - - nvinfer1::IExecutionContext& getContext(SizeType32 contextIndex) const - { - return *mContexts.at(contextIndex); - } - - SizeType32 getNbProfiles() const - { - return static_cast(mEngine->getNbOptimizationProfiles()); - } - - /// @brief If multiple TensorRT optimization profiles are built in the engine, this function selects the - /// corresponding profile that is going to be used based on the runtime shape, for now, TensorRT LLM only split - /// multiple profiles on the num_tokens dimension, hence the profile index is selected based on which profile - /// handles the actual num_tokens - /// @return The index of the selected TensorRT optimization profile - [[nodiscard]] SizeType32 getOptProfileId(int numTokens, std::vector const& splitPoints) const - { - if (getNbProfiles() == 1) - { - return 0; - } - auto const it = std::lower_bound(splitPoints.begin(), splitPoints.end(), numTokens); - auto const optProfileId = std::distance(splitPoints.begin(), it); - return optProfileId; - } - - nvinfer1::IExecutionContext& addContext(std::int32_t profileIndex); - - void clearContexts(); - - /// @brief Set input tensors from tensorMap for all contexts. - /// @details The function can be used to set static input tensors for all iterations. If a tensor was set this way, - /// it doesn't need to included in calls to setInputTensors anymore. - void setStaticInputTensors(TensorMap const& tensorMap); - - /// @brief Set input tensors from tensorMap for context at contextIndex. - /// @details The function expects that all input tensors (excluding the ones set by setStaticInputTensors) are - /// contained in the tensorMap. If a tensor is missing, has a bad shape or type, it will throw. - void setInputTensors(SizeType32 contextIndex, TensorMap const& tensorMap); - - /// @brief Set output tensors from tensorMap for context at contextIndex. - /// @details The function expects that all output tensors are contained in the tensorMap. If a tensor is missing and - /// shape inference is enabled, it will allocate the tensor on GPU and insert it into the tensorMap. Otherwise it - /// will throw. - void setOutputTensors(SizeType32 contextIndex, TensorMap& tensorMap); - - bool executeContext(SizeType32 contextIndex) const; - - CudaStream const& getStream() const; - - BufferManager::CudaStreamPtr getStreamPtr() - { - return mStream; - } - - nvinfer1::ICudaEngine& getEngine() - { - return *mEngine; - } - - nvinfer1::ICudaEngine const& getEngine() const - { - return *mEngine; - } - - nvinfer1::IEngineInspector& getEngineInspector() - { - return *mEngineInspector; - } - - nvinfer1::IEngineInspector const& getEngineInspector() const - { - return *mEngineInspector; - } - - BufferManager& getBufferManager() - { - return mBufferManager; - } - - BufferManager const& getBufferManager() const - { - return mBufferManager; - } - - void setLayerProfiler(); - bool hasLayerProfiler(SizeType32 contextId) const; - std::string getLayerProfileInfo() const; - void reportToProfiler(SizeType32 contextId); - void loadManagedWeights(RawEngine const& rawEngine, int localRank); - void initializeUserBuffer(tensorrt_llm::runtime::WorldConfig const& world_config, SizeType32 maxBatchSize, - SizeType32 maxBeamWidth, SizeType32 maxSequenceLength, SizeType32 hiddenSize, - std::optional maxNumTokens); - - bool isUserBufferEnabled() const - { - return mUserBufferEnabled; - } - - void setCurrentBeamWidths(std::vector const& beamWidth) noexcept - { - mCurrentBeamWidths = beamWidth; - } - - [[nodiscard]] SizeType32 const& getCurrentBeamWidth() const noexcept - { - // 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 widths in `mCurrentBeamWidths` are the same. - // 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. - TLLM_CHECK_WITH_INFO(mCurrentBeamWidths.size() > 0, "`mCurrentBeamWidths` is empty."); - bool const isEqual = std::all_of(mCurrentBeamWidths.begin(), mCurrentBeamWidths.end(), - [&](int elem) { return elem == mCurrentBeamWidths.front(); }); - TLLM_CHECK_WITH_INFO(isEqual, "beam widths in `mCurrentBeamWidths` are not all equal."); - return mCurrentBeamWidths.front(); - } - -private: - void cacheTensorNames(); - - void setInputTensorsImpl(SizeType32 contextIndex, TensorMap const& tensorMap, bool throwOnMiss); - - void setUserBufferTensors(SizeType32 contextIndex, TensorMap& tensorMap); - - void printEngineInfo(); - - void printContextInfo(SizeType32 contextIndex); - - // Tool functions for `printEngineInfo()`. - static std::string shapeToString(nvinfer1::Dims64 const& dim) - { - std::string output("("); - if (dim.nbDims == 0) - { - return output + ")"; - } - for (int i = 0; i < dim.nbDims - 1; ++i) - { - output += std::to_string(dim.d[i]) + ", "; - } - output += std::to_string(dim.d[dim.nbDims - 1]) + ")"; - return output; - } - - static std::string dataTypeToString(nvinfer1::DataType type) - { - switch (type) - { - case nvinfer1::DataType::kINT64: return "INT64"; - case nvinfer1::DataType::kINT32: return "INT32"; - case nvinfer1::DataType::kFLOAT: return "FP32"; - case nvinfer1::DataType::kBF16: return "BF16"; - case nvinfer1::DataType::kHALF: return "FP16"; - case nvinfer1::DataType::kBOOL: return "BOOL"; - case nvinfer1::DataType::kUINT8: return "UINT8"; - case nvinfer1::DataType::kINT8: return "INT8"; - case nvinfer1::DataType::kFP8: return "FP8"; - case nvinfer1::DataType::kINT4: return "INT4"; - case nvinfer1::DataType::kFP4: return "FP4"; - default: return "UNKNOWN"; - } - return ""; - } - - static std::string alignText( - std::string const& text, int const width, bool const bCenter = true, char const blank = ' ') - { - int textLen = text.size(); - int padLeft = 0; - int padRight = 0; - padLeft = bCenter ? (width - textLen) / 2 : 0; - padRight = width - padLeft - textLen; - return std::string(padLeft, blank) + text + std::string(padRight, blank); - } - - BufferManager::CudaStreamPtr mStream; - BufferManager mBufferManager; - std::unique_ptr mRuntime; - std::unique_ptr mEngine; - BufferManager::IBufferPtr mEngineBuffer; - std::vector> mContexts; - std::unique_ptr mDummyTensor; - std::unique_ptr mEngineInspector; - std::unique_ptr mLayerProfiler; - bool mUseShapeInference; - TensorMap mManagedWeightsMap; - // List of input tensor names. - // Names of static tensors are removed from this list when setStaticInputTensors is called. - std::vector mInputTensorNames; - std::vector mOutputTensorNames; - - bool mUserBufferEnabled; - // For Variable-Beam-Width-Search - std::vector mCurrentBeamWidths; -}; -} // namespace tensorrt_llm::runtime diff --git a/cpp/tensorrt_llm/runtime/tllmStreamReaders.cpp b/cpp/tensorrt_llm/runtime/tllmStreamReaders.cpp deleted file mode 100644 index 55440bbe714f..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmStreamReaders.cpp +++ /dev/null @@ -1,217 +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 "tllmStreamReaders.h" -#include "tensorrt_llm/common/assert.h" -#include "tensorrt_llm/common/logger.h" - -#include -#include -#include -#include -#include -#include -#include - -// Non-GDS StreamReader - -StreamReader::StreamReader(std::filesystem::path fp) -{ - mFile.open(fp.string(), std::ios::binary | std::ios::in); - TLLM_CHECK_WITH_INFO(mFile.good(), std::string("Error opening engine file: " + fp.string())); -} - -StreamReader::~StreamReader() -{ - if (mFile.is_open()) - { - mFile.close(); - } -} - -int64_t StreamReader::read(void* destination, int64_t nbBytes) -{ - if (!mFile.good()) - { - return -1; - } - - mFile.read(static_cast(destination), nbBytes); - - return mFile.gcount(); -} - -// StreamReader using GDS - -GDSStreamReader::GDSStreamReader(std::filesystem::path const& filePath) -{ - auto const start_time = std::chrono::high_resolution_clock::now(); - initializeDriver(); - auto const elapsed_ms - = std::chrono::duration_cast(std::chrono::high_resolution_clock::now() - start_time); - - TLLM_LOG_INFO("GDS driver initialization time %lld ms", elapsed_ms); - - open(filePath); -} - -bool GDSStreamReader::open(std::string const& filepath) -{ - if (!initializeDriver()) - { - TLLM_LOG_INFO("Failed to initialize cuFile driver"); - return false; - } - - int32_t const ret = ::open(filepath.c_str(), O_CREAT | O_RDWR | O_DIRECT, 0664); - - if (ret < 0) - { - TLLM_LOG_INFO("Failed to open engine file"); - return false; - } - - mFd = ret; - mFileSize = lseek(mFd, 0, SEEK_END); - lseek(mFd, 0, SEEK_SET); - - CUfileDescr_t fileDescr; - memset((void*) &fileDescr, 0, sizeof(fileDescr)); - fileDescr.handle.fd = mFd; - fileDescr.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD; - - CUfileError_t gdsStatus = cuFileHandleRegister(&mFileHandle, &fileDescr); - - if (gdsStatus.err != CU_FILE_SUCCESS) - { - TLLM_LOG_INFO("Failed to cuFileHandleRegister"); - ::close(mFd); - return false; - } - return true; -} - -void GDSStreamReader::close() -{ - if (mFd >= 0) - { - ::close(mFd); - mFd = -1; - } -} - -GDSStreamReader::~GDSStreamReader() -{ - if (mFileHandle) - { - cuFileHandleDeregister(mFileHandle); - mFileHandle = nullptr; - } - - if (mDriverInitialized) - { - cuFileDriverClose(); - } -} - -bool GDSStreamReader::seek(int64_t offset, nvinfer1::SeekPosition where) noexcept -{ - switch (where) - { - case nvinfer1::SeekPosition::kSET: mCursor = offset; return true; - case nvinfer1::SeekPosition::kCUR: mCursor += offset; return true; - case nvinfer1::SeekPosition::kEND: mCursor = -offset; return true; - default: return false; - } - return true; -} - -int64_t GDSStreamReader::read(void* dest, int64_t bytes, cudaStream_t stream) noexcept -{ - cudaPointerAttributes attributes{}; - if (cudaPointerGetAttributes(&attributes, dest) != cudaSuccess) - { - TLLM_LOG_INFO("cudaPointerGetAttributes failed"); - } - - off_t destOffset = 0; - void* destBase = dest; - - if (attributes.type == cudaMemoryTypeDevice) - { - CUdeviceptr cuDest = reinterpret_cast(dest); - CUdeviceptr cuBufBase = 0; - size_t cuBufSize = 0; - - cuMemGetAddressRange(&cuBufBase, &cuBufSize, cuDest); - destOffset += cuDest - cuBufBase; - destBase = reinterpret_cast(cuBufBase); - } - cuFileRead(this->mFileHandle, destBase, bytes, mCursor, destOffset); - - mCursor += bytes; - return bytes; -} - -void GDSStreamReader::reset() -{ - lseek(mFd, 0, SEEK_SET); - mCursor = 0; -} - -[[nodiscard]] bool GDSStreamReader::isOpen() const -{ - bool open = mFd >= 0; - return open; -} - -bool GDSStreamReader::initializeDriver() -{ - if (mDriverInitialized) - { - return true; - } - - mCuFileLibHandle = dlopen("libcufile.so", RTLD_LAZY | RTLD_GLOBAL); - if (!mCuFileLibHandle) - { - TLLM_LOG_INFO("Failed to dlopen libcufile.so"); - return false; - } - - // Load the required functions - *reinterpret_cast(&cuFileDriverOpen) = dlsym(mCuFileLibHandle, "cuFileDriverOpen"); - *reinterpret_cast(&cuFileHandleRegister) = dlsym(mCuFileLibHandle, "cuFileHandleRegister"); - *reinterpret_cast(&cuFileHandleDeregister) = dlsym(mCuFileLibHandle, "cuFileHandleDeregister"); - *reinterpret_cast(&cuFileDriverClose) = dlsym(mCuFileLibHandle, "cuFileDriverClose"); - *reinterpret_cast(&cuFileRead) = dlsym(mCuFileLibHandle, "cuFileRead"); - - if (!cuFileDriverOpen || !cuFileHandleRegister || !cuFileHandleDeregister || !cuFileDriverClose || !cuFileRead) - { - TLLM_LOG_INFO("Failed to dlsym libcufile.so"); - return false; - } - - CUfileError_t gdsStatus = cuFileDriverOpen(); - if (gdsStatus.err != CU_FILE_SUCCESS) - { - TLLM_LOG_INFO("cuFileDriverOpen failed"); - return false; - } - - mDriverInitialized = true; - return true; -} diff --git a/cpp/tensorrt_llm/runtime/tllmStreamReaders.h b/cpp/tensorrt_llm/runtime/tllmStreamReaders.h deleted file mode 100644 index 943f0bb3e32e..000000000000 --- a/cpp/tensorrt_llm/runtime/tllmStreamReaders.h +++ /dev/null @@ -1,71 +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 -#include - -class StreamReader final : public nvinfer1::IStreamReader -{ -public: - StreamReader(std::filesystem::path fp); - - virtual ~StreamReader(); - - int64_t read(void* destination, int64_t nbBytes) final; - -private: - std::ifstream mFile; -}; - -class GDSStreamReader final : public nvinfer1::IStreamReaderV2 -{ -public: - explicit GDSStreamReader(std::filesystem::path const& filePath); - - virtual ~GDSStreamReader(); - - void close(); - - [[nodiscard]] bool isOpen() const; - - bool open(std::string const& filepath); - - int64_t read(void* dest, int64_t bytes, cudaStream_t stream) noexcept final; - - void reset(); - - bool seek(int64_t offset, nvinfer1::SeekPosition where) noexcept final; - -private: - bool initializeDriver(); - - void* mCuFileLibHandle{}; - CUfileHandle_t mFileHandle{nullptr}; - bool mDriverInitialized{false}; - int32_t mFd{-1}; - int64_t mCursor{0}; - int64_t mFileSize{0}; - - CUfileError_t (*cuFileDriverOpen)(){}; - CUfileError_t (*cuFileHandleRegister)(CUfileHandle_t*, CUfileDescr_t*){}; - CUfileError_t (*cuFileHandleDeregister)(CUfileHandle_t){}; - CUfileError_t (*cuFileDriverClose)(){}; - ssize_t (*cuFileRead)(CUfileHandle_t, void*, size_t, int64_t, int64_t){}; -}; diff --git a/cpp/tensorrt_llm/runtime/utils/debugUtils.cu b/cpp/tensorrt_llm/runtime/utils/debugUtils.cu index 661dacd9a7ac..46f549835cf9 100644 --- a/cpp/tensorrt_llm/runtime/utils/debugUtils.cu +++ b/cpp/tensorrt_llm/runtime/utils/debugUtils.cu @@ -167,7 +167,7 @@ template bool tensorHasInvalid(ITensor const& tensor, BufferManager const& manager, std::string const& infoStr) { printLogitsKeyInfo(tensor, infoStr); - auto foundInvalid = BufferManager::pinnedPool(ITensor::makeShape({1}), nvinfer1::DataType::kINT32); + auto foundInvalid = BufferManager::pinnedPool(ITensor::makeShape({1}), tensorrt_llm::DataType::kINT32); auto foundInvalidPtr = bufferCast(*foundInvalid); foundInvalidPtr[0] = 0; auto const size = tensor.getSize(); @@ -184,24 +184,24 @@ template bool tensorHasInvalid<__nv_fp8_e4m3>( ITensor const& tensor, BufferManager const& manager, std::string const& infoStr); bool tensorHasInvalid( - size_t M, size_t K, nvinfer1::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr) + size_t M, size_t K, tensorrt_llm::DataType type, void const* data, cudaStream_t stream, std::string const& infoStr) { auto tensorView = ITensor::wrap( const_cast(data), type, ITensor::makeShape({static_cast(M), static_cast(K)})); auto manager = BufferManager(std::make_shared(stream)); - if (type == nvinfer1::DataType::kFLOAT) + if (type == tensorrt_llm::DataType::kFLOAT) { return tensorHasInvalid(*tensorView, manager, infoStr); } - else if (type == nvinfer1::DataType::kHALF) + else if (type == tensorrt_llm::DataType::kHALF) { return tensorHasInvalid(*tensorView, manager, infoStr); } - else if (type == nvinfer1::DataType::kBF16) + else if (type == tensorrt_llm::DataType::kBF16) { return tensorHasInvalid<__nv_bfloat16>(*tensorView, manager, infoStr); } - else if (type == nvinfer1::DataType::kFP8) + else if (type == tensorrt_llm::DataType::kFP8) { return tensorHasInvalid<__nv_fp8_e4m3>(*tensorView, manager, infoStr); } diff --git a/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp b/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp index 6f95704455d3..2ee970da83dc 100644 --- a/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp +++ b/cpp/tensorrt_llm/runtime/utils/numpyUtils.cpp @@ -22,7 +22,7 @@ #include "tensorrt_llm/common/stringUtils.h" #include "tensorrt_llm/runtime/bufferManager.h" #include "tensorrt_llm/runtime/iTensor.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include @@ -34,9 +34,9 @@ namespace tc = tensorrt_llm::common; namespace tensorrt_llm::runtime::utils { -std::string getNumpyTypeDesc(nvinfer1::DataType type) +std::string getNumpyTypeDesc(tensorrt_llm::DataType type) { - using dt = nvinfer1::DataType; + using dt = tensorrt_llm::DataType; static std::unordered_map const type_map{{dt::kBOOL, "?"}, {dt::kUINT8, "u1"}, {dt::kINT8, "i1"}, {dt::kINT32, "i4"}, {dt::kINT64, "i8"}, {dt::kHALF, "f2"}, {dt::kFLOAT, "f4"}}; @@ -51,11 +51,11 @@ std::string getNumpyTypeDesc(nvinfer1::DataType type) return type_map.count(type) > 0 ? type_map.at(type) : "x"; } -nvinfer1::DataType typeFromNumpyDesc(std::string const& type) +tensorrt_llm::DataType typeFromNumpyDesc(std::string const& type) { TLLM_LOG_DEBUG("numpy type: %s", type.c_str()); - using dt = nvinfer1::DataType; + using dt = tensorrt_llm::DataType; static std::unordered_map const type_map{{"?", dt::kBOOL}, {"u1", dt::kUINT8}, {"i1", dt::kINT8}, {"i4", dt::kINT32}, {"i8", dt::kINT64}, {"f2", dt::kHALF}, {"f4", dt::kFLOAT}}; TLLM_CHECK_WITH_INFO(type_map.count(type) > 0, "numpy data type '" + type + "' not supported"); @@ -102,7 +102,7 @@ void parseNpyIntro(FILE*& f_ptr, uint32_t& header_len, uint32_t& start_data) start_data = 8 + 2 * npy_major + header_len; } -int parseNpyHeader(FILE*& f_ptr, uint32_t header_len, nvinfer1::DataType& type, std::vector& shapeVec) +int parseNpyHeader(FILE*& f_ptr, uint32_t header_len, tensorrt_llm::DataType& type, std::vector& shapeVec) { char* header_c = (char*) malloc(header_len * sizeof(char)); TLLM_CHECK_WITH_INFO(header_c != nullptr, "Failed to allocate memory for npy header"); @@ -168,11 +168,11 @@ int parseNpyHeader(FILE*& f_ptr, uint32_t header_len, nvinfer1::DataType& type, uint32_t header_len, start_data; utils::parseNpyIntro(f_ptr, header_len, start_data); - nvinfer1::DataType type; + tensorrt_llm::DataType type; std::vector shape; utils::parseNpyHeader(f_ptr, header_len, type, shape); - nvinfer1::Dims dims; + tensorrt_llm::Dims dims; dims.nbDims = shape.size(); std::copy(shape.begin(), shape.end(), dims.d); @@ -203,10 +203,10 @@ void saveNpy(BufferManager const& manager, ITensor const& tensor, std::string co auto const dtype = tensor.getDataType(); #ifdef ENABLE_BF16 - if (dtype == nvinfer1::DataType::kBF16) + if (dtype == tensorrt_llm::DataType::kBF16) { TLLM_CHECK(where == MemoryType::kGPU); - auto tensorFp32 = manager.gpu(shape, nvinfer1::DataType::kFLOAT); + auto tensorFp32 = manager.gpu(shape, tensorrt_llm::DataType::kFLOAT); auto dataFp32 = bufferCast(*tensorFp32); auto dataBf16 = bufferCast<__nv_bfloat16 const>(tensor); tc::invokeCudaD2DcpyConvert(dataFp32, dataBf16, tensorSize); diff --git a/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h b/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h index f131ab3419bf..da94419eff64 100644 --- a/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h +++ b/cpp/tensorrt_llm/runtime/utils/runtimeUtils.h @@ -25,7 +25,6 @@ namespace tensorrt_llm::runtime { -class TllmRuntime; namespace utils { diff --git a/cpp/tensorrt_llm/runtime/virtualMemory.cpp b/cpp/tensorrt_llm/runtime/virtualMemory.cpp index 9b23866d6281..38342aa4a671 100644 --- a/cpp/tensorrt_llm/runtime/virtualMemory.cpp +++ b/cpp/tensorrt_llm/runtime/virtualMemory.cpp @@ -141,8 +141,8 @@ void OffloadConfigurator::teardown(CUmemGenericAllocationHandle, bool destructin { switch (mBackType) { - case MemoryType::kCPU: mBackedStorage = BufferManager::cpu(mSize, nvinfer1::DataType::kINT8); break; - case MemoryType::kPINNED: mBackedStorage = BufferManager::pinned(mSize, nvinfer1::DataType::kINT8); break; + case MemoryType::kCPU: mBackedStorage = BufferManager::cpu(mSize, tensorrt_llm::DataType::kINT8); break; + case MemoryType::kPINNED: mBackedStorage = BufferManager::pinned(mSize, tensorrt_llm::DataType::kINT8); break; default: TLLM_THROW("Unknown memory type: %d", static_cast(mBackType)); } } diff --git a/cpp/tensorrt_llm/testing/modelSpec.h b/cpp/tensorrt_llm/testing/modelSpec.h index 5b6f88dcd135..105058419ae8 100644 --- a/cpp/tensorrt_llm/testing/modelSpec.h +++ b/cpp/tensorrt_llm/testing/modelSpec.h @@ -17,7 +17,7 @@ #pragma once -#include "NvInfer.h" +#include "tensorrt_llm/common/tllmDataType.h" #include "tensorrt_llm/runtime/common.h" #include "tensorrt_llm/runtime/modelConfig.h" #include "tensorrt_llm/runtime/speculativeDecodingMode.h" @@ -50,7 +50,7 @@ enum class OutputContentType class ModelSpec { public: - ModelSpec(std::string const& inputFile, nvinfer1::DataType dtype, + ModelSpec(std::string const& inputFile, tensorrt_llm::DataType dtype, std::shared_ptr otherModelSpecToCompare = nullptr) : mInputFile{std::move(inputFile)} , mDataType{dtype} @@ -294,14 +294,14 @@ class ModelSpec static ModelSpec getDefaultModelSpec() { - static ModelSpec modelSpec{"input_tokens.npy", nvinfer1::DataType::kHALF}; + static ModelSpec modelSpec{"input_tokens.npy", tensorrt_llm::DataType::kHALF}; modelSpec.useGptAttentionPlugin().setKVCacheType(KVCacheType::kPAGED).usePackedInput(); return modelSpec; } std::string mInputFile; - nvinfer1::DataType mDataType; + tensorrt_llm::DataType mDataType; bool mUseGptAttentionPlugin{false}; bool mUsePackedInput{false}; diff --git a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp index 696b8d471043..f780faf8e258 100644 --- a/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp +++ b/cpp/tensorrt_llm/thop/IndexerTopKOp.cpp @@ -19,7 +19,7 @@ #include "tensorrt_llm/kernels/IndexerTopK.h" -// #include +// #include "tensorrt_llm/common/tllmDataType.h" // #include // #include // #include diff --git a/cpp/tensorrt_llm/thop/allgatherOp.cpp b/cpp/tensorrt_llm/thop/allgatherOp.cpp index 0d92aa966901..66ed8d141d83 100644 --- a/cpp/tensorrt_llm/thop/allgatherOp.cpp +++ b/cpp/tensorrt_llm/thop/allgatherOp.cpp @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #include diff --git a/cpp/tensorrt_llm/thop/allreduceOp.cpp b/cpp/tensorrt_llm/thop/allreduceOp.cpp index 9f4de156ea55..5083283179f9 100644 --- a/cpp/tensorrt_llm/thop/allreduceOp.cpp +++ b/cpp/tensorrt_llm/thop/allreduceOp.cpp @@ -235,7 +235,7 @@ class AllreduceOp { public: AllreduceOp( - std::set group, nvinfer1::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, float eps) + std::set group, tensorrt_llm::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, float eps) : mGroup(std::move(group)) , mIsNVLINKSupported(false) , mIsP2PSupported(false) @@ -248,7 +248,7 @@ class AllreduceOp } AllreduceOp(std::set group, c10::intrusive_ptr const& process_group_, - nvinfer1::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, float eps) + tensorrt_llm::DataType type, AllReduceStrategyType strategy, AllReduceFusionOp op, float eps) : mGroup(std::move(group)) , mIsNVLINKSupported(false) , mIsP2PSupported(false) @@ -348,7 +348,7 @@ class AllreduceOp { TORCH_CHECK(norm_weight, "norm_weight is required for residual rms norm allreduce"); TORCH_CHECK(!bias, "bias is not supported for residual rms norm allreduce"); - TORCH_CHECK(mType == nvinfer1::DataType::kHALF || mType == nvinfer1::DataType::kBF16); + TORCH_CHECK(mType == tensorrt_llm::DataType::kHALF || mType == tensorrt_llm::DataType::kBF16); auto [norm_out, ub_buffer1] = torch_ext::create_userbuffers_tensor(input.sizes(), input.scalar_type()); tensorrt_llm::kernels::ub::allreduce2_userbuff_rmsnorm_launcher(ub_buffer0.handle, 0, ub_buffer1.handle, 0, size, hidden_size, nullptr, norm_weight.value().data_ptr(), mEps, residual.value().data_ptr(), @@ -1461,7 +1461,7 @@ class AllreduceOp bool mIsNVLINKSupported; bool mIsP2PSupported; bool mIsMNNVLSupported; - nvinfer1::DataType mType; + tensorrt_llm::DataType mType; AllReduceStrategyType mStrategy; AllReduceFusionOp mOp; float mEps; diff --git a/cpp/tensorrt_llm/thop/attentionOp.cpp b/cpp/tensorrt_llm/thop/attentionOp.cpp index 64b205e54085..dba3a3f87795 100644 --- a/cpp/tensorrt_llm/thop/attentionOp.cpp +++ b/cpp/tensorrt_llm/thop/attentionOp.cpp @@ -1067,7 +1067,7 @@ void attention(torch::Tensor q, std::optional k, std::optional k, std::optional>(); } } - else if (dtype == nvinfer1::DataType::kFLOAT) + else if (dtype == tensorrt_llm::DataType::kFLOAT) { TLLM_CHECK(out_dtype == torch::kFloat32); runner = std::make_shared>(); } #ifdef ENABLE_BF16 - else if (dtype == nvinfer1::DataType::kBF16) + else if (dtype == tensorrt_llm::DataType::kBF16) { if (is_fp8_out) { @@ -1112,7 +1112,7 @@ void attention(torch::Tensor q, std::optional k, std::optional(); op->mType = dtype; - op->mFMHAForceFP32Acc = dtype == nvinfer1::DataType::kBF16; + op->mFMHAForceFP32Acc = dtype == tensorrt_llm::DataType::kBF16; op->mLayerIdx = local_layer_idx; op->mNumHeads = num_heads; op->mNumKVHeads = num_kv_heads; @@ -1372,7 +1372,7 @@ bool attention_supports_nvfp4_output(int64_t const num_heads, int64_t const num_ } auto op = std::make_shared(); - op->mType = nvinfer1::DataType::kHALF; + op->mType = tensorrt_llm::DataType::kHALF; op->mNumHeads = num_heads; op->mNumKVHeads = num_kv_heads; op->mHeadSize = head_size; diff --git a/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp index a9ad46ad8f04..8e4da99dbadc 100644 --- a/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasFp4ScaledMM.cpp @@ -16,7 +16,6 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/cudaUtils.h" -#include "tensorrt_llm/plugins/common/plugin.h" #include "tensorrt_llm/thop/outputTensor.h" #include "tensorrt_llm/thop/thUtils.h" #include "userbuffersTensor.h" diff --git a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp index 4951db12cfd2..9cef9a19137f 100644 --- a/cpp/tensorrt_llm/thop/cublasScaledMM.cpp +++ b/cpp/tensorrt_llm/thop/cublasScaledMM.cpp @@ -18,8 +18,6 @@ #include "tensorrt_llm/common/cublasMMWrapper.h" #include "tensorrt_llm/common/cudaUtils.h" #include "tensorrt_llm/kernels/userbuffers/ub_interface.h" -#include "tensorrt_llm/plugins/common/plugin.h" -#include "tensorrt_llm/plugins/gemmPlugin/gemmPlugin.h" #include "tensorrt_llm/runtime/torchUtils.h" #include "tensorrt_llm/thop/outputTensor.h" #include "tensorrt_llm/thop/thUtils.h" diff --git a/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp b/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp index 8e9e817bbb51..646736d914fa 100644 --- a/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp +++ b/cpp/tensorrt_llm/thop/dynamicDecodeOp.cpp @@ -54,7 +54,7 @@ FtDynamicDecode::FtDynamicDecode(size_t const maxBatchSize, size_t const maxB auto bufferManager = std::make_shared(cudaStreamPtr); mFinishedSum = bufferManager->pinnedPool( - tr::ITensor::makeShape({static_cast(maxBatchSize)}), nvinfer1::DataType::kINT32); + tr::ITensor::makeShape({static_cast(maxBatchSize)}), tensorrt_llm::DataType::kINT32); mDynamicDecodeLayer = std::make_shared>(tle::DecodingMode::Auto(), decodingDomain, bufferManager); mBatchSlots = tr::getDefaultBatchSlots(maxBatchSize); diff --git a/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp b/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp index c408a8c286fb..989ef348712d 100644 --- a/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp +++ b/cpp/tensorrt_llm/thop/groupRmsNormOp.cpp @@ -100,9 +100,9 @@ void groupRMSNormBase(torch::TensorList const& inputs, torch::TensorList const& /* Handle dtype conversion */ \ switch (dtype) \ { \ - case torch::ScalarType::Half: params.dtype = nvinfer1::DataType::kHALF; break; \ - case torch::ScalarType::BFloat16: params.dtype = nvinfer1::DataType::kBF16; break; \ - case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break; \ + case torch::ScalarType::Half: params.dtype = tensorrt_llm::DataType::kHALF; break; \ + case torch::ScalarType::BFloat16: params.dtype = tensorrt_llm::DataType::kBF16; break; \ + case torch::ScalarType::Float: params.dtype = tensorrt_llm::DataType::kFLOAT; break; \ default: TORCH_CHECK(false, "Unsupported data type"); \ } \ tensorrt_llm::kernels::group_rms_norm::GroupRMSNormBaseKernelLauncher(params); \ @@ -181,9 +181,9 @@ void groupRMSNormLargeBatch(torch::TensorList const& inputs, torch::TensorList c // Handle dtype conversion switch (dtype) { - case torch::ScalarType::Half: params.dtype = nvinfer1::DataType::kHALF; break; - case torch::ScalarType::BFloat16: params.dtype = nvinfer1::DataType::kBF16; break; - case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break; + case torch::ScalarType::Half: params.dtype = tensorrt_llm::DataType::kHALF; break; + case torch::ScalarType::BFloat16: params.dtype = tensorrt_llm::DataType::kBF16; break; + case torch::ScalarType::Float: params.dtype = tensorrt_llm::DataType::kFLOAT; break; default: TORCH_CHECK(false, "Unsupported data type"); } @@ -260,9 +260,9 @@ void groupRMSNormHeuristic(torch::TensorList const& inputs, torch::TensorList co /* Handle dtype conversion */ \ switch (dtype) \ { \ - case torch::ScalarType::Half: params.dtype = nvinfer1::DataType::kHALF; break; \ - case torch::ScalarType::BFloat16: params.dtype = nvinfer1::DataType::kBF16; break; \ - case torch::ScalarType::Float: params.dtype = nvinfer1::DataType::kFLOAT; break; \ + case torch::ScalarType::Half: params.dtype = tensorrt_llm::DataType::kHALF; break; \ + case torch::ScalarType::BFloat16: params.dtype = tensorrt_llm::DataType::kBF16; break; \ + case torch::ScalarType::Float: params.dtype = tensorrt_llm::DataType::kFLOAT; break; \ default: TORCH_CHECK(false, "Unsupported data type"); \ } \ \ diff --git a/cpp/tensorrt_llm/thop/loraOp.cpp b/cpp/tensorrt_llm/thop/loraOp.cpp index b35ca2608625..3c996f40bf51 100644 --- a/cpp/tensorrt_llm/thop/loraOp.cpp +++ b/cpp/tensorrt_llm/thop/loraOp.cpp @@ -151,11 +151,11 @@ std::vector lora_grouped_gemm(th::Tensor const& input, th::Tensor co { outHiddenSizes[i] = output_hidden_sizes[i]; } - nvinfer1::DataType loraRuntimeDataType; + tensorrt_llm::DataType loraRuntimeDataType; switch (input.scalar_type()) { - case torch::kFloat16: loraRuntimeDataType = nvinfer1::DataType::kHALF; break; - case torch::kBFloat16: loraRuntimeDataType = nvinfer1::DataType::kBF16; break; + case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; + case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; default: throw std::invalid_argument("Invalid dtype, only supports float16, bfloat16"); } @@ -221,11 +221,11 @@ void lora_grouped_gemm_cuda_graph(th::Tensor const& lora_in_sizes, // [layer_mod auto* splitk_offsets_gpu = reinterpret_cast(const_cast(splitk_offsets.data_ptr())); // Get data type - nvinfer1::DataType loraRuntimeDataType; + tensorrt_llm::DataType loraRuntimeDataType; switch (dtype) { - case torch::kFloat16: loraRuntimeDataType = nvinfer1::DataType::kHALF; break; - case torch::kBFloat16: loraRuntimeDataType = nvinfer1::DataType::kBF16; break; + case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; + case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; default: TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, got %s", c10::toString(dtype)); } @@ -301,11 +301,11 @@ void lora_group_gemm_param_fill_row_reorder_fusion(th::Tensor const& in_sizes, / int32_t const module_count = static_cast(in_sizes.size(0)); // Get data type info - nvinfer1::DataType loraRuntimeDataType; + tensorrt_llm::DataType loraRuntimeDataType; switch (dtype) { - case torch::kFloat16: loraRuntimeDataType = nvinfer1::DataType::kHALF; break; - case torch::kBFloat16: loraRuntimeDataType = nvinfer1::DataType::kBF16; break; + case torch::kFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kHALF; break; + case torch::kBFloat16: loraRuntimeDataType = tensorrt_llm::DataType::kBF16; break; default: TORCH_CHECK(false, "Invalid dtype, only supports float16, bfloat16, got %s", c10::toString(dtype)); } diff --git a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp index 7a767976dabd..c75aed44b479 100644 --- a/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp +++ b/cpp/tensorrt_llm/thop/moeAlltoAllOp.cpp @@ -434,20 +434,20 @@ torch::Tensor moeA2ACombineOp(torch::Tensor const& payload, int64_t localNumToke TORCH_CHECK(epRank >= 0 && epRank < epSize, "epRank must be in the range [0, epSize)"); TORCH_CHECK(topK > 0 && topK <= kMaxTopK, "topK must be in the range (0, kMaxTopK]"); - // Map torch dtype to nvinfer1::DataType - nvinfer1::DataType nvDtype = nvinfer1::DataType::kFLOAT; + // Map torch dtype to tensorrt_llm::DataType + tensorrt_llm::DataType nvDtype = tensorrt_llm::DataType::kFLOAT; auto scalarType = payload.scalar_type(); if (scalarType == at::kHalf) { - nvDtype = nvinfer1::DataType::kHALF; + nvDtype = tensorrt_llm::DataType::kHALF; } else if (scalarType == at::kBFloat16) { - nvDtype = nvinfer1::DataType::kBF16; + nvDtype = tensorrt_llm::DataType::kBF16; } else if (scalarType == at::kFloat) { - nvDtype = nvinfer1::DataType::kFLOAT; + nvDtype = tensorrt_llm::DataType::kFLOAT; } else { diff --git a/cpp/tensorrt_llm/thop/moeOp.cpp b/cpp/tensorrt_llm/thop/moeOp.cpp index feede19bf71a..e1b002f41017 100644 --- a/cpp/tensorrt_llm/thop/moeOp.cpp +++ b/cpp/tensorrt_llm/thop/moeOp.cpp @@ -89,7 +89,7 @@ enum class MoeLoraRequestType : int32_t // --------------------------------------------------------------------------- inline void moeLoraDeviceRunImpl(::tensorrt_llm::kernels::cutlass_kernels::MoeLoraDevicePathModule const& mod, int64_t num_permuted_tokens, int64_t in_hidden_size, int64_t max_lora_rank, int64_t dtype_bytes, - int64_t splitk_slices, void const* input_base, void* output_base, nvinfer1::DataType data_type, cudaStream_t stream) + int64_t splitk_slices, void const* input_base, void* output_base, tensorrt_llm::DataType data_type, cudaStream_t stream) { TLLM_CHECK_WITH_INFO(mod.permuted_ranks_dev != nullptr, "Device-path LoRA module is missing permuted ranks buffer (forgot to populate device_path?)."); @@ -1259,15 +1259,15 @@ class FusedMoeRunner : public torch::CustomClassHolder } } - // Map a torch dtype to the TRT-LLM nvinfer1::DataType expected by LoraImpl. - static nvinfer1::DataType loraTypeFromActDtype(c10::ScalarType dtype) + // Map a torch dtype to the TRT-LLM tensorrt_llm::DataType expected by LoraImpl. + static tensorrt_llm::DataType loraTypeFromActDtype(c10::ScalarType dtype) { switch (dtype) { - case c10::ScalarType::Half: return nvinfer1::DataType::kHALF; - case c10::ScalarType::Float: return nvinfer1::DataType::kFLOAT; + case c10::ScalarType::Half: return tensorrt_llm::DataType::kHALF; + case c10::ScalarType::Float: return tensorrt_llm::DataType::kFLOAT; #ifdef ENABLE_BF16 - case c10::ScalarType::BFloat16: return nvinfer1::DataType::kBF16; + case c10::ScalarType::BFloat16: return tensorrt_llm::DataType::kBF16; #endif default: C10_THROW_ERROR_FORMATTED(Error, "MoE LoRA only supports fp16/bf16/fp32 activation dtype."); } @@ -2073,7 +2073,7 @@ class FusedMoeRunner : public torch::CustomClassHolder // Compute cuBLAS LoraImpl scratch size for the current call. Returns 0 when LoRA inactive. size_t computeLoraWorkspaceSize(LoraImplPtr const& fc1_impl, LoraImplPtr const& fc2_impl, int64_t num_tokens, - int64_t experts_per_token, int64_t num_seqs, nvinfer1::DataType lora_dtype) const + int64_t experts_per_token, int64_t num_seqs, tensorrt_llm::DataType lora_dtype) const { int64_t const num_lora_tokens = num_tokens * experts_per_token; // num_reqs upper-bounded by num_lora_tokens; mirrors plugin convention. diff --git a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp index e445206e1d78..f8ff3ed18c50 100644 --- a/cpp/tensorrt_llm/thop/noAuxTcOp.cpp +++ b/cpp/tensorrt_llm/thop/noAuxTcOp.cpp @@ -20,7 +20,7 @@ #include "tensorrt_llm/kernels/noAuxTcKernels.h" -// #include +// #include "tensorrt_llm/common/tllmDataType.h" // #include // #include // #include diff --git a/cpp/tensorrt_llm/thop/reducescatterOp.cpp b/cpp/tensorrt_llm/thop/reducescatterOp.cpp index 40f89e40ff75..7f9a24194d86 100644 --- a/cpp/tensorrt_llm/thop/reducescatterOp.cpp +++ b/cpp/tensorrt_llm/thop/reducescatterOp.cpp @@ -20,7 +20,7 @@ #include "tensorrt_llm/runtime/utils/mpiUtils.h" #include "tensorrt_llm/runtime/utils/pgUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include #include #if ENABLE_MULTI_DEVICE diff --git a/cpp/tensorrt_llm/thop/thUtils.cpp b/cpp/tensorrt_llm/thop/thUtils.cpp index 97fe6acaab7b..c151414127fa 100644 --- a/cpp/tensorrt_llm/thop/thUtils.cpp +++ b/cpp/tensorrt_llm/thop/thUtils.cpp @@ -15,7 +15,7 @@ */ #include "tensorrt_llm/thop/thUtils.h" -#include +#include "tensorrt_llm/common/tllmDataType.h" #include TRTLLM_NAMESPACE_BEGIN @@ -25,12 +25,12 @@ namespace torch_ext tensorrt_llm::runtime::ITensor::Shape convert_shape(torch::Tensor tensor) { - constexpr auto trtMaxDims = nvinfer1::Dims::MAX_DIMS; + constexpr auto trtMaxDims = tensorrt_llm::Dims::MAX_DIMS; auto const torchTensorNumDims = tensor.dim(); TLLM_CHECK_WITH_INFO(torchTensorNumDims <= trtMaxDims, "TensorRT supports at most %i tensor dimensions. Found a Torch tensor with %li dimensions.", trtMaxDims, torchTensorNumDims); - auto result = nvinfer1::Dims{}; + auto result = tensorrt_llm::Dims{}; result.nbDims = static_cast(torchTensorNumDims); for (int i = 0; i < torchTensorNumDims; i++) { diff --git a/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp b/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp index 9302e2a4978f..bab8b094d417 100644 --- a/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp +++ b/cpp/tensorrt_llm/thop/trtllmGenQKVProcessOp.cpp @@ -394,16 +394,16 @@ trtllmGenContextPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, tor switch (qkvDtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; @@ -540,16 +540,16 @@ void trtllmGenContextPostprocess(torch::Tensor qkv_input, torch::Tensor workspac switch (qkvDtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: tensorrt_llm::kernels::invokeKvCachePostprocessing( reinterpret_cast&>(qkvParams), stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: tensorrt_llm::kernels::invokeKvCachePostprocessing( reinterpret_cast&>(qkvParams), stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: tensorrt_llm::kernels::invokeKvCachePostprocessing( reinterpret_cast&>(qkvParams), stream); break; @@ -717,16 +717,16 @@ trtllmGenGenerationPreprocess(torch::Tensor qkv_input, torch::Tensor workspace, switch (qkvDtype) { - case nvinfer1::DataType::kFLOAT: + case tensorrt_llm::DataType::kFLOAT: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; - case nvinfer1::DataType::kHALF: + case tensorrt_llm::DataType::kHALF: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; #ifdef ENABLE_BF16 - case nvinfer1::DataType::kBF16: + case tensorrt_llm::DataType::kBF16: tensorrt_llm::kernels::invokeQKVPreprocessing( reinterpret_cast&>(qkvParams), stream); break; diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index 0a06d40ee85e..73a3fe5e731d 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -38,7 +38,7 @@ function(add_gtest test_name test_src) ${ARGN}) add_executable(${test_name} ${test_src}) - target_link_libraries(${test_name} PUBLIC gmock_main TensorRT::OnnxParser) + target_link_libraries(${test_name} PUBLIC gmock_main) if(NOT ARGS_NO_GTEST_MAIN) target_link_libraries(${test_name} PUBLIC gtest_main) endif() diff --git a/cpp/tests/e2e_tests/executor/CMakeLists.txt b/cpp/tests/e2e_tests/executor/CMakeLists.txt index 4813c92584fc..10e9545ed7b8 100644 --- a/cpp/tests/e2e_tests/executor/CMakeLists.txt +++ b/cpp/tests/e2e_tests/executor/CMakeLists.txt @@ -18,5 +18,3 @@ add_gtest(executorTest executorTest.cpp) target_link_libraries(executorTest PRIVATE testingUtils) add_gtest(encDecTest encDecTest.cpp) target_link_libraries(encDecTest PRIVATE testingUtils) -add_gtest(disaggExecutorTest disaggExecutorTest.cpp) -target_link_libraries(disaggExecutorTest PRIVATE testingUtils) diff --git a/cpp/tests/e2e_tests/executor/disaggExecutorTest.cpp b/cpp/tests/e2e_tests/executor/disaggExecutorTest.cpp deleted file mode 100644 index 0eb05d2cc807..000000000000 --- a/cpp/tests/e2e_tests/executor/disaggExecutorTest.cpp +++ /dev/null @@ -1,1437 +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. - */ - -#include "disaggExecutor.h" -#include "executorTest.h" -#include "tensorrt_llm/common/envUtils.h" -#include "tensorrt_llm/runtime/utils/numpyUtils.h" -#include "tests/utils/common.h" - -#include -#include - -namespace tr = tensorrt_llm::runtime; - -using namespace tensorrt_llm::testing; - -namespace -{ -auto constexpr LLAMA_INPUT_FILE = "input_tokens_llama.npy"; -auto constexpr LLAMA_VOCAB_SIZE_PADDED = 128256; -auto constexpr LLAMA_END_ID = 128001; -auto constexpr LLAMA_PAD_ID = 128001; - -using CondDisaggParamsType = std::tuple; // modelName - -enum class InstanceRole : int -{ - kCONTEXT = 1, - kGENERATION = 0, - kMIXED = 2 -}; - -using DisaggParamsType = std::tuple< // - int, // processNum - std::vector, // modelNames - std::vector>, // participantIdsEachInstance - std::vector>, // participantDeviceIdsEachInstance - std::vector, // instanceRoles - int // controllerRank - >; - -std::string convertToString(std::vector> const& vec) -{ - std::ostringstream oss; - oss << "XX"; - - for (size_t i = 0; i < vec.size(); ++i) - { - for (size_t j = 0; j < vec[i].size(); ++j) - { - oss << vec[i][j]; - if (j < vec[i].size() - 1) - { - oss << "_"; - } - } - if (i < vec.size() - 1) - { - oss << "X_X"; - } - } - - oss << "XX"; - return oss.str(); -}; - -std::string convertToString(std::vector const& vec) -{ - std::ostringstream oss; - oss << "XX"; - - for (size_t j = 0; j < vec.size(); ++j) - { - oss << static_cast(vec[j]); - if (j < vec.size() - 1) - { - oss << "_"; - } - } - - oss << "XX"; - return oss.str(); -}; - -std::string generateTestNameDisaggParams(testing::TestParamInfo const& info) -{ - auto const processNum = std::get<0>(info.param); - auto const modelNames = std::get<1>(info.param); - auto const participantIdsEachInstance = std::get<2>(info.param); // std::vector> - auto const participantDeviceIdsEachInstance = std::get<3>(info.param); // std::vector>; - auto const instanceRoles = std::get<4>(info.param); // std::vector ; //1 is context , 0 is generation - auto const controllerRank = std::get<5>(info.param); - - std::string name = "DisaggExecutorTest_"; - - name.append("ProcessNum_" + std::to_string(processNum)); - // name.append("_contextModel_" + contextModel + "_genModel_" + genModel); - name.append("_modelNames_"); - for (auto&& modelName : modelNames) - { - name.append(modelName).append("_"); - } - - name.append("_controllerRank_" + std::to_string(controllerRank)); - - name.append("_ranks_").append(convertToString(participantIdsEachInstance)); - name.append("_devices_").append(convertToString(participantDeviceIdsEachInstance)); - name.append("_roles_").append(convertToString(instanceRoles)); - name.append("_controllerRank_" + std::to_string(controllerRank)); - - return name; -} - -std::string generateTestNameCondDisaggParams(testing::TestParamInfo const& info) -{ - auto const modelName = std::get<0>(info.param); - return "Model_" + modelName; -} - -class DisaggParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class DisaggOrchestratorParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -class ConditionalDisaggParamsTest : public GptExecutorTest, public ::testing::WithParamInterface -{ -}; - -void verifyGenerateDistStats(std::deque const& iterationStats) -{ - for (auto const& iteration : iterationStats) - { - for (auto const& requestStats : iteration.requestStats) - { - // exclude context only requests for mixed server - if (requestStats.stage == RequestStage::kGENERATION_COMPLETE && requestStats.numGeneratedTokens > 1) - { - EXPECT_TRUE(requestStats.disServingStats.has_value()); - EXPECT_GT(requestStats.disServingStats.value().kvCacheTransferMS, 0.0); - } - if (requestStats.stage != RequestStage::kQUEUED) - { - EXPECT_TRUE(requestStats.disServingStats.has_value()); - } - else - { - EXPECT_FALSE(requestStats.disServingStats.has_value()); - } - } - } -} -} // namespace - -void runDisaggTest(tensorrt_llm::testing::disaggexecutor::DisaggExecutorLeader& executor, - tensorrt_llm::runtime::BufferManager& manager, ITensor const& givenInput, ModelIds const& modelIds, - FlakyTestInfo const& flakyTestInfo, bool streaming, SizeType32 const vocabSizePadded, BeamResult const& beamResult, - OutputConfig const& outConfig, bool isSpeculativeDecoding, int maxWaitMs, BatchingType batchingType, - bool returnAllGeneratedTokens) -{ - - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - auto const worldSize = comm.getSize(); - auto const beamWidth = beamResult.beamWidth; - - std::unordered_map reqIdToBatchId; - std::unordered_map> tokens; - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(givenInput, modelIds.padId); - auto const* const givenInputData = tr::bufferCast(givenInput); - - auto const& inputShape = givenInput.getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - SizeType32 numRequests = static_cast(givenInputLengths.size()); - SizeType32 maxRequests = numRequests; - std::vector requests; - std::vector reqMaxNewTokens; - SizeType32 const numReturnSequences = 1; - - for (SizeType32 req = 0; req < maxRequests; ++req) - { - SizeType32 inputLen = givenInputLengths.at(req); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - SizeType32 endId = -1; - auto const* const seqBegin = givenInputData + req * maxInputLength; - VecTokens tokens(seqBegin, seqBegin + inputLen); - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - samplingConfig.setNumReturnSequences(numReturnSequences); - auto request = Request( - VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, samplingConfig, outConfig, endId); - request.setReturnAllGeneratedTokens(returnAllGeneratedTokens); - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_ONLY); - requests.emplace_back(std::move(request)); - } - - if (executor.isControllerRank()) - { - std::vector reqIds; - - for (int i = 0; i < requests.size(); ++i) - { - std::vector resultTokens; - resultTokens.reserve(numReturnSequences); - for (SizeType32 seqIdx = 0; seqIdx < numReturnSequences; ++seqIdx) - { - resultTokens.emplace_back(beamWidth); - } - auto retReqId = executor.enqueueRequests({requests[i]}); - reqIds.push_back(retReqId.front()); - tokens[i] = std::move(resultTokens); - reqIdToBatchId[retReqId.front()] = i; - } - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < maxRequests && iter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto batchId = reqIdToBatchId.at(response.getRequestId()); - auto seqIdx = result.sequenceIndex; - - auto& contextLogits = result.contextLogits; - auto& genLogits = result.generationLogits; - auto& outputTokenIds = result.outputTokenIds; - - EXPECT_EQ(result.finishReasons.size(), beamWidth); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto& newTokens = outputTokenIds.at(beam); - auto& reqTokens = tokens.at(batchId).at(seqIdx).at(beam); - - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - // FinishReason is only supported for bw=1 and inflight batching. - if (beamWidth == 1 && batchingType == BatchingType::kINFLIGHT) - { - EXPECT_EQ(result.finishReasons.at(beam), - result.isFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - - auto& cumLogProbs = result.cumLogProbs; - auto& logProbs = result.logProbs; - auto& beamTokens = tokens.at(batchId).at(seqIdx); - testData.verifyLogProbs(outConfig.returnLogProbs, streaming, outConfig.excludeInputFromOutput, - givenInputLengths.at(batchId), beamWidth, beamTokens, cumLogProbs, logProbs, batchId, - flakyTestInfo); - - testData.validateContextLogits(outConfig.returnContextLogits, givenInputLengths.at(batchId), - beamWidth, contextLogits, vocabSizePadded, batchId); - testData.validateGenerationLogits(outConfig.returnGenerationLogits, result.isFinal, streaming, - outConfig.excludeInputFromOutput, givenInputLengths.at(batchId), reqMaxNewTokens.at(batchId), - beamWidth, beamTokens, genLogits, vocabSizePadded, batchId, returnAllGeneratedTokens); - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, maxWaitMs); - testData.verifyOutput(tokens, givenInputLengths, streaming, outConfig.excludeInputFromOutput, flakyTestInfo, - isSpeculativeDecoding, beamWidth, numReturnSequences, false); - } - comm.barrier(); - if (executor.isGenerationRank()) - { - verifyGenerateDistStats(executor.getLatestRequestStats()); - } -} - -void runDisaggTest(DisaggExecutorOrchestrator& executor, tensorrt_llm::runtime::BufferManager& manager, - ITensor const& givenInput, ModelIds const& modelIds, FlakyTestInfo const& flakyTestInfo, bool streaming, - SizeType32 const vocabSizePadded, BeamResult const& beamResult, OutputConfig const& outConfig, - bool isSpeculativeDecoding, int maxWaitMs, BatchingType batchingType, bool returnAllGeneratedTokens) -{ - - auto& comm = tensorrt_llm::mpi::MpiComm::world(); - auto const worldRank = comm.getRank(); - auto const worldSize = comm.getSize(); - auto const beamWidth = beamResult.beamWidth; - - std::unordered_map reqIdToBatchId; - std::unordered_map> tokens; - // std::unordered_map gGenIdIdTogContextId; - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(givenInput, modelIds.padId); - auto const* const givenInputData = tr::bufferCast(givenInput); - - auto const& inputShape = givenInput.getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - SizeType32 numRequests = static_cast(givenInputLengths.size()); - SizeType32 maxRequests = numRequests; - std::vector requests; - std::vector reqMaxNewTokens; - SizeType32 const numReturnSequences = 1; - - for (SizeType32 req = 0; req < maxRequests; ++req) - { - SizeType32 inputLen = givenInputLengths.at(req); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - SizeType32 endId = -1; - auto const* const seqBegin = givenInputData + req * maxInputLength; - VecTokens tokens(seqBegin, seqBegin + inputLen); - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - samplingConfig.setNumReturnSequences(numReturnSequences); - auto request = Request( - VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, samplingConfig, outConfig, endId); - request.setReturnAllGeneratedTokens(returnAllGeneratedTokens); - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_ONLY); - requests.emplace_back(std::move(request)); - } - - if (worldRank == 0) - { - std::vector reqIds; - - for (int i = 0; i < requests.size(); ++i) - { - std::vector resultTokens; - resultTokens.reserve(numReturnSequences); - for (SizeType32 seqIdx = 0; seqIdx < numReturnSequences; ++seqIdx) - { - resultTokens.emplace_back(beamWidth); - } - auto retReqId = executor.enqueueContext({requests[i]}, std::nullopt); - reqIds.push_back(retReqId.front()); - tokens[i] = std::move(resultTokens); - reqIdToBatchId[retReqId.front()] = i; - } - - int32_t numContextFinished = 0; - int contextIter = 0; - while (numContextFinished < maxRequests && contextIter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - - auto contextResponses = executor.awaitContextResponses(waitTime); - contextIter++; - numContextFinished += contextResponses.size(); - - for (auto&& responseWithId : contextResponses) - { - auto contextGid = responseWithId.gid; - int batchId = reqIdToBatchId[contextGid]; - auto&& request = requests[batchId]; - request.setRequestType(RequestType::REQUEST_TYPE_GENERATION_ONLY); - request.setContextPhaseParams(responseWithId.response.getResult().contextPhaseParams.value()); - executor.enqueueGeneration({request}, {responseWithId.gid}, std::nullopt); - } - } - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < maxRequests && iter < maxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitGenerationResponses(waitTime); - for (auto& responseWithId : responses) - { - numResponses++; - if (!responseWithId.response.hasError()) - { - auto result = responseWithId.response.getResult(); - numFinished += result.isFinal; - auto batchId = reqIdToBatchId.at(responseWithId.gid); - auto seqIdx = result.sequenceIndex; - - auto& contextLogits = result.contextLogits; - auto& genLogits = result.generationLogits; - auto& outputTokenIds = result.outputTokenIds; - - EXPECT_EQ(result.finishReasons.size(), beamWidth); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto& newTokens = outputTokenIds.at(beam); - auto& reqTokens = tokens.at(batchId).at(seqIdx).at(beam); - - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - // FinishReason is only supported for bw=1 and inflight batching. - if (beamWidth == 1 && batchingType == BatchingType::kINFLIGHT) - { - EXPECT_EQ(result.finishReasons.at(beam), - result.isFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - - auto& cumLogProbs = result.cumLogProbs; - auto& logProbs = result.logProbs; - auto& beamTokens = tokens.at(batchId).at(seqIdx); - testData.verifyLogProbs(outConfig.returnLogProbs, streaming, outConfig.excludeInputFromOutput, - givenInputLengths.at(batchId), beamWidth, beamTokens, cumLogProbs, logProbs, batchId, - flakyTestInfo); - - testData.validateContextLogits(outConfig.returnContextLogits, givenInputLengths.at(batchId), - beamWidth, contextLogits, vocabSizePadded, batchId); - testData.validateGenerationLogits(outConfig.returnGenerationLogits, result.isFinal, streaming, - outConfig.excludeInputFromOutput, givenInputLengths.at(batchId), reqMaxNewTokens.at(batchId), - beamWidth, beamTokens, genLogits, vocabSizePadded, batchId, returnAllGeneratedTokens); - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(responseWithId.gid) - + " has already been processed and was terminated."; - EXPECT_EQ(responseWithId.response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, maxWaitMs); - testData.verifyOutput(tokens, givenInputLengths, streaming, outConfig.excludeInputFromOutput, flakyTestInfo, - isSpeculativeDecoding, beamWidth, numReturnSequences, false); - } - comm.barrier(); -} - -TEST_P(DisaggParamsTest, DisaggTokenComparison) -{ - -#if ENABLE_MULTI_DEVICE - - if (!(tensorrt_llm::common::getEnvUseUCXKvCache())) - { - setenv("UCX_TLS", "^cuda_ipc", 1); // disable cuda_ipc for testing for mpi - } - else - { - setenv("UCX_TCP_CM_REUSEADDR", "y", - 1); // tests creates and destroies ucxCacheCommunicatoers frequently, so listener ports must be reused - } - auto const processNum = std::get<0>(GetParam()); - auto const modelNames = std::get<1>(GetParam()); - auto const participantIdsEachInstance = std::get<2>(GetParam()); // std::vector> - auto const participantDeviceIdsEachInstance = std::get<3>(GetParam()); // std::vector>; - auto const instanceRoles - = std::get<4>(GetParam()); // std::vector ; //1 is context , 0 is generation, 2 is mixed - auto const controllerRank = std::get<5>(GetParam()); - - // params_check - auto const& world_comm = tensorrt_llm::mpi::MpiComm::world(); - int const commRank = world_comm.getRank(); - int const commSize = world_comm.getSize(); - if (commSize != processNum) - { - GTEST_SKIP() << " need " << processNum << " processes but got " << commSize << " mpi processes, skip test."; - } - ASSERT_EQ(participantIdsEachInstance.size(), participantDeviceIdsEachInstance.size()); - SizeType32 instanceNum = participantIdsEachInstance.size(); - ASSERT_EQ(instanceNum, instanceRoles.size()); - ASSERT_EQ(instanceNum, modelNames.size()); - - std::unordered_set deviceIdsSet; - for (auto const& ids : participantDeviceIdsEachInstance) - { - for (auto const& id : ids) - { - deviceIdsSet.insert(id); - } - } - if (mDeviceCount < deviceIdsSet.size()) - { - GTEST_SKIP() << " need " << deviceIdsSet.size() << " devices but got " << mDeviceCount - << " devices, skip test."; - } - - ASSERT_GE(controllerRank, 0); - ASSERT_LT(controllerRank, commSize); - int ranksNum = 0; - std::unordered_map rankCounter; - std::unordered_map deviceCounter; - SizeType32 deviceRuseNum = 1; - bool isContext = false; - bool isGeneration = false; - std::vector participatntIds; - std::vector deviceIds; - std::string modelName; - bool isController = (commRank == controllerRank); - for (SizeType32 i = 0; i < instanceNum; i++) - { - auto const& ranksThisInstance = participantIdsEachInstance[i]; - auto const& devicesThisInstance = participantDeviceIdsEachInstance[i]; - - ASSERT_EQ(ranksThisInstance.size(), devicesThisInstance.size()); - SizeType32 rankNumThisInstance = ranksThisInstance.size(); - ASSERT_GT(rankNumThisInstance, 0); - ranksNum += rankNumThisInstance; - for (SizeType32 j = 0; j < rankNumThisInstance; j++) - { - rankCounter[ranksThisInstance[j]]++; - deviceCounter[devicesThisInstance[j]]++; - ASSERT_GE(rankCounter[ranksThisInstance[j]], 1); - deviceRuseNum = std::max(deviceCounter[devicesThisInstance[j]], deviceRuseNum); - ASSERT_GE(ranksThisInstance[j], 0); - ASSERT_LT(ranksThisInstance[j], commSize); - - if (commRank == ranksThisInstance[j]) - { - participatntIds = ranksThisInstance; - deviceIds = devicesThisInstance; - isContext = instanceRoles[i] == InstanceRole::kCONTEXT || instanceRoles[i] == InstanceRole::kMIXED; - isGeneration - = instanceRoles[i] == InstanceRole::kGENERATION || instanceRoles[i] == InstanceRole::kMIXED; - // modelName = isContext ? contextModel : genModel; - modelName = modelNames[i]; - } - } - } - ASSERT_GE(ranksNum, commSize); - - OutputConfig outConfig; - int const beamWidth = 1; - BeamResult beamResult{beamWidth}; - - bool streaming = false; - int const maxBeamWidth = 1; - ASSERT_TRUE(fs::exists(DATA_PATH)); - - fs::path modelPath; - // set defaults and adjust if needed by different models - fs::path inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded - bool isSpeculativeDecoding{false}; - - // NOTE: This can be used to disable checks for certain prompt batch entries - FlakyTestInfo flakyTestInfo; - - if (modelName == "gpt") - { - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - if (outConfig.returnContextLogits || outConfig.returnGenerationLogits) - { - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_GATHER_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_RESULT_FILE(); - beamResult.contextLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CONTEXT_LOGITS_FILE(); - beamResult.genLogitsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GENERATION_LOGITS_FILE(); - if (outConfig.returnLogProbs) - { - beamResult.cumLogProbsFile - = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_CUM_LOG_PROBS_FILE(); - beamResult.logProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_GATHER_LOG_PROBS_FILE(); - } - } - else - { - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_FILE(); - if (outConfig.returnLogProbs) - { - beamResult.cumLogProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_CUM_LOG_PROBS_FILE(); - beamResult.logProbsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_LOG_PROBS_FILE(); - } - } - } - else if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1" - || modelName == "llama_tp1_pp2_cp1" || modelName == "llama_tp2_pp1_cp1" || modelName == "llama_tp1_pp1_cp1") - { - inputPath = DATA_PATH / LLAMA_INPUT_FILE; - vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; - - auto const resultsPath - = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelIds.padId = LLAMA_PAD_ID; - modelIds.endId = LLAMA_END_ID; - if (modelName == "llama_tp4_pp1_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP4_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP4_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp2_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp1_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp2_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp1_cp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - } - else - { - TLLM_THROW("Unrecognized modelName"); - } - - // Warning: This should be the last check before running the test. - // It will initialize MPI which can take significant time. - if (modelName == "llama_tp4_pp1_cp1" || modelName == "llama_tp1_pp4_cp1" || modelName == "llama_tp2_pp2_cp1" - || modelName == "llama_tp1_pp2_cp1" || modelName == "llama_tp2_pp1_cp1") - { - if (outConfig.returnLogProbs || outConfig.returnContextLogits || outConfig.returnGenerationLogits) - { - GTEST_SKIP() << "Skipping logits and log probs tests for mpi runs"; - } - } - - // Returning logits will bring higher latency - if (streaming && (outConfig.returnContextLogits || outConfig.returnGenerationLogits)) - { - mMaxWaitMs = 20000; - } - - auto executorConfig = ExecutorConfig(maxBeamWidth); - FloatType freeGpuMemoryFraction = 0.9f / (deviceRuseNum); // context and gen instance run on same device - KvCacheConfig kvCacheConfig{true, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction}; - executorConfig.setKvCacheConfig(kvCacheConfig); - executorConfig.setRequestStatsMaxIterations(1000); - executorConfig.setCacheTransceiverConfig( - texec::CacheTransceiverConfig(texec::CacheTransceiverConfig::BackendType::DEFAULT)); - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - world_comm.barrier(); - auto disaggExecutor = tensorrt_llm::testing::disaggexecutor::DisaggExecutorLeader(modelPath, - ModelType::kDECODER_ONLY, executorConfig, isController, isContext, isGeneration, givenInputLengths.size(), - participatntIds, deviceIds, commRank); - - runDisaggTest(disaggExecutor, manager, *givenInput, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, - outConfig, isSpeculativeDecoding, mMaxWaitMs, executorConfig.getBatchingType(), false); - -#else - - GTEST_SKIP() << "Skipping DisaggExecutor Test"; - -#endif -} - -TEST_P(DisaggOrchestratorParamsTest, DisaggTokenComparison) -{ - -#if ENABLE_MULTI_DEVICE - - if (!(tensorrt_llm::common::getEnvUseUCXKvCache())) - { - setenv("UCX_TLS", "^cuda_ipc", 1); // disable cuda_ipc for testing for mpi - } - else - { - setenv("UCX_TCP_CM_REUSEADDR", "y", - 1); // tests creates and destroies ucxCacheCommunicatoers frequently, so listener ports must be reused - } - auto const processNum = std::get<0>(GetParam()); - auto const modelNames = std::get<1>(GetParam()); - auto const participantIdsEachInstance = std::get<2>(GetParam()); // std::vector> - auto const participantDeviceIdsEachInstance = std::get<3>(GetParam()); // std::vector>; - auto const instanceRoles = std::get<4>(GetParam()); // std::vector ; //1 is context , 0 is generation - auto const controllerRank = std::get<5>(GetParam()); - - // params_check - auto const& world_comm = tensorrt_llm::mpi::MpiComm::world(); - int const commRank = world_comm.getRank(); - int const commSize = world_comm.getSize(); - if (commSize != processNum) - { - GTEST_SKIP() << " need " << processNum << " processes but got " << commSize << " mpi processes, skip test."; - } - - bool spawnProcess = false; - if (commSize == 1) - { - spawnProcess = true; - if (mDeviceCount < 4) - { - GTEST_SKIP() << "DisaggExecutorTest requires at least 4 GPUs"; - } - ASSERT_TRUE(tensorrt_llm::common::getEnvUseUCXKvCache() || tensorrt_llm::common::getEnvUseNixlKvCache()); - } - - ASSERT_EQ(participantIdsEachInstance.size(), participantDeviceIdsEachInstance.size()); - SizeType32 instanceNum = participantIdsEachInstance.size(); - ASSERT_EQ(instanceNum, instanceRoles.size()); - ASSERT_EQ(instanceNum, modelNames.size()); - - std::unordered_set deviceIdsSet; - for (auto const& ids : participantDeviceIdsEachInstance) - { - for (auto const& id : ids) - { - deviceIdsSet.insert(id); - } - } - if (mDeviceCount < deviceIdsSet.size()) - { - GTEST_SKIP() << " need " << deviceIdsSet.size() << " devices but got " << mDeviceCount - << " devices, skip test."; - } - - ASSERT_GE(controllerRank, 0); - ASSERT_LT(controllerRank, commSize); - std::string modelName = modelNames[0]; - bool isController = (commRank == controllerRank); - std::vector contextModels; - std::vector genModels; - - auto getModelPath = [=](std::string modelNN) - { - fs::path retPath; - if (modelNN == "llama_tp4_pp1") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelNN == "llama_tp1_pp4") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - } - else if (modelNN == "llama_tp1_pp2") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp2-cp1-gpu"; - } - else if (modelNN == "llama_tp2_pp1") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp1-cp1-gpu"; - } - else if (modelNN == "llama_tp2_pp2") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - } - else if (modelNN == "llama_tp1_pp1") - { - retPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - return retPath; - }; - for (SizeType32 i = 0; i < instanceNum; i++) - { - if (instanceRoles[i] == InstanceRole::kCONTEXT) - { - contextModels.push_back(getModelPath(modelNames[i])); - } - else - { - genModels.push_back(getModelPath(modelNames[i])); - } - } - - OutputConfig outConfig; - int const beamWidth = 1; - BeamResult beamResult{beamWidth}; - - bool streaming = false; - int const maxBeamWidth = 1; - ASSERT_TRUE(fs::exists(DATA_PATH)); - - fs::path modelPath; - // set defaults and adjust if needed by different models - fs::path inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded - bool isSpeculativeDecoding{false}; - - // NOTE: This can be used to disable checks for certain prompt batch entries - FlakyTestInfo flakyTestInfo; - if (modelName == "llama_tp4_pp1" || modelName == "llama_tp1_pp4" || modelName == "llama_tp2_pp2" - || modelName == "llama_tp1_pp2" || modelName == "llama_tp2_pp1" || modelName == "llama_tp1_pp1") - { - inputPath = DATA_PATH / LLAMA_INPUT_FILE; - vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; - - auto const resultsPath - = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelIds.padId = LLAMA_PAD_ID; - modelIds.endId = LLAMA_END_ID; - if (modelName == "llama_tp4_pp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP4_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp4-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp4") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP4_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp4-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp2") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp1-cp1-gpu"; - } - else if (modelName == "llama_tp2_pp2") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp2-pp2-cp1-gpu"; - } - else if (modelName == "llama_tp1_pp1") - { - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP2_PP2_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - } - - else - { - TLLM_THROW("Unrecognized modelName"); - } - - // Warning: This should be the last check before running the test. - // It will initialize MPI which can take significant time. - if (modelName == "llama_tp4_pp1" || modelName == "llama_tp1_pp4" || modelName == "llama_tp2_pp2" - || modelName == "llama_tp1_pp2" || modelName == "llama_tp2_pp1") - { - if (outConfig.returnLogProbs || outConfig.returnContextLogits || outConfig.returnGenerationLogits) - { - GTEST_SKIP() << "Skipping logits and log probs tests for mpi runs"; - } - } - - // Returning logits will bring higher latency - if (streaming && (outConfig.returnContextLogits || outConfig.returnGenerationLogits)) - { - mMaxWaitMs = 20000; - } - - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - world_comm.barrier(); - auto contextNum = contextModels.size(); - auto genNum = genModels.size(); - // int deviceCount = -1; - // TLLM_CUDA_CHECK(cudaGetDeviceCount(&deviceCount)); - bool isOrchestrator = commRank == 0; - std::vector ctxExecutorConfigs; - std::vector genExecutorConfigs; - for (int in = 0; in < instanceNum; in++) - { - tensorrt_llm::executor::SchedulerConfig schedulerConfig(CapacitySchedulerPolicy::kMAX_UTILIZATION); - KvCacheConfig kvCacheConfig{true, std::nullopt, std::nullopt, std::nullopt, 0.2}; - - tensorrt_llm::executor::ExecutorConfig executorConfig(maxBeamWidth, schedulerConfig, kvCacheConfig); - tensorrt_llm::executor::OrchestratorConfig orchestratorConfig{ - isOrchestrator, PathUtil::EXECUTOR_WORKER_PATH(), nullptr, spawnProcess}; - - tensorrt_llm::executor::ParallelConfig parallelConfig{tensorrt_llm::executor::CommunicationType::kMPI, - tensorrt_llm::executor::CommunicationMode::kORCHESTRATOR, participantDeviceIdsEachInstance.at(in), - spawnProcess ? std::nullopt : std::optional>(participantIdsEachInstance.at(in)), - orchestratorConfig}; - executorConfig.setParallelConfig(parallelConfig); - executorConfig.setCacheTransceiverConfig( - texec::CacheTransceiverConfig(texec::CacheTransceiverConfig::BackendType::DEFAULT)); - if (in < contextNum) - { - ctxExecutorConfigs.push_back(executorConfig); - } - else - { - genExecutorConfigs.push_back(executorConfig); - } - } - auto disaggExecutor - = DisaggExecutorOrchestrator(contextModels, genModels, ctxExecutorConfigs, genExecutorConfigs, true, true); - - runDisaggTest(disaggExecutor, manager, *givenInput, modelIds, flakyTestInfo, streaming, vocabSizePadded, beamResult, - outConfig, isSpeculativeDecoding, mMaxWaitMs, BatchingType::kINFLIGHT, false); - -#else - - GTEST_SKIP() << "Skipping DisaggExecutor Test"; - -#endif -} - -TEST_P(ConditionalDisaggParamsTest, DisaggTokenComparison) -{ -#if ENABLE_MULTI_DEVICE - if (!tensorrt_llm::common::getEnvUseUCXKvCache()) - { - setenv("UCX_TLS", "^cuda_ipc", 1); // disable cuda_ipc for testing for mpi - } - auto constexpr processNum = 2; - auto constexpr deviceNum = 2; - auto const& modelName = std::get<0>(GetParam()); - auto constexpr controllerRank = 0; - - // params_check - auto const& world_comm = tensorrt_llm::mpi::MpiComm::world(); - int const commRank = world_comm.getRank(); - int const commSize = world_comm.getSize(); - if (commSize != processNum) - { - GTEST_SKIP() << " need " << processNum << " processes but got " << commSize << " mpi processes, skip test."; - } - if (mDeviceCount < deviceNum) - { - GTEST_SKIP() << " need " << deviceNum << " devices but got " << mDeviceCount << " devices, skip test."; - } - - bool isContext = commRank == 0; - bool isGeneration = commRank == 1; - std::vector participatntIds = {commRank}; - std::vector deviceIds = {commRank}; - bool isController = (commRank == controllerRank); - - OutputConfig outConfig(false, false, false, false, false, false); - int const beamWidth = 1; - BeamResult beamResult{beamWidth}; - - bool streaming = false; - int const maxBeamWidth = 1; - ASSERT_TRUE(fs::exists(DATA_PATH)); - - fs::path modelPath; - // set defaults and adjust if needed by different models - fs::path inputPath = DATA_PATH / "input_tokens.npy"; - ModelIds modelIds{50256, 50256}; - SizeType32 vocabSizePadded{50257}; // gpt vocabSizePadded - bool isSpeculativeDecoding{false}; - - // NOTE: This can be used to disable checks for certain prompt batch entries - FlakyTestInfo flakyTestInfo; - - if (modelName == "gpt") - { - auto const resultsPath - = GPT_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelPath = GPT_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_FILE(); - } - else if (modelName == "llama_tp1_pp1_cp1") - { - inputPath = DATA_PATH / LLAMA_INPUT_FILE; - vocabSizePadded = LLAMA_VOCAB_SIZE_PADDED; - - auto const resultsPath - = LLAMA_DATA_PATH / ((beamWidth == 1) ? "sampling" : "beam_search_" + std::to_string(beamWidth)); - modelIds.padId = LLAMA_PAD_ID; - modelIds.endId = LLAMA_END_ID; - beamResult.resultsFile = resultsPath / PathUtil::FP16_PLUGIN_PACKED_PAGED_RESULT_TP1_PP1_FILE(); - modelPath = LLAMA_MODEL_PATH / PathUtil::FP16_GPT_ATTENTION_PACKED_PAGED_DIR() / "tp1-pp1-cp1-gpu"; - } - else - { - TLLM_THROW("Unrecognized modelName"); - } - - auto executorConfig = ExecutorConfig(maxBeamWidth); - FloatType freeGpuMemoryFraction = 0.9f; - KvCacheConfig kvCacheConfig{true, std::nullopt, std::nullopt, std::nullopt, freeGpuMemoryFraction}; - executorConfig.setKvCacheConfig(kvCacheConfig); - executorConfig.setRequestStatsMaxIterations(1000); - executorConfig.setCacheTransceiverConfig( - texec::CacheTransceiverConfig(CacheTransceiverConfig::BackendType::DEFAULT)); - auto manager = tr::BufferManager(std::make_shared()); - auto const& givenInput = tr::utils::loadNpy(manager, inputPath.string(), tr::MemoryType::kCPU); - auto [givenInputLengths, nbGivenInputs, maxInputLength] = getGivenInputLengths(*givenInput, modelIds.padId); - world_comm.barrier(); - auto executor = tensorrt_llm::testing::disaggexecutor::DisaggExecutorLeader(modelPath, ModelType::kDECODER_ONLY, - executorConfig, isController, isContext, isGeneration, givenInputLengths.size(), participatntIds, deviceIds, - commRank); - - std::unordered_map reqIdToBatchId; - std::unordered_map> tokens; - auto const* const givenInputData = tr::bufferCast(*givenInput); - - auto const& inputShape = givenInput->getShape(); - ASSERT_EQ(inputShape.nbDims, 2); - ASSERT_GT(inputShape.d[0], 0); - - // Load expected outputs for each beam width value - auto testData = TestData::loadTestData(beamResult, *givenInput, beamWidth, manager, outConfig, modelIds); - auto const maxSeqLen = testData.maxSeqLen; - - // Load expected outputs and inputs - SizeType32 numRequests = static_cast(givenInputLengths.size()); - SizeType32 maxRequests = numRequests; - std::vector requests; - std::vector reqMaxNewTokens; - SizeType32 const numReturnSequences = 1; - - for (SizeType32 req = 0; req < maxRequests; ++req) - { - SizeType32 inputLen = givenInputLengths.at(req); - auto maxNewTokens = maxSeqLen - maxInputLength; - reqMaxNewTokens.push_back(maxNewTokens); - SizeType32 endId = -1; - auto const* const seqBegin = givenInputData + req * maxInputLength; - VecTokens tokens(seqBegin, seqBegin + inputLen); - auto samplingConfig = tensorrt_llm::executor::SamplingConfig(beamWidth); - samplingConfig.setNumReturnSequences(numReturnSequences); - auto request = Request( - VecTokens(seqBegin, seqBegin + inputLen), maxNewTokens, streaming, samplingConfig, outConfig, endId); - request.setReturnAllGeneratedTokens(false); - // setting request type to context/full by condition - if (req % 2 == 0) - { - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_ONLY); - } - else - { - request.setRequestType(RequestType::REQUEST_TYPE_CONTEXT_AND_GENERATION); - } - requests.emplace_back(std::move(request)); - } - - if (isController) - { - std::vector reqIds; - - for (int i = 0; i < requests.size(); ++i) - { - std::vector resultTokens; - resultTokens.reserve(numReturnSequences); - for (SizeType32 seqIdx = 0; seqIdx < numReturnSequences; ++seqIdx) - { - resultTokens.emplace_back(beamWidth); - } - auto retReqId = executor.enqueueRequests({requests[i]}); - reqIds.push_back(retReqId.front()); - tokens[i] = std::move(resultTokens); - reqIdToBatchId[retReqId.front()] = i; - } - - // Get the new tokens for each requests - int32_t numFinished = 0; - int iter = 0; - SizeType32 numResponses = 0; - while (numFinished < maxRequests && iter < mMaxWaitMs) - { - std::chrono::milliseconds waitTime(1); - auto responses = executor.awaitResponses(waitTime); - for (auto& response : responses) - { - numResponses++; - if (!response.hasError()) - { - auto result = response.getResult(); - numFinished += result.isFinal; - auto batchId = reqIdToBatchId.at(response.getRequestId()); - auto seqIdx = result.sequenceIndex; - - auto& outputTokenIds = result.outputTokenIds; - - EXPECT_EQ(result.finishReasons.size(), beamWidth); - for (SizeType32 beam = 0; beam < beamWidth; ++beam) - { - auto& newTokens = outputTokenIds.at(beam); - auto& reqTokens = tokens.at(batchId).at(seqIdx).at(beam); - - reqTokens.insert(reqTokens.end(), newTokens.begin(), newTokens.end()); - // FinishReason is only supported for bw=1 and inflight batching. - if (beamWidth == 1 && executorConfig.getBatchingType() == BatchingType::kINFLIGHT) - { - EXPECT_EQ(result.finishReasons.at(beam), - result.isFinal ? FinishReason::kLENGTH : FinishReason::kNOT_FINISHED); - } - } - } - else - { - // Allow response with error only if awaitResponse processed a terminated request id - std::string err = "ReqId " + std::to_string(response.getRequestId()) - + " has already been processed and was terminated."; - EXPECT_EQ(response.getErrorMsg(), err); - } - } - ++iter; - } - EXPECT_LT(iter, mMaxWaitMs); - testData.verifyOutput(tokens, givenInputLengths, streaming, outConfig.excludeInputFromOutput, flakyTestInfo, - isSpeculativeDecoding, beamWidth, numReturnSequences, false); - } - world_comm.barrier(); -#else - GTEST_SKIP() << "Skipping DisaggExecutor Test"; -#endif -} - -INSTANTIATE_TEST_SUITE_P(GptDisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0, 1) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptDisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED, InstanceRole::kMIXED}), // instanceRoles - testing::Values(1) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptSingleDeviceDisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptSingleDeviceDisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"gpt", "gpt"}), // modelNames - testing::Values(std::vector>{{0}, {1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED, InstanceRole::kMIXED}), // instanceRoles - testing::Values(1) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(GptConditionalDisaggSymmetricExecutorTest, ConditionalDisaggParamsTest, - testing::Combine(testing::Values("gpt")), generateTestNameCondDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConditionalDisaggSymmetricExecutorTest, ConditionalDisaggParamsTest, - testing::Combine(testing::Values("llama_tp1_pp1_cp1")), generateTestNameCondDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaTP2DisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp2_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaPP2DisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp1_pp2_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // participantIdsEachInstance - testing::Values(std::vector>{{1, 0}, {3, 2}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaTP2DisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaPP2DisaggSymmetricExecutorMixedTest, DisaggParamsTest, - testing::Combine( // - testing::Values(2), // processNum - testing::Values(std::vector{"llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kMIXED}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaTP2PP2DisaggSymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(8), // processNum - testing::Values(std::vector{"llama_tp2_pp2_cp1", "llama_tp2_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1, 2, 3}, {4, 5, 6, 7}}), // participantIdsEachInstance - testing::Values(std::vector>{{2, 3, 0, 1}, {2, 3, 0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConPP2GenTP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp1_pp2_cp1", "llama_tp2_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // (1,0) (2,3) // participantIdsEachInstance - testing::Values(std::vector>{{1, 0}, {2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2GenPP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(4), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0, 1}, {2, 3}}), // (0,1), (3,2)// participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {3, 2}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2PP2GenPP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp2_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values( - std::vector>{{0, 1, 2, 3}, {4, 5}}), // (2,3,0,1) , (5,4)// participantIdsEachInstance - testing::Values(std::vector>{{2, 3, 0, 1}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2PP2GenTP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp2_cp1", "llama_tp2_pp1_cp1"}), // modelNames - testing::Values( - std::vector>{{0, 1, 2, 3}, {4, 5}}), // (2,3,0,1), (4,5)// participantIdsEachInstance - testing::Values(std::vector>{{2, 3, 0, 1}, {0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); -INSTANTIATE_TEST_SUITE_P(LlamaConTP2PP1GenTP2PP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp2_pp2_cp1"}), // modelNames - testing::Values( - std::vector>{{0, 1}, {2, 3, 4, 5}}), // (0,1) , (4,5,2,3)%4// participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {0, 1, 2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaConTP2GenPP4DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{"llama_tp2_pp1_cp1", "llama_tp1_pp4_cp1"}), // modelNames - testing::Values( - std::vector>{{4, 5}, {0, 1, 2, 3}}), // (4,5) ,(3,2,1,0)// participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {3, 2, 1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon4TP1Gen1TP4DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(8), // processNum - testing::Values(std::vector{"llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", - "llama_tp1_pp1_cp1", "llama_tp4_pp1_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2}, {3}, {4, 5, 6, 7}}), // participantIdsEachInstance - testing::Values( - std::vector>{{0}, {1}, {2}, {3}, {0, 1, 2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(4) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2TP2AndPP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{ - "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp2_pp1_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2, 3}, {4, 5}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {2, 3}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2PP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(6), // processNum - testing::Values(std::vector{ - "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp1_pp2_cp1", "llama_tp1_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2, 3}, {4, 5}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {3, 2}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon4TP1Gen1TP2PP2DisaggAsymmetricExecutorTest, DisaggParamsTest, - testing::Combine( // - testing::Values(8), // processNum - testing::Values(std::vector{"llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", "llama_tp1_pp1_cp1", - "llama_tp1_pp1_cp1", "llama_tp2_pp2_cp1"}), // modelNames - testing::Values(std::vector>{{0}, {1}, {2}, {3}, {4, 5, 6, 7}}), // participantIdsEachInstance - testing::Values( - std::vector>{{0}, {1}, {2}, {3}, {2, 3, 0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(4) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2TP2DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values( - std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp2_pp1", "llama_tp2_pp1"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4}, {5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {2, 3}, {0, 1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); -// for disaggOrchestrator 1->0, 2->1, 3->2, 4->3, 5->0, 6->1 - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP2Gen2TP1DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values( - std::vector{"llama_tp2_pp1", "llama_tp2_pp1", "llama_tp1_pp1", "llama_tp1_pp1"}), // modelNames - testing::Values(std::vector>{{1, 2}, {3, 4}, {5}, {6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {2, 3}, {0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2PP2DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values( - std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp1_pp2", "llama_tp1_pp2"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4}, {5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {3, 2}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen1TP2PP2DisaaggOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(7), // processNum - testing::Values(std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp2_pp2"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4, 5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {0, 1, 2, 3}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{ - InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP2Gen2TP1DisaggSpawnOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(1), // processNum - testing::Values( - std::vector{"llama_tp2_pp1", "llama_tp2_pp1", "llama_tp1_pp1", "llama_tp1_pp1"}), // modelNames - testing::Values(std::vector>{{1, 2}, {3, 4}, {5}, {6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0, 1}, {2, 3}, {0}, {1}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); - -INSTANTIATE_TEST_SUITE_P(LlamaCon2TP1Gen2PP2DisaggSpawnOrchestrator, DisaggOrchestratorParamsTest, - testing::Combine( // - testing::Values(1), // processNum - testing::Values( - std::vector{"llama_tp1_pp1", "llama_tp1_pp1", "llama_tp1_pp2", "llama_tp1_pp2"}), // modelNames - testing::Values(std::vector>{{1}, {2}, {3, 4}, {5, 6}}), // participantIdsEachInstance - testing::Values(std::vector>{{0}, {1}, {3, 2}, {1, 0}}), // participantDeviceIdsEachInstance - testing::Values(std::vector{InstanceRole::kCONTEXT, InstanceRole::kCONTEXT, - InstanceRole::kGENERATION, InstanceRole::kGENERATION}), // instanceRoles - testing::Values(0) // controllerRank - ), - generateTestNameDisaggParams); diff --git a/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md b/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md index 22d4688c503d..18c5ce06c0f2 100644 --- a/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md +++ b/docs/source/blogs/Best_perf_practice_on_DeepSeek-R1_in_TensorRT-LLM.md @@ -414,13 +414,13 @@ Note that, `max_batch_size` and `max_num_tokens` can easily affect the performan Generally, you should make sure that `max_batch_size` is not too low to bottleneck the throughput, and `max_num_tokens` needs to be large enough so that it covers the max input sequence length of the samples in dataset, as mentioned in below section "WIP: Chunked context support on DeepSeek models". -For more details on `max_batch_size` and `max_num_tokens`, refer to [Tuning Max Batch Size and Max Num Tokens](../legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md). +For more details on `max_batch_size` and `max_num_tokens`, refer to the performance tuning guidance in the developer documentation. ### MLA chunked context MLA currently supports the chunked context feature on both Hopper and Blackwell GPUs. You can use `--enable_chunked_context` to enable it. This feature is primarily designed to reduce TPOT (Time Per Output Token). The default chunk size is set to `max_num_tokens`. If you want to achieve a lower TPOT, you can appropriately reduce the chunk size. However, please note that this will also decrease overall throughput. Therefore, a trade-off needs to be considered. -For more details on `max_num_tokens`, refer to [Tuning Max Batch Size and Max Num Tokens](../legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md). +For more details on `max_num_tokens`, refer to the performance tuning guidance in the developer documentation. ### Out of memory issues diff --git a/docs/source/blogs/XQA-kernel.md b/docs/source/blogs/XQA-kernel.md index d4018966cd2d..456024a10ed5 100644 --- a/docs/source/blogs/XQA-kernel.md +++ b/docs/source/blogs/XQA-kernel.md @@ -2,7 +2,7 @@ XQA kernel provides optimization for [MQA](https://arxiv.org/abs/1911.02150) and [GQA](https://arxiv.org/abs/2305.13245v3) during generation phase. It also provides optimization for beam search. Using tensor cores for acceleration, reducing data loading and conversion, it delivers increased throughput within the same latency budget. Increased throughput allows serving greater number of user requests while providing the same experience. -Support matrix and usage flags are described in [docs/source/legacy/advanced/gpt-attention](../legacy/advanced/gpt-attention.md#xqa-optimization). +Support matrix and usage flags are described in the [Attention](../features/attention.md) documentation. **Increased Throughput:** Looking at the Throughput-Latency curves below, we see that the enabling of XQA optimization increases throughput. Higher throughput equates to serving more users, and we can see that TPOT on the Y-axis flattens out when XQA gets enabled. diff --git a/docs/source/commands/trtllm-build.rst b/docs/source/commands/trtllm-build.rst deleted file mode 100644 index d209d3eac9ca..000000000000 --- a/docs/source/commands/trtllm-build.rst +++ /dev/null @@ -1,7 +0,0 @@ -trtllm-build -=========================== - -.. argparse:: - :module: tensorrt_llm.commands.build - :func: parse_arguments - :prog: trtllm-build diff --git a/docs/source/index.rst b/docs/source/index.rst index 9d92a8770148..07fde255c0c9 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -127,10 +127,9 @@ Welcome to TensorRT LLM's Documentation! .. toctree:: :maxdepth: 2 - :caption: Use TensorRT Engine - :hidden: + :caption: Migration - legacy/tensorrt_quickstart.md + legacy/tensorrt-backend-removal.md Indices and tables ================== diff --git a/docs/source/legacy/advanced/disaggregated-service.md b/docs/source/legacy/advanced/disaggregated-service.md deleted file mode 100644 index ac6cbad472e9..000000000000 --- a/docs/source/legacy/advanced/disaggregated-service.md +++ /dev/null @@ -1,86 +0,0 @@ -(disaggregated-service)= - -# Disaggregated-Service (Prototype) - -```{note} -Note: -This feature is currently in prototype, and the related API is subject to change in future versions. -``` -Currently TRT-LLM supports `disaggregated-service`, where the context and generation phases of a request can run on different executors. TRT-LLM's disaggregated service relies on the executor API, please make sure to read the [executor page](executor.md) before reading the document. - -For more information on disaggregated service in LLM inference, one can refer to papers such as [DistServe](https://arxiv.org/abs/2401.09670), [SplitWise](https://arxiv.org/abs/2311.18677). - -An [architectural and performance overview](../../../docs/source/blogs/tech_blog/blog05_Disaggregated_Serving_in_TensorRT-LLM.md), as well as [usage examples](../../../examples/disaggregated/README.md), are provided. - -## Environment Variables - -TRT-LLM uses some environment variables to control the behavior of disaggregated service. - - -* `TRTLLM_DISABLE_KV_CACHE_TRANSFER_OVERLAP`: If set to `1`, generationExecutor will not overlap KV cache transfer with model inference. The default value is `0`. - -* `TRTLLM_ENABLE_KVCACHE_RECEIVE_PARALLEL`: When the generation rank receives KV cache from multiple context ranks within a single context instance, it will receive KV cache from each rank sequentially. If set to `1`, the generation rank will receive KV cache from each rank within one context instance in parallel. The default value is `0`. - -* `TRTLLM_REQUEST_KV_CACHE_CONCURRENT`: If set to `1`, generationExecutor prepares independent resources for each context executor to receive KV cache, requests whose KV cache are received from different context executors will be processed concurrently. If set to `0`, the generation executor will reuse the same resource to process KV cache transfer for each request sequentially, reducing the resources used by KV cache transmission and thereby lowering the risk of running out of memory. The default value is `0`. - -* `TRTLLM_TRY_ZCOPY_FOR_KVCACHE_TRANSFER`: TRT-LLM typically copies non-contiguous data into a temporary buffer before sending KV cache. If set to `1`, TRT-LLM will attempt to directly transmit each KV cache block, eliminating extra copies. The default value is `0`. - -* `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE`: By default, TRT-LLM uses a `stream-ordered memory allocator` to allocate temporary buffers. If this environment variable is set to #Size, TRT-LLM will use `cudaMalloc` to allocate buffer of size #Size for KV cache transmission. The default value is `512MB`. Users can set `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE=1GB` to allocate a 1 GB buffer with `cudaMalloc` for KV cache transmission. - -* `TRTLLM_KVCACHE_TRANSFER_USE_ASYNC_BUFFER`: If set to `1`, TRT-LLM will use `cudaMallocAsync` to allocate buffers for KV cache transmission. The default value is `0`. This environment variable only takes effect when `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE` is greater than 0. - -* `TRTLLM_KVCACHE_SEND_MAX_CONCURRENCY_NUM`: The maximum number of concurrent KV cache sends. The default value is `1`. This environment variable only takes effect when `TRTLLM_KVCACHE_TRANSFER_BUFFER_SIZE` is greater than 0. - -There are some other useful environment variables that may help when encountering failures or performance issues. - -* `NCCL_GRAPH_MIXING_SUPPORT`: With the default value `1`, the CUDA driver may create too many CUDA streams while working with one CUDA graph, leading to performance drop. Setting it to `0` will reduce the number of CUDA streams, but please make sure there are no other NCCL ops outside the one CUDA graph, otherwise it's unsafe. - -* `UCX_MAX_RNDV_RAILS`: With the default value `2`, UCX attempts to use two InfiniBand (IB) NIC devices per GPU for Rendezvous (RNDV) transfers. When both the context and generation instances enable tensor- and expert-parallel (TEP), multiple TP ranks may transfer KV cache concurrently. Because each TP rank can use up to two NIC devices, some NIC devices can be shared across GPUs, causing contention and reduced throughput. Setting `UCX_MAX_RNDV_RAILS=1` can reduce contention in this case. - -## Troubleshooting and FAQ - -### General FAQs - -*Q. What are the limitations of disaggregated-service in TRT-LLM?* - -A. Currently, only `decoder-only engine` and `beamWidth=1` are supported, and the KV cache at each layer of the model is required to be homogeneous, with the same data type and the same number of attention heads. - -*Q. Is the engine used by disaggregated-service different from other engines?* - -A. No. There are no special requirements for the arguments to build engine. - -*Q. Do the engines used by the context executor and generation executor need to be the same?* - -A. No. The engines used by context executor and generation executor can be different, and their parallelism can be heterogeneous, i.e., TP,PP can be different, and TRT-LLM will handle the heterogeneity of KV cache. - -*Q. Does TRT-LLM support running multiple context executor instances and generation executor instances?* - -A. Yes. TRT-LLM supports running multiple context executors and generation executors at the same time, and each executor can use different engine, but it is the user's responsibility to route requests to different executors and manage `requestId`. - -*Q. Can an executor handle both context-only requests and generation-only requests?* - -A. Yes, but it's not recommended, TRT-LLM does not implement proper scheduling for the case where the executor handles mixed context-only requests and generation-only requests, it's better to run context-only requests and generation-only requests on different executors. - -*Q. Does disaggregated-service in TRT-LLM support multi-gpu and multi-node?* - -A. Yes, it's recommended that different executor use different GPUs. We support context-only executor and generation-only executor run on same node or different nodes. The `participantIds` and `deviceIds` used by each executor need to be explicitly set by the user, and the `participantIds` of each executor must not be intersecting. - -### Debugging FAQs - -*Q. Does TRT-LLM support using GPU direct RDMA for inter-node KV Cache transfer?* - -A. Yes, TRT-LLM supports using GPU direct RDMA for inter-node KV cache transfer. - -*Q. What causes the substantial bandwidth fluctuations in kvCache transfers, especially during the first few requests following service initialization?* - -A. The communication for kvCache transfer between executors are established dynamically. The connection establishment process incurs significant overhead, which explains the apparently lower kvCache transfer bandwidth observed during the initial requests after service startup. This lower bandwidth reflects the inclusion of connection establishment overhead. When conducting benchmarks, it is recommended to perform a warm-up phase to ensure accurate performance measurements. - -*Q. When my servers are running on different NVLink domains, some servers hang or have a lower performance. How to fix that?* - -A. NVLink domain can be found with `nvidia-smi -q` in the `Fabric.ClusterUUID` field. A few UCX environment variables can be adjusted when your servers have different NVLink domains: - -* `UCX_CUDA_IPC_ENABLE_MNNVL`: Set to `n`. This also can reduce UCX timeout error messages like `UCX ERROR cuMemImportFromShareableHandle failed: invalid resource handle`, although these errors don't necessarily cause your trtllm-serve to fail. - -* `UCX_NET_DEVICES`: Check if this is set correctly, or unset this variable to allow UCX to use all possible devices. - -* `UCX_RNDV_SCHEME`: Set to `get_zcopy` or `put_zcopy` on GB200 for better performance. The default value is `auto`. diff --git a/docs/source/legacy/advanced/executor.md b/docs/source/legacy/advanced/executor.md deleted file mode 100644 index 1f59d0802a16..000000000000 --- a/docs/source/legacy/advanced/executor.md +++ /dev/null @@ -1,153 +0,0 @@ -(executor)= - -# Executor API - -TensorRT-LLM includes a high-level C++ API called the Executor API which allows you to execute requests -asynchronously, with in-flight batching, and without the need to define callbacks. - -A software component (referred to as "the client" in the text that follows) can interact -with the executor using the API defined in the [`executor.h`](source:cpp/include/tensorrt_llm/executor/executor.h) file. -For details about the API, refer to the {ref}`_cpp_gen/executor.rst`. - -The following sections provide an overview of the main classes defined in the Executor API. - -## API - -### The Executor Class - -The `Executor` class is responsible for receiving requests from the client, and providing responses for those requests. The executor is constructed by providing a path to a directory containing the TensorRT-LLM engine or buffers containing the engine and the model JSON configuration. The client can create requests and enqueue those requests for execution using the `enqueueRequest` or `enqueueRequests` methods of the `Executor` class. Enqueued requests will be scheduled for execution by the executor, and multiple independent requests can be batched together at every iteration of the main execution loop (a process often referred to as continuous batching or iteration-level batching). Responses for a particular request can be awaited for by calling the `awaitResponses` method, and by providing the request id. Alternatively, responses for any requests can be awaited for by omitting to provide the request id when calling `awaitResponses`. The `Executor` class also allows to cancel requests using the `cancelRequest` method and to obtain per-iteration and per-request statistics using the `getLatestIterationStats`. - -### The Request Class - -The `Request` class is used to define properties of the request, such as the input token ids and the maximum number of tokens to generate. The `streaming` parameter can be used to indicate if the request should generate a response for each new generated tokens (`streaming = true`) or only after all tokens have been generated (`streaming = false`). Other mandatory parameters of the request include the sampling configuration (defined by the `SamplingConfig` class) which contains parameters controlling the decoding process and the output configuration (defined by the `OutputConfig` class) which controls what information should be included in the `Result` for a particular response. - -Optional parameters can also be provided when constructing a request such as a list of bad words, a list of stop words, a client id, or configurations objects for prompt tuning, LoRA, or speculative decoding, or a number of sequences to generate for example. - -### The Response Class - -The `awaitResponses` method of the `Executor` class returns a vector of responses. Each response contains the request id associated with this response, and also contains either an error or a `Result`. Check if the response has an error by using the `hasError` method before trying to obtain the `Result` associated with this response using the `getResult` method. - -### The Result Class - -The `Result` class holds the result for a given request. It contains a Boolean parameter called `isFinal` that indicates if this is the last `Result` that will be returned for the given request id. It also contains the generated tokens. If the request is configured with `streaming = false` and `numReturnSequences = 1`, a single response will be returned, the `isFinal` Boolean will be set to `true` and all generated tokens will be included in the `outputTokenIds`. If `streaming = true` and `numReturnSequences = 1` is used, a `Result` will include one or more tokens (depending on the request `returnAllGeneratedTokens` parameter) except the last result and the `isFinal` flag will be set to `true` for the last result associated with this request. - -The request `numReturnSequences` parameter controls the number of output sequences to generate for each prompt. When this option is used, the Executor will return at least `numReturnSequences` responses for each request, each containing one Result. In beam search (`beamWidth > 1`), the number of beams to be returned will be limited by `numReturnSequences` and the `sequenceIndex` attribute of the `Result` class will always be zero. Otherwise, in sampling (`beamWidth = 1`), the `sequenceIndex` attribute indicates the index of the generated sequence in the result (`0 <= sequenceIndex < numReturnSequences`). It contains a Boolean parameter called `isSequenceFinal` that indicates if this is the last result for the sequence and also contains a Boolean parameter `isFinal` that indicates when all sequences for the request have been generated. When `numReturnSequences = 1`, `isFinal` is identical to `isSequenceFinal`. - -Here is an example that shows how a subset of 3 responses might look like for `numReturnSequences = 3`: - -``` -Response 1: requestId = 1, Result with sequenceIndex = 0, isSequenceFinal = false, isFinal = false -Response 2: requestId = 1, Result with sequenceIndex = 1, isSequenceFinal = true, isFinal = false -Response 3: requestId = 1, Result with sequenceIndex = 2, isSequenceFinal = false, isFinal = false -``` - -In this example, each response contains one result for different sequences. The `isSequenceFinal` flag of the second Result is set to true, indicating that it is the last result for `sequenceIndex = 1`, however, the isFinal flag of each Response is set to false because sequences 0 and 2 are not completed. - -### Sending Requests with Different Beam Widths - -The executor can process requests with different beam widths if the following conditions are met: - -- The model was built with a `max_beam_width > 1`. -- The executor is configured with a `maxBeamWidth > 1` (the configured `maxBeamWidth` must be less than or equal to the model's `max_beam_width`). -- The requested beam widths are less than or equal to the configured `maxBeamWidth`. - -The executor may schedule successive requests with the same beam width at the same time. For successive requests with two different beam widths, `x` and `y`, requests with beam width `y` are not scheduled until all requests with beam width `x` have been processed. -This allows the runtime to reconfigure itself for a new beam width when no requests are in flight. The reconfiguration happens automatically each time requests with a different beam width than currently configured are detected. Waiting for previous requests to finish and reconfiguring the runtime may cause significant overhead and reduce overall throughput. - -### Controlling output with Logits Post-Processor - -Optionally, you can alter the logits produced by the network by providing an instance of `Executor::LogitsPostProcessorConfig`. For instance, this feature can be used to generate JSON formatted output. {cpp:class}`Executor::LogitsPostProcessorConfig ` specifies a map of named callbacks in the following form - -```cpp -std::unordered_map)>> -``` - -The map key is the name associated with that logits post-processing callback. Each request can then specify the name of the logits post-processor to use for that particular request, if any. - -The first argument to the callback is the request id, second is the logits tensor, third are the tokens produced by the request so far, fourth is the operation stream used by the logits tensor, and last one is an optional client id. The callback returns a modified tensor of logits. Multiple requests can share same client id and callback can use different logic based on client id. - -You must use the stream to access the logits tensor. For example, to perform an addition with a bias tensor, the addition operation is enqueued on that stream. Alternatively, you can call `stream->synchronize()`, however, that will slow down the entire execution pipeline. - -The executor also includes a {cpp:class}`LogitsPostProcessorBatched ` method that enables altering logits of multiple requests in a batch. The batched method allows further optimizations and reduces callback overheads. - -```cpp -std::function const&, std::vector&, std::vector> const&, StreamPtr const&, std::vector> const&)> -``` - -A single batched callback can be specified in `LogitsPostProcessorConfig`. Each request can opt to apply this callback by specifying the name of the logits post-processor as `Request::kBatchedPostProcessorName`. - -Note: Neither callback variant is supported with the `STATIC` batching type for the moment. - -In a multi-GPU run, the callback is invoked on all ranks in the first tensor-parallel group, by default. To ensure correct execution, replicate the client-side state that is accessed by the callback on these ranks. If replication is expensive or infeasible, use `LogitsPostProcessorConfig::setReplicate(false)` to invoke the callback only on rank 0. The executor broadcasts the sampled tokens internally to ensure correct execution. - -### Structured output with guided decoding -Guided decoding controls the generation outputs to be amenable to pre-defined structured formats, e.g., JSON or XML. Currently, guided decoding is supported with the [XGrammar](https://github.com/mlc-ai/xgrammar) backend. - -To enable guided decoding, a valid instance of `GuidedDecodingConfig` must be provided when constructing `Executor`. `GuidedDecodingConfig` should be constructed with some tokenizer information, including `encodedVocab`, `tokenizerStr` (optional) and `stopTokenIds` (optional). Given a Hugging Face tokenizer, these can be extracted by: - -```python -encoded_vocab = tokenizer.get_vocab() -encoded_vocab = [token for token, _ in sorted(encoded_vocab.items(), key=lambda x: x[1])] -tokenizer_str = tokenizer.backend_tokenizer.to_str() -stop_token_ids = [tokenizer.eos_token_id] -``` - -Refer to [`tensorrt_llm/llmapi/tokenizer.py`](source:tensorrt_llm/llmapi/tokenizer.py) for more details. You may dump these materials to disk, and reload them to C++ runtime for use. - -Each request can be optionally specified with a `GuidedDecodingParams`, which defines the desired structured format. Currently, it supports four types: -* `GuidedDecodingParams::GuideType::kJSON`: The generated text is amenable to JSON format; -* `GuidedDecodingParams::GuideType::kJSON_SCHEMA`: The generated text is amenable to JSON format with additional restrictions; -* `GuidedDecodingParams::GuideType::kREGEX`: The generated text is amenable to regular expression; -* `GuidedDecodingParams::GuideType::kEBNF_GRAMMAR`: The generated text is amenable to the extended Backus-Naur form (EBNF) grammar. - -The latter three types should be used with the schema/regex/grammar provided to `GuidedDecodingParams`. - -### Obtaining Arbitrary Output Tensors -The executor API gives the user the possibility to read the arbitrary outputs from the model. For example, it is possible to obtain hidden states or logits. - -#### Mark Tensors As Output -For a tensor to be obtainable using this feature, it needs to be marked as an output in the model definition (e.g. add `topk_logits.mark_output("TopKLogits")`) before building the TRT engine. - -#### Configure The Executor -Assuming the TensorRT engine you are planning to use has a tensor named `TopKLogits` marked as output, you should then configure the `Executor` to read from this output tensor by passing its name to the `ExecutorConfig` configuration object: -```cpp -auto const executorConfig = ExecutorConfig{}; - -std::vector additionalOutputs{ - executor::AdditionalModelOutput{"TopKLogits", /*whether or not to get the output for the context too */ true}}; -executorConfig.setAdditionalModelOutputs(additionalOutputs); - -// ... set more configuration options if needed -// ... create the `Executor` instance -``` - -### Request Additional Output -Construct a request to enqueue in the executor to query this tensor output: -```cpp -executor::Request request{requestTokens, parameters.maxOutputLength, true, executor::SamplingConfig{}, - executor::OutputConfig{false, false, false, true, false, false, additionalOutputs}}; -executor.enqueueRequest(request); -``` - -The output can be found at the `additionalOutputs` property of each response. - -#### Note on context outputs - -If KV cache reuse is enabled, context outputs will not contain outputs for the part of the context that has been reused. This part of the outputs can only be obtained from the prior request with the same prefix that generated this part of the KV cache. - -## C++ Executor API Example - -Two C++ examples are provided that show how to use the Executor API and can be found in the [`examples/cpp/executor`](source:examples/cpp/executor/) folder. - -## Python Bindings for the Executor API - -Python bindings for the Executor API are also available to use the Executor API from Python. The Python bindings are defined in [bindings.cpp](source:cpp/tensorrt_llm/nanobind/executor/bindings.cpp) and once built, are available in package `tensorrt_llm.bindings.executor`. Running `help('tensorrt_llm.bindings.executor')` in a Python interpreter will provide an overview of the classes available. - -In addition, three Python examples are provided to demonstrate how to use the Python bindings to the Executor API for single and multi-GPU models. They can be found in [`examples/bindings`](source:examples/bindings). - -## In-flight Batching with the Triton Inference Server - -A Triton Inference Server C++ [backend](https://github.com/triton-inference-server/tensorrtllm_backend) is provided with TensorRT-LLM that -includes the mechanisms needed to serve models using in-flight batching. That -backend is also a good starting example of how to implement in-flight batching using -the TensorRT-LLM C++ Executor API. diff --git a/docs/source/legacy/advanced/expert-parallelism.md b/docs/source/legacy/advanced/expert-parallelism.md deleted file mode 100644 index 55d948a80cae..000000000000 --- a/docs/source/legacy/advanced/expert-parallelism.md +++ /dev/null @@ -1,30 +0,0 @@ -(expert-parallelism)= - -# Expert Parallelism in TensorRT-LLM - -## Mixture of Experts (MoE) - -Mixture of Experts (MoE) architectures have become widespread, with models such as [Mistral Mixtral 8×7B](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1). Specifically, MoE’s structure supports multiple parallel feed-forward neural-network (FFN) layers (called experts) in place of the single FFN layer in a dense model. When tokens arrive, the router layer selects the top-k experts for each token, and the corresponding hidden state of each token is dispatched to those experts. As a result, there are multiple tokens’ hidden states that are dispatched to each expert. - -moe_structure - -the MOE structure in Switch Transformer: [https://arxiv.org/pdf/2101.03961.pdf](https://arxiv.org/pdf/2101.03961.pdf) - -## Tensor Parallel vs Expert Parallel - -Parallelism on multi-GPUs is necessary if the MoE model cannot be accommodated by a single GPU’s memory. We support three kinds of parallel patterns for MoE structure: Tensor Parallel (default pattern), Expert Parallel, and a hybrid of the two. - -tensor parallel vs expert parallel - -Tensor Parallel evenly splits each expert’s weight and distributes them to different GPUs, which means each GPU holds partial weight of all experts, While Expert Parallel evenly distributes some of the experts’ full weight to different GPUs, which means each GPU holds part of the experts’ full weight. As a result, each GPU rank in the Tensor Parallel group receives all tokens’ hidden states for all experts, then computes using the partial weights, while for Expert Parallel, each GPU rank only receives part of tokens’ hidden states for experts on this rank, then computes using the full weights. - -When both Tensor Parallel and Expert Parallel are enabled, each GPU handles a portion of the expert weights matrices (as in EP mode) and these weights are further sliced across multiple GPUs (as in TP mode). This hybrid approach aims to balance the workload more evenly across GPUs, enhancing efficiency and reducing the likelihood of bottlenecks associated with EP mode alone. - - -## How to Enable - -The default parallel pattern is Tensor Parallel. You can enable Expert Parallel or hybrid parallel by setting `--moe_tp_size` and `--moe_ep_size` when calling `convert_checkpoint.py`. If only `--moe_tp_size` is provided, TRT-LLM will use Tensor Parallel for the MoE model; if only `--moe_ep_size` is provided, TRT-LLM will use Expert Parallel; if both are provided, the hybrid parallel will be used. - -Ensure the product of `moe_tp_size` and `moe_ep_size` is equal to `tp_size`, since the total number of MoE parallelism across all GPUs must match the total number of parallelism in other parts of the model. - -The other parameters related to the MoE structure, such as `num_experts_per_tok` (TopK in previous context) and `num_local_experts`, can be found in the model’s configuration file, such as the one for [Mixtral 8x7B model](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1/blob/main/config.json). diff --git a/docs/source/legacy/advanced/gpt-attention.md b/docs/source/legacy/advanced/gpt-attention.md deleted file mode 100644 index d5b4d29718bf..000000000000 --- a/docs/source/legacy/advanced/gpt-attention.md +++ /dev/null @@ -1,343 +0,0 @@ -(gpt-attention)= - -# Multi-Head, Multi-Query, and Group-Query Attention - -This document details the implementation of multi-head attention (MHA), -multi-query attention (MQA) and group-query attention (GQA) for auto-regressive -GPT-like models in TensorRT-LLM. As a quick reminder, the multi-head attention -is the sequence of a batched matmul, a softmax and another batched matmul -described in the -[Attention Is All You Need](https://arxiv.org/abs/1706.03762) article. [Multi-query Attention (MQA)](https://arxiv.org/abs/1911.02150) and [Group-query Attention (GQA)](https://arxiv.org/abs/2307.09288) are variants of MHA that use fewer, so-called, K/V head than the number of query heads. TensorRT-LLM, MHA, MQA and GQA are implemented by the operator [`tensorrt_llm.functional.gpt_attention`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/functional.py). - -## Important Note - -As discussed below, the current implementation supports two input modes: Padded -and packed (non-padded). As the packed mode is always more memory-efficient and -faster than the padded mode, ***support for padded mode may be removed in the -future***. - -## Padded and Packed Tensors - -In TensorRT-LLM, the GPT attention operator supports two different types -of QKV inputs: Padded and packed (i.e. non padded) inputs. The mode is -determined by the global configuration parameter `remove_input_padding` defined -in [`tensorrt_llm.plugin`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/plugin/plugin.py). - -When padding is enabled (that is, `remove_input_padding` is `False`), the sequences -that are shorter than the `max_sequence_length` are padded to that maximum -length. It may result in excessive memory consumption as well as unneeded -computations on padding tokens (in the various matrix multiplications that -surround the MHA block). - -To overcome that problem, TensorRT-LLM supports a mode without padding where -the different tokens are packed together and the user provides the operator -with a 1D tensor containing the lengths of the different sequences. It is -recommended that users to always use packed mode (and support for the padded -mode may be removed in the future). - -## Context and Generation Phases - -The GPT attention operator encapsulates different implementations for both -context and generation phases in auto-regressive models like GPT. - -### Context Phase - -If the `context_fmha_type` is set to `disabled` (refer to -[`tensorrt_llm.plugin`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/tensorrt_llm/plugin/plugin.py)), -the implementation maps to a sequence of GPU kernels that will store the -intermediate `Q*K^T` tensor in memory before calling the softmax operator. It -is the slowest method and the memory footprint is significant (quadratically -depends on the sequence length). - -Otherwise, if `context_fmha_type` is set to a `enabled` or -`enabled_with_fp32_acc` (accumulation in the first batched matmul is forced to -FP32), that function will trigger a kernel that performs the MHA/MQA block -using a single kernel. For short sequences, that kernel uses a vanilla -implementation of MHA/MQA. For larger sequences, this kernel uses the Flash -Attention algorithm as described in -[FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://arxiv.org/abs/2205.14135) -and -[FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://arxiv.org/abs/2307.08691). - -Currently, the implementation triggers extra kernels that apply pre-processing -to the elements (like RoPE) and populate the KV cache (see below). In a future -release, the number of such kernels is planned on being reduced in order to -improve the overall performance. - -#### FP8 Context FMHA - -When FP8 quantization is activated, the attention can be further accelerated by -enabling FP8 Context FMHA (`use_fp8_context_fmha = enable`). - -FP8 Paged Context FMHA is also supported with the fp8 quantization workflow. -You need to specify `use_fp8_context_fmha = enable` and -`use_paged_context_fmha = enable` at the same time. - -Please be aware that this feature is only supported on Ada and Hopper. - -### Generation Phase - -The generation phase is implemented using a single kernel called the masked -multi-head attention in TensorRT-LLM. That kernel is able to apply -pre-processing on the Q, K, and V elements on-the-fly: adds the QKV bias, applies -RoPE, and performs dequantization and quantization. TensorRT-LLM will continue to add (or -enable) additional features in future releases. For example, enable the support -for IA3. - -_The masked MHA kernel has a special version that distributes the work across -multiple CUDA thread-blocks on the GPU for cases where the GPU occupancy is -low. That mode called multi-block is turned on by default starting from TRT-LLM 0.13, -and can be disabled using `--multi_block_mode=False` during runtime. -Users are recommended to test that mode in scenarios where both the batch -size and the number of heads in the model are relatively small. The exact -definition of small in that context will depend on the model of the GPU and is -hard to predict but to provide with a rule of thumb, it is worth testing that -mode when `batch_size * num_heads` is less than the number of multi-processors -on the GPU (that suggestion may evolve in the future as more research is -conducted and the software improves)_. - -_Note that even if the multi-block mode is enabled, the attention operator will -not immediately trigger the multi-block version of the GPU kernel. There is a -minimum number of tokens (input + generated) that are required for the -multi-block version to become more efficient than the "vanilla" implementation -that uses a single CUDA thread-block per head. It is controlled by an internal -heuristic._ - -Another note is that as the masked MHA kernels use shared memory size -proportional to sequence length, so there can be some cases that GPU's shared -memory is not enough when multi-block mode is not enabled. To get masked MHA -kernel work in these cases, multi-block mode is forced on and a warning log is -printed. - -#### XQA Optimization - -Another optimization for MQA/GQA in the generation phase is called XQA optimization. - -Support matrix of the XQA optimization: - - FP16 / BF16 compute data type. - - FP16 / BF16 / FP8 / INT8 KV cache data type. - - Paged KV cache (8 / 16 / 32 / 64 / 128 tokens per block). - -By default, this is enabled. Note that a heuristic algorithm -is also used to decide whether to use XQA kernel or masked MHA kernel to get -better performance. -If you want to use that kernel whenever possible, set `TRTLLM_FORCE_XQA=1` to force use of the XQA kernel when the model config is supported. -Supported configurations can be found using the `shouldUse` function of the `DecoderXQARunner` class in -`cpp/tensorrt_llm/kernels/decoderMaskedMultiheadAttention/decoderXQARunner.h`. - - -(inflight-batching)= - -## In-flight Batching - -TensorRT-LLM supports in-flight batching of requests (also known as continuous -batching or iteration-level batching) for higher serving throughput. With this feature, -sequences in context phase can be processed together with sequences in -generation phase. The purpose of that technique is to better interleave -requests to reduce latency as well as make better use of the GPUs. -For efficiency reasons (1), the support for inflight batching ***requires the -input tensors to be packed (no padding)***. - -***In the current implementation, the sequences that are going through the -context phase must be before the sequences in the generation phase in the input -tensor. For example, for sequences `S0`, `S1` and `S2`, if `S0` and `S2` are in -context phase (and `S1` in generation), tokens from `S0` and `S2` must appear -before the tokens of `S1` in the input tensor***. The constraint may or may not -be relaxed in a future version. - -_(1) Padding sequences in the generation phase, that contain a single token, to -the length of the maximum input sequence is inefficient use of resources_. - - - -## Chunked Context - -In the original state, the common behavior was to process all context tokens at -once. This feature splits the context into several chunks. In this way, the -context chunks can be batched with more tokens during the generation phase, -which is expected to increase the total throughput. Chunking contexts also removes -constraints on input length. To enable this feature, the FMHA paged kv-cache also -needs to be enabled. Except for the last one, the size of the context chunk needs -to be an integer multiple of the kv-cache block size. Refer to -[the chunked context feature](../../features/long-sequence.md) for usage. - -## KV Cache - -In the generation phase, a common optimization is to provide the MHA kernel -with a cache containing the values of the past K and V elements that have -already been computed. That cache is known as the KV cache. TensorRT-LLM uses -that technique to accelerate its generation phase. In TensorRT-LLM, there is -one KV cache per Transformer layer, which means that there are as many KV -caches as layers in a model. The current version of TensorRT-LLM supports two -different types of KV caches: **contiguous** and **paged** KV caches. - -### Contiguous KV Cache - -The contiguous KV cache is a monolithic tensor. Its shape is: -``` -[max_batch_size * max_beam_width, 2, num_heads, max_seqlen, hidden_dim_per_head]. -``` - -That implementation uses a lot more memory than needed when the sequences are -shorter than the maximum sequence length (even if they end up close to the -limit after the generation of many output tokens, it may take a lot of steps to -reach that point). - -### Paged KV Cache - -The paged KV cache decomposes the KV cache into blocks that are distributed to -the different requests by a cache manager during processing. That cache manager -keeps track of the sequences, allocate new blocks from a pool and recycle those -blocks when required. See the simplified implementation of -[`tensorrt_llm.runtime.KVCacheManager`](source:tensorrt_llm/runtime/kv_cache_manager.py). -A more efficient C++ implementation is included in the -[Batch Manager](source:cpp/include/tensorrt_llm/batch_manager). - -## INT8/FP8 KV Caches - -In its current implementation, even if the rest of the network runs in INT8 or -FP8, the GPT attention operator works with FP32, FP16, and BFloat16 inputs and -outputs. However, TensorRT-LLM supports INT8 and FP8 -(`kv_cache_quant_mode=QuantMode.INT8_KV_CACHE` and -`kv_cache_quant_mode=QuantMode.FP8_KV_CACHE`) KV caches. - -The GPT attention operator populates the KV cache. When INT8 or FP8 KV caches -are enabled, the input values have to be quantized to 8 bits using a scaling -factor. For quantization, the scaling factor is stored in the -`kv_cache_scaling_factor` tensor. Its shape is `[1]` and only per-tensor -quantization is supported in the current version. Quantization uses inversed scale -since it does multiply as `fp_value * (1.0 / kv_cache_scaling_factor)` in plugin. - -During generation, the values read from the cache are dequantized on-the-fly in -the MHA/MQA kernel, dequantization can be described as -`quantized_value * kv_cache_scaling_factor`. - - -## Sliding Window Attention, Cyclic (Rolling Buffer) KV Cache - -TensorRT-LLM has a feature called `Cyclic KV Cache`, which treats the kv cache -as a circular buffer. This means that it only stores the kv cache for the last N -tokens, where N is determined by the `max_attention_window_size` parameter in -`GenerationSession.setup`. You can see examples of this in the `run.py` or -`summarize.py` files. When the cache is full, new tokens’ kv cache will -overwrite the "least recently used" caches. - -In the context phase, if the input length surpasses the `max_attention_window_size`, -`Sliding Window Attention` will be activated. This serves the same function as -the `sliding window_size`. - -This feature helps to reduce the memory footprint of the kv cache when -dealing with very long sequences. - -The feature, which allows different `max_attention_window_size` values -for each layer, is also supported. To utilize this feature, simply provide an -`int32 torch.Tensor` or `list` to the `GenerationSession.setup` when using python -runtime session, or provide a vector to the `KvCacheConfig` when using cpp runtime. -If the number of the provided elements is less than the number of layers, the provided -tensor/list/vector will be repeated multiple times to the number of layers and then be -saved as a new tensor. This tensor will serve as the buffer for `max_attention_window_size`, -setting unique values for each layer. However, it’s important to note that the -memory allocation for the kv cache still relies on the buffer’s maximum value. - -_Note that the cyclic kv cache feature doesn't work with beam searching currently as -the context kv cache are shared across beams. - -## StreamingLLM - -The StreamingLLM feature uses a window attention to perform efficient and stable LLM -on long texts, which means that only `N` tokens need to be stored in the KV cache. -Similar to the cyclic KV cache feature in TensorRT-LLM, `max_attention_window_size` -parameter is used to determine `N`. Different from the cyclic KV cache feature, -the first `S` tokens, called sink tokens, are always kept in the attention window, -where `S` is determined by `sink_token_length` parameter in `GenerationSession.setup`. -But in context phase, the self-attention is dense in the official implementation of -StreamingLLM, and it uses all of the tokens for computation and only saves `N` tokens -to the KV cache. - -In addition, the relative position embedding is also changed in StreamingLLM. -When determining the relative distance and adding positional information to tokens, -StreamingLLM uses the positions within the cache rather than those in the original text. - -`streamingllm` flag is used to enable this feature. - -## Beam-Search - -The GPT attention operator supports beam-search. In the context phase, a single -beam is computed per input sequence. In the generation phase, the MHA/MQA/GQA -kernel uses an additional tensor to reconstruct the correct path for each beam. -That tensor is called the `cache_indirection`. Its shape is `[batch_size, -beam_width, max_seqlen]`. - -For a sequence `si`, a beam `bi` and a token `ti`, the element -`cache_indirection[si][bi][ti]` is an integer between `0` and `beam_width-1` -that indicates which path in the beam to read the K and V elements from in the -KV cache. This tensor is populated in the sampling stage. - -## Input QKV tensor - -The input QKV tensor packs the Q, K and V tensors (concatenated along the last -dimension) after the projection of the hidden states. It is a 3D tensor. RoPE -and quantization to INT8 or FP8 (when needed) are performed by the GPT -attention operator. - -In padded mode, its shape is `[batch_beam_size, max_seqlen, 3 * hidden_dim]` -where `batch_beam_size` is the batch size (number of sequences) for the context -phase and the batch size multiplied by the beam width for the generation phase. -Having different beam widths per sequence in padded mode is not supported. - -In packed mode, its shape is `[num_tokens, 3 * hidden_dim]` where -`num_tokens` is the total number of tokens in the batch. For the sequences in -context phase, the number of tokens of a sequence corresponds to its input -length (even if the beam width is greater than `1` for beam search). For the -sequences in generation phase, there are `beam_width` tokens per sequence. The -beam width can be different for each sequence. - -In other words, the pseudo-code to compute the number of tokens is: - -```python -num_tokens = 0 - -# Add the length of each sequence in context phase. -for seq in context_phase: - num_tokens += seq.length - -# Add the width of the beam for each sequence in generation phase. -for seq in generation_phase: - num_tokens += seq.beam_width -``` - -### Rotary Positional Embedding (RoPE) - -The GPT attention operation can perform the computation of the Rotary -Positional Embedding (RoPE). When that operation is enabled, -`rotary_embedding_dim` is set to a value greater than 0, it is fused with other -operations. The GPT operator supports GPT-NeoX and GPT-J forms of RoPE by -setting `position_embedding_type` to `PositionEmbeddingType.rope_gpt_neox` -or `PositionEmbeddingType.rope_gptj`. - -### ALiBi - -The GPT attention operator can apply ALiBi to the result of the `Q*K^T` -product. The bias is computed on-the-fly from the ALiBi slopes in the optimized -kernel. - -### Scaling factor(s) - -In MHA, the output of the `Q*K^T` product is scaled by a constant value that -is computed as: - -``` -norm_factor = 1.f / (q_scaling * sqrt(head_size)). -``` - -### Cross Attention - -On top of the MHA as self attention needed by GPT-style decoder-only models, `gpt_attention` also supports cross attention. - -This enables using `gpt_attention` in a broader aspect as a generic decoder component. For example, the Encoder-Decoder model uses `gpt_attention` to issue both the self attention and cross attention modules in its Decoder. - -### Relative Attention Bias (RAB) - -Relative attention bias (RAB) is a kind of relative position modeling, adding an attention bias (`Q*K^T+bias`) according to relative positions. RAB is a lightweight method to include the information of relative positions, and is used in the popular Encoder-Decoder model [T5](https://huggingface.co/docs/transformers/model_doc/t5) and also other models in the T5 family. - -RAB is supported in two modes: i) regular mode which user passes in relative attention bias computed ahead of MHA. ii) implicit mode which computes the relative attention bias on the fly in MHA. The implicit mode suits the case when the relative attention bias is too large to fit in memory and can be turned on by passing in `max_distance`. diff --git a/docs/source/legacy/advanced/gpt-runtime.md b/docs/source/legacy/advanced/gpt-runtime.md deleted file mode 100644 index e384c45175e4..000000000000 --- a/docs/source/legacy/advanced/gpt-runtime.md +++ /dev/null @@ -1,202 +0,0 @@ -(gpt-runtime)= - -# C++ GPT Runtime - -TensorRT-LLM includes a C++ component to execute TensorRT engines built with -the Python API as described in the {ref}`architecture-overview` section. -That component is called the C++ runtime. - -The API of the C++ runtime is composed of the classes declared in -[`cpp/include/tensorrt_llm/runtime`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/include/tensorrt_llm/runtime) and -implemented in [`cpp/tensorrt_llm/runtime`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/runtime). - -Even if the different components described in that document mention GPT in -their name, they are not restricted to this specific model. Those classes can -be used to implement auto-regressive models like BLOOM, GPT-J, GPT-NeoX or -LLaMA, for example. - -Complete support of encoder-decoder models, like T5, will be added to -TensorRT-LLM in a future release. An experimental version, only in Python for -now, can be found in the [`examples/models/core/enc_dec`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) folder. - -## Overview - -Runtime models are described by an instance of the -[`ModelConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime//modelConfig.h) -class and a pointer to the TensorRT engine that must be -executed to perform the inference. -The environment is configured through the -[`WorldConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/worldConfig.h) -(that name comes from -[MPI](https://en.wikipedia.org/wiki/Message_Passing_Interface) and its "famous" -`MPI_COMM_WORLD` default communicator). -The [`SamplingConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/samplingConfig.h) -class encapsulates parameters that control the -[generation](https://huggingface.co/blog/how-to-generate) of new tokens. - -### Model Configuration - -The model configuration is an instance of the -[`ModelConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime//modelConfig.h) class. -That class encapsulates the following parameters (they are declared as private -member variables and exposed through getters and setters): - - * `vocabSize`, the size of the vocabulary, - * `numLayers`, the number of layers in the model, - * `numHeads`, the number of heads in the attention block, - * `numKvHeads`, the number of heads for K and V in the attention component. - When the number of K/V heads is the same as the number of (Q) heads, the - model uses multi-head attention. When the number of K/V heads is 1, it uses - multi-query attention. Otherwise, it uses group-query attention. Refer to {ref}`gpt-attention` for more information, - * `hiddenSize`, the size of the hidden dimension, - * `dataType`, the datatype that was used to build the TensorRT engine and that - must be used to run the model during inference, - * `useGptAttentionPlugin`, indicates if the {ref}`gpt-attention` operator was compiled using the - [GPT Attention plugin](https://github.com/NVIDIA/TensorRT-LLM/tree/main/cpp/tensorrt_llm/plugins/gptAttentionPlugin), - * `inputPacked`, indicates that the input must be packed (or padded when set - to `false`). For performance reasons, it is recommended to always use packed, - even if its default is set to `false` (will be changed in a future release). - Refer to {ref}`gpt-attention` for more information, - * `pagedKvCache`, indicates if the K/V cache uses paging. - Refer to {ref}`gpt-attention` for more information, - * `tokensPerBlock`, is the number of tokens in each block of the K/V cache. - It's relevant when the paged K/V cache is enabled. By default, the value is - 64. Refer to {ref}`gpt-attention` for more information, - * `quantMode`, controls the quantization method. Refer to {ref}`precision` for more information. - * `maxBatchSize`, indicates the maximum batch size that the TensorRT engine - was built for, - * `maxInputLen`, the maximum size of the input sequences, - * `maxSequenceLen`, the maximum total size (input+output) of the sequences. - -### World Configuration - -Familiarity with -[MPI](https://en.wikipedia.org/wiki/Message_Passing_Interface), is not required -to utilize the TensorRT-LMM C++ runtime. There are two main things -you need to know: -* The C++ Runtime in TensorRT-LLM uses -[processes](https://en.wikipedia.org/wiki/Process_(computing)) to execute -TensorRT engines on the different GPUs. Those GPUs can be located on a single -node as well as on different nodes in a cluster. Each process is called a -*rank* in MPI. -* The ranks are grouped in communication groups. The -TensorRT-LLM C++ Runtime calls that group the *world*. - -The world configuration is an instance of the -[`WorldConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/worldConfig.h) -class, which encapsulates the following parameters: - -* `tensorParallelism`, the number of ranks that collaborate together to - implement Tensor Parallelism (TP). With TP, each GPU performs computations for - all the layers of the model. Some of those computations are distributed - across the GPU. TP is more balanced than Pipeline Parallelism (PP), in most cases, but - requires higher bandwidth between the GPUs. It is the recommended setting in - the presence of NVLINK between GPUs, -* `pipelineParallelism`, the number of ranks that collaborate together to - implement Pipeline Parallelism (PP). With PP, each GPU works on a subset of - consecutive layers. Communications between the GPUs happen only at the - boundaries of the subsets of layers. It is harder to guarantee the full - utilization of the GPUs with PP but it requires less memory bandwidth. It - is the recommended setting in the absence of NVLINK between GPUs, -* `rank`, the unique identifier of the rank, -* `gpusPerNode`, indicates the number of GPUs on each node. Having that - information allows the C++ runtime to optimize communications between GPUs in - a node (like taking advantage of the - [NVLINK](https://www.nvidia.com/en-us/data-center/nvlink/) - interconnect between GPUs of an A100 - [DGX](https://www.nvidia.com/en-us/data-center/dgx-platform/) - node). - -### Sampling Parameters - -The [`SamplingConfig`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/samplingConfig.h) -class encapsulates parameters that control the -[generation](https://huggingface.co/blog/how-to-generate) of new tokens. -A comparison of selecting decoding method is listed as the table below (`X` means it is not supported yet). -Except for the `beamWidth` parameter, all the fields are optional and the -runtime will use a default value if no values are provided by the user. For -vector fields, the TensorRT-LLM runtime supports one value per sequence (that is, -the vector contains `batchSize` values). If all the sequences use the same -value for a given parameter, the vector can be limited to a single element -(that is, `size() == 1`). - -| Method name in HF | Condition in HF | Method name in TRT-LLM | Condition in TRT-LLM | -| :------------------------------: | :---------------------------------------------------: | :--------------------: | :--------------------------------------------: | -| assisted decoding | `assistant_model` or `prompt_lookup_num_tokens!=None` | X | | -| beam-search decoding | `num_beams>1` and `do_sample=False` | beam search | `beamWidth > 1` | -| beam-search multinomial sampling | `num_beams>1` and `do_sample=True` | X | | -| constrained beam-search decoding | `constraints!=None` or `force_words_ids!=None` | X | | -| contrastive search | `penalty_alpha>0` and `top_k>1` | X | | -| diverse beam-search decoding | `num_beams>1` and `num_beam_groups>1` | X | | -| greedy decoding | `num_beams=1` and `do_sample=False` | sampling | `beamWidth == 1` and `topK=0` and `topP=0.0f` | -| multinomial sampling | `num_beams=1` and `do_sample=True` | sampling | `beamWidth == 1` and (`topK>0` or `topP>0.0f`) | - -***General*** - -| Name in TRT-LLM | Description | Data type | Range of value | Default value | Name in HF | -| :-----------------: | :-------------------------------------------------------------------------------: | :-----------: | :---------------------------------------------------------------------------------------: | :---------------------------------------------------: | :--------------------: | -| `temperature` | modulation of logits in sampling workflow | List\[Float\] | \[0.0f, $+\infty$\) | `1.0f` (no modulation) | `temperature` | -| `minLength` | lower-bound on the number of tokens generated | List\[Int\] | \[0, $+\infty$\) | `0` (no effect (the first generated token can be EOS) | `min_length` | -| `repetitionPenalty` | penalize repetitive tokens
multiplicative, irrespective of appearances count | List\[Float\] | \[0.0f, $+\infty$\)
`< 1.0f` encourages repetition
`> 1.0f` discourages it | `1.0f` (no effect) | `repetition_penalty` | -| `presencePenalty` | penalize existed tokens
additive, irrespective of appearances count | List\[Float\] | \($-\infty$, $+\infty$\)
`< 0.0f` encourages repetition
`> 0.0f` discourages it | `0.0f` (no effect) | no | -| `frequencyPenalty` | penalize existed tokens
additive, dependent on appearances count | List\[Float\] | \($-\infty$, $+\infty$\)
`< 0.0f` encourages repetition
`> 0.0f` discourages it | `0.0f` (no effect) | no | -| `noRepeatNgramSize` | | List\[Int\] | \[0, $+\infty$\)
`> 0` all ngrams of that size can only occur once | `0` (no effect) | `no_repeat_ngram_size` | - -* The tokens of input prompt are included during adopting `repetitionPenalty`, `presencePenalty`, and `frequencyPenalty` onto logits. - -* The parameters `repetitionPenalty`, `presencePenalty`, and `frequencyPenalty` are not mutually exclusive. - -***Sampling*** - -| Name in TRT-LLM | Description | Data type | Range of value | Default value | Name in HF | -| :-------------: | :---------------------------------------------------------------------: | :-----------: | :---------------: | :--------------: | :--------: | -| `randomSeed` | random seed for random number generator | Int64 | \[0, 2^64-1\] | `0` | no | -| `topK` | the number of logits to sample from | List\[Int\] | \[0, 1024\] | `0` | `top_k` | -| `topP` | the top-P probability to sample from | List\[Float\] | \[0.0f, 1.0f\] | `0.0f` | `top_p` | -| `topPDecay` | the decay in the `topP` algorithm | List\[Float\] | \(0.0f, 1.0f\] | `1.0f` | no | -| `topPMin` | the decay in the `topP` algorithm | List\[Float\] | \(0.0f, 1.0f\] | `1.0e-6,f` | no | -| `topPResetIds` | the decay in the `topP` algorithm | List\[Int\] | \[-1, $+\infty$\) | `-1` (no effect) | no | -| `minP` | scale the most likely token to determine the minimum token probability. | List\[Float\] | \[0.0f, 1.0f\] | `0.0` (no effect) | `min_p` | - - * If setting `topK = 0` and `topP = 0.0f`, greedy search is performed. - * If setting `topK > 0` and `topP = 0.0f`, `topK` tokens of highest probabilities will become the candidates of sampling (named `TopK sampling` in TRT-LLM). - * If setting `topK = 0` and `topP > 0.0f`, tokens will be sorted with probability descendly, then the tokens with highest probabilities which the accumulated probability larger than `topP` will become the candidates of sampling (named `TopP sampling` in TRT-LLM). - * If setting `topK > 0` and `topP > 0.0f`, `topK` tokens of highest probabilities will be selected, then those selected tokens will be sorted with probability descendly and their probability will be normalized, then the tokens with highest normalized probabilities which the accumulated probability larger than `topP` will become the candidates of sampling (named `TopKTopP sampling` in TRT-LLM) - - * If different `topK` values are provided for the different sequences in the batch, the performance of the implementation will depend on the largest value. For efficiency reasons, we recommend to batch requests with similar `topK` values together. - - * `topPDecay`, `topPMin` and `topPResetIds` are explained in - [_Factuality Enhanced Language Models for Open-Ended Text Generation_](https://arxiv.org/abs/2206.04624). - `topPDecay` is the decay, `topPMin` is the lower-bound and `topPResetIds` indicates where to reset the decay. - - * `minP` is explained in [_Turning Up the Heat: Min-p Sampling for Creative and Coherent LLM Outputs_](https://arxiv.org/abs/2407.01082). - - * TensorRT-LLM does not generate all possible tokenizations of a word. Therefore, stop words may appear in the output if there are multiple ways to tokenize a stop word and the token sequence in the output differs from the one in `stopWords`. - -***Beam-search*** - -| Name in TRT-LLM | Description | Data type | Range of value | Default value | Name in HF | -| :-----------------------: | :-----------------------------: | :-----------------: | :----------------------: | :-----------------------: | :-----------------: | -| `beamWidth` | width for beam-search algorithm | Int | \[0, 1024\] | `0` (disable beam search) | `beam_width` | -| `beamSearchDiversityRate` | diversity of generated tokens | List\[Float\] | \[0, $+\infty$\) | `0.0f` | `diversity_penalty` | -| `lengthPenalty` | penalize longer sequences | List\[Float\] | \[0, $+\infty$\) | `0.0f` | `length_penalty` | -| `earlyStopping` | see description below | List\[Int\] | \($-\infty$, $+\infty$\) | `0` | `early_stopping` | -| `beamWidthArray` | see description below | List\[List\[Int\]\] | \[0, 1024\] | `` | no | - - * Beam-search algorithm: [beam search](https://en.wikipedia.org/wiki/Beam_search). - * Parameter `diversity_penalty` in HF is only used for `diverse beam-search decoding` (or named `Group-Beam-Search`), which is not supported by TRT-LLM yet. - * If setting `earlyStopping = 1`, decoding will stop once `beamWidth` finished sentences are generated. - * If setting `earlyStopping = 0`, decoding will keep going until no better sentences (with better score) can be generated. - * If setting `earlyStopping` to other values, decoding will stop only depending on `lengthlengthPenalty`. - * `beamWidthArray` is a list of beam width for each step. Using `beamWidthArray = [20,40,80]` as an example, -beam width will be 20 for the first step, 40 for second step, 80 for the later all steps. - * The `beamWidth` parameter is a scalar value. It means that in this release of -TensorRT-LLM, it is not possible to specify a different width for each input -sequence. This limitation is likely to be removed in a future release. - -### Internal Components - -The [`TllmRuntime`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/tensorrt_llm/runtime/tllmRuntime.h) is in charge of the execution of the TensorRT engine. -The `TllmRuntime` class is an internal component and you are not expected to use that class directly. -The [`GptDecoder`](https://github.com/NVIDIA/TensorRT-LLM/blob/main/cpp/include/tensorrt_llm/runtime/gptDecoder.h) generates tokens from the logits. -The `GptDecoder` can be used directly to implement a custom generation loop and for use cases that cannot be satisfied by the TRT-LLM implementation. diff --git a/docs/source/legacy/advanced/graph-rewriting.md b/docs/source/legacy/advanced/graph-rewriting.md deleted file mode 100644 index 40552a0497d9..000000000000 --- a/docs/source/legacy/advanced/graph-rewriting.md +++ /dev/null @@ -1,197 +0,0 @@ -(graph-rewriting)= - - -# Graph Rewriting Module - -TensorRT-LLM uses a declarative approach to define neural networks and contains -techniques to optimize the underlying graph. It provides a wrapper similar to PyTorch's Module. When a user invokes the `forward` method, the layers are lowered to TensorRT's `ILayer`s and become part of an `INetworkDefinition`. The Graph Rewriting (GW) module can be used to manipulate the network at the `ILayer`/`INetworkDefinition` level. - -## When to Use Graph Rewriting? - -For network manipulation, there are two options in TensorRT-LLM: - -1. **Module Rewriting:** This method modifies the members of `Module` instances before triggering the `forward` method (that is, creating the TensorRT graph). It works on the highest level of the network representation and facilitates the modification of sequences of operations (like modifying the GEMM + activation for SmoothQuant), - -2. **Graph Rewriting:** Graph Rewriting manipulates TensorRT's `INetworkDefinition` after the `forward` method is triggered. It operates at a finer-grained `ILayer` level and can alter the structure across multiple Module instances. It is typically used for layer fusion. - -Graph Rewriting (GW) is ideally used in the following conditions: - -1. When only `ILayer`/`INetworkDefinition` is available, -2. When Module Rewriting would lead to nested control flow or scattered functionality. - -## Graph Rewriting APIs - -Several core APIs are provided for Graph Rewriting: - -### Tensor-Related Methods - -- `Tensor.get_parent`: Get the `ILayer` that produces this tensor, -- `Tensor.get_users`: Get the consumer `ILayer`s of this tensor, -- `replace_all_uses_with`: Replace this tensor with another tensor in all consumer `ILayer`s. - -### FLayerInfo for Retrieving High-Level Information for a Functional - -For all the layers located in `functional.py`, the original input information is missing once lowered to `INetworkDefinition`, especially for TensorRT plugins, which are opaque in the Python world. `FLayerInfo` holds their original information as a high-level signature containing inputs like `Tensor`, Python attributes, and more. There is a Network-wise singleton called `FLayerInfoMemo` to map each `ILayer` to its corresponding `FLayerInfo`. - -For `FLayerInfo`: - -- `FLayerInfo.replace_input_with`: Replace some input tensor with another tensor, -- `FLayerInfo.replace_output_uses_with`: Redirect the usage of the original output tensors to a set of new tensors. - -For `FLayerInfoMemo`: - -- `FLayerInfoMemo.instance()`: Get the singleton instance, -- `FLayerInfoMemo.get`: Get the corresponding `FLayerInfo` for an `ILayer`. - -`FLayerInfo` remains consistent with the actual `ILayer` during GW, making it safe to use. - -### Pattern and Pattern Manager - -There are two kinds of patterns: - -- `PatternRewriter`: Used for defining a rewriting pattern, which actually alters the network. - - `match`: Match the pattern; returns true if a layer is matched, - - `rewrite`: Manipulate a layer, - - `match_and_rewrite`: Combines both `match` and `rewrite`, used for complex states that need to pass from `match` to `rewrite`. - -- `PatternAnalyzer`: Used for defining an analysis pattern, which collects information from the network. - - `match`: Match the pattern, - - `analyze`: Perform analysis on a list of layers. - -There are two managers for managing multiple `PatternRewriter` or `PatternAnalyzer`: - -- `RewritePatternManager`: - - `add`: Add a pattern with its label and benefit; the benefit specifies its privilege, - - `get`: Get a pattern by label, - - `rewrite`: Apply the rewriting patterns contained to a network. - -- `AnalysisPatternManager`: - - `add`: Add a pattern with its label and benefit; the benefit specifies its privilege, - - `get`: Get a pattern by label, - - `analyze`: Apply the analysis patterns contained to a network. - -### @record_signature to Decorate Functionals Requiring FLayerInfo - -The `@record_signature` decorator is used to record the `FLayerInfo` for a functional. While FLayerInfo is vital for GW when analyzing or rewriting certain functionals, it is used in an "add as needed" manner. If you are adding GW patterns, ensure that the functional requires the `@record_signature` decorator. - -## Classical Workflow - -There are specific routines for defining a GW pattern. Let's start with a simple example: replacing a sum layer with a subtract layer, which can also be found in the `test_graph_rewriting.py` file. - -```python -class NaivePatternRewriter_ReplaceAddWithSub(PatternRewriter): - - def __init__(self): - super().__init__('replace_add_with_sub', - root_layer={trt.LayerType.ELEMENTWISE}, - separate_match_rewrite=True) - - def match(self, layer: Layer): - # The rewriter will stop at the first matched layer, and then the Rewriter will enter the rewrite() to do the rewriting. - return layer.as_layer().op == trt.ElementWiseOperation.SUM - - def rewrite(self, layer: Layer) -> None: - # The layer here should be an Elementwise_SUM layer. - with net_guard(layer.network): - # There are several stages to replace some subgraph with another subgraph: - - # Stage 1: Get the input tensors and output tensors of the subgraph to replace. - # - For Elementwise_SUM, there are two inputs and one output. - a, b = layer.get_inputs(0, 1) - o = layer.get_outputs(0)[0] - - # Stage 2: Create a new subgraph that takes the old one's inputs. - # - Here we insert an Elementwise_SUB layer, and 'c' is the output. - c = a - b - - # Stage 3: Redirect all the layers depending on the outputs of the old subgraph to the new subgraph's. - # - After this, the SUM becomes dangling and will be pruned by TensorRT when building the engine. - # - Note that there is no API in TensorRT python to remove a layer explicitly; `replace_all_uses_with` is the only way to "remove" a layer. - o.replace_all_uses_with(c) - - # Stage 4: Mark all the layers in the old subgraph as removed. - # - This helps the PatternRewriter to skip the removed layers. - layer.mark_as_removed() -``` - -In this example, we deal with `ILayer` rather than Plugins, so `FLayerInfo` is unnecessary. As illustrated in the `rewrite` method, there are four stages that are shared across nearly all rewrite patterns. - -Note that in GW, we **NEVER** rewrite a layer directly. Instead, we do it in two steps: first, create another layer with the same input and deprive all the users of the original outputs, redirecting them to the outputs of the new layers. In this way, the old layer will be dangling and pruned automatically by TensorRT during the engine building phase. This is a limitation of TensorRT since remove-layer-like APIs are not available in Python. - -In Stage 2, we rely on operators and layers commonly used during the network building phase. Ideally, you can replace them with any network structure during GW. - -For the usage of `FLayerInfo`, let's rewrite the `gpt_attention` to enable the `remove-padding` feature. `gpt_attention` is actually - - a TensorRT plugin, so we need `FLayerInfo` to hold the original Tensor-wise inputs to help create new `gpt_attention` layers. - -```python -class GPTAttentionPluginRemovePaddingRewritePass(PatternRewriter): - - def __init__(self): - super().__init__('gpt_attention_plugin_remove_padding', - root_layer={trt.LayerType.PLUGIN_V2}) - - def match_and_rewrite(self, layer: Layer) -> bool: - if layer.as_layer().type != trt.LayerType.PLUGIN_V2 or \ - layer.as_layer().plugin.plugin_namespace != 'tensorrt_llm' or \ - layer.as_layer().plugin.plugin_type != 'GPTAttention': - return False - - # Retrieve the FLayerInfo - flayer = FLayerInfoMemo.instance().get(layer.name) - assert flayer - # Although the layer is a plugin, which is a black box, we get some high-level input information from the FLayerInfo. - tensor_input: Tensor = flayer.get_input('qkv') - if tensor_input.shape[0] == 1: # Already in remove-padding mode - return False - - # Some information could be passed in from external - assert self.args is not None, "args should be passed in from RewritePatternManager.rewrite()" - batch_size, in_len, hidden_size = self.args['batch_size'], self.args['in_len'], self.args['hidden_size'] - - with net_guard(layer.network): - new_inputs = flayer.clone_inputs() - - # Step 1: Create new inputs and replace the original arglist. - input = Tensor( - name='qkv', - dtype=trt.float16, - shape=(1, batch_size * in_len, hidden_size), - ) - new_inputs['qkv'] = input - - # Step 2: Create a new plugin instance. - new_outs = gpt_attention(**new_inputs) - - # Step 3: Deprive all the users of the old plugin instance. - flayer.replace_outputs_uses_with(layer.network, new_outs) - - # Step 4: Remove the old plugin instance. - layer.mark_as_removed() - - return True -``` - -This is quite similar to the first example, with the focus on the `FLayerInfo` part. Through the code below, we can get the original inputs of this layer, enabling us to alter the inputs related to remove-padding and create a new layer to replace it. - -```python -flayer = FLayerInfoMemo.instance().get(layer.name) -assert flayer -``` - -```python -new_inputs = flayer.clone_inputs() - -# Step 1: Create new inputs and replace the original arglist. -input = Tensor( - name='tensor', - dtype=trt.float16, - shape=(1, batch_size * in_len, hidden_size), -) -new_inputs['tensor'] = input - -# Step 2: Create a new plugin instance. -new_outs = gpt_attention(**new_inputs) -``` - -For real examples, please refer to the `FuseAttentionWithBiasPass` in the `graph_rewriting.py`. diff --git a/docs/source/legacy/advanced/images/8x_l20_L40S_node_architecture.png b/docs/source/legacy/advanced/images/8x_l20_L40S_node_architecture.png deleted file mode 100644 index 725427f163c4..000000000000 Binary files a/docs/source/legacy/advanced/images/8x_l20_L40S_node_architecture.png and /dev/null differ diff --git a/docs/source/legacy/advanced/images/disaggregated-service_usage.png b/docs/source/legacy/advanced/images/disaggregated-service_usage.png deleted file mode 100644 index 6b98a2233227..000000000000 Binary files a/docs/source/legacy/advanced/images/disaggregated-service_usage.png and /dev/null differ diff --git a/docs/source/legacy/advanced/kv-cache-management.md b/docs/source/legacy/advanced/kv-cache-management.md deleted file mode 100644 index f4506d6ee996..000000000000 --- a/docs/source/legacy/advanced/kv-cache-management.md +++ /dev/null @@ -1,75 +0,0 @@ -(kv-cache-management)= - -# KV Cache Management: Pools, Blocks, and Events - -This document provides an overview of the internal hierarchy and event system for paged KV cache management, as implemented in the TensorRT-LLM codebase. - -For more information on KV cache reuse see [KV cache reuse](kv-cache-reuse.md). - ---- - -## Hierarchy: Pool, Block, and Page - -### **Block** -- **Definition:** The smallest unit of KV cache allocation. A `KVCacheBlock` holds metadata (not the actual data) for a chunk of KV cache. -- **Purpose:** Each block represents a fixed number of tokens' worth of KV data (can be specified by `tokens_per_block` parameter). -- **Usage:** Blocks are allocated, reused, or evicted as sequences are processed. - -### **Page** -- **Definition:** In this codebase, "page" is often used interchangeably with "block" (as in "paged KV cache"), but technically, a page could refer to a memory page (hardware-level), while a block is a logical unit for the cache. -- **In Practice:** The code uses "block" as the main unit; "page" is not a distinct class or struct. - -### **Pool** -- **Definition:** A pool is a contiguous memory buffer (or set of buffers) that holds the actual KV data for one or more layers. -- **Types:** There are primary pools (fast GPU memory) and secondary pools (slower, e.g., CPU or offload memory). -- **Organization:** Each pool can serve multiple layers that share the same KV head configuration. Pools are managed by `KVCacheBlockPool` and tracked in vectors in `WindowBlockManager`. -- **Block ↔ Pool:** Each block is an index into a pool; the pool provides the actual storage, while the block is the metadata handle. - -### **WindowBlockManager/BlockManager** - -TRT-LLM supports 2 complex features related to KV cache management: -1. **Variable Group-Query Attention (VGQA)** - i.e. a different `num_kv_heads` value for different layers. -2. **Variable Sliding Window Attention (VSWA)** - i.e. a different `attention_window_size` value for different layers. - -In order to support both of these features, the pool management works as described below. - -But in the simple, *most common case*, for most models, where -1. [MHA/MQA/Non-variable GQA](gpt-attention.md#multi-head-multi-query-and-group-query-attention), i.e., same `num_kv_heads` value for all layers, -2. Global attention/[SWA](gpt-attention.md#sliding-window-attention-cyclic-rolling-buffer-kv-cache), i.e., same `attention_window_size` value for all layers, - -only a *single* pool will be created within the structure described below. - -#### KV Cache Pool Management - -- **WindowBlockManager:** Manages blocks and pools for a specific attention window size. Within a `WindowBlockManager`, there can be multiple pools - each corresponding a unique number of KV heads - i.e., to support VGQA. -- **BlockManager:** Manages all `WindowBlockManager` instances, one per unique window size. - -**Hierarchy Summary:** -- **Pool** (memory buffer for KV data) - - Contains many blocks. -- **Blocks** (metadata for a chunk of the pool, each block = `tokens_per_block` tokens) - - (Optionally, blocks can be swapped between primary/secondary pools.) -- **BlockManager/WindowBlockManager**: Manage pools and blocks, handle allocation, reuse, and eviction. - ---- - -## Events in `KVCacheEventManager` - -The `KVCacheEventManager` is responsible for tracking and reporting significant changes in the state of the KV cache. Events are used for logging, debugging, or possibly for external monitoring. - -### **Types of Events** -- **Created Event:** When pools or blocks are created/allocated. -- **Updated Event:** When a block's state changes (e.g., moved between primary/secondary, priority updated). -- **Removed Event:** When a block is removed from the cache (evicted or released). -- **Stored Event:** When blocks are stored for potential reuse (e.g., after a sequence finishes and its blocks are reusable). - -### **What Triggers an Event?** -- **Allocation/Deallocation:** Creating or freeing memory pools or blocks. -- **Eviction/Reuse:** When a block is evicted, reused, or its priority changes. -- **Block Movement:** When a block is moved between memory levels (primary ↔ secondary). -- **Block Storage:** When blocks are stored for future reuse (e.g., after a sequence completes). - -**In summary:** -An "event" is any significant change in the lifecycle or state of a KV cache block or pool, tracked for monitoring, debugging, or optimization purposes. - ---- diff --git a/docs/source/legacy/advanced/kv-cache-reuse.md b/docs/source/legacy/advanced/kv-cache-reuse.md deleted file mode 100644 index 5f3a5d73cf34..000000000000 --- a/docs/source/legacy/advanced/kv-cache-reuse.md +++ /dev/null @@ -1,94 +0,0 @@ -(kv-cache-reuse)= - -# KV cache reuse - -This document describes how kv cache pages can be shared and reused by requests that start with the same prompt. This can greatly lower first token latency, the time it takes before the first output token is generated. Many use cases can benefit from this, including multi-turn requests and system prompts. - -## How to enable kv cache reuse - -There are two steps to enabling kv cache reuse. - -1. Model must support it - -KV cache reuse requires the model to be built for paged context attention. This is done with `trtllm-build`: - -```trtllm-build --use_paged_context_fmha enable``` - -2. KV cache reuse is enabled by default in KVCacheManager - -If you are running gptManagerBenchmark application, you can disable kv cache reuse with a command-line switch: - -```gptManagerBenchmark --enable_kv_cache_reuse enable=false``` - -If you are running a Triton server, you can enable kv cache reuse with a parameter: - -``` -parameters: { - key: "enable_kv_cache_reuse" - value: { - string_value: "true" - } -} -``` - -If you are writing your own application using Executor API, you can enable kv cache reuse by including `enableBlockReuse=true` when you create the `KvCacheConfig` object. Note that this is the default, if you wish to disable kv cache reuse, pass `enableBlockReuse=false` instead. - -### Enable kv cache reuse for p-tuning - -When using p-tuning, different requests may use same fake input ids (i.e. prompt ids whose values are larger than vocabulary size). That may lead to incorrect kv cache reuse, since TRT-LLM could not distinguish these requests only by input ids. To enable kv cache reuse for p-tuning correctly, users should provide an extra id (uint64) for each input id. Extra ids for normal input ids (i.e. text token ids) should always be 0, while fake input ids should have extra ids which are larger than 0. Requests using same prompt embeddings should use same extra ids, while requests using different prompt embeddings should use different extra ids. - -Example: -Assume vocabulary size is 100, which means normal text token ids are in range [0, 99] and prompt ids start from 100. - -```python -# Request 1 uses prompt embedding table 1 -input_ids = [100, 101, 102, 103, 1, 2, 3, 4] -extra_ids = [1, 1, 1, 1, 0, 0, 0, 0] - -# Request 2 uses prompt embedding table 2 -input_ids = [100, 101, 102, 103, 1, 2, 3, 4] -extra_ids = [2, 2, 2, 2, 0, 0, 0, 0] - -# Request 3 uses prompt embedding table 1 and different text tokens -input_ids = [100, 101, 102, 103, 5, 6, 7, 8] -extra_ids = [1, 1, 1, 1, 0, 0, 0, 0] -``` - -## Performance expectations - -KV cache state can be reused when two requests start with the same partial prompt. This reduces first token latency, the time it takes until the first output token is generated. Bigger savings are realized when the shared prompt is longer, relative to the overall prompt length. The biggest saving is realized when two identical requests are run back-to-back, in which case the latency for the first output token approaches latency for subsequent tokens. - -## Situations that can prevent kv cache reuse - -There are a few pitfalls that can prevent kv cache reuse when that seems possible. KV cache state only becomes reusable after the request that computed the state terminates. If you have a shared system prompt, the first request will compute kv cache state for the system prompt, the second request will reuse it, but only if the second request launches after the first request completed. If you run with a large batch-size, it is likely that many requests that share a common system prompt will be launched before the first request has terminated. No reuse will occur until one of the requests terminate, then subsequently scheduled requests can reuse. - -Kv cache state for system prompts will remain reusable until memory is needed for launching a new request or propagating an existing one. When this happens, reusable blocks are evicted based on LRU. System prompts that are frequently used have a better chance of remaining reusable, but there is no guarantee since launching new requests take priority over possible reuse. Running with a larger batch size, or larger output sequence lengths for example will reduce the probability of kv cache blocks being reused, since it increases memory needs. - -KV cache state is stored in blocks, each block holds multiple tokens. Only full blocks can be shared by multiple requests, thus the block size matters. Partially matched blocks can also be reused, but that creates a new copy of the block for each sequence. The block size is a trade-off, larger block size may improve efficiency of compute kernels, but it reduces the likelihood of kv cache state reuse. The block defaults to 128 tokens, this can be changed when the model is built with the trtllm-build command, for example - -```trtllm-build --tokens_per_block 32 ...``` - -will create a model where one KV cache block can hold 32 tokens. Note that tokens_per_block must be a power of 2. - -## Offloading to host memory - -Offloading to host memory increases likelihood of kv cache reuse. Reusable blocks that are needed for higher priority tasks, like propagating an already running request, are copied to a buffer in host memory instead of being evicted. This greatly extends the amount of memory available for reuse, allowing blocks to remain reusable much longer. On the other hand, offloading of blocks (and subsequent onboarding when a block is reused) has some cost since the blocks must be copied from CPU to GPU memory and vice versa. This cost is negligible on Grace-Hopper machines, and small enough to yield a net benefit for many use cases on x86 machines with Hopper GPUs. Offloading is unlikely to yield benefits on older architectures because of the (relatively) slow link between GPU and host memory. - -If you are running gptManagerBenchmark, you can enable offloading with a command-line switch. For example, - -```gptManagerBenchmark --kv_host_cache_bytes 45000000000``` - -will create a 45 GiB offloading buffer in host memory. Note that this buffer is pinned memory, allocating a lot of pinned memory on x86 machines can take a substantial amount of time (10s of seconds). This is a one-time cost. - -If you are running a Triton server, you can enable offloading to host memory with the kv_cache_host_memory_bytes parameter. For example, adding this to your model config file will create a 45 GiB offloading buffer in host memory. - -``` -parameters: { - key: "kv_cache_host_memory_bytes" - value: { - string_value: "45000000000" - } -} -``` - -If you are writing your own application using Executor API, you can enable offloading to host by including `hostCacheSize=45000000000` when you create the `KvCacheConfig` object. This will create a 45 GiB offloading buffer in host memory. diff --git a/docs/source/legacy/advanced/lora.md b/docs/source/legacy/advanced/lora.md deleted file mode 100644 index 5ac3aa7c5797..000000000000 --- a/docs/source/legacy/advanced/lora.md +++ /dev/null @@ -1,142 +0,0 @@ -(lora)= - -## Run gpt-2b + LoRA using Executor / cpp runtime - -First build a model with LoRA and inflight-batching enabled. - -```bash -git-lfs clone https://huggingface.co/qychen/luotuo-lora-7b-0.1 -git-lfs clone https://huggingface.co/kunishou/Japanese-Alpaca-LoRA-7b-v0 -BASE_MODEL=llama-7b-hf - -python examples/models/core/llama/convert_checkpoint.py --model_dir ${BASE_MODEL} \ - --output_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --dtype float16 - -trtllm-build --checkpoint_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --output_dir /tmp/llama_7b_with_lora_qkv/trt_engines/fp16/1-gpu/ \ - --remove_input_padding enable \ - --gpt_attention_plugin float16 \ - --context_fmha enable \ - --paged_kv_cache enable \ - --gemm_plugin float16 \ - --lora_plugin float16 \ - --max_batch_size 128 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir Japanese-Alpaca-LoRA-7b-v0 \ - --max_lora_rank 8 \ - --lora_target_modules "attn_q" "attn_k" "attn_v" -``` - -To pass LoRAs into the cpp runtime they must be converted to the format below. -The script below will convert a Hugging Face LoRA model to the correct NumPy tensor. - -```bash -python3 tensorrt_llm/examples/hf_lora_convert.py -i Japanese-Alpaca-LoRA-7b-v0 -o Japanese-Alpaca-LoRA-7b-v0-weights --storage-type float16 -python3 tensorrt_llm/examples/hf_lora_convert.py -i luotuo-lora-7b-0.1 -o luotuo-lora-7b-0.1-weights --storage-type float16 -``` - -Refer to the [tensorrtllm_backend documentation](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/lora.md) for a Multi-LoRA example using Triton. - -### LoRA tensor format details - -To run inference using `Executor`, a `Request` must have a `LoraConfig` that contains a `task_id`, `weights` and `config` parameters. - -`task_id` the unique task ID for the given LoRA. - -To perform inference with a specific LoRA for the first time, `task_id`, `weights`, and `config` must all be given. The LoRA will be cached, so that subsequent requests for the same task only require `task_id`. -If the cache is full, the oldest LoRA will be evicted to make space for new ones. An error is returned if `task_id` is not cached. - -`weights` contains the weights for all the LoRAs. Currently, this should include weights for all TP and PP ranks. -The weights tensor has the shape `[num_lora_modules_layers, D x Hi + Ho x D ]`. The last dimension holds the in / out adapter weights for the associated module (for example, `attn_qkv`) and model layer. - -Each of the in / out tensors are first flattened and then concatenated together in the format above. -The first dimension (of size `num_lora_module_layers`) has an entry for each module-layer (that is, there is an entry for `attn_q layer1` and another for `attn_k layer1`). - -`D=adapter_size (i.e. R value), Hi=hidden_size_in, Ho=hidden_size_out.` - -`config` is a configuration tensor which identifies the moduleId, layerId, and adapter size of each element of `LoraWeights`. It has the shape `[num_lora_modules_layers, 3]`. The last dimension holds `[module_id, layer_idx, adapter_size D (i.e. R value)]`. - -This feature supports LoRAs as described in https://arxiv.org/pdf/2106.09685.pdf - -#### Example LoRA tensors - -Here is an example of `LoraWeights` and `LoraConfig` tensors for a model with tp=1, pp=1, 4 layers, and a hidden size of 4. -The following tensors are for a LoRA which has a `q` and `k` adapter. - -``` -# loraConfig -[ - [1, 0, 2] - [2, 0, 4] - [1, 1, 2] - [2, 1, 4] - [1, 2, 2] # Note that the final 2 layers only adapt `q` - [1, 3, 8] -] -# Note: The loraConfig tensor configures the loraWeights tensor. -# The contents of each row of loraWeights is specified be the corresponding row in loraConfig - -# loraWeights -# Note: that 'in weights' and 'out weights' are 'A' and 'B' in the LoRA paper. -[ - [ <2 x 4 in weights>, <4 x 2 out weights> ] # `q` adapter for layer 0 - [ <4 x 4 in weights>, <4 x 4 out weights> ] # `k` adapter for layer 0 - [ <2 x 4 in weights>, <4 x 2 out weights> ] # `q` adapter for layer 1 - [ <4 x 4 in weights>, <4 x 4 out weights> ] # `k` adapter for layer 1 - [ <2 x 4 in weights>, <4 x 2 out weights> ] # `q` adapter for layer 2 - [ <8 x 4 in weights>, <4 x 8 out weights> ] # `q` adapter for layer 3. Note the final layer has a adapter size of 8 -] - -``` - -#### LoRA Module id mapping - -| module name (as specified in `convert_checkpoint.py` scripts) | module id | description | -| --------------------------------------------- | --------- | ----------- | -| attn_qkv | 0 | compbined qkv adapter | -| attn_q | 1 | q adapter | -| attn_k | 2 | k adapter | -| attn_v | 3 | v adapter | -| attn_dense | 4 | adapter for the dense layer in attention | -| mlp_h_to_4h | 5 | for llama2 adapter for gated mlp layer after attention / RMSNorm: up projection | -| mlp_4h_to_h | 6 | for llama2 adapter for gated mlp layer after attention / RMSNorm: down projection | -| mlp_gate | 7 | for llama2 adapter for gated mlp later after attention / RMSNorm: gate | -| cross_attn_qkv | 8 | compbined qkv adapter for cross attention | -| cross_attn_q | 9 | q adapter for cross attention | -| cross_attn_k | 10 | k adapter for cross attention | -| cross_attn_v | 11 | v adapter for cross attention | -| cross_attn_dense | 12 | adapter for the dense layer in cross attention | -| moe_h_to_4h | 13 | for mixtral adapter for expert mlp layer: up projection | -| moe_4h_to_h | 14 | for mixtral adapter for expert mlp layer: down projection | -| moe_gate | 15 | for mixtral adapter for expert mlp layer: gate | -| moe_router | 16 | for mixtral adapter for expert router layer | -| mlp_router | 17 | for qwen2-moe adapter for shared expert gate layer | -| mlp_gate_up | 18 | adapter for gated mlp layer after attention / RMSNorm: gate + up projection | - -#### LoraCache configuration - -The core idea is that we will have a fixed size, 2-level LoRA cache in TRT-LLM. The higher level cache resides on the host and the lower level is on GPU (distinct from the existing KV cache). Sizes of both are user configurable. - -The CPU cache is configured to be a max size. The GPU cache is configured to a percentage of free GPU memory after engine load. As requests come in LoRAs are stored in the host cache. - -As requests are scheduled for execution LoRAs are loaded into the GPU cache. - -#### LoRA with tensor parallel - -The partition of tensor parallel for LoRA is special. There are two cases: `RowLinear` and `ColumnLinear`. Assume we have a linear layer and the input feature size is `K` and the output feature size is `N`. Then, the shape of the weight is `[K, N]`. - -First, consider this linear layer is a `ColumnLinear` layer. When we partition the weight, we split the weight by column with `tp_size`. Then, there are `tp_size` split weights and the shapes of these weights are `[K, N // tp_size]`. When we apply LoRA adapter on such `ColumnLinear` layer, the shapes of original two weights are `[K, lora_rank]` and `[lora_rank, N]`. So, we only partition the second weight and get `tp_size` split weights with shapes `[lora_rank, N // tp_size]`. For the first weight, each GPU maintains the same entire weight (with shape `[K, lora_rank]`). - -Next, consider this linear layer is a `RowLinear` layer. When we partition the weight, we split the weight by row with `tp_size`. Then, there are `tp_size` split weights and the shapes of these weights are `[K // tp_size, N]`. When we apply LoRA adapter on such `RowLinear` layer, the shapes of original two weights are `[K, lora_rank]` and `[lora_rank, N]`. So, we only partition the first weight and get `tp_size` split weights with shapes `[K // tp_size, lora_rank]`. For the second weight, each GPU maintains the same entire weight (with shape `[lora_rank, N]`). - -#### DoRA - -TensorRT LLM supports DoRA as described in https://arxiv.org/abs/2402.09353 . To enable DoRA, you must add the additional `--dora_plugin enable` flag to the `trtllm-build` command. - -The DoRA scales must be normalized before they are submitted to TensorRT LLM in an inference request. The normalization requires the base model weights. To normalize your adapter you may use the script provided in `tensorrt_llm/examples/dora/normalize_weights.py`. - -When using DoRA, the format of `LoraWeights` and `LoraConfig` changes slightly. -The shape of `LoraConfig` becomes `[module_id, layer_idx, adapter_size D (i.e. R value), is_dora]`, with `is_dora` a boolean flag that determines whether the supplied adapter contains DoRA scales or not. If the old config shape is used, it is assumed the adapter does not have DoRA scales. -The shape of `LoraWeights` becomes `[num_lora_modules_layers, D x Hi + Ho x D + Ho]`, and the last `Ho` values are the DoRA scale vector. diff --git a/docs/source/legacy/advanced/lowprecision-pcie-allreduce.md b/docs/source/legacy/advanced/lowprecision-pcie-allreduce.md deleted file mode 100644 index b7ab50703701..000000000000 --- a/docs/source/legacy/advanced/lowprecision-pcie-allreduce.md +++ /dev/null @@ -1,61 +0,0 @@ -# Low-Precision-AllReduce - -```{note} -Note: -This feature is optimized for PCIe-based GPU topologies and may affect model accuracy. Please evaluate precision impact for your specific workload. -``` - - -TRT-LLM supports `low-precision-allreduce`, a communication optimization that accelerates AllReduce operations in PCIe-based GPU environments. This feature quantizes FP16/BF16 data to FP8 during network transmission, reducing communication volume and improving performance. - -## Algorithm - -The Low-Precision-AllReduce algorithm works by: -1. Quantizing input FP16/BF16 tensors to FP8 format before network transmission - - - **Quantization details**: We use a "per-warp" quantization approach where each CUDA warp (32 threads) processes a batch of data. In each warp, 31 threads quantize FP16/BF16 values to FP8 e4m3 format (16 bytes per thread), while the last thread transmits a scalar value. This results in each warp collectively quantizing 496 elements plus one scalar at a time. - -2. Transmitting the quantized data through the network -3. Dequantizing received data back to the original precision -4. Performing the reduction operation - -In 8-GPU scenarios, this approach shifts the communication bottleneck from cross-NUMA QPI to the PCIe switch, resulting in better overall performance. - -## Topology Requirements - -![8x L20/L40s Node Architecture](images/8x_l20_L40S_node_architecture.png) - -Low-Precision-AllReduce is specifically designed for the topology shown above, where: -- Each node contains 2 NUMA domains -- Each NUMA domain has 4 GPUs connected via PCIe switch -- GPUs within the same NUMA node communicate via the PCIe switch - -**Important:** This optimization will not accelerate performance in different topologies (e.g., where each GPU is in a separate NUMA domain). - -## Usage - -The Low-Precision-AllReduce algorithm can be enabled in two ways: - -1. **Direct specification** in your code: -``` -AllReduce allreduce(mapping=mapping, strategy=AllReduceStrategy.LOWPRECISION); -``` - -2. Enable by LlmArgs -``` -Set allreduce_strategy field in LlmArgs. -Candidates of strategies are "AUTO", "NCCL", "UB", "MINLATENCY", "ONESHOT", "TWOSHOT", "LOWPRECISION" and "MNNVL". -If no strategy is set, AUTO will be set. -``` - -## Performance and Accuracy Considerations - -Low-Precision-AllReduce reduces communication volume by using FP8 data format for transmission. This optimization: -- Improves performance for large message sizes in PCIe-based topologies -- May slightly reduce numerical precision -- Automatically falls back to other strategies when no performance benefit is expected (e.g., with NVLink or small messages) - -Users should evaluate the precision impact on their specific models and workloads. - -**Note**: When compiling TensorRT-LLM without enabling the `ENABLE_FP8` option, setting Low Precision allreduce will not take effect. diff --git a/docs/source/legacy/advanced/open-sourced-cutlass-kernels.md b/docs/source/legacy/advanced/open-sourced-cutlass-kernels.md deleted file mode 100644 index 80e9c066b642..000000000000 --- a/docs/source/legacy/advanced/open-sourced-cutlass-kernels.md +++ /dev/null @@ -1,26 +0,0 @@ -We have recently open-sourced a set of Cutlass kernels that were previously known as "internal_cutlass_kernels". Due to internal dependencies, these kernels were previously only available to users as static libraries. We have now decoupled these internal dependencies, making the kernels available as source code. - -The open-sourced Cutlass kernels are on the path `cpp/tensorrt_llm/kernels/cutlass_kernels`, including: -- `low_latency_gemm` -- `moe_gemm` -- `fp4_gemm` -- `allreduce_gemm` - -To ensure stability and provide an optimized performance experience, we have maintained the previous method of calling these kernels via static libraries as an alternative option. You can switch between open-sourced Cutlass kernels and static library Cutlass kernels through the `USING_OSS_CUTLASS_*` macro (where * represents the specific kernel name), enabling kernel-level control. By default, the open-source Cutlass kernels are used. -Note that support for these static libraries will be gradually deprioritized in the future and may eventually be deprecated. - -**Default Configuration (Using open-sourced Cutlass Kernels)** - -To build using the open-source Cutlass kernels (default setting), run: - -```bash -python3 ./scripts/build_wheel.py --cuda_architectures "90-real;100-real" -``` - -**Using Static Library Cutlass Kernels** - -If you prefer to use the Cutlass kernels from the static library, you can control this during compilation by setting the `USING_OSS_CUTLASS_*` macro to `OFF`. For example, to use the static library implementation specifically for `low_latency_gemm` and `moe_gemm` while keeping other kernels as OSS, use the following compilation command: - -```bash -python3 ./scripts/build_wheel.py --cuda_architectures "90-real;100-real" -D "USING_OSS_CUTLASS_MOE_GEMM=OFF;USING_OSS_CUTLASS_LOW_LATENCY_GEMM=OFF" -``` diff --git a/docs/source/legacy/advanced/speculative-decoding.md b/docs/source/legacy/advanced/speculative-decoding.md deleted file mode 100644 index 2faf885f8d10..000000000000 --- a/docs/source/legacy/advanced/speculative-decoding.md +++ /dev/null @@ -1,182 +0,0 @@ -# Speculative Sampling - -- [About Speculative Sampling](#about-speculative-sampling) -- [Performance Improvements](#Performance-improvements) -- [Draft-Target-Model](#Draft-Target-Model) -- [NGram](#ngram) -- [Medusa](#medusa) - - [Medusa Tree](#medusa-tree) - - [Using Medusa with TensorRT-LLM](#using-medusa-with-tensorrt-llm) - - [Limitations](#limitations) -- [ReDrafter](#redrafter) -- [EAGLE](#eagle) - - [Disaggregated Serving](#disaggregated-serving) -- [Lookahead decoding](#lookahead-decoding) - -## About Speculative Sampling - -Speculative Sampling (also referred to as Speculative Decoding) is a set of techniques designed to allow generation of more than one token per forward pass iteration. This can lead to a reduction in the average per-token latency **in situations where the GPU -is underutilized due to small batch sizes.** - -Speculative Sampling involves predicting a sequence of future tokens, referred to as draft tokens, using a method -that is substantially more efficient than repeatedly executing the target Large Language Model (LLM). -These draft tokens are then collectively validated by processing them through the target LLM in a single forward pass. -The underlying assumptions are twofold: - -1. processing multiple draft tokens concurrently will be as rapid as processing a single token -2. multiple draft tokens will be validated successfully over the course of the full generation - -If the first assumption holds true, the latency of speculative decoding will no worse than the standard approach. If the second holds, output token generation advances by statistically more than one token per forward pass. -The combination of both these allows speculative decoding to result in reduced latency. - -TensorRT-LLM supports several approaches for generating draft tokens, including: - -1. Utilizing a smaller, auxiliary model, known as the draft model approach. For more information, refer to the [Fast Inference from Transformers via Speculative Decoding paper](https://arxiv.org/pdf/2211.17192.pdf). -2. Implementing additional language model heads that predict tokens for future positions: - 1. [Medusa: Simple LLM Inference Acceleration Framework with Multiple Decoding Heads paper](https://arxiv.org/abs/2401.10774). - 2. [Recurrent Drafter for Fast Speculative Decoding in Large Language Models](https://arxiv.org/html/2403.09919v1). - 3. [EAGLE: Speculative Sampling Requires Rethinking Feature Uncertainty](https://arxiv.org/pdf/2401.15077). -3. Utilizing prompt tokens as draft tokens. For more information, refer to [NGram](https://github.com/apoorvumang/prompt-lookup-decoding/). -4. Utilizing Jacobi-like decoding to predict and verify draft tokens using the same model which does not need additional fine-tuning. Refer to [Break the Sequential Dependency of LLM Inference Using Lookahead Decoding](https://arxiv.org/pdf/2402.02057). - - -## Performance Improvements - -It's important to note that the effectiveness of speculative decoding techniques is highly dependent -on the specific task at hand. For instance, forecasting subsequent tokens in a code-completion scenario -may prove simpler than generating a summary for an article. - -Furthermore, when integrating Medusa with a standard PyTorch model implementation which may not be as finely -tuned as TensorRT-LLM, the potential time savings are more pronounced. - -## Draft-Target-Model - -The Draft-Target-Model involves the use of two distinct models (a smaller Draft model and a larger Target model) trained independently but sharing the same vocabulary. For example, GPT 125M / 6.7B models serve as the Draft / Target model. - -The management of Draft and Target models is facilitated through two separate `Executor` instances. -It is essential that you to coordinate the interactions between the Draft and Target models effectively. -Initially, the Draft model is queried to generate up to `K` draft tokens. -These tokens are then forwarded to the Target model for verification. -Upon verification, the Target model may return up to `K+1` tokens. -Subsequently, the prompt, now updated with the accepted tokens, is sent back to the Draft model to initiate the generation of new draft tokens. -This iterative process continues until a predefined stop conditions are met. -An example orchestration script is available in the Triton backend repository’s -[draft-target-model client example](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/client/python/draft_target_model_client.py). - -We provide two styles of running Draft-Target-Model now: using TensorRT-LLM-BLS in Triton Inference Server, or using TensorRT-LLM directly. Detailed steps of running can be found in [examples/draft_target_model/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/draft_target_model/README.md) and the code can be found in [examples/ngram/run_dtm_ngram.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/run_dtm_ngram.py). - -## NGram - -The NGram speculative decoding directly copies from the input prompt and previous generated output as draft tokens while generating the later output. It works like Draft-Target-Model but involves only one Target LLM model without further fine-tuning. The NGram profit from the scenarios which have high n-gram overlap between input prompt and output, such as summarization, document QA, multi-turn chat, code editing, etc. - -See document in [examples/ngram/README.md](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/README.md) and the code can be found in [examples/ngram/run_dtm_ngram.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/ngram/run_dtm_ngram.py). - -## Medusa - -This approach leverages a single model to both generate and verify draft tokens. -It enhances the existing model by adding multiple extra language model heads, known as Medusa heads. -These additional heads are trained to predict future tokens while the base model remains unchanged. -Specifically, the first Medusa head is tasked with predicting the immediate next token, -the second head predicts the token after that, and so on. -With `K` Medusa heads, the model can forecast up to `K` tokens ahead. -The draft tokens generated by the Medusa heads during iteration `i` -are then verified and potentially accepted in the subsequent iteration, `i+1`. - -The true potential of the Medusa strategy is realized when more than one token per head is used, -employing a TopK approach to create multiple potential paths, essentially forming a tree, rather than -a single linear path as seen in the Draft model approach. To reduce redundant computations, many of these paths, -which often share common prefixes, are consolidated into a single path. -This is achieved by applying attention with a sparse mask that represents the various paths. Sparse mask formed by Medusa tree is described in detail later. - -By validating multiple paths simultaneously, there is an increased likelihood of accepting more than one token per iteration, -albeit at the expense of additional computational effort. - -It is crucial to recognize that as the number of potential paths grows exponentially with `K`, -it is not necessary to explore or validate all of them. A recommended strategy for managing this complexity is to prune the tree -by focusing only on the paths with higher-probability tokens. - -You must strike a balance between the breadth and depth of the tree you want to explore and the impact of a larger tree on the overall -performance for your specific application. - -In the TensorRT-LLM implementation of Medusa, the configuration of the tree is a runtime parameter. -This flexibility allows you to experiment and identify the optimal tree structure for your use case, -which can then be utilized in a production environment. - -### Medusa Tree - -Consider the following diagram, which illustrates how the hidden states from the last layer of the base model -are passed to the base model's language model (LM) head and to four Medusa heads (MHs). - -

- Example Medusa Tree -

- -In this example: - -1. The token l0 represents the actual token generated by the model. -All other tokens, denoted as phk, are predictions from the MHs, -where `h` indicates the Medusa head index (1-based) and `k` represents the TopK choice index (0-based). -1. Four MHs are used, which means the model is predicting four future tokens. -2. The first two MHs utilize Top-2 predictions, while the last two use Top-1. -For instance, p10 and p11 are the top and -second top predictions from the first Medusa Head (MH1). -1. A total of four paths are explored, which is fewer than the 16 that would be examined -if a complete binary tree were used (assuming Top-2 predictions for all MHs). -1. As some of these paths may be accepted, there are ten potential candidates, referred to as `medusa_choices`. -The number of tokens that can be accepted at each step, including the true token, -ranges from 1 (if all Medusa predictions are incorrect) to 5 (if all are correct). - -During the generation phase, the model receives an input of 10 tokens, -which corresponds to the last tokens of each candidate path, rather than just one. - -In TensorRT-LLM, you have the option to define such trees by providing all the Medusa choices -or by simply specifying the unique paths. - -- Since each candidate/path begins with the true token (l0), -there is no need to specify it separately. For the predicted tokens, only the TopK indices are required. -- For example, to specify the path l0p10p21p30, -one would use `[0,1,0]`. And -to specify the path l0p11p20, -one would use `[1,0]`. -- To specify all 4 paths in the example, use `medusa_choices=[[0,0,0,0], [0,1,0], [1,0], [1,1]]`. -- It's also possible to specify all candidates explicitly, similar to the Medusa repository. -For instance, `medusa_choices=[[0], [0,0], [0,0,0], [0,0,0,0], [0,1], -[0,1,0], [1], [1,0], [1,1]]`. Note that when specifying all the candidates explicitly, **we don't include -the empty `[]` candidate** for the case where only the true token is accepted, that is, all the predictions from MHs are wrong. -So, only `9` candidates are specified. - -**Specifying paths-only instead of all choices is currently supported only in the Python runtime.** - -### Using Medusa with TensorRT-LLM - -For guidance on constructing and executing Medusa with the Python runtime, consult the [Medusa README](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/medusa/README.md). When utilizing the Inflight Fused Batching (IFB) with the C++ API, it is necessary to define the `medusa_choices` explicitly within the model configuration. For detailed instructions, refer to the [model configuration in TensorRT-LLM backend](https://github.com/triton-inference-server/tensorrtllm_backend?tab=readme-ov-file#modify-the-model-configuration) for more details. - -#### Limitations - -- TensorRT-LLM supports Medusa only for Vicuna (fine tuned LLaMA). -However, similar to any new model, you can follow the same approach to define your own Medusa model and deploy with TensorRT-LLM. -- We match only tokens during the validation phase that is `medusa_temperature=0`. -- Beam search is **not** compatible with Medusa. - - -## ReDrafter - -The ReDrafter approach enhances the single-model Medusa method by predicting and verifying tokens using the same model. However, unlike Medusa, it predicts draft tokens using a recurrent predictor, where each draft token depends on the previous one. This method also allows the use of beam search to identify more prominent draft tokens. For more details, please read [the ReDrafter paper](https://arxiv.org/html/2403.09919v1). - -TensorRT-LLM implements the ReDrafter model such that logits prediction, beam search, and draft token acceptance are performed inside the TensorRT engine. This contrasts with standard model inference, which only predicts logits and performs decoding outside the engine. Since the engine predicts explicit draft tokens instead of implicit tokens decoded from logits, we categorize this speculative decoding method as `explicit_draft_tokens`. Please, visit the [ReDrafter README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/redrafter/README.md) for information about building and running the model. ReDrafter supports both Inflight Fused Batching runtime and Python static batching runtime. - -## EAGLE - -The EAGLE approach enhances the single-model Medusa method by predicting and verifying tokens using the same model. Similarly to ReDrafter, it predicts draft tokens using a recurrent predictor where each draft token depends on the previous one. However, unlike ReDrafter, it uses a single-layer transformer model to predict draft tokens from previous hidden states and decoded tokens. In the EAGLE-1 decoding tree needs to be known during the decoding. In the EAGLE-2 this tree is asssembled during the execution by searching for the most probable hypothesis along the beam. - -Similarly to ReDrafter, TensorRT-LLM implements the EAGLE model such that logits prediction, draft tokens acceptance and draft token generation are performed inside of the TensorRT engine(EAGLE-1 and EAGLE-2 are both supported). Please, visit the [EAGLE README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/eagle/README.md) for information about building and running the model. - -> **EAGLE3 note.** If the EAGLE3 draft head config omits `draft_vocab_size`, TensorRT-LLM assumes it matches `vocab_size` and emits a warning. Set `draft_vocab_size` explicitly if the draft head uses a different vocabulary. - -### Disaggregated Serving - -[Disaggregated Serving](https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/features/disaggregated-service.md) with EAGLE3 using the two model approach is supported in the Pytorch backend. Please refer to the following [Dynamo example](https://github.com/ai-dynamo/dynamo/blob/main/examples/tensorrt_llm/llama4_plus_eagle.md) on how to run EAGLE3 with Disaggregated Serving for Llama 4 Maverick. - -## Lookahead Decoding - -Lookahead decoding algorithm operates through two parallel computation branches within the same model: a lookahead branch that generates n-grams using a fixed-sized 2D window, and a verification branch that validates promising n-gram candidates. This approach eliminates the necessity for additional model training or fine-tuning and can be enabled for any autoregressive model. Refer to the [Lookahead decoding README](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/lookahead/README.md) for information about building and running the model. diff --git a/docs/source/legacy/advanced/weight-streaming.md b/docs/source/legacy/advanced/weight-streaming.md deleted file mode 100644 index 0cbe4b031f7b..000000000000 --- a/docs/source/legacy/advanced/weight-streaming.md +++ /dev/null @@ -1,58 +0,0 @@ -(weight-streaming)= - -## Running With Weight Streaming to Reduce GPU Memory Consumption - -TensorRT Weight Streaming can offload some weights to the CPU memory and stream them to the GPU memory during runtime. -This can reduce the weights size in GPU memory, therefore, we can run larger models or larger batch sizes in the same GPU memory budget. - - -During build time, build the engine with `--weight-streaming --gemm_plugin disable` since Weight Streaming only supports non-plugin weights. During runtime, run with `--gpu_weights_percent x` to config the percent of weights that remained on the GPU. `x` can be a value from `0.0` to `1.0`. - -Here is an example to run llama-7b with Weight Streaming: -```bash - -# Convert model as normal. Assume hugging face model is in llama-7b-hf/ -python3 examples/models/core/llama/convert_checkpoint.py \ - --model_dir llama-7b-hf/ \ - --output_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --dtype float16 - -# Build engine that enabled Weight Streaming. -trtllm-build \ - --checkpoint_dir /tmp/llama_7b/trt_ckpt/fp16/1-gpu/ \ - --output_dir /tmp/llama_7b/trt_engines/fp16/1-gpu/ \ - --weight_streaming \ - --gemm_plugin disable \ - --max_batch_size 128 \ - --max_input_len 512 \ - --max_seq_len 562 - -# Run the engine with 20% weights in GPU memory. -python3 examples/summarize.py \ - --engine_dir /tmp/llama_7b/trt_engines/fp16/1-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir llama-7b-hf/ \ - --data_type fp16 \ - --gpu_weights_percent 0.2 - -``` - - -### API Changes - -To build engines with Weight Streaming enabled, some API changes are needed for the builder: -- Added a new bool member `weight_streaming` to class `BuildConfig`. -- Added a new bool parameter `weight_streaming` to method `create_builder_config` of class `Builder`. - -To run with Weight Streaming with `Executor`, there are some API change to its config `ExecutorConfig`: -- Added a new float parameter `gpuWeightsPercent` to the constructor of `ExecutorConfig`. -- Added two member functions `setGpuWeightsPercent` and `getGpuWeightsPercent` to set and get the GPU weights percentage. - -Here is an example to create an `Executor` with Weight Streaming: -```c++ -... -auto executorConfig = tle::ExecutorConfig(gpuWeightsPercent=0.5); -auto executor = tle::Executor("model_path", tensorrt_llm::executor::ModelType::kDECODER_ONLY, executorConfig); -... -``` diff --git a/docs/source/legacy/architecture/add-model.md b/docs/source/legacy/architecture/add-model.md deleted file mode 100644 index 7235ee07fbd5..000000000000 --- a/docs/source/legacy/architecture/add-model.md +++ /dev/null @@ -1,95 +0,0 @@ -(add-model)= - -# Adding a Model - -This document describes how to add a typical decoder-only model in TensorRT LLM. - -## Step 1. Write Modeling Part - -TensorRT LLM provides different levels of APIs: - -- Low-level functions, for example, `concat`, `add`, and `sum`. -- Basic layers, such as, `Linear` and `LayerNorm`. -- High-level layers, such as, `MLP` and `Attention`. -- Base class for typical decoder-only models, such as, `DecoderModelForCausalLM`. - -1. Create a model directory in `tensorrt_llm/models`, for example `my_model`. -2. Write a `model.py` with TensorRT LLM's APIs - -```python -class MyDecoderLayer(Module): - def __init__(self, config: PretrainedConfig, layer_idx: int): - self.layer_idx = layer_idx - self.config = config - self.input_layernorm = LayerNorm(...) - self.attention = Attention(...) - self.post_layernorm = LayerNorm(...) - self.mlp = MLP(...) - - def forward(self, hidden_states, ...): - # decoder layer forward - return hidden_states - -class MyModel(Module): - def __init__(self, config: PretrainedConfig): - self.config = config - self.vocab_embedding = Embedding(...) - self.layers = DecoderLayerList(MyDecoderLayer, config) - self.ln_f = LayerNorm(...) - - def forward(self, input_ids, ...): - # model forward - return hidden_states - - -class MyModelForCausalLM(DecoderModelForCausalLM): - def __init__(self, config: PretrainedConfig): - transformer = MyModel(config) - lm_head = ColumnLinear(...) - super().__init__(config, transformer, lm_head) -``` - - -## Step 2. Implement Weight Conversion - -The weights from source framework need to be converted and bound to the new added TensorRT LLM model. Here is an example of converting HuggingFace weights: - -```python -class MyModelForCausalLM(DecoderModelForCausalLM): - @classmethod - def from_hugging_face( - cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None) -> MyModelForCausalLM - # create a TensorRT LLM MyModelForCausalLM model object - # convert HuggingFace checkpoint to TensorRT LLM expected weights dict - # load the weights to MyModelForCausalLM object -``` - -It's optional to develop a `convert_checkpoint.py` script in the `examples/my_model/` directory for the convenience of offline weights conversion. - -## Step 3. Register New Model - -Please register the new model class `MyModelForCausalLM` in `tensorrt_llm/models/__init__.py`. - -## Step 4. Verify New Model - -At last, let's verify the new model. The typical commands are as following: - -```bash -cd examples/my_model/ - -python convert_checkpoint.py --model_dir hf_model_dir --output_dir tllm_ckpt_dir - -trtllm-build --checkpoint_dir tllm_ckpt_dir --output_dir tllm_engine_dir - -# try the model with a single prompt -python ../run.py --engine_dir tllm_engine_dir --tokenizer_dir hf_model_dir --input_text "Born in north-east France, Soyer trained as a" -# run summarization task -python ../summarize.py --engine_dir tllm_engine_dir --hf_model_dir hf_model_dir --test_trt_llm -``` - -## Reference - -It's recommended to read the workflow[./workflow.md] and checkpoint[./checkpoint.md] documents for more details. diff --git a/docs/source/legacy/architecture/checkpoint.md b/docs/source/legacy/architecture/checkpoint.md deleted file mode 100644 index 81d0955e1eb9..000000000000 --- a/docs/source/legacy/architecture/checkpoint.md +++ /dev/null @@ -1,251 +0,0 @@ -# TensorRT LLM Checkpoint - -> [!WARNING] -> This page describes the **legacy** TensorRT engine-build workflow. -> For new projects, use [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -## Overview - -The earlier versions (pre-0.8 version) of TensorRT LLM were developed with a very aggressive timeline. For those versions, emphasis was not put on defining a unified workflow. Now that TensorRT LLM has reached some level of feature richness, the development team has decided to put more effort into unifying the APIs and workflow of TensorRT LLM. This file documents the workflow around TensorRT LLM checkpoint and the set of CLI tools to generate checkpoint, build engines, and evaluate engines. - -There are three steps in the workflow: - -1. Convert weights from different source frameworks into TensorRT LLM checkpoint. -2. Build the TensorRT LLM checkpoint into TensorRT engines with a unified build command. -3. Load the engines to TensorRT LLM model runner and evaluate with different evaluation tasks. - -``` -NeMo ------------- - | -HuggingFace ------ - | convert build load -Modelopt --------- ----------> TensorRT LLM Checkpoint --------> TensorRT Engine ------> TensorRT LLM ModelRunner - | -JAX -------------- - | -DeepSpeed -------- -``` - -## Prepare the TensorRT LLM Checkpoint - -TensorRT LLM aims at supporting different sources: - -1. Trained models from NVIDIA NeMo, Microsoft DeepSpeed, and JAX -2. Quantized models from NVIDIA Modelopt -3. Popular models from HuggingFace - -TensorRT LLM defines its own checkpoint format. A checkpoint directory includes: - -1. One config `json` file, which contains several model hyper-parameters. -2. One or several rank weights files, each file contains a dictionary of tensors (weights). -The different files are loaded by different ranks in a multi-GPU (multi-process) scenario. - -### Config - -| Field | Type | Default Value | -| :------------------------------------- | :--------- | :------------------ | -| architecture | string | mandatory | -| dtype | string | mandatory | -| logits_dtype | string | 'float32' | -| vocab_size | int | mandatory | -| max_position_embeddings | int | null | -| hidden_size | int | mandatory | -| num_hidden_layers | int | mandatory | -| num_attention_heads | int | mandatory | -| num_key_value_heads | int | num_attention_heads | -| hidden_act | string | mandatory | -| intermediate_size | int | null | -| norm_epsilon | float | 1e-5 | -| position_embedding_type | string | 'learned_absolute' | -| mapping.world_size | int | 1 | -| mapping.tp_size | int | 1 | -| mapping.pp_size | int | 1 | -| quantization.quant_algo | str | null | -| quantization.kv_cache_quant_algo | str | null | -| quantization.group_size | int | 64 | -| quantization.has_zero_point | bool | False | -| quantization.pre_quant_scale | bool | False | -| quantization.exclude_modules | list | null | - -`mapping.world_size` means `mapping` is a dictionary containing the `world_size` sub field. - -```json -{ - "architecture": "OPTForCausalLM", - "mapping": { - "world_size": 1 - } -} -``` - -Supported quantization algorithm list: - -- W8A16 -- W4A16 -- W4A16_AWQ -- W4A8_AWQ -- W4A16_GPTQ -- FP8 -- W8A8_SQ_PER_CHANNEL - -Supported KV cache quantization algorithm list: - -- FP8 -- INT8 - -The config field is extensible, a model could add its own specific config fields. -For example, OPT model has a `do_layer_norm_before` field. - -Here is the model specific config list: - -| Field | Type | Default Value | -| :------------------------------------- | :--------- | :------------------ | -| OPT | | | -| do_layer_norm_before | bool | False | -| | | | -| Falcon | | | -| bias | bool | True | -| new_decoder_architecture | bool | False | -| parallel_attention | bool | False | - -### Rank Weights - -Like PyTorch, the tensor (weight) name is a string containing hierarchical information, -which is uniquely mapped to a certain parameter of a TensorRT LLM model. - -For example, each transformer layer of the OPT model contains an `Attention` layer, an `MLP` layer, and two `LayerNorm` layers. - -#### Attention Weights - -The `Attention` layer contains two `Linear` layers, qkv and dense; each `Linear` layer contains one weight and one bias. -There are four tensors (weights) in total, whose names are: - -- `transformer.layers.0.attention.qkv.weight` -- `transformer.layers.0.attention.qkv.bias` -- `transformer.layers.0.attention.dense.weight` -- `transformer.layers.0.attention.dense.bias` - -where `transformer.layers.0.attention` is the prefix name, indicating that the weights/biases are in the Attention module of the 0-th transformer layer. - -#### MLP Weights - -The `MLP` layer also contains two `Linear` layers, fc and proj; each `Linear` layer contains one weight and one bias. -There are four tensors (weights) in total, whose names are: - -- `transformer.layers.0.mlp.fc.weight` -- `transformer.layers.0.mlp.fc.bias` -- `transformer.layers.0.mlp.proj.weight` -- `transformer.layers.0.mlp.proj.bias` - -where `transformer.layers.0.mlp` is the prefix name, indicating that the weights/biases are in the MLP module of the 0-th transformer layer. - -#### LayerNorm Weights - -Each of the two `LayerNorm` layers, namely `input_layernorm` and `post_layernorm`, contains one weight and one bias. -There are four tensors (weights) in total, whose names are: - -- `transformer.layers.0.input_layernorm.weight` -- `transformer.layers.0.input_layernorm.bias` -- `transformer.layers.0.post_layernorm.weight` -- `transformer.layers.0.post_layernorm.bias` - -where `transformer.layers.0.input_layernorm` and `transformer.layers.0.post_layernorm` are prefix names for the two `layernorm` modules. - -#### KV Cache Quantization Scaling Factors - -If we quantize the model, there will be different tensors (depending on the quantization method applied). -For example, if we quantize the KV cache, the `Attention` layer will have this extra scaling factor: - -- `transformer.layers.0.attention.kv_cache_scaling_factor` - -#### FP8 Quantization Scaling Factors - -Here is the FP8 scaling factors of `attention.qkv` linear layer: - -- `transformer.layers.0.attention.qkv.activation_scaling_factor` -- `transformer.layers.0.attention.qkv.weights_scaling_factor` - -#### AWQ Quantization Scaling Factors - -Here is the AWQ scaling factors of `mlp.fc` linear layer: - -- `transformer.layers.0.mlp.fc.weights_scaling_factor` -- `transformer.layers.0.mlp.fc.prequant_scaling_factor` - - ```{note} - The linear weights in TensorRT LLM checkpoint always follows (`out_feature`, `in_feature`) shape, whereas some quantized linear in TensorRT LLM implemented by plugin may use (`in_feature`, `out_feature`) shape. The `trtllm-build` command adds a transpose operation to post-process it. - -### Example - -Let's take OPT as an example and deploy the model with tensor parallelism 2: - -```bash -cd examples/opt -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --tp_size 2 \ - --output_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ -``` - -Here is the checkpoint directory: - -``` -./opt/125M/trt_ckpt/fp16/1-gpu/ - config.json - rank0.safetensors - rank1.safetensors -``` - -Here is the `config.json`: - -```json -{ - "architecture": "OPTForCausalLM", - "dtype": "float16", - "logits_dtype": "float32", - "num_hidden_layers": 12, - "num_attention_heads": 12, - "hidden_size": 768, - "vocab_size": 50272, - "position_embedding_type": "learned_absolute", - "max_position_embeddings": 2048, - "hidden_act": "relu", - "mapping": { - "world_size": 2, - "tp_size": 2 - }, - "use_parallel_embedding": false, - "embedding_sharding_dim": 0, - "do_layer_norm_before": true, -} -``` - -## Build Checkpoint into TensorRT Engine - -TensorRT LLM provides a unified build command: `trtllm-build`. Before using it, -you may need to add it to the `PATH`. - -```bash -export PATH=/usr/local/bin:$PATH - -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/2-gpu/ -``` - -## Make Evaluation - -```bash -mpirun -n 2 --allow-run-as-root \ - python3 ../summarize.py --engine_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-125m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 -``` diff --git a/docs/source/legacy/architecture/core-concepts.md b/docs/source/legacy/architecture/core-concepts.md deleted file mode 100644 index 26b40dbef8fa..000000000000 --- a/docs/source/legacy/architecture/core-concepts.md +++ /dev/null @@ -1,389 +0,0 @@ -(core-concepts)= - -# Model Definition - -> [!WARNING] -> This page describes the **legacy** TensorRT engine-build workflow. -> For new projects, use [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -TensorRT-LLM has a Model Definition API that can be used to define -Large Language Models. This API is built on top of the powerful -[TensorRT Python API](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/index.html) -to create graph representations of deep neural networks in TensorRT. To become -familiar with the core concepts of the TensorRT API, refer to the -[Core Concepts](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/coreConcepts.html) -section of the TensorRT documentation before proceeding further. - -In TensorRT-LLM, the [`tensorrt_llm.Builder`](source:tensorrt_llm/builder.py) class -contains a -[`tensorrt.Builder`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#id1) -object. That instance is used in the `tensorrt_llm.Builder.create_network` -method to create an instance of the -[`tensorrt.INetworkDefinition`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Network.html#tensorrt.INetworkDefinition) -class. The `INetworkDefinition` object can then be populated using the free -functions defined in the -[`tensorrt_llm.functional`](source:tensorrt_llm/functional.py). - -A simple example of such a free function is `tensorrt_llm.activation` that inserts a -[`tensorrt.IActivationLayer`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Layers.html#tensorrt.IActivationLayer) -node in the graph of the model: - -```python -# In tensorrt_llm.functional: - -def activation(input: Tensor, act_type: trt.ActivationType) -> Tensor: - layer = default_trtnet().add_activation(input.trt_tensor, act_type) # default_trtnet() -> INetworkDefinition - return _create_tensor(layer.get_output(0), layer) -``` - -To make it even easier for users, a few of the most standard activation -functions found in LLMs are derived from that function: - -```python -# In tensorrt_llm.functional: - -relu = partial(activation, act_type=trt.ActivationType.RELU) -sigmoid = partial(activation, act_type=trt.ActivationType.SIGMOID) - -``` - -Specialized activation functions can be used to assemble more advanced -functions such as the `silu` activation: - -```python -# In tensorrt_llm.functional: - -def silu(input: Tensor) -> Tensor: - return input * sigmoid(input) -``` - -When the TensorRT-LLM's Model Definition API is utilized, a graph of the network is -assembled. The graph can later be traversed or transformed using the graph -traversal API exposed by the -[`tensorrt.ILayer`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/LayerBase.html#tensorrt.ILayer) -class. That graph will also be optimized by TensorRT during the compilation of -the engine, as explained in the next section. - -# Compilation - -Once populated, the instance of the -[`tensorrt.INetworkDefinition`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Graph/Network.html#tensorrt.INetworkDefinition), -can be compiled into an efficient engine by the -[`tensorrt.Builder`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#id1) -In TensorRT-LLM, it is done through the `build_engine` member function of the -`tensorrt_llm.Builder` class that calls the -[`build_serialized_network`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#tensorrt.Builder.build_serialized_network -method of the -[`tensorrt.Builder`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Core/Builder.html#id1) -object. That call, if everything works as expected, produces an instance of the -[`tensorrt.IHostMemory`](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/FoundationalTypes/HostMemory.html#tensorrt.IHostMemory) -class. That object is an optimized TensorRT engine that can be stored as a -binary file. - -## TensorRT Compiler - -The TensorRT compiler can sweep through the graph to choose the best kernel for each operation and available GPU. Crucially, it can also identify patterns in the graph where multiple operations are good candidates for being fused into a single kernel. This reduces the required amount of memory movement and the overhead of launching multiple GPU kernels. - -TensorRT also compiles the graph of operations into a single [CUDA Graph](https://developer.nvidia.com/blog/cuda-graphs/) that can be launched all at one time, further reducing the kernel launch overhead. - -The TensorRT compiler is extremely powerful for fusing layers and increasing execution speed, but there are some complex layer fusions—like [FlashAttention](https://arxiv.org/abs/2307.08691) — that involve interleaving many operations together and which can’t be automatically discovered. For those, you can explicitly replace parts of the graph with [plugins](#plugins) at compile time. - -## Model Engine - -The engine file contains the information that you need for executing the model, but LLM usage in practice requires much more than a single forward pass through the model. TensorRT-LLM includes a highly optimized C++ runtime for executing built LLM engines and managing processes like sampling tokens from the model output, managing the KV cache, and batching requests together. - -You can use that runtime directly to execute the model locally, or you can use the TensorRT-LLM runtime backend for NVIDIA Triton Inference Server to serve the model for multiple users. - -## Weight Bindings - -TensorRT engines embed the network weights, that must be known for compilation. -For that reason, the weights must be bound to parameters in the model -definition before calling `tensorrt_llm.Builder.build_engine`. It leads to code like: - -```python -# The Linear operator exposes two parameters (see tensorrt_llm/layers/linear.py): -class Linear(Module): - def __init__(self, ...): - self.weight = Parameter(shape=(self.out_features, self.in_features), dtype=dtype) - self.bias = Parameter(shape=(self.out_features, ), dtype=dtype) - -# The parameters are bound to the weights before compiling the model. See examples/models/core/gpt/weight.py: -tensorrt_llm_gpt.layers[i].mlp.fc.weight.value = fromfile(...) -tensorrt_llm_gpt.layers[i].mlp.fc.bias.value = fromfile(...) -``` - -Note that TensorRT can also -[refit](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#refitting-engine-c) -engines to update the weights after compilation. This feature is available to -TensorRT-LLM users through the `refit_engine` method in the -`tensorrt_llm.Builder` class. - -## Pattern-Matching and Fusion - -One of the key steps performed by TensorRT when it compiles the network graph -is the fusion of operations. Fusion is a well-known technique to improve the -efficiency when executing LLMs. It helps reduce the amount of data transferred -between the memory (DRAM) and the compute cores (CUDA cores as well as Tensor -Cores located on the [Streaming -Multiprocessors](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#introduction) -of a GPU). It also removes kernel launch overhead (each time a kernel is -launched on the GPU, there is a small additional CPU cost that is called the -launch overhead). A classical example is the fusion of the activation function -with the matrix multiplication (matmul) that usually precedes it in the -network. - -In TensorRT-LLM, when defining the model, such a sequence can be written as: - -```python -c = tensorrt_llm.functional.matmul(a, b) -c = tensorrt_llm.functional.relu(c) -``` - -During inference, if the above sequence is executed without fusion, the `c` -tensor has to be written to global memory at the end of the `matmul`, read from -that same memory in `relu` and written again after `relu`. If no other -operation uses the intermediate values between `matmul` and `relu`, it is -suboptimal. That is why, during compilation, TensorRT will identify that -pattern and automatically produce a GPU kernel that applies `relu` at the end -of `matmul` without an intermediate step through global memory. With that -optimization, the `c` tensor is written only once (after `relu`) instead of -twice, and is not read between the two operations. - -The process of identifying the sequences of operations that can be fused is -called _pattern-matching_. TensorRT has a powerful pattern-matching algorithm -that can identify a lot of possible fusions. All the identified patterns are -converted into more efficient kernels by an advanced kernel compiler. - -## Plugins - -The number of possible fusions is almost infinite and some useful fusions -involve very advanced modifications of the graph. A well-known example -is the [Flash-Attention](https://arxiv.org/abs/2205.14135) technique to -optimize the [Multihead-Attention](https://arxiv.org/abs/1706.03762) block -found in many LLMs. Flash-Attention requires modifications to the arithmetic -performed in the sequence `BMM-Softmax-BMM` (where `BMM` stands for Batched -Matrix-Matrix product) and the interleaving of the `for`-loops of the two -batched matrix products. That's non-trivial and not necessarily something -you can expect a compiler to "discover" on its own (or it might require the -support for a [polyhedral -model](https://en.wikipedia.org/wiki/Polytope_model)). - -As a result, even if TensorRT has a powerful pattern-matching algorithm and -supports a lot of possible fusions, there is always the risk that it cannot -identify uncommon and/or very advanced patterns. To overcome that inevitable -limitation, TensorRT offers a powerful mechanism known as -[plugins](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/infer/Plugin/pyPlugin.html). - -The plugins are nodes inserted in the network graph definition that map to user-defined -GPU kernels. TensorRT-LLM uses a number of such plugins. They can be found in -the [`cpp/tensorrt_llm/plugins`](source:/cpp/tensorrt_llm/plugins) directory. - -Plugins are written in C++ and follow a well-defined interface described in the -[Extending TensorRT with Custom Layers](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/extending-custom-layers.html) -section of the TensorRT -[Developer Guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/index.html). -When executed within a TensorRT engine, plugins trigger the execution of -their encapsulated GPU kernels. A fairly simple example of plugins is the -[`QuantizeTensorPlugin`](source:/cpp/tensorrt_llm/plugins/quantizeTensorPlugin) that -triggers a CUDA kernel in the `QuantizeTensorPlugin::enqueue` member function: - -```cpp -// In cpp/tensorrt_llm/plugins/quantizeTensorPlugin/quantizeTensorPlugin.cpp: - -int QuantizeTensorPlugin::enqueue(...) { - if (inputDesc[0].type == DataType::kFLOAT) { - invokeQuantization(...); - } else { - invokeQuantization(...); - } - return 0; -} - -// In cpp/tensorrt_llm/kernels/quantization.cu: - -template -void invokeQuantization(...) { - // The standard <<< >>> construct to launch CUDA kernels - quantizedKernel<<>>(...); -} -``` - -For more details on how TensorRT-LLM implements the GPT Attention operator, see -the [Multi-head, Multi-query and Group-query Attention](../advanced/gpt-attention.md) document. - -# Runtime - -TensorRT-LLM includes an API to implement Python and C++ runtimes. The role of -the runtime components is to load the TensorRT engines and drive their -execution. Typically, for an auto-regressive model like GPT, the runtime is in -charge of loading the engine that implements both the processing of the input -sequence as well as the body of the generation loop. See the [GPT C++ -Runtime](../advanced/gpt-runtime.md) document for details on the C++ Runtime. - -(multi-gpu-multi-node)= - -# Multi-GPU and Multi-Node Support - -Even if TensorRT is designed for single-GPU systems, TensorRT-LLM adds the -support for systems with multiple GPUs and nodes. It is enabled -using TensorRT plugins that wrap communication primitives from the -[NCCL](https://developer.nvidia.com/nccl) library as well as a custom -plugin that optimize the All-Reduce primitive in the presence of All-to-all -connections between GPUs (through NVSwitch in DGX systems). - -The communication plugins can be found in -[cpp/tensorrt_llm/plugins/ncclPlugin](source:cpp/tensorrt_llm/plugins/ncclPlugin) -and the multi-GPU functions are exposed in the TensorRT-LLM Model Definition API -as: - -```python -# In tensorrt_llm/functional.py: - -# Collectives. -def allreduce(tensor: Tensor, group: List[int]) -> Tensor -def allgather(tensor: Tensor, group: List[int], gather_dim: int = 0) -> Tensor - -# Point-to-point communication primitives. -def send(tensor: Tensor, tgt: int) -> Tensor -def recv(tensor: Tensor, src: int) -> Tensor -``` - -The multi-GPU support can be enabled through two different modes of model -parallelism: Tensor Parallelism and Pipeline Parallelism. The former mode -splits the different layers of a model across the GPUs. Each GPU runs the -entire network and synchronizes with its siblings when needed. The Pipeline -Parallelism distributes the different layers to the GPUs. Each GPU runs a -subset of the entire model and communications happen at the boundary of those -subsets of layers. Tensor Parallelism usually leads to more balanced executions -but requires more memory bandwidth between the GPUs. Pipeline Parallelism -reduces the need for high-bandwidth communication but may incur load-balancing -issues and may be less efficient in terms of GPU utilization. - -## Examples - -Here are examples of Llama 3.1 70B and Llama 3.1 405B showing how to perform multi-GPU and multi-node inference in TensorRT-LLM. The example of Llama 3.1 70B performs multi-GPU inference on a single node, while the example of Llama 3.1 405B performs multi-node inference. - -### Llama 3.1 70B - -The following sample commands build an engine for running the Llama 3.1 70B model with tensor parallelism (TP=4) using 4 GPUs on a single node. - -```bash -folder_trt_llm=../TensorRT-LLM -model_dir=Llama-3.1-70B -ckpt_dir=ckpt_llama_3.1_70b -engine_dir=engine_llama_3.1_70b -dtype=bfloat16 -tp_size=4 -pp_size=1 -kv_cache_type=paged -max_input_len=128 -max_output_len=128 -max_batch_size=4 -workers=$(( tp_size * pp_size )) - -python ${folder_trt_llm}/examples/models/core/llama/convert_checkpoint.py \ - --output_dir ${ckpt_dir} \ - --model_dir ${model_dir} \ - --dtype ${dtype} \ - --tp_size ${tp_size} \ - --pp_size ${pp_size} \ - --workers ${workers} \ - --use_parallel_embedding - -trtllm-build \ - --output_dir ${engine_dir} \ - --checkpoint_dir ${ckpt_dir} \ - --gemm_plugin ${dtype} \ - --gpt_attention_plugin ${dtype} \ - --kv_cache_type ${kv_cache_type} \ - --max_input_len ${max_input_len} \ - --max_seq_len $(( max_input_len + max_output_len )) \ - --max_batch_size ${max_batch_size} \ - --workers ${workers} -``` - -The following sample commands perform inference using 4 GPUs on a single node by running `examples/run.py`. - -```bash -input_text="Born in north-east France, Soyer trained as a" - -mpirun -n $(( tp_size * pp_size )) \ - python ${folder_trt_llm}/examples/run.py \ - --engine_dir ${engine_dir} \ - --tokenizer_dir ${model_dir} \ - --input_text "${input_text}" \ - --max_output_len ${max_output_len} -``` - -### Llama 3.1 405B - -The following sample commands build an engine for running the Llama 3.1 405B model with tensor parallelism (TP=16) on 2 nodes that each have 8 GPUs. Although the model runs on multiple nodes, you can build the engine on a single node. - -```bash -folder_trt_llm=../TensorRT-LLM -model_dir=Llama-3.1-405B -ckpt_dir=ckpt_llama_3.1_405b -engine_dir=engine_llama_3.1_405b -dtype=bfloat16 -tp_size=16 -pp_size=1 -kv_cache_type=paged -max_input_len=128 -max_output_len=128 -max_batch_size=4 -workers=8 - -python ${folder_trt_llm}/examples/models/core/llama/convert_checkpoint.py \ - --output_dir ${ckpt_dir} \ - --model_dir ${model_dir} \ - --dtype ${dtype} \ - --tp_size ${tp_size} \ - --pp_size ${pp_size} \ - --workers ${workers} \ - --use_parallel_embedding - -trtllm-build \ - --output_dir ${engine_dir} \ - --checkpoint_dir ${ckpt_dir} \ - --gemm_plugin ${dtype} \ - --gpt_attention_plugin ${dtype} \ - --kv_cache_type ${kv_cache_type} \ - --max_input_len ${max_input_len} \ - --max_seq_len $(( max_input_len + max_output_len )) \ - --max_batch_size ${max_batch_size} \ - --workers ${workers} -``` - -The following sample script, `launch_llama_3.1_405b.sh`, shows how to perform inference with Slurm on 2 nodes that each have 8 GPUs. If you use a different workload management software, the key concern is to run the `examples/run.py` command. - -```bash -#!/bin/bash -#SBATCH --account account -#SBATCH --partition partition -#SBATCH --job-name job-name -#SBATCH --time 1:00:00 -#SBATCH --nodes 2 - -folder_trt_llm=../TensorRT-LLM -engine_dir=engine_llama_3.1_405b -model_dir=Llama-3.1-405B -max_output_len=128 - -input_text="Born in north-east France, Soyer trained as a" - -srun \ - --ntasks-per-node 8 \ - --mpi pmix \ - python ${folder_trt_llm}/examples/run.py \ - --engine_dir ${engine_dir} \ - --tokenizer_dir ${model_dir} \ - --input_text "${input_text}" \ - --max_output_len ${max_output_len} -``` - -You can perform inference by running the script on the Slurm cluster. - -```bash -sbatch launch_llama_3.1_405b.sh -``` diff --git a/docs/source/legacy/architecture/model-weights-loader.md b/docs/source/legacy/architecture/model-weights-loader.md deleted file mode 100644 index d41191364f1e..000000000000 --- a/docs/source/legacy/architecture/model-weights-loader.md +++ /dev/null @@ -1,256 +0,0 @@ -# TensorRT-LLM Model Weights Loader - -## Overview - -The weights loader is designed for easily converting and loading external weight checkpoints into TensorRT-LLM models. - -## Workflow - -Weight checkpoints can be generated from all sources, and may have different naming and data layouts compared to TRT-LLM's requirements. E.g.: - -```bash -# HuggingFace LLaMA checkpoints -{ - "model.embed_tokens.weight": torch.Tensor([vocab_size, hidden_size]) - "model.layers.0.input_layernorm.weight": torch.Tensor([hidden_size]), - "model.layers.0.mlp.down_proj.weight": torch.Tensor([hidden_size, inter_size]), - "model.layers.0.mlp.gate_proj.weight": torch.Tensor([inter_size, hidden_size]), - "model.layers.0.mlp.up_proj.weight": torch.Tensor([inter_size, hidden_size]), - "model.layers.0.post_attention_layernorm.weight": torch.Tensor([hidden_size]), - "model.layers.0.self_attn.q_proj.weight": torch.Tensor([hidden_size, hidden_size]), - "model.layers.0.self_attn.k_proj.weight": torch.Tensor([hidden_size, hidden_size]), - "model.layers.0.self_attn.v_proj.weight": torch.Tensor([hidden_size, hidden_size]), - "model.layers.0.self_attn.o_proj.weight": torch.Tensor([hidden_size, hidden_size]), - ..., -} -# TensorRT-LLM expected weights -{ - "transformer.vocab_embedding.weight": torch.Tensor([vocab_size, hidden_size]) - "transformer.layers.0.input_layernorm.weight": torch.Tensor([hidden_size]), - "transformer.layers.0.mlp.down_proj.weight": torch.Tensor([hidden_size, inter_size]), - "transformer.layers.0.mlp.gate_proj.weight": torch.Tensor([inter_size, hidden_size]), - "transformer.layers.0.mlp.up_proj.weight": torch.Tensor([inter_size, hidden_size]), - "transformer.layers.0.post_layernorm.weight": torch.Tensor([hidden_size]), - "transformer.layers.0.attention.qkv.weight": torch.Tensor([hidden_size * 3, hidden_size]), # Different layout - "transformer.layers.0.attention.dense.weight": torch.Tensor([hidden_size, hidden_size]), - ..., -} -``` - -Conversion means converting the dictionary of `{external_keys:external_weights}` into `{tllm_keys:tllm_weights}`, it includes changing the naming logic and data layouts, and consists of the following parts: - -1. Translate a TRT-LLM parameter name into external-format name(s). -2. Loading tensor slice(s) according to the translated names. -3. Postprocess the tensor(s) into target layout. - -### Translator - -TRT-LLM parameter names are translated in units of sections divided by dots. E.g.: - -| TensorRT-LLM key | `transformer` |.| `layers` |.| `0` |.| `attention` |.| `dense` |.| `weight` | -| :---------------------: | :-----------: |-| :------: |-|:---:|-| :---------: |-| :------: |-| :------: | -| Translated external key | `model` |.| `layers` |.| `0` |.| `self_attn` |.| `o_proj` |.| `weight` | - -The mapping between TRT-LLM keywords and HF keywords are described in `tllm_to_externel_key_dict` of `ModelWeightsLoader` class object. \ -If any of the mappings has one-to-multiple corresponding, the translated key will get multiplied accordingly. E.g.: - -| TensorRT-LLM key and related keyword mapping | Translated external keys | -| :----------------------------------------------------------: | :----------------------: | -| `transformer.layers.0.attention.qkv.weight`
`{"qkv":[q_proj, k_proj, v_proj]}` | `model.layers.0.self_attn.q_proj.weights`
`model.layers.0.self_attn.k_proj.weights`
`model.layers.0.self_attn.v_proj.weights`| -| `transformer.layers.0.mlp.fc.weight`
`{"weight":[qweight, qzeros, scales]}` | `model.layers.0.mlp.gate_proj.qweight`
`model.layers.0.mlp.gate_proj.qzeros`
`model.layers.0.mlp.gate_proj.scales`| - -The default `tllm_to_externel_key_dict` is based on HF LLaMA as: - -```python -class ModelWeightsLoader: - def __init__(self, model_dir, customized_key_dict: dict = {}) -> None: - ... - self.tllm_to_externel_key_dict = { - "transformer": "model", - "vocab_embedding": "embed_tokens", - "lm_head": "lm_head", - "ln_f": "norm", - "attention": "self_attn", - "qkv": ["q_proj", "k_proj", "v_proj"], - "dense": "o_proj", - "gate": "up_proj", - "proj": "down_proj", - "fc": "gate_proj", - "input_layernorm": "input_layernorm", - "post_layernorm": "post_attention_layernorm", - } - self.tllm_to_externel_key_dict.update(customized_key_dict) - ... -``` - -It can be updated through passing `customized_key_dict` when initializing `ModelWeightsLoader`. - -The dictionary will also get updated according to the layer classes. When iterating over parameters, -if the layer class has attribute `tllm_to_externel_key_dict`, for keywords exist both in the default one and the layer-specified one, -the weight loader will translate according to the layer attribute with higher priority. -This can enable the support for different quantization precisions automatically. - - -### Loading function - -The loading function can load an arbitrary tensor slice according to its `key`, `tp_size`, `tp_dim` and `tp_rank`. - -The template for loading function is as following. - -```python -def load_tensor(self, key, tp_size, tp_dim, tp_rank): - # Retrieve file pointer index - if key in self.shard_map: - ptr_idx = self.shard_map[key] - else: - return None - - # Load tensor from the corresponding shard - if self.format == ModelWeightsFormat.SAFETENSORS: - tensor = self.shards[ptr_idx].get_slice(key) - tensor_shape = tensor.get_shape() - else: - ... - - # Shard and return a tensor slice - slice_shape = ... - return tensor[slice_shape] -``` - -When initializing the `ModelWeightsLoader` object, the file format will be derived from `model_dir` through `detect_format`. The following formats are supported for now: - - * Directory contains or file named `*.safetensors` (Recommended, has better performance) - * Directory contains or file named `*.bin` - * Directory contains or file named `*.pth` - -To support other formats or in-memory loaded models, the format need to be claimed in `ModelWeightsFormat`, `detect_format()`, `preload()` and `load_tensor()`. - -### Postprocessing functions - -After translation and loading, a TRT-LLM key will become a tensor or a list of tensors, which is the input of postprocessing functions. \ -Operations including QKV concatenating, MoE weight stacking and weight-only quantization can be handled here. -The template of postprocessing function is: - -```python -# Example for 1-1 weights mapping -class CustomizedModuleA(Module): - def __init__(...): - super().__init__(...) - ... - self.tp_dim = 0 # Need to set or inherit from parent class - - def postprocess(self, tllm_key, weights, **kwargs): - weights = proc(weights) - return {tllm_key: weights} - -# Example for multiple-multiple weights mapping -class CustomizedModuleB(Module): - def __init__(...): - super().__init__(...) - ... - self.tp_dim = 0 # Need to set or inherit from parent class - # The default value of "weights" in tllm_to_externel_key_dict will be override - self.tllm_to_externel_key_dict = {"weight": ["qweight", "qzeros", "scales"]} - - def postprocess(self, tllm_key, weights, **kwargs): - # Skipped the postprocess of zeros and weights_scaling_factor - # They are loaded in the postprocess of weight - config = kwargs.get("config", None) # Passed through kwargs by default - if not tllm_key.endswith("weight"): - return {} - # The order in weights is defined in tllm_to_externel_key_dict - qweight, qzeros, scales = weights - proccessed_weight, proccessed_zeros = proc(qweight, qzeros, config.num_heads) - return { - tllm_key: proccessed_weight, - tllm_key.replace("weight", "zeros"): proccessed_zeros, - tllm_key.replace("weight", "weights_scaling_factor"): scales, - } -``` - -## Examples - -The `ModelWeightsLoader` class can support different models with the following levels: - -### Natively supported models -For models with native support, users can call the default weight loader without any other operations. -```python -# Using the model weights loader for LLaMA -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -loader = ModelWeightsLoader(external_checkpoint_dir) -loader.generate_tllm_weights(trtllm_model) -``` -For calibration-free quantization precisions, passing a properly quantized `trtllm_model` will let the weight loader load at the given precision accordingly. The configurations will be read from `trtllm_model.config` automatically. For now, LLaMA family models using the default `tllm_to_externel_key_dict` is supported natively. - -### Models with customized key names -For models with different naming logic, users can still call the default weight loader with `customized_key_dict` specified. -```python -# Using the model weights loader for the LLM part of LLaVA -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -llava_dict = { - "transformer": "language_model.model", - "lm_head": "language_model.lm_head" -} -loader = ModelWeightsLoader(external_checkpoint_dir, llava_dict) -loader.generate_tllm_weights(trtllm_model) -``` -Users need to specify the different part from the default `tllm_to_externel_key_dict`. The loader still have support across different precisions. -The support for LLaVA and Exaone is in `LLaMAForCausalLM.from_hugging_face()` of [model.py](../../../tensorrt_llm/models/llama/model.py), and can also be taken as examples. - -### Models with customized weight layout -For models with different weight layout, users can write the conversion loop explicitly and do customized operations. -```python -# Using the model weights loader for BLOOM -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -bloom_dict = { - "transformer": "", - "layers": "h", - "ln_f": "ln_f", - "lm_head": "word_embeddings", - "ln_embed": "word_embeddings_layernorm", - "vocab_embedding": "word_embeddings", - "attention": "self_attention", - "qkv": "query_key_value", - "dense": "dense", - "fc": "dense_h_to_4h", - "proj": "dense_4h_to_h", - "post_layernorm": "post_attention_layernorm", -} -loader = ModelWeightsLoader(external_checkpoint_dir, bloom_dict) -# See ModelWeightsLoader.generate_tllm_weights() -loader.update_key_mapping(trtllm_model) -tllm_weights = {} -for tllm_key, _ in tqdm(trtllm_model.named_parameters()): - if tllm_key.endswith("qkv"): - # Passing the callable handle - tllm_weights.update(loader.load(tllm_key, preprocess=customized_preprocess)) - else: - tllm_weights.update(loader.load(tllm_key)) -loader.fill(tllm_weights) -``` -This will apply `preprocess` after `load_tensor()` and before `postprocess`, and demonstrates how to convert the loaded shard into default HF layout. The loader still have support for precisions quantized from FP16/BF16 (e.g. INT8-wo/INT4-wo), the other precisions may require special operations, and can be addressed inside the `preprocess` function. -The support for Qwen-1 is in `QWenForCausalLM.from_hugging_face()` of [model.py](../../../tensorrt_llm/models/qwen/model.py), and can also be taken as example. - -### Fully customized -If the model weights loader cannot satisfy the requirements, users can write the conversion loop totally on their own. -```python -tllm_weights = {} -for tllm_key, param in tqdm(trtllm_model.named_parameters()): - # Load from external checkpoints - # The load_tensor() function can also be called here - tensor = ... - # Convert tensor and set the values according to the config - if trtllm_model.config.quantization.quant_algo == xxx: - ... - else: - ... - param.value = tensor -``` -In this mode, every precision require user's own support. - -## Troubleshooting -The weights loader is enabled for LLaMA family models and Qwen models by default with TensorRT flow only. - -If users encounter failure caused by `ModelWeightsLoader`, a workaround is passing environmental variable `TRTLLM_DISABLE_UNIFIED_CONVERTER=1` to disable the model weights loader and fallback to the legacy path. - -This workaround will be removed in future version after the LLaMA/Qwen weights conversion is stable. diff --git a/docs/source/legacy/architecture/workflow.md b/docs/source/legacy/architecture/workflow.md deleted file mode 100644 index 3c880e9723fb..000000000000 --- a/docs/source/legacy/architecture/workflow.md +++ /dev/null @@ -1,241 +0,0 @@ -# TensorRT-LLM Build Workflow - -> [!WARNING] -> This page describes the **legacy** TensorRT engine-build workflow. -> For new projects, use [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -## Overview - - -The build workflow contains two major steps. - -1. Create TensorRT-LLM models from existing model checkpoints exported by the training framework. -2. Build the TensorRT-LLM models to TensorRT-LLM engines. - -To generalize the TensorRT-LLM optimization features to all models, and to share the same workflow between different models for TensorRT-LLM users, TensorRT-LLM has conventions about how the models shall be defined and how the models shall be imported. - -TensorRT-LLM checkpoint convention is documented in [](checkpoint.md) and all decoder-only models had been migrated to adopt the convention. Model-specific convert_checkpoint.py scripts are shipped as source code in example directories, and a trtllm-build CLI tool had been added. However, there are some disadvantages of providing convert checkpoint scripts outside the core TensorRT-LLM lib as example: - -1. TensorRT-LLM evolves so quickly that the model's definition code might have changed for better performance; which means the `convert_checkpoint.py` is out of date. - - -2. TensorRT-LLM is creating a new set of high-level APIs which handle model conversion, engine building, and inference in one class for easier-of-use. Thus, the high-level APIs need to call the weights conversion code, which shall be part of TensorRT-LLM core lib, not the example. And the conversion code of different models shall have same interface such that the high-level APIs do not need to add many ad-hoc code for different models. - -To mitigate these issues, the model specific `convert_checkpoint.py` scripts are being refactored. Most of the conversion code will be moved into core lib, sitting next to the model definition. Refer to `tensorrt_llm/models/llama/` as an example. There is a new set of APIs for importing models and converting weights. The 0.9 release refactored the LLaMA model class to adopt the new APIs, others models' refactor work is ongoing. - - -## Conversion APIs - -The API for weight conversion of the LLaMA model looks like this. A `TopModelMixin` class is introduced, `from_hugging_face()` interface is declared, the `LLaMAForCausalLM` class inherits `TopModelMixin` (not direct parent class, but in its base class hierarchy), and implements the interface. - -```python -class TopModelMixin - @classmethod - def from_hugging_face(cls, - hf_model_dir: str, - dtype: Optional[str] = 'float16', - mapping: Optional[Mapping] = None, - **kwargs): - raise NotImplementedError("Subclass shall override this") - -# TopModelMixin is in the part of base class hierarchy -class LLaMAForCausalLM (DecoderModelForCausalLM): - @classmethod - def from_hugging_face(cls, - hf_model_dir, - dtype='float16', - mapping: Optional[Mapping] = None) -> LLaMAForCausalLM: - # creating a TensorRT-LLM llama model object - # converting HuggingFace checkpoint to TensorRT-LLM expected weights dict - # Load the weights to llama model object -``` - - -Then, in the convert_checkpoint.py script in the -[`examples/models/core/llama/`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama/) directory of the GitHub repo, -the logic can be greatly simplified. Even if the model definition code of TensorRT-LLM LLaMA class is changed due to some reason, the `from_hugging_face` API will keep the same, thus the existing workflow using this interface will not be affected. - - -```python -#other args omitted for simplicity here. -llama = LLaMAForCausalLM.from_hugging_face(model_dir, dtype, mapping=mapping) -llama.save_checkpoint(output_dir, save_config=(rank==0)) -``` - -The `from_hugging_face` API does not save the checkpoint into disk intentionally, instead it returns an in-memory object. Call `save_checkpoint` to save the models. This keeps the flexibility and makes the flow of convert->build in one process faster. Typically, saving and loading disk for large models are slower and thus should be avoided. - - -Since LLaMA models were also released with different formats, such as the Meta checkpoint, the `LLaMAForCausalLM` class has a `from_meta_ckpt` function for that. This function is not declared in the `TopModelMixin` class due to it being LLaMA specific, and therefore, other models don't use it. - - -In the 0.9 release, only LLaMA is refactored. Since popular LLaMA (and its variants) models are released by Hugging Face and Meta checkpoint formats, only these two functions are implemented. - - -In future releases, there might be `from_jax`, `from_nemo`, `from_keras` or other factory methods for different training checkpoints added. -For example, the Gemma 2B model and the convert_checkpoint.py file in the [`examples/models/core/gemma`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gemma/) -directory support JAX and Keras formats in addition to Hugging Face. The model developers can choose to implement **any subset** of these factory methods for the models they contributed to TensorRT-LLM. - - -For some formats which are not supported by TensorRT-LLM model developers, you still have the freedom to implement your own weights conversion outside the core lib; the flow will look like this: - - -```python -config = read_config_from_the_custom_training_checkpoint(model_dir) -llama = LLaMAForCausalLM(config) - -# option 1: -# Create a weights dict and then calls LLaMAForCausalLM.load -weights_dict = convert_weights_from_custom_training_checkpoint(model_dir) -llama.load(weights_dict) - -# option 2: -# Internally assign the model parameters directly -convert_and_load_weights_into_trtllm_llama(llama, model_dir) -# Use the llama object as usual, to save the checkpoint or build engines -``` - -Though there are some limitations and pitfalls of doing these custom weights loading, if the model definition is inside TensorRT-LLM core lib, and the weights loading/conversion are outside the core lib, the conversion code might need to be updated when new TensorRT-LLM is released. - - -## Quantization APIs - -TensorRT-LLM relies on NVIDIA Modelopt toolkit to support some of the quantization like: FP8, W4A16_AWQ, W4A8_AWQ, while it also has some its own quantization implementation for Smooth Quant, INT8 KV cache, and INT4/INT8 weight only. - - -In TensorRT-LLM 0.8 version: - -* For Modelopt-supported quantization algorithms, a standalone script, - [example/quantization/quantize.py](https://github.com/NVIDIA/TensorRT-LLM/blob/main/examples/quantization/quantize.py) - can export TensorRT-LLM checkpoints, and the trtllm-build command needs to be executed to build the checkpoints to engines. - -* For the non-Modelopt quantization algorithms, users need to use the per-model convert_checkpoint.py scripts to export TensorRT-LLM checkpoints. - -Use the `quantize()` interface to unify the different quantization flows. The default implementation is added in the `PretrainedModel` class. - - -```python -class PretrainedModel: - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantConfig, - mapping: Optional[Mapping] = None): #some args are omitted here - # Internally quantize the given hugging face models using Modelopt - # and save the checkpoint to output_dir -``` - -* The default implementation only handles the Modelopt supported quantization. The LLaMA class then inherits this `PretrainedModel` and dispatches the Modelopt quantization to the super class's default implementation. -* The model developer raises errors in the sub-class implementation if the new model is not supported by Modelopt yet. - - -```python -class LLaMAForCausalLM: - @classmethod - def quantize( - cls, - hf_model_dir, - output_dir, - quant_config: QuantiConfig, - mapping: Optional[Mapping] = None): #some args are omitted here - use_modelopt_quantization = ... # determine if to use Modelopt or use native - if use_modelopt_quantization: - super().quantize(hf_model_dir, - output_dir, - quant_config) - else: - # handles TensorRT-LLM native model specific quantization - # or raise exceptions if not supported -``` - - -The `quantize` API is designed to take multi-GPU resources internally to make quantization. For example, a LLaMA 70B BF16 takes 140G memory, if we make FP8 quantization, then, another 70G is needed. So, we need at least 210G, 4 * A100(H100) is needed to quantize the LLaMA 70B model. If you want to call `quantize` API inside a MPI program, be cautious and ensure the quantize API is only called by rank 0. - - -Usage of the `quantize` API in an MPI program looks like this, only rank 0 calls it. In an non-MPI program, the `if rank == 0` and the `mpi_barrier()` are not needed. - -```python -quant_config = QuantConfig() -quant_config.quant_algo = quant_mode.W4A16_AWQ -mapping = Mapping(world_size=tp_size, tp_size=tp_size) -if rank == 0: - LLaMAForCausalLM.quantize(hf_model_dir, - checkpoint_dir, - quant_config=quant_config) -mpi_barrier() # wait for rank-0 to finish the quantization -llama = LLaMAForCausalLM.from_checkpoint(checkpoint_dir, rank) -engine = build(llama, build_config) -engine.save(engine_dir) -``` - - -The `examples/quantization/quantize.py` is kept for backward compatibility. - - -## Build APIs - - -The `tensorrt_llm.build` API builds the TensorRT-LLM model object to TensorRT-LLM engine. This new API replaced the older flow: creating a builder, creating a network object, tracing the model to the network, and building TensorRT engines. -The usage of this API looks like this: - -```python -llama = ... # create LLaMAForCausalLM object -build_config = BuildConfig(max_batch_size=1) -engine = tensorrt_llm.build(llama, build_config) -engine.save(engine_dir) -``` - - -The Llama object can be created by any method mentioned in the [](#conversion-apis) or [](#quantization-apis) sections. - - -The `trtllm-build` CLI tool is a thin wrapper around this `tensorrt_llm.build` API. The flags of the CLI tool are kept close to the fields of the `BuildConfig` class. - - -If a model were to be saved into disk and then built to the engine later, TensorRT-LLM provides a `from_checkpoint` API to deserialize the checkpoint. - -```python -## TensorRT-LLM code -class PretrainedModel: - @classmethod - def from_checkpoint(cls, - ckpt_dir: str, - rank: int = 0, - config: PretrainedConfig = None): - # Internally load the model weights from a given checkpoint directory -``` - - -The `from_checkpoint` API is called to deserialize the checkpoint to a model object. The `tensorrt_llm.build` API can be called to build the engine. - - -```python -llama = LLaMAForCausalLM.from_checkpoint(checkpoint_dir) -engine = build(llama, build_config) -engine.save(engine_dir) -``` - -## CLI Tools - -All the weights conversion, quantization, and build APIs mentioned above have corresponding CLI tools for convenience. - -* Model specific `convert_checkpoint.py` scripts are inside the `examples//` folder. -* A unified quantization script is inside the `examples/quantization/quantize.py` and can be shared by all **supported** models. -* A `trtllm-build` CLI tool builds all models from TensorRT-LLM checkpoint. - -Refer to the following considerations for the CLI tools: - -* These scripts and tools should be used for scripting. Do not import the Python functions/class defined in these tools. TensorRT-LLM does not promise the content of these scripts can be compatible with previous versions. The options of these tools may also be changed when it’s not avoidable. - -* These scripts in the example folder may use TensorRT-LLM internal/unstable APIs, which is not guaranteed to work if the examples’ version and the TensorRT-LLM install version are mismatched. There are some GitHub issues caused by version mismatch. - - https://github.com/NVIDIA/TensorRT-LLM/issues/1293 - - https://github.com/NVIDIA/TensorRT-LLM/issues/1252 - - https://github.com/NVIDIA/TensorRT-LLM/issues/1079 - - You should always install the same TensorRT-LLM version specified in `examples//requirements.txt`. - -* In the future, the per-model conversion script may or may not be unified to one single script shared by models, given the nature of different models’ attributes may be different. However, the TensorRT-LLM team will try to make sure the flags for the same feature are consistent between different scripts. - -* The TensorRT-LLM team encourages use of the new low-level conversion/quantization/build API instead of these scripts. The conversion APIs will be added model-by-model gradually, which may span a few releases. diff --git a/docs/source/legacy/dev-on-cloud/build-image-to-dockerhub.md b/docs/source/legacy/dev-on-cloud/build-image-to-dockerhub.md deleted file mode 100644 index a07fbfdbc819..000000000000 --- a/docs/source/legacy/dev-on-cloud/build-image-to-dockerhub.md +++ /dev/null @@ -1,56 +0,0 @@ -(build-image-to-dockerhub)= - -# Build the TensorRT LLM Docker Image -When you develop trt-llm on cloud platform such as runpod, you may need to provide a docker image for the platform. So you firstly need to upload the image to dockerhub. - -## Build the TensorRT LLM Docker Image and Upload to DockerHub - -```bash -make -C docker build -``` -Then we can get the docker image named `tensorrt_llm/devel:latest` - -### Enable ssh access to the container -Since the default docker image doesn’t have ssh support, we can’t ssh into it. We need to add ssh support to the container. -Let’s first create a new Dockerfile with below content: - -```Dockerfile -FROM tensorrt_llm/devel:latest - -RUN apt update && apt install openssh-server -y -RUN mkdir -p /run/sshd && chmod 755 /run/sshd -RUN mkdir -p /root/.ssh && chmod 700 /root/.ssh && touch /root/.ssh/authorized_keys && chmod 600 /root/.ssh/authorized_keys -# add sshd to entrypoint script -RUN echo "sshd -E /opt/sshd.log" >> /opt/nvidia/entrypoint.d/99-start-sshd.sh -``` - -If we save this Dockerfile as `Dockerfile.ssh`. Then we can build the docker image with below command: - -```bash -docker build -t tensorrt_llm/devel:with_ssh -f Dockerfile.ssh . -``` - -Then we can get the docker image named `tensorrt_llm/devel:with_ssh` - -## Upload the Docker Image to DockerHub - -You need to register a [dockerhub](https://hub.docker.com) account first if you don't have one. - -Then you can click 'Personal Access Tokens' in the user menu and create a new token. - -With the token, you can login to dockerhub with below command: - -```bash -docker login -u -``` - -Enter the token to the console. - -After login, you can tag and push the docker image to dockerhub with below command: - -```bash -docker tag tensorrt_llm/devel:with_ssh /tensorrt_llm:devel -docker push /tensorrt_llm:devel -``` - -Finally, you can see the docker image in your dockerhub repository and can use it with the link such as `docker.io//tensorrt_llm:devel`. diff --git a/docs/source/legacy/dev-on-cloud/dev-on-runpod.md b/docs/source/legacy/dev-on-cloud/dev-on-runpod.md deleted file mode 100644 index 254f8fe14895..000000000000 --- a/docs/source/legacy/dev-on-cloud/dev-on-runpod.md +++ /dev/null @@ -1,47 +0,0 @@ -(dev-on-runpod)= - -# Develop TensorRT LLM on Runpod -[Runpod](https://runpod.io) is a popular cloud platform among many researchers. This doc describes how to develop TensorRT LLM on Runpod. - -## Prepare - -### Create a Runpod account -Please refer to the [Runpod Getting Started](https://docs.runpod.io/get-started/). - -### Configure SSH Key -Please refer to the [Configure SSH Key](https://docs.runpod.io/pods/configuration/use-ssh). - -Note that we can skip the step of "Start your Pod. Make sure of the following things" here as we will introduce it below. - -## Build the TensorRT LLM Docker Image and Upload to DockerHub -Please refer to the [Build Image to DockerHub](build-image-to-dockerhub.md). - -Note that the docker image must enable ssh access. See on [Enable ssh access to the container](build-image-to-dockerhub.md#enable-ssh-access-to-the-container). - -## Create a Pod Template -Click the "Template" button in the menu, then click "Create Template". - -Fill the docker image link of DockerHub such as `docker.io//tensorrt_llm:devel` on "Docker Image" field. - -Fill "22" into "Expose TCP Ports" field. - -Fill -```bash -sleep infinity -``` -into 'Container Start Command' field. - -## Connect to the Pod -Please refer to the [Connect to the Pod](https://docs.runpod.io/pods/connect-to-a-pod). - -You can connect the pod with SSH or Web Terminal. - -If you want to connect the pod with SSH, you can copy the command from "SSH over exposed TCP" field and run it on your host. - -In some scenarios such as using a team account, your public key has not been added to the pod successfully. You can directly add this command to the 'Container Start Command' field as: - -```bash -bash -c 'echo "" >> ~/.ssh/authorized_keys;sleep infinity' -``` - -Enjoy your development! diff --git a/docs/source/legacy/key-features.md b/docs/source/legacy/key-features.md deleted file mode 100644 index 6fac19262df4..000000000000 --- a/docs/source/legacy/key-features.md +++ /dev/null @@ -1,10 +0,0 @@ -# Key Features - -This document lists key features supported in TensorRT-LLM. - -- [Quantization](reference/precision.md) -- [Inflight Batching](advanced/gpt-attention.md#in-flight-batching) -- [Chunked Context](advanced/gpt-attention.md#chunked-context) -- [LoRA](advanced/lora.md) -- [KV Cache Reuse](advanced/kv-cache-reuse.md) -- [Speculative Sampling](advanced/speculative-decoding.md) diff --git a/docs/source/legacy/performance/perf-analysis.md b/docs/source/legacy/performance/perf-analysis.md deleted file mode 100644 index 71f5320ee3a5..000000000000 --- a/docs/source/legacy/performance/perf-analysis.md +++ /dev/null @@ -1,91 +0,0 @@ -(perf-analysis)= - -# Performance Analysis - -NVIDIA Nsight Systems reports at the application level are highly informative. Metric sampling capabilities have increased over generations and provide a clean middle-ground between timing analysis and kernel-level deep dives with NVIDIA Nsight Compute. - -Given the potential long runtimes of Large Languages Models (LLMs) and the diversity of workloads a model may experience during a single inference pass or binary execution, we have added features to TensorRT-LLM to get the most out of Nsight Systems capabilities. This document outlines those features as well as provides examples of how to best utilize them to understand your application. - - -## Feature Descriptions - -The main functionality here: - * Relies on toggling the CUDA profiler runtime API on and off. - * (PyTorch workflow only) Toggling the PyTorch profiler on and off. - * Provides a means to understand which regions a user may want to focus on. - -Toggling the CUDA profiler runtime API on and off: - * Allows users to know specifically what the profiled region corresponds to. - * Results in smaller files to post-process (for metric extraction or similar). - -(PyTorch workflow only) Toggling the PyTorch profiler on and off: - * Help users to analyze the performance breakdown in the model. - * Results in smaller files to post-process (for metric extraction or similar). - - -## Coordinating with NVIDIA Nsight Systems Launch - -Consult the Nsight Systems User Guide for full overview of options. - -On the PyTorch workflow, basic NVTX markers are by default provided. On the C++/TensorRT workflow, append `--nvtx` when calling `scripts/build_wheel.py` script to compile, and clean build the code. - -### Only collect specific iterations - -To reduce the Nsight Systems profile size, and to control that only specific iterations are collected, set environment variable `TLLM_PROFILE_START_STOP=A-B`, and append `-c cudaProfilerApi` to `nsys profile` command. - - -### Enable more NVTX markers for debugging -Set environment variable `TLLM_NVTX_DEBUG=1`. - -### Enable garbage collection (GC) NVTX markers -Set environment variable `TLLM_PROFILE_RECORD_GC=1`. - - -### Enable GIL information in NVTX markers -Append “python-gil” to Nsys “-t” option. - - -## Coordinating with PyTorch profiler (PyTorch workflow only) - -### Collect PyTorch profiler results -1. Set environment variable `TLLM_PROFILE_START_STOP=A-B` to specify the range of the iterations to be collected. -2. Set environment variable `TLLM_TORCH_PROFILE_TRACE=`, and the results will be saved to ``. - -### Visualize the PyTorch profiler results -Use to inspect the saved profile. - - -## Examples -Consult the Nsight Systems User Guide for full overview of MPI-related options. - -### Profiling specific iterations on a trtllm-bench/trtllm-serve run - -Say we want to profile iterations 100 to 150 on a trtllm-bench/trtllm-serve run, we want to collect as much information as possible for debugging, such as GIL, debugging NVTX markers, etc: - -```bash -#!/bin/bash - -# Prepare dataset for the benchmark -trtllm-bench \ - --model=${MODEL_PATH} prepare-dataset \ - --output /tmp/dataset.txt token-norm-dist --num-requests=${NUM_SAMPLES} \ - --input-mean=1000 --output-mean=1000 --input-stdev=0 --output-stdev=0 - -# Benchmark and profile -TLLM_PROFILE_START_STOP=100-150 nsys profile \ - -o trace -f true \ - -t 'cuda,nvtx,python-gil' -c cudaProfilerApi \ - --cuda-graph-trace node \ - -e TLLM_PROFILE_RECORD_GC=1,TLLM_LLMAPI_ENABLE_NVTX=1,TLLM_TORCH_PROFILE_TRACE=trace.json \ - --trace-fork-before-exec=true \ - trtllm-bench \ # or trtllm-serve command - --model deepseek-ai/DeepSeek-V3 \ - --model_path ${MODEL_PATH} \ - throughput \ - --dataset /tmp/dataset.txt --warmup 0 \ - --streaming -``` - -The Nsight Systems reports will be saved to `trace.nsys-rep`. Use NVIDIA Nsight Systems application to open it. - -The PyTorch profiler results will be saved to `trace.json`. Use to inspect the saved profile. diff --git a/docs/source/legacy/performance/perf-benchmarking.md b/docs/source/legacy/performance/perf-benchmarking.md deleted file mode 100644 index 4fc460596f22..000000000000 --- a/docs/source/legacy/performance/perf-benchmarking.md +++ /dev/null @@ -1,871 +0,0 @@ -(perf-benchmarking)= - -# TensorRT-LLM Benchmarking - -```{important} -This benchmarking suite is a work in progress. -Expect breaking API changes. -``` - -TensorRT-LLM provides the `trtllm-bench` CLI, a packaged benchmarking utility that aims to make it -easier for users to reproduce our officially published [performance overview](../../developer-guide/perf-overview.md#throughput-measurements). `trtllm-bench` provides the follows: - -- A streamlined way to build tuned engines for benchmarking for a variety of models and platforms. -- An entirely Python workflow for benchmarking. -- Ability to benchmark various flows and features within TensorRT-LLM. - -`trtllm-bench` executes all benchmarks using [in-flight batching] -- for more information see -the [this section](../advanced/gpt-attention.md#in-flight-batching) that describes the concept -in further detail. - -## Before Benchmarking - -For rigorous benchmarking where consistent and reproducible results are critical, proper GPU configuration is essential. These settings help maximize GPU utilization, eliminate performance variability, and ensure optimal conditions for accurate measurements. While not strictly required for normal operation, we recommend applying these configurations when conducting performance comparisons or publishing benchmark results. - -### Persistence mode -Ensure persistence mode is enabled to maintain consistent GPU state: -```shell -sudo nvidia-smi -pm 1 -``` - -### GPU Clock Management -Allow the GPU to dynamically adjust its clock speeds based on workload and temperature. While locking clocks at maximum frequency might seem beneficial, it can sometimes lead to thermal throttling and reduced performance. Reset GPU clocks using: -```shell -sudo nvidia-smi -rgc -``` - -### Set power limits -First query the maximum power limit: -```shell -nvidia-smi -q -d POWER -``` -Then configure the GPU to operate at its maximum power limit for consistent performance: -```shell -sudo nvidia-smi -pl -``` - -### Boost settings -Potentially a GPU may support boost levels. First query available boost levels: -```shell -sudo nvidia-smi boost-slider -l -``` -If supported, enable the boost slider using one of the available levels for maximum performance: -```shell -sudo nvidia-smi boost-slider --vboost -``` - - -## Throughput Benchmarking - -### Limitations and Caveats - -#### Validated Networks for Benchmarking - -While `trtllm-bench` should be able to run any network that TensorRT-LLM supports, the following are the list -that have been validated extensively and is the same listing as seen on the -[Performance Overview](../../developer-guide/perf-overview.md) page. - -- [meta-llama/Llama-2-7b-hf](https://huggingface.co/meta-llama/Llama-2-7b-hf) -- [meta-llama/Llama-2-70b-hf](https://huggingface.co/meta-llama/Llama-2-70b-hf) -- [tiiuae/falcon-180B](https://huggingface.co/tiiuae/falcon-180B) -- [EleutherAI/gpt-j-6b](https://huggingface.co/EleutherAI/gpt-j-6b) -- [meta-llama/Meta-Llama-3-8B](https://huggingface.co/meta-llama/Meta-Llama-3-8B) -- [meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B) -- [meta-llama/Meta-Llama-3-70B](https://huggingface.co/meta-llama/Meta-Llama-3-70B) -- [meta-llama/Llama-3.1-70B](https://huggingface.co/meta-llama/Llama-3.1-70B) -- [meta-llama/Llama-3.1-405B](https://huggingface.co/meta-llama/Llama-3.1-405B) -- [mistralai/Mixtral-8x7B-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-v0.1) -- [mistralai/Mistral-7B-v0.1](https://huggingface.co/mistralai/Mistral-7B-v0.1) -- [meta-llama/Llama-3.1-8B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct) -- [meta-llama/Llama-3.1-70B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-70B-Instruct) -- [meta-llama/Llama-3.1-405B-Instruct](https://huggingface.co/meta-llama/Llama-3.1-405B-Instruct) -- [mistralai/Mixtral-8x7B-Instruct-v0.1](https://huggingface.co/mistralai/Mixtral-8x7B-Instruct-v0.1) - -```{tip} -`trtllm-bench` can automatically download the model from Hugging Face Model Hub. -Export your token in the `HF_TOKEN` environment variable. -``` - -#### Supported Quantization Modes - -`trtllm-bench` supports the following quantization modes: - -- None (no quantization applied) -- `FP8` -- `NVFP4` - -For more information about quantization, refer to [](../reference/precision.md) and -the [support matrix](../reference/precision.md#support-matrix) of the supported quantization methods for each network. - -```{tip} -Although TensorRT-LLM supports more quantization modes than listed above, `trtllm-bench` currently only configures for -a smaller subset. -``` - -### Quickstart - -This quick start focuses on running a short max throughput benchmark on -`meta-llama/Llama-3.1-8B` on a synthetic dataset with a uniform distribution of prompts with ISL:OSL -of 128:128. -To run the benchmark from start to finish, run the following commands: - -```shell -trtllm-bench --tokenizer meta-llama/Llama-3.1-8B prepare-dataset --output /tmp/synthetic_128_128.txt token-norm-dist --input-mean 128 --output-mean 128 --input-stdev 0 --output-stdev 0 --num-requests 3000 -trtllm-bench --model meta-llama/Llama-3.1-8B build --dataset /tmp/synthetic_128_128.txt --quantization FP8 -trtllm-bench --model meta-llama/Llama-3.1-8B throughput --dataset /tmp/synthetic_128_128.txt --engine_dir /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -``` - -After the benchmark completes, `trtllm-bench` prints a summary with summary metrics. - -```shell -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: meta-llama/Llama-3.1-8B -Engine Directory: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -TensorRT-LLM Version: 0.17.0 -Dtype: bfloat16 -KV Cache Dtype: FP8 -Quantization: FP8 -Max Input Length: 256 -Max Sequence Length: 256 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -Max Runtime Batch Size: 4096 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 5.0689E+14 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 3000 -Average Input Length (tokens): 128.0000 -Average Output Length (tokens): 128.0000 -Token Throughput (tokens/sec): 28390.4265 -Request Throughput (req/sec): 221.8002 -Total Latency (ms): 13525.6862 - -=========================================================== -``` - -### Workflow - -The workflow for `trtllm-bench` is composed of the following steps: - -1. Prepare a dataset to drive the inflight batching benchmark. -2. Build a benchmark engine using `trtllm-bench build` subcommand (not required for [PyTorch flow](#running-with-the-pytorch-workflow)). -3. Run the max throughput benchmark using the `trtllm-bench throughput` subcommand or low latency benchmark using the `trtllm-bench latency` subcommand. - - -#### Preparing a Dataset - -The throughput benchmark utilizes a fixed JSON schema to specify requests. The schema is defined as follows: - -| Key | Required | Type | Description | -| :-------------- | :------: | :-----------: | :---------------------------------------------- | -| `task_id` | Y | String | Unique identifier for the request. | -| `prompt` | N* | String | Input text for a generation request. | -| `input_ids` | Y* | List[Integer] | List of token IDs that make up the request prompt. | -| `output_tokens` | Y | Integer | Number of generated tokens for this request. | - -```{tip} -\* Specifying `prompt` or `input_ids` is required. However, you cannot have both prompts and token IDs (`input_ids`) -defined at the same time. If you specify `input_ids`, the `prompt` entry is ignored for request generation. -``` - -Refer to the following examples of valid entries for the benchmark: - -- Entries with a human-readable prompt and no logits. - - ```json - {"task_id": 1, "prompt": "Generate an infinite response to the following: This is the song that never ends, it goes on and on my friend.", "output_tokens": 1000} - {"task_id": 2, "prompt": "Generate an infinite response to the following: Na, na, na, na", "output_tokens": 1000} - ``` - -- Entries which contain logits. - - ```json - {"task_id":0,"input_ids":[863,22056,25603,11943,8932,13195,3132,25032,21747,22213],"output_tokens":128} - {"task_id":1,"input_ids":[14480,13598,15585,6591,1252,8259,30990,26778,7063,30065,21764,11023,1418],"output_tokens":128} - ``` - -```{tip} -Specify each entry on one line. -To simplify passing the data, a complete JSON entry is on each line so that the benchmarker -can simply read a line and assume a complete entry. When creating a dataset, be sure that a complete -JSON entry is on every line. -``` - -In order to prepare a synthetic dataset, you can use the provided script in the `benchmarks/cpp` -directory. For example, to generate a synthetic dataset of 1000 requests with a uniform ISL/OSL of -128/128 for [meta-llama/Llama-3.1-8B](https://huggingface.co/meta-llama/Llama-3.1-8B), run: - -```shell -trtllm-bench --tokenizer meta-llama/Llama-3.1-8B prepare-dataset --output /tmp/synthetic_128_128.txt token-norm-dist --input-mean 128 --output-mean 128 --input-stdev 0 --output-stdev 0 --num-requests 3000 -``` - -### Building a Benchmark Engine - -#### Default Build Behavior -The `trtllm-bench` CLI tool provides the `build` subcommand to build the TRT-LLM engines for max throughput benchmark. -To build an engine for benchmarking, you can specify the dataset generated with `prepare_dataset.py` through `--dataset` option. -By default, `trtllm-bench`'s tuning heuristic uses the high-level statistics of the dataset (average ISL/OSL, max sequence length) -to optimize engine build settings. The following command builds an FP8 quantized engine optimized using the dataset's ISL/OSL. - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B build --quantization FP8 --dataset /tmp/synthetic_128_128.txt -``` - -#### Other Build Modes - -The build subcommand also provides other ways to build the engine where users have larger control over the tuning values. - -- Build engine with self-defined tuning values: -You specify the tuning values to build the engine with by setting `--max_batch_size` and `--max_num_tokens` directly. -`max_batch_size` and `max_num_tokens` control the maximum number of requests and tokens that can be scheduled in each iteration. -If no value is specified, the default `max_batch_size` and `max_num_tokens` values of `2048` and `8192` are used. -The following command builds an FP8 quantized engine by specifying the engine tuning values. - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B build --quantization FP8 --max_seq_len 4096 --max_batch_size 1024 --max_num_tokens 2048 -``` - - -#### Parallelism Mapping Support -The `trtllm-bench build` subcommand supports combinations of tensor-parallel (TP) and pipeline-parallel (PP) mappings as long as the world size (`tp_size x pp_size`) `<=` `8`. The parallelism mapping in build subcommad is controlled by `--tp_size` and `--pp_size` options. The following command builds an engine with TP2-PP2 mapping. - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B build --quantization FP8 --dataset /tmp/synthetic_128_128.txt --tp_size 2 --pp_size 2 -``` - - -#### Example of Build Subcommand Output: -The output of the `build` subcommand looks similar to the snippet below (for `meta-llama/Llama-3.1-8B`): - -```shell -user@387b12598a9e:/scratch/code/trt-llm/tekit_2025$ trtllm-bench --model meta-llama/Llama-3.1-8B build --dataset /tmp/synthetic_128_128.txt --quantization FP8 -[TensorRT-LLM] TensorRT-LLM version: 0.17.0 -[01/18/2025-00:55:14] [TRT-LLM] [I] Found dataset. -[01/18/2025-00:55:14] [TRT-LLM] [I] -=========================================================== -= DATASET DETAILS -=========================================================== -Max Input Sequence Length: 128 -Max Output Sequence Length: 128 -Max Sequence Length: 256 -Target (Average) Input Sequence Length: 128 -Target (Average) Output Sequence Length: 128 -Number of Sequences: 3000 -=========================================================== - - -[01/18/2025-00:55:14] [TRT-LLM] [I] Max batch size and max num tokens are not provided, use tuning heuristics or pre-defined setting from trtllm-bench. -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated total available memory for KV cache: 132.37 GB -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated total KV cache memory: 125.75 GB -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated max number of requests in KV cache memory: 8048.16 -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated max batch size (after fine-tune): 4096 -[01/18/2025-00:55:14] [TRT-LLM] [I] Estimated max num tokens (after fine-tune): 8192 -[01/18/2025-00:55:14] [TRT-LLM] [I] Set dtype to bfloat16. -[01/18/2025-00:55:14] [TRT-LLM] [I] Set multiple_profiles to True. -[01/18/2025-00:55:14] [TRT-LLM] [I] Set use_paged_context_fmha to True. -[01/18/2025-00:55:14] [TRT-LLM] [I] Set use_fp8_context_fmha to True. -[01/18/2025-00:55:14] [TRT-LLM] [I] -=========================================================== -= ENGINE BUILD INFO -=========================================================== -Model Name: meta-llama/Llama-3.1-8B -Model Path: None -Workspace Directory: /tmp -Engine Directory: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 - -=========================================================== -= ENGINE CONFIGURATION DETAILS -=========================================================== -Max Sequence Length: 256 -Max Batch Size: 4096 -Max Num Tokens: 8192 -Quantization: FP8 -KV Cache Dtype: FP8 -=========================================================== - -Loading Model: [1/3] Downloading HF model -Downloaded model to /data/models--meta-llama--Llama-3.1-8B/snapshots/d04e592bb4f6aa9cfee91e2e20afa771667e1d4b -Time: 0.321s -Loading Model: [2/3] Loading HF model to memory -Loading checkpoint shards: 100%|█████████████████████████████████████████████████████████████████████████████████████████████████████| 4/4 [00:59<00:00, 14.79s/it] -Generating train split: 100%|████████████████████████████████████████████████████████████████████████████████████| 287113/287113 [00:06<00:00, 41375.57 examples/s] -Generating validation split: 100%|█████████████████████████████████████████████████████████████████████████████████| 13368/13368 [00:00<00:00, 41020.63 examples/s] -Generating test split: 100%|███████████████████████████████████████████████████████████████████████████████████████| 11490/11490 [00:00<00:00, 41607.11 examples/s] -Inserted 675 quantizers -/usr/local/lib/python3.12/dist-packages/modelopt/torch/quantization/model_quant.py:71: DeprecationWarning: forward_loop should take model as argument, but got forward_loop without any arguments. This usage will be deprecated in future versions. - warnings.warn( -Disable lm_head quantization for TRT-LLM export due to deployment limitations. -current rank: 0, tp rank: 0, pp rank: 0 -Time: 122.568s -Loading Model: [3/3] Building TRT-LLM engine -/usr/local/lib/python3.12/dist-packages/tensorrt/__init__.py:85: DeprecationWarning: Context managers for TensorRT types are deprecated. Memory will be freed automatically when the reference count reaches 0. - warnings.warn( -Time: 53.820s -Loading model done. -Total latency: 176.709s - - - -=========================================================== -ENGINE SAVED: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -=========================================================== -``` - -The engine in this case will be written to `/tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1` (the end of the log). - - -### Max Throughput Benchmark - -The `trtllm-bench` command line tool provides a max throughput benchmark that is accessible via the -`throughput` subcommand. This benchmark tests a TensorRT-LLM engine or PyTorch backend under maximum load to provide an -upper bound throughput number. - -#### How the Benchmarker Works - -The benchmarker reads a data file where a single line contains -a complete JSON request entry as specified in [](#preparing-a-dataset). -The process that the benchmarker follows is: - -1. Iterate over all input requests. If `logits` is specified, construct the request using the specified -list of logits. Otherwise, tokenize the `prompt` as specified by `--model $HF_MODEL_NAME`. -1. Submit the dataset to the TensorRT-LLM `Executor` API as fast as possible (offline mode). -1. Wait for all requests to return, compute statistics, and then report results. - -To run the benchmarker, run the following commands with the [engine](#building-a-benchmark-engine) and -[dataset](#preparing-a-dataset) generated from previous steps: - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B throughput --dataset /tmp/synthetic_128_128.txt --engine_dir /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -[TensorRT-LLM] TensorRT-LLM version: 0.17.0 -[01/18/2025-01:01:13] [TRT-LLM] [I] Preparing to run throughput benchmark... -[01/18/2025-01:01:13] [TRT-LLM] [I] Setting up throughput benchmark. - - - -[01/18/2025-01:01:26] [TRT-LLM] [I] Setting up for warmup... -[01/18/2025-01:01:26] [TRT-LLM] [I] Running warmup. -[01/18/2025-01:01:26] [TRT-LLM] [I] Starting benchmarking async task. -[01/18/2025-01:01:26] [TRT-LLM] [I] Starting benchmark... -[01/18/2025-01:01:26] [TRT-LLM] [I] Request submission complete. [count=2, time=0.0000s, rate=121847.20 req/s] -[01/18/2025-01:01:28] [TRT-LLM] [I] Benchmark complete. -[01/18/2025-01:01:28] [TRT-LLM] [I] Stopping LLM backend. -[01/18/2025-01:01:28] [TRT-LLM] [I] Cancelling all 0 tasks to complete. -[01/18/2025-01:01:28] [TRT-LLM] [I] All tasks cancelled. -[01/18/2025-01:01:28] [TRT-LLM] [I] LLM Backend stopped. -[01/18/2025-01:01:28] [TRT-LLM] [I] Warmup done. -[01/18/2025-01:01:28] [TRT-LLM] [I] Starting benchmarking async task. -[01/18/2025-01:01:28] [TRT-LLM] [I] Starting benchmark... -[01/18/2025-01:01:28] [TRT-LLM] [I] Request submission complete. [count=3000, time=0.0012s, rate=2590780.97 req/s] -[01/18/2025-01:01:42] [TRT-LLM] [I] Benchmark complete. -[01/18/2025-01:01:42] [TRT-LLM] [I] Stopping LLM backend. -[01/18/2025-01:01:42] [TRT-LLM] [I] Cancelling all 0 tasks to complete. -[01/18/2025-01:01:42] [TRT-LLM] [I] All tasks cancelled. -[01/18/2025-01:01:42] [TRT-LLM] [I] LLM Backend stopped. -[01/18/2025-01:01:42] [TRT-LLM] [I] - -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: meta-llama/Llama-3.1-8B -Engine Directory: /tmp/meta-llama/Llama-3.1-8B/tp_1_pp_1 -TensorRT-LLM Version: 0.17.0 -Dtype: bfloat16 -KV Cache Dtype: FP8 -Quantization: FP8 -Max Input Length: 256 -Max Sequence Length: 256 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -Max Runtime Batch Size: 4096 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 5.0689E+14 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 3000 -Average Input Length (tokens): 128.0000 -Average Output Length (tokens): 128.0000 -Token Throughput (tokens/sec): 28390.4265 -Request Throughput (req/sec): 221.8002 -Total Latency (ms): 13525.6862 - -=========================================================== - -[01/18/2025-01:01:42] [TRT-LLM] [I] Thread proxy_dispatch_result_thread stopped. -[TensorRT-LLM][INFO] Refreshed the MPI local session -``` - -### Running with the PyTorch Workflow - -```{eval-rst} -.. include:: ../../_includes/note_sections.rst - :start-after: .. start-note-config-flag-alias - :end-before: .. end-note-config-flag-alias -``` - -To benchmark the PyTorch backend (`tensorrt_llm._torch`), use the following command with [dataset](#preparing-a-dataset) generated from previous steps. With the PyTorch flow, you will not need to -run `trtllm-bench build`; the `throughput` benchmark initializes the backend by tuning against the -dataset provided via `--dataset` (or the other build mode settings described [above](#other-build-modes)). -Note that CUDA graph is enabled by default. You can add additional pytorch config with -`--config` followed by the path to a YAML file. For more details, please refer to the -help text by running the command with `--help`. - -```{tip} -The command below specifies the `--model_path` option. The model path is optional and used only when you want to run a locally -stored checkpoint. When using `--model_path`, the `--model` is still required for reporting reasons and in order to look up parameters -for build heuristics. -``` - -```shell -trtllm-bench --model meta-llama/Llama-3.1-8B --model_path /Ckpt/Path/To/Llama-3.1-8B throughput --dataset /tmp/synthetic_128_128.txt - -# Example output - -=========================================================== -= PyTorch backend -=========================================================== -Model: meta-llama/Llama-3.1-8B -Model Path: /Ckpt/Path/To/Llama-3.1-8B -TensorRT-LLM Version: 0.17.0 -Dtype: bfloat16 -KV Cache Dtype: None -Quantization: FP8 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -Max Runtime Batch Size: 2048 -Max Runtime Tokens: 4096 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 7.6753E+14 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 3000 -Average Input Length (tokens): 128.0000 -Average Output Length (tokens): 128.0000 -Token Throughput (tokens/sec): 20685.5510 -Request Throughput (req/sec): 161.6059 -Total Latency (ms): 18563.6825 - -``` - -#### Benchmarking with LoRA Adapters in PyTorch workflow - -The PyTorch workflow supports benchmarking with LoRA (Low-Rank Adaptation) adapters. This requires preparing a dataset with LoRA metadata and configuring the LoRA settings. - -**Preparing LoRA Dataset** - -Use `prepare_dataset.py` with LoRA-specific options to generate requests with LoRA metadata: - -```shell -python3 benchmarks/cpp/prepare_dataset.py \ - --stdout \ - --rand-task-id 0 1 \ - --tokenizer /path/to/tokenizer \ - --lora-dir /path/to/loras \ - token-norm-dist \ - --num-requests 100 \ - --input-mean 128 \ - --output-mean 128 \ - --input-stdev 16 \ - --output-stdev 24 \ - > synthetic_lora_data.json -``` - -Key LoRA options: -- `--lora-dir`: Parent directory containing LoRA adapter subdirectories named by their task IDs (e.g., `0/`, `1/`, etc.) -- `--rand-task-id`: Range of LoRA task IDs to randomly assign to requests -- `--task-id`: Fixed LoRA task ID for all requests (alternative to `--rand-task-id`) - -The generated dataset will include LoRA request metadata. Below is an example of a single such request data entry: - -```json -{ - "task_id": 0, - "input_ids": [3452, 88226, 102415, ...], - "output_tokens": 152, - "lora_request": { - "lora_name": "lora_0", - "lora_int_id": 0, - "lora_path": "/path/to/loras/0" - } -} -``` - -**LoRA Configuration** - -Create a `config.yaml` file with LoRA configuration: - -```yaml -lora_config: - lora_dir: - - /path/to/loras/0 - - /path/to/loras/1 - max_lora_rank: 64 - lora_target_modules: - - attn_q - - attn_k - - attn_v - trtllm_modules_to_hf_modules: - attn_q: q_proj - attn_k: k_proj - attn_v: v_proj -``` - -**Running LoRA Benchmark** - -```shell -trtllm-bench --model /path/to/base/model \ - throughput \ - --dataset synthetic_lora_data.json \ - --config config.yaml -``` - -```{note} -The LoRA directory structure should have task-specific subdirectories named by their task IDs (e.g., `loras/0/`, `loras/1/`). -Each subdirectory should contain the LoRA adapter files for that specific task. -``` - -#### Running multi-modal models in the PyTorch Workflow - -To benchmark multi-modal models with PyTorch workflow, you can follow the similar approach as above. - -First, prepare the dataset: -``` -python ./benchmarks/cpp/prepare_dataset.py \ - --tokenizer Qwen/Qwen2-VL-2B-Instruct \ - --stdout \ - dataset \ - --dataset-name lmms-lab/MMMU \ - --dataset-split test \ - --dataset-image-key image \ - --dataset-prompt-key question \ - --num-requests 10 \ - --output-len-dist 128,5 > mm_data.jsonl -``` -It will download the media files to `/tmp` directory and prepare the dataset with their paths. Note that the `prompt` fields are texts and not tokenized ids. This is due to the fact that -the `prompt` and the media (image/video) are processed by a preprocessor for multimodal files. - -Sample dataset for multimodal: -``` -{"task_id":0,"prompt":"Brahma Industries sells vinyl replacement windows to home improvement retailers nationwide. The national sales manager believes that if they invest an additional $25,000 in advertising, they would increase sales volume by 10,000 units. What is the total contribution margin?","media_paths":["/tmp/tmp9so41y3r.jpg"],"output_tokens":126} -{"task_id":1,"prompt":"Let us compute for the missing amounts under work in process inventory, what is the cost of goods manufactured? ","media_paths":["/tmp/tmpowsrb_f4.jpg"],"output_tokens":119} -{"task_id":2,"prompt":"Tsuji is reviewing the price of a 3-month Japanese yen/U.S. dollar currency futures contract, using the currency and interest rate data shown below. Because the 3-month Japanese interest rate has just increased to .50%, Itsuji recognizes that an arbitrage opportunity exists nd decides to borrow $1 million U.S. dollars to purchase Japanese yen. Calculate the yen arbitrage profit from Itsuji's strategy, using the following data: ","media_paths":["/tmp/tmpxhdvasex.jpg"],"output_tokens":126} -... -``` - -Run the benchmark: -``` -trtllm-bench --model Qwen/Qwen2-VL-2B-Instruct \ - throughput \ - --dataset mm_data.jsonl \ - --num_requests 10 \ - --max_batch_size 4 \ - --modality image -``` - - -Sample output: -``` -=========================================================== -= REQUEST DETAILS -=========================================================== -Number of requests: 10 -Number of concurrent requests: 5.3019 -Average Input Length (tokens): 411.6000 -Average Output Length (tokens): 128.7000 -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 1 -PP Size: 1 -EP Size: None -Max Runtime Batch Size: 4 -Max Runtime Tokens: 12288 -Scheduling Policy: GUARANTEED_NO_EVICT -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 1.4117E+17 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Request Throughput (req/sec): 1.4439 -Total Output Throughput (tokens/sec): 185.8351 -Per User Output Throughput (tokens/sec/user): 38.1959 -Per GPU Output Throughput (tokens/sec/gpu): 185.8351 -Total Token Throughput (tokens/sec): 780.1607 -Total Latency (ms): 6925.4963 -Average request latency (ms): 3671.8441 - --- Request Latency Breakdown (ms) ----------------------- - -[Latency] P50 : 3936.3022 -[Latency] P90 : 5514.4701 -[Latency] P95 : 5514.4701 -[Latency] P99 : 5514.4701 -[Latency] MINIMUM: 2397.1047 -[Latency] MAXIMUM: 5514.4701 -[Latency] AVERAGE: 3671.8441 - -=========================================================== -= DATASET DETAILS -=========================================================== -Dataset Path: /workspaces/tensorrt_llm/mm_data.jsonl -Number of Sequences: 10 - --- Percentiles statistics --------------------------------- - - Input Output Seq. Length ------------------------------------------------------------ -MIN: 167.0000 119.0000 300.0000 -MAX: 1059.0000 137.0000 1178.0000 -AVG: 411.6000 128.7000 540.3000 -P50: 299.0000 128.0000 427.0000 -P90: 1059.0000 137.0000 1178.0000 -P95: 1059.0000 137.0000 1178.0000 -P99: 1059.0000 137.0000 1178.0000 -=========================================================== -``` - -**Notes and Limitations**: -- Only image datasets are supported for now. -- `--output-len-dist` is a required argument for multimodal datasets. -- Tokenizer is unused during the prepare step but it is still a required argument. -- Since the images are converted to tokens when the model is run, `trtllm-bench` uses a default large value for the maximum input sequence length when setting up the execution settings. - You can also modify the behavior by specifying a different value with the flag `--max_input_len` that suits your use-case. - -#### Quantization in the PyTorch Flow - -In order to run a quantized run with `trtllm-bench` utilizing the PyTorch flow, you will need to use a pre-quantized -To run a quantized benchmark with `trtllm-bench` utilizing the PyTorch flow, you will need to use a pre-quantized -checkpoint. For the Llama-3.1 models, TensorRT-LLM provides the following checkpoints via HuggingFace: - -- [`nvidia/Llama-3.1-8B-Instruct-FP8`](https://huggingface.co/nvidia/Llama-3.1-8B-Instruct-FP8) -- [`nvidia/Llama-3.1-70B-Instruct-FP8`](https://huggingface.co/nvidia/Llama-3.1-70B-Instruct-FP8) -- [`nvidia/Llama-3.1-405B-Instruct-FP8`](https://huggingface.co/nvidia/Llama-3.1-405B-Instruct-FP8) - -`trtllm-bench` utilizes the `hf_quant_config.json` file present in the pre-quantized checkpoints above. The configuration -file is present in checkpoints quantized with [Model Optimizer](https://github.com/NVIDIA/Model-Optimizer) -and describes the compute and KV cache quantization that checkpoint was compiled with. For example, from the checkpoints -above: - -```json -{ - "producer": { - "name": "modelopt", - "version": "0.23.0rc1" - }, - "quantization": { - "quant_algo": "FP8", - "kv_cache_quant_algo": null - } -} -``` - -The checkpoints above are quantized to run with a compute precision of `FP8` and default to no KV cache quantization (full -`FP16` cache). When running `trtllm-bench throughput`. The benchmark will select a KV cache quantization that is best suited -for the compute precision in the checkpoint automatically if `kv_cache_quant_algo` is specified as `null`, otherwise it will -be forced to match the specified non-null KV cache quantization. The following are the mappings that `trtllm-bench` will -follow when a checkpoint does not specify a KV cache quantization algorithm: - -| Checkpoint Compute Quant | Checkpoint KV Cache Quant | `trtllm-bench` | Note | -| - | - | - | - | -| `null` | `null` | `null` | In this case, a quantization config doesn't exist. | -| `FP8` | `FP8` | `FP8` | Matches the checkpoint | -| `FP8` | `null` | `FP8` | Set to `FP8` via benchmark | -| `NVFP4` | `null` | `FP8` | Set to `FP8` via benchmark | - -If you would like to force the KV cache quantization, you can specify the following in the YAML file to force the precision -when the checkpoint precision is `null`: - -```yaml -kv_cache_dtype: "fp8" -``` - -```{tip} -The valid values for `kv_cache_dtype` are `auto`, `fp8`, and `nvfp4`. -``` - -## Low Latency Benchmark - -The low latency benchmark follows a similar workflow to the [throughput benchmark](#max-throughput-benchmark) -but requires building the engine separately from `trtllm-bench`. Low latency benchmarks has the following modes: - -- A single-request low-latency engine -- A Medusa-enabled speculative-decoding engine - -### Low Latency TensorRT-LLM Engine for Llama-3 70B - -To build a low-latency engine for the latency benchmark, run the following quantize and build commands. -The `$checkpoint_dir` is the path to the [meta-llama/Meta-Llama-3-70B](https://huggingface.co/meta-llama/Meta-Llama-3-70B) Hugging Face checkpoint in your cache or downloaded to a specific location with the [huggingface-cli](https://huggingface.co/docs/huggingface_hub/en/guides/cli). -To prepare a dataset, follow the same process as specified in [](#preparing-a-dataset). - -#### Benchmarking a non-Medusa Low Latency Engine - -To quantize the checkpoint: - -```shell -cd tensorrt_llm/examples/models/core/llama -python ../quantization/quantize.py \ - --model_dir $checkpoint_dir \ - --dtype bfloat16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir /tmp/meta-llama/Meta-Llama-3-70B/checkpoint \ - --calib_size 512 \ - --tp_size $tp_size -``` - -then build, - -```shell -trtllm-build \ - --checkpoint_dir /tmp/meta-llama/Meta-Llama-3-70B/checkpoint \ - --use_fused_mlp enable \ - --gpt_attention_plugin bfloat16 \ - --output_dir /tmp/meta-llama/Meta-Llama-3-70B/engine \ - --max_batch_size 1 \ - --max_seq_len $(($isl+$osl)) \ - --reduce_fusion enable \ - --gemm_plugin fp8 \ - --workers $tp_size \ - --use_fp8_context_fmha enable \ - --max_num_tokens $isl \ - --use_paged_context_fmha disable \ - --multiple_profiles enable -``` - -After the engine is built, run the low-latency benchmark: - -```shell -env TRTLLM_ENABLE_MMHA_MULTI_BLOCK_DEBUG=1 \ - TRTLLM_MMHA_KERNEL_BLOCK_SIZE=256 \ - TRTLLM_MMHA_BLOCKS_PER_SEQUENCE=32 \ - FORCE_MULTI_BLOCK_MODE=ON \ - TRTLLM_ENABLE_PDL=1 \ - trtllm-bench --model meta-llama/Meta-Llama-3-70B \ - latency \ - --dataset $DATASET_PATH \ - --engine_dir /tmp/meta-llama/Meta-Llama-3-70B/engine -``` - -### Building a Medusa Low-Latency Engine - -To build a Medusa-enabled engine requires checkpoints that contain Medusa heads. -NVIDIA provides TensorRT-LLM checkpoints on the [NVIDIA](https://huggingface.co/nvidia) page on Hugging Face. -The checkpoints are pre-quantized and can be directly built after downloading them with the -[huggingface-cli](https://huggingface.co/docs/huggingface_hub/en/guides/cli). -After you download the checkpoints, run the following command. Make sure to -specify the `$tp_size` supported by your Medusa checkpoint and the path to its stored location `$checkpoint_dir`. -Additionally, `$max_seq_len` should be set to the model's maximum position embedding. - -Using Llama-3.1 70B as an example, for a tensor parallel 8 and bfloat16 dtype: - -```shell -tp_size=8 -max_seq_len=131072 -trtllm-build --checkpoint_dir $checkpoint_dir \ - --speculative_decoding_mode medusa \ - --max_batch_size 1 \ - --gpt_attention_plugin bfloat16 \ - --max_seq_len $max_seq_len \ - --output_dir /tmp/meta-llama/Meta-Llama-3.1-70B/medusa/engine \ - --use_fused_mlp enable \ - --paged_kv_cache enable \ - --use_paged_context_fmha disable \ - --multiple_profiles enable \ - --reduce_fusion enable \ - --use_fp8_context_fmha enable \ - --workers $tp_size \ - --low_latency_gemm_plugin fp8 -``` - -After the engine is built, you need to define the Medusa choices. -The choices are specified with a YAML file like the following example (`medusa.yaml`): - -```yaml -- [0] -- [0, 0] -- [1] -- [0, 1] -- [2] -- [0, 0, 0] -- [1, 0] -- [0, 2] -- [3] -- [0, 3] -- [4] -- [0, 4] -- [2, 0] -- [0, 5] -- [0, 0, 1] -``` - -To run the Medusa-enabled engine, run the following command: - -```shell -env TRTLLM_ENABLE_PDL=1 \ - UB_ONESHOT=1 \ - UB_TP_SIZE=$tp_size \ - TRTLLM_ENABLE_PDL=1 \ - TRTLLM_PDL_OVERLAP_RATIO=0.15 \ - TRTLLM_PREFETCH_RATIO=-1 \ - trtllm-bench --model meta-llama/Meta-Llama-3-70B \ - latency \ - --dataset $DATASET_PATH \ - --engine_dir /tmp/meta-llama/Meta-Llama-3-70B/medusa/engine \ - --medusa_choices medusa.yml -``` - -## Summary - -The following table summarizes the commands needed for running benchmarks: - -| Scenario | Phase | Command | -| - | - | - | -| Dataset | Preparation | `python benchmarks/cpp/prepare_dataset.py --stdout --tokenizer $HF_MODEL token-norm-dist --input-mean $ISL --output-mean $OSL --input-stdev 0 --output-stdev 0 --num-requests $NUM_REQUESTS > $DATASET_PATH` | -| Throughput | Build | `trtllm-bench --model $HF_MODEL build --dataset $DATASET_PATH` | -| Throughput | Benchmark | `trtllm-bench --model $HF_MODEL throughput --dataset $DATASET_PATH --engine_dir $ENGINE_DIR` | -| Latency | Build | See [section about building low latency engines](#low-latency-tensorrt-llm-engine-for-llama-3-70b) | -| Non-Medusa Latency | Benchmark | `trtllm-bench --model $HF_MODEL latency --dataset $DATASET_PATH --engine_dir $ENGINE_DIR` | -| Medusa Latency | Benchmark | `trtllm-bench --model $HF_MODEL latency --dataset $DATASET_PATH --engine_dir $ENGINE_DIR --medusa_choices $MEDUSA_CHOICES` | - -where, - -`$HF_MODEL` -: The Hugging Face name of a model. - -`$NUM_REQUESTS` -: The number of requests to generate. - -`$DATASET_PATH` -: The path where the dataset was written when preparing the dataset. - -`$ENGINE_DIR` -: The engine directory as printed by `trtllm-bench build`. - -`$MEDUSA_CHOICES` -: A YAML config representing the Medusa tree for the benchmark. diff --git a/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md b/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md deleted file mode 100644 index 17cb9aef45ba..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/benchmarking-default-performance.md +++ /dev/null @@ -1,243 +0,0 @@ -(benchmarking-default-performance)= - -# Benchmarking Default Performance - -This section discusses how to build an engine for the model using the LLM-API and benchmark it using TRTLLM-Bench. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Before You Begin: TensorRT-LLM LLM-API - -TensorRT-LLM's LLM-API aims to make getting started with TensorRT-LLM quick and easy. For example, the following script instantiates `Llama-3.3-70B-Instruct` and runs inference on a small set of prompts. For those familiar with TensorRT-LLM's [CLI workflow](./benchmarking-default-performance.md#building-and-saving-engines-via-cli), the call to `LLM()` handles converting the model checkpoint and building the engine in one line. - -```python -#quickstart.py -from tensorrt_llm import LLM, SamplingParams - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", #HuggingFace model name, no need to download the checkpoint beforehand - tensor_parallel_size=4 - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` -### Troubleshooting Tips and Pitfalls To Avoid - -Since we are running on multiple GPUs, MPI is used to spawn processes for each GPU. This raises the following requirements - -1. The entrypoint to the script should be guarded via `if __name__ == '__main__'`. This requirement comes from mpi4py. -2. Depending on your environment, it might be required to wrap the `python` command with `mpirun`. For example the command to run the script above could be `mpirun -n 1 --oversubscribe --allow-run-as-root python quickstart.py`. For running on multiple GPUs on one node like the example is attempting to do it is usually not required to prefix with `mpirun` but if you are getting MPI errors then you should add it. Additionally, the `-n 1` which says just to launch one process is intentional as TensorRT-LLM handles spawning the processes for the remaining GPUs -3. If you get a HuggingFace access error when loading the Llama weights, this is likely because the model is gated. Request access on the HuggingFace page for the model. Then follow the instructions on [Huggingface's quickstart guide](https://huggingface.co/docs/huggingface_hub/en/quick-start#authentication) to authenticate in your environment. - - -## Building and Saving the Engine - -Save the engine using `.save()`. Just like the previous example, this script and all subsequent scripts might need to be run via `mpirun`. - -```python -from tensorrt_llm import LLM - -def main(): - llm = LLM( - model="/scratch/Llama-3.3-70B-Instruct", - tensor_parallel_size=4 - ) - - llm.save("baseline") - -if __name__ == '__main__': - main() -``` - -### Building and Saving Engines via CLI - -TensorRT-LLM also has a command line interface for building and saving engines. This workflow consists of two steps - -1. Convert model checkpoint (HuggingFace, Nemo) to TensorRT-LLM checkpoint via `convert_checkpoint.py`. Each supported model has a `convert_checkpoint.py` associated it with it and can be found in the examples folder. For example, the `convert_checkpoint.py` script for Llama models can be found [here](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama/convert_checkpoint.py) -2. Build engine by passing TensorRT-LLM checkpoint to `trtllm-build` command. The `trtllm-build` command is installed automatically when the `tensorrt_llm` package is installed. - -The README in the examples folder for supported models walks through building engines using this flow for a wide variety of situations. The examples folder for Llama models can be found at [https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama). - -## Benchmarking with `trtllm-bench` - -`trtllm-bench` provides a command line interface for benchmarking the throughput and latency of saved engines. - -### Prepare Dataset - -`trtllm-bench` expects to be passed in a dataset of requests to run through the model. This guide creates a dummy dataset of 1000 requests with every request having input and output sequence length of 2048. TensorRT-LLM provides the `prepare_dataset.py` script to produce the dataset. To use it clone the TensorRT-LLM Repo and run the following command: - -`python benchmarks/cpp/prepare_dataset.py --stdout --tokenizer /path/to/hf/Llama-3.3-70B-Instruct/ token-norm-dist --input-mean 2048 --output-mean 2048 --input-stdev 0 --output-stdev 0 --num-requests 1000 > synthetic_2048_2048.txt` - -`trtllm-bench` can also take in real data, see [`trtllm-bench` documentation](../perf-benchmarking.md) for more details on the required format. - -### Running Throughput and Latency Benchmarks - - To benchmark the baseline engine built in the previous script, run the following commands. Again, due to the multi-gpu nature of the workload you may need prefix the `trtllm-bench` command with `mpirun -n 1 --oversubscribe --allow-run-as-root`. - -**Throughput** - -```bash -trtllm-bench \ ---model /path/to/hf/Llama-3.3-70B-Instruct/ \ -throughput \ ---dataset /path/to/dataset/synthetic_2048_2048_1000.txt \ ---engine_dir /path/to/engines/baseline #replace baseline with name used in llm.save() -``` - -This command will send all 1000 requests to the model immediately. Run `trtllm-bench throughput -h` to see a list of options that help you control the request rate and cap the total number of requests if the benchmark is taking too long. For reference, internal testing of the above command took around 20 minutes on a 4 NVLink connected H100-sxm-80GB. - -Running this command will provide a throughput overview like this: - -```bash -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: /scratch/Llama-3.3-70B-Instruct/ -Engine Directory: /scratch/grid_search_engines/baseline -TensorRT-LLM Version: 0.16.0 -Dtype: bfloat16 -KV Cache Dtype: None -Quantization: None -Max Sequence Length: 131072 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 4 -PP Size: 1 -Max Runtime Batch Size: 2048 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% -Issue Rate (req/sec): 7.9353E+13 - -=========================================================== -= PERFORMANCE OVERVIEW -=========================================================== -Number of requests: 1000 -Average Input Length (tokens): 2048.0000 -Average Output Length (tokens): 2048.0000 -Token Throughput (tokens/sec): 1585.7480 -Request Throughput (req/sec): 0.7743 -Total Latency (ms): 1291504.1051 - -=========================================================== -``` - -**Latency** - -```bash -trtllm-bench \ ---model /path/to/hf/Llama-3.3-70B-Instruct/ \ -latency \ ---dataset /path/to/dataset/synthetic_2048_2048_1000.txt \ ---num-requests 100 \ ---warmup 10 \ ---engine_dir /path/to/engines/baseline #replace baseline with name used in llm.save() -``` -The latency benchmark enforces a batch size of 1 to accurately measure latency, which can significantly increase testing duration. In the example above the total number of requests is limited to 100 via `--num-requests` to make the test duration more manageable. This example benchmark was designed to produce very stable numbers, but in real scenarios even 100 requests is likely more than you need and can take a long time to complete (in the case-study it took about an hour and a half). Reducing the number of requests to 10 would still provide accurate data and enable faster development iterations. In general you should adjust the number of requests per your needs. Run `trtllm-bench latency -h` to see other configurable options. - -Running this command will provide a latency overview like this: - -```bash -=========================================================== -= ENGINE DETAILS -=========================================================== -Model: /scratch/Llama-3.3-70B-Instruct/ -Engine Directory: /scratch/grid_search_engines/baseline -TensorRT-LLM Version: 0.16.0 -Dtype: bfloat16 -KV Cache Dtype: None -Quantization: None -Max Input Length: 1024 -Max Sequence Length: 131072 - -=========================================================== -= WORLD + RUNTIME INFORMATION -=========================================================== -TP Size: 4 -PP Size: 1 -Max Runtime Batch Size: 1 -Max Runtime Tokens: 8192 -Scheduling Policy: Guaranteed No Evict -KV Memory Percentage: 90.00% - -=========================================================== -= GENERAL OVERVIEW -=========================================================== -Number of requests: 100 -Average Input Length (tokens): 2048.0000 -Average Output Length (tokens): 2048.0000 -Average request latency (ms): 63456.0704 - -=========================================================== -= THROUGHPUT OVERVIEW -=========================================================== -Request Throughput (req/sec): 0.0158 -Total Token Throughput (tokens/sec): 32.2742 -Generation Token Throughput (tokens/sec): 32.3338 - -=========================================================== -= LATENCY OVERVIEW -=========================================================== -Total Latency (ms): 6345624.0554 -Average time-to-first-token (ms): 147.7502 -Average inter-token latency (ms): 30.9274 -Acceptance Rate (Speculative): 1.00 - -=========================================================== -= GENERATION LATENCY BREAKDOWN -=========================================================== -MIN (ms): 63266.8804 -MAX (ms): 63374.7770 -AVG (ms): 63308.3201 -P90 (ms): 63307.1885 -P95 (ms): 63331.7136 -P99 (ms): 63374.7770 - -=========================================================== -= ACCEPTANCE BREAKDOWN -=========================================================== -MIN: 1.00 -MAX: 1.00 -AVG: 1.00 -P90: 1.00 -P95: 1.00 -P99: 1.00 - -=========================================================== -``` - -## Results - -The baseline engine achieves the following performance for token throughput, request throughput, average time to first token, and average inter-token latency. These metrics will be analyzed throughout the guide. - - -| Metric | Value | -|-------------------------------|---------------| -| Token Throughput (tokens/sec) | 1564.3040 | -| Request Throughput (req/sec) | 0.7638 | -| Average Time To First Token (ms) | 147.6976 | -| Average Inter-Token Latency (ms) | 31.3276 | - -The following sections show ways you can improve these metrics using different configuration options. diff --git a/docs/source/legacy/performance/performance-tuning-guide/deciding-model-sharding-strategy.md b/docs/source/legacy/performance/performance-tuning-guide/deciding-model-sharding-strategy.md deleted file mode 100644 index 2970d9e7e444..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/deciding-model-sharding-strategy.md +++ /dev/null @@ -1,50 +0,0 @@ -(deciding-model-sharding-strategy)= - -# Deciding Model Sharding Strategy - -Large models often can't fit on one GPU and need to be sharded across multiple GPUs. The sharding strategies used to accomplish this can have significant impacts on performance. This guide walks through how to determine if tensor parallelism, pipeline parallelism, or a mix of both are the best strategy for you. If you are not familiar with tensor parallelism and pipeline parallelism please refer to [Mastering LLM Techniques - Inference Optimization](https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/) - - -## How to Think about Model Sharding: Communication is Key - -Splitting your model weights between multiple GPUs requires them to communicate activations between each other, adding additional overhead. How expensive this overhead is on your system is the key factor in determining the best strategy for you. - -In pipeline parallelism, the model is split into sets of contiguous layers and each GPU houses one of these sets. In this setup, the only required communication is for each GPU to send the outputs of its set to the GPU with the next set. - -![Pipeline Parallel Visualization](../../media/Pipeline_Parallel_Vis.svg) - - - On the other hand, tensor parallelism takes each layer of the model and splits it between the GPUs. This means that every GPU houses a portion of every layer. However since each layer needs the full outputs of the previous layer as an input, each GPU has to perform the heavier All-Reduce communication operation to share its results with all other GPUs before it can begin processing the next layer. While this seems disadvantageous, because each GPU only holds partial layers, it also performs smaller matrix multiplications, allowing it to compute its outputs quicker. - - ![Tensor Parallel Visualization](../../media/Tensor_Parallelism_Vis.svg) - - - Ultimately deciding the best strategy comes down to whether the extra overhead from the All-Reduce operation overshadows the gains from the smaller matrix multiplications. If the interconnects between the GPUs are sufficiently fast, the gains from the reduced computation burden per layer can outweigh the additional communication cost. Consequently, a general rule of thumb is that if your GPUs have fast connections between them like NVLink then tensor parallel is likely a good choice. However if the communication will go over slow connections (across nodes for example) pipeline parallel is likely better. Overall we provide the following guidelines: - -**If your model fits on one gpu:** Unless you have a very specific reason, don't shard your model. The best communication overhead is no communication overhead. - -**If your model fits in one node:** Tensor parallel is likely the best option here, especially if you have fast connections between the GPUs like NVLink. If you don't, then pipeline parallel might be needed. Start with tensor parallel and sanity check if pipeline parallel is better. - -**If your model is sharded across multiple nodes:** Inter-node connections are typically significantly slower than intra-node connections, so if you have tensor parallelism across nodes it will be bottlenecked by the slow interconects. Consequently, a good starting point is having tensor parallelism within the node and pipeline parallelism between nodes. An exception is if you are running on NVL36 or NVL72 Blackwell systems. These have multinode NVLink so as long as you stay within the 36 or 72 GPUs, tensor parallel won't be bottlenecked by inter-node connections. - -## How to set Tensor Parallelism and Pipeline Parallelism - -The `LLM` class takes `tensor_parallel_size` and `pipeline_parallel_size` as parameters. `tensor_parallel_size * pipeline_parallel_size` should be equal to the total number of GPUs you are sharding the model over, referred to as the world size. For example, if you were sharding a model over 2 nodes, each with 16 GPUs, you might set tensor parallel to 8 (for tensor parallelism within the node) and pipeline parallel to 2 (pipeline parallel between nodes) like this: - -```python - llm = LLM( - model="/scratch/Llama-3.1-405B-Instruct", - tensor_parallel_size=8, - pipeline_parallel_size=2 - ) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) you can specify tensor parallelism and pipeline parallelism by providing the `--tp_size` and `--tp_size` arguments to `convert_checkpoint.py` - -``` -python examples/models/core/llama/convert_checkpoint.py --model_dir ./tmp/llama/405B/ \ - --output_dir ./tllm_checkpoint_16gpu_tp8_pp2 \ - --dtype float16 \ - --tp_size 8 - --pp_size 2 -``` diff --git a/docs/source/legacy/performance/performance-tuning-guide/fp8-quantization.md b/docs/source/legacy/performance/performance-tuning-guide/fp8-quantization.md deleted file mode 100644 index 3a5121b7c173..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/fp8-quantization.md +++ /dev/null @@ -1,217 +0,0 @@ -(fp8-quantization)= - -# FP8 Quantization - -Quantization is a technique that allows models to run in lower precisions like int8 and fp8 while maintaining acceptable output quality. Running in lower precisions can greatly boost performance, significantly increasing throughput and decreasing latency. The tradeoff is a drop in output quality, but in many cases the output quality is still acceptable and many real world deployments utilize quantization. If you want to learn more about quantization refer to [Mastering LLM Techniques - Inference Optimization](https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/) - -This section walks through enabling fp8 quantization and highlight some fp8 quantization specific configuration options for boosting performance. It also continues the case study of Llama-3.3-70B split across 4 H100-sxm-80GB GPUs via tensor parallelism and showcase the effects of enabling these configuration options on performance. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Enabling Quantization - -To enable quantization you need to configure the `QuantConfig` class and pass it to the `quant_config` parameter of the LLM class. At a minimum the `quant_algo` parameter, which sets the quantization algorithm (fp8, fp8 per token, int8awq, etc.) must be specified. You can find all supported quantization algorithms and other configurable options for `QuantConfig` in the LLM-API->Reference section of the docs. While it is not required if you are using weights/checkpoints from that are already quantized, if you are using an fp16 checkpoint then you also need to specify the calibration dataset that will be used to determine the quantization scales via `CalibConfig`. `CalibConfig` provides several options for setting the calibration dataset that can also be referenced in the LLM-API->Reference section of the docs. Although TensorRT-LLM supports several other types of quantization, this guide focuses on fp8. - - -Here is an example of building and saving an fp8 engine from a bf16 checkpoint (Note that fp8 is supported only on devices with compute capability > 8.9 - Ada, Hopper, Blackwell, and beyond): -```python -from tensorrt_llm import LLM, BuildConfig -from tensorrt_llm.llmapi import QuantConfig, QuantAlgo, CalibConfig - -def main(): - - quant_config = QuantConfig(quant_algo=QuantAlgo.FP8) - - calib_config = CalibConfig( - calib_batches=512, - calib_batch_size=1, - calib_max_seq_length=2048, - tokenizer_max_seq_length=4096 - ) - - build_config = BuildConfig( - max_num_tokens=2048, - max_batch_size=512, - ) - - build_config.plugin_config.use_paged_context_fmha = True - build_config.plugin_config.multiple_profiles = True - - llm = LLM( - model="/path/to/Llama-3.3-70B", - tensor_parallel_size=4, - pipeline_parallel_size=1, - build_config=build_config, - quant_config=quant_config, - calib_config=calib_config - ) - - llm.save("baseline_fp8_engine") - -if __name__ == '__main__': - main() -``` - -For an example of how to build an fp8 engine using the [TensorRT-LLM CLI workflow](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) flow see [TensorRT-LLM LLaMA examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama). In short you first run [`examples/quantization/quantize.py`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/quantization) to quantize and convert the model checkpoint to TensorRT-LLM format and then use `trtllm-build`. - -> ***Note: While quantization aims to preserve model accuracy this is not guaranteed and it is extremely important you check that the quality of outputs remains sufficient after quantization.*** - -## FP8 "Baseline" Performance - -Benchmarking the engine produced by the example above yielded the following performance results. Note that we enabled some of the build flags we mentioned [earlier](./useful-build-time-flags.md) (multiple profiles, paged_context_fmha) and also tuned max batch size and max num tokens. This is done to give a sense of what performance is achievable if you tune an fp8 engine but exclude options that have been tailored for quantization. We recommend disabling the gemm plugin for quantized engines which is why it is not included here (it is off by default). Reduce fusion has a quantization specific optimization that will be covered later. For the remainder of this page we will refer to this setup as the "baseline" numbers for fp8. - - -| Metric | Value | -| -------------------------------- | --------- | -| Token Throughput (tokens/sec) | 3389.5305 | -| Request Throughput (req/sec) | 1.6550 | -| Average Time To First Token (ms) | 96.1597 | -| Average Inter-Token Latency (ms) | 12.4248 | - - -## Quantized KV-Cache - -By default the KV-Cache is not quantized but TensorRT-LLM supports quantizing the KV-Cache to further improve performance. However, quantizing the model more aggressively also increases the risk of model output quality degrading so it is important to check that when using this feature. - -### Enabling Quantized KV Cache - -The LLM-API exposes the quantization algorithm to be used for kv cache via the `kv_cache_quant_algo` field in `QuantConfig`. To enable fp8 kv cache, you would modify `QuantConfig` as such: - -```python -quant_config = QuantConfig(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--kv_cache_dtype fp8` to [`examples/quantization/quantize.py`](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/quantization). - -### Performance with Quantized KV Cache - -| Metric | Baseline | FP8 KV-Cache ON | -| -------------------------------- | --------- | --------------- | -| Token Throughput (tokens/sec) | 3389.5305 | 5299.6372 | -| Request Throughput (req/sec) | 1.6550 | 2.5877 | -| Average Time To First Token (ms) | 96.1597 | 97.1287 | -| Average Inter-Token Latency (ms) | 12.4248 | 12.5496 | - -## Reduce Norm Fusion with User Buffers for Llama Models - -The [Reduce Norm Fusion](./useful-build-time-flags.md#reduce-norm-fusion-plugin-for-llama-models) feature is supported for fp8. An additional optimization called "User Buffers" is also supported for fp8 models. The user buffer feature aims to eliminate extra copies from the local buffer to the shared buffer in the communication kernel, leading to improved end-to-end performance. - - -### Enabling Reduce Norm Fusion with User Buffers - - -To enable reduce norm fusion with user buffers, add the following lines below `BuildConfig`'s initialization - -```python -build_config.plugin_config.reduce_fusion = True -build_config.plugin_config.user_buffer = True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--reduce_fusion enable` and `--user_buffer enable` to `trtllm-build` to enable the feature. - -> Note: You must have enabled `reduce_fusion` in order to enable `user_buffer` - -### Performance with Reduce Norm Fusion + User Buffers: - -Reduce Norm Fusion + User Buffer ON: Same engine previously referred to as FP8 KV-Cache ON. - -Reduce Norm Fusion + User Buffer ON: Previous example with reduce fusion and user buffers enabled. Max-num tokens set to 16384 and max-batch size set to 512 after tuning. - - -| Metric | Reduce Norm Fusion + User Buffer OFF | Reduce Norm Fusion + User Buffer ON | -| -------------------------------- | ------------------------------------ | ----------------------------------- | -| Token Throughput (tokens/sec) | 5299.6372 | 5980.7842 | -| Request Throughput (req/sec) | 2.5877 | 2.9203 | -| Average Time To First Token (ms) | 97.1287 | 82.2679 | -| Average Inter-Token Latency (ms) | 12.5496 | 12.6975 | - -## GEMM + SwiGLU Fusion in Gated-MLP - -The GEMM + SwiGLU fusion in Gated-MLP combines two Matmul operations and one SwiGLU operation into a single kernel. Currently this is only supported for FP8 precision on Hopper. While this fusion improves performance, it can slightly reduce accuracy in FP8 PTQ because one quantization scaling factor is discarded. - -We recommend enabling this feature for large models running on Hopper with FP8 precision. We do not recommend enabling this feature for very small workloads or if the -accuracy loss is unacceptable. - -### Enabling GEMM + SwiGLU Fusion - -To enable the GEMM + SwiGLU fusion, add the following lines below `BuildConfig`'s initialization - -```python -build_config.plugin_config.gemm_swiglu_plugin = 'fp8' -``` -For small batch size cases where latency is important, you can replace the above line with - -```python -build_config.plugin_config.low_latency_gemm_swiglu_plugin = 'fp8' -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--gemm_swiglu_plugin=fp8` or `--low_latency_gemm_swiglu_plugin=fp8` for the low latency case (only include one or the other) to `trtllm-build`. - -### Performance with GEMM + SwiGLU Fusion - - -| Metric | GEMM + SwiGLU fusion OFF | GEMM + SwiGLU fusion ON | -| -------------------------------- | ------------------------ | ----------------------- | -| Token Throughput (tokens/sec) | 5980.7842 | 5976.7977 | -| Request Throughput (req/sec) | 2.9203 | 2.9184 | -| Average Time To First Token (ms) | 82.2679 | 81.8841 | -| Average Inter-Token Latency (ms) | 12.6975 | 11.7031 | - -In this case, the GEMM + SwiGLU plugin performs almost equivalently to when it was disabled. The throughput drop is within run to run variance and the TTFT and ITL improvements are slight. However, we found that when paired with the low latency gemm plugin discussed next, enabling this feature was necessary for getting the maximum throughput. - -## Low Latency GEMM Plugin - -Previously we mentioned the [GEMM Plugin](./useful-build-time-flags.md#gemm-plugin) feature. Although it has fp8 support we recommend disabling it (by default it is disabled). However for low-latency scenarios in fp8 we recommend trying the low latency GEMM plugin to see if it is effective for your workload. - -### Enabling Low Latency GEMM plugin - -To enable the low latency GEMM plugin, add the following lines below `BuildConfig`'s initialization - -```python -build_config.plugin_config.low_latency_gemm_plugin = 'fp8' -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--low_latency_gemm_plugin=fp8` to `trtllm-build` to enable the feature. Again, **we recommend disabling the gemm plugin for fp8** so if you are passing `--gemm_plugin=fp8` to `trtllm-build` we recommend removing that. - -### Performance with Low Latency GEMM plugin - -Low Latency GEMM ON: Same configuration as previous example but with low latency GEMM plugin enabled. Max num tokens was set to 16384 and max-batch size was set to 512 after tuning. - -| Metric | Low Latency GEMM OFF | Low Latency GEMM ON | -| -------------------------------- | -------------------- | ------------------- | -| Token Throughput (tokens/sec) | 5976.7977 | 6049.1625 | -| Request Throughput (req/sec) | 2.9184 | 2.9537 | -| Average Time To First Token (ms) | 81.8841 | 88.0162 | -| Average Inter-Token Latency (ms) | 11.7031 | 10.8225 | - -In this case, enabling the low-latency gemm plugin actually provided a meaningful boost to throughput. Additionally it also improved ITL but at the expense of TTFT. Furthermore, when used without the gemm+swiglu fusion, performance was actually worse than with out the plugin turned on. This suggests that for this workload the low-latency gemm plugin was choosing a worse kernel for the gemm right before the swiglu, but once that was handled by the gemm+swiglu fusion custom kernel, the rest of the kernels the low-latency gemm plugin was choosing was better than the baseline, resulting in improved performance. This underscores the importance of benchmarking different settings as the impact of this plugin is highly workload dependent. If possible some grid searching can be useful for extremely performance sensitive workloads - -## Conclusion - -Overall leveraging quantization can provide significant uplifts in performance. Here are the performance uplifts from our tuned fp8 model as compared to the tuned fp16 numbers we reached in the [previous page of guide](./tuning-max-batch-size-and-max-num-tokens.md) - -| Metric | Tuned FP16 Model | Tuned FP8 Model | % Improvement | -| -------------------------------- | ---------------- | --------------- | ------------- | -| Token Throughput (tokens/sec) | 2474.2581 | 6049.1625 | 144.48 | -| Request Throughput (req/sec) | 1.2081 | 2.9537 | 144.49 | -| Average Time To First Token (ms) | 147.5742 | 88.0162 | 40.36 | -| Average Inter-Token Latency (ms) | 14.6852 | 10.8225 | 26.30 | - -Additionally, compared to the fp8 baseline numbers (the baseline numbers had some degree of tuning, see [Baseline Performance](./fp8-quantization.md#fp8-baseline-performance) for details), we received the following performance uplifts from enabling the flags discussed above: - -| Metric | Baseline FP8 Model | Tuned FP8 Model | % Improvement | -| -------------------------------- | ------------------ | --------------- | ------------- | -| Token Throughput (tokens/sec) | 3389.5305 | 6049.1625 | 78.47 | -| Request Throughput (req/sec) | 1.6550 | 2.9537 | 78.47 | -| Average Time To First Token (ms) | 96.1597 | 88.0162 | 8.47 | -| Average Inter-Token Latency (ms) | 12.4248 | 10.8225 | 12.90 | - -As mentioned previously, the caveat with leveraging quantization are potential drops in accuracy, and we strongly recommend having a way to test whether model output quality is acceptable before attempting to use quantization. That said, many real world cases successfully use quantization and the significant performance boosts it enables are often worth the effort to see if it is a fit. - -### Summary of Configuration Option Recommendations: - -1. Quantized KV-cache: Typically provides significant throughput boost. We recommend turning it on as long as output quality is still acceptable with the feature enabled. -2. Reduce fusion + user buffers: This feature is only supported on fp8 Llama and Mistral/Mixtral models. Effectiveness is workload dependent so we recommend turning it on and benchmarking to check. -3. Gemm + Swiglu Plugin: This feature is only supported on fp8 models with Swiglu operators like Llama, Mixtral etc. Like reduce fusion effectiveness is workload dependent and we recommend sanity checking effectiveness. Has increased risk of affecting accuracy since it drops a quantization scale. -4. Low-Latency GEMM plugin: Effectiveness is workload dependent so we recommend turning it on and benchmarking. Effectiveness can be affected by other flags as we saw in our case study, so if possible benchmarking various combinations of configuration options is ideal. diff --git a/docs/source/legacy/performance/performance-tuning-guide/index.rst b/docs/source/legacy/performance/performance-tuning-guide/index.rst deleted file mode 100644 index f642a4ba2a88..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/index.rst +++ /dev/null @@ -1,15 +0,0 @@ -Performance Tuning Guide -======================= - -.. include:: introduction.md - :parser: myst_parser.sphinx_ - -.. toctree:: - :maxdepth: 1 - - benchmarking-default-performance - useful-build-time-flags - tuning-max-batch-size-and-max-num-tokens - deciding-model-sharding-strategy - fp8-quantization - useful-runtime-flags diff --git a/docs/source/legacy/performance/performance-tuning-guide/introduction.md b/docs/source/legacy/performance/performance-tuning-guide/introduction.md deleted file mode 100644 index 800ae65f37bf..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/introduction.md +++ /dev/null @@ -1,17 +0,0 @@ -While defaults are expected to provide solid performance, TensorRT-LLM has several configurable options that can improve performance for your particular workload. This guide is meant to help you tune TensorRT-LLM to extract the best performance for your use case. It covers several of the most helpful tunable parameters and provides intuition for thinking about them. This guide also doubles as an example of how to work with TensorRT-LLM's LLM-API and its TRTLLM-Bench benchmarking workflow. - -This guide uses Llama-3.3-70b on 4 H100-sxm-80GB connected via NVLink as a case study and focuses on optimizing performance on input sequence length/output sequence length of 2048/2048. Case study sections throughout this guide reference internal performance testing and results to help reinforce the conclusions and recommendations given. - -## Prerequisite Knowledge - -This guide expects you have some familiarity with the following concepts - -- Phases of Inference: Context (Prefill) Phase and Generation Phase -- Inflight Batching -- Tensor Parallelism and Pipeline Parallelism -- Quantization - - Please refer to [Mastering LLM Techniques - Inference Optimization](https://developer.nvidia.com/blog/mastering-llm-techniques-inference-optimization/) for an introduction to these concepts. - -## Table of Contents - diff --git a/docs/source/legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md b/docs/source/legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md deleted file mode 100644 index 6ed3a5c5c205..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/tuning-max-batch-size-and-max-num-tokens.md +++ /dev/null @@ -1,126 +0,0 @@ -(tuning-max-batch-size-and-max-num-tokens)= - -# Tuning Max Batch Size and Max Num Tokens - -One of TensorRT-LLM's key features is its inflight batching scheduler and runtime, which simultaneously schedules and executes requests in both context and generation phases. Max num tokens, together with max batch size, dictate how and when the scheduler admits new requests and continues current ones, and understanding how to tune them can provide significant performance benefits. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Understanding the TensorRT-LLM scheduler - -This section visualizes how TensorRT-LLM schedules requests based on max-batch size and max-num tokens. The example starts out with a newly initialized engine as well as a few unscheduled requests that have come in. For the sake of this example, toy values are set to `max batch size = 4` and `max num tokens = 12`. Each square block represents a token, and its color represents which request it belongs to. - -![TRT-LLM Scheduler Visualization 1](../../media/TRTLLM_Scheduler_Vis_1.svg) - - -Now the scheduler takes the first two requests, Request 1 and Request 2, and schedules them to execute the context phase. However, it cannot schedule any more requests because the prompts of the first two requests had 5 tokens each, leaving a budget of 2 tokens due to the max num tokens limit. Since all remaining requests have more than 2 prompt tokens none of them can be scheduled (context chunking can help in this situation, see the paged context attention section below). The tokens are marked with a "C" on them to represent that they are prompt tokens that were processed in the context phase. - -> Note: The tokens for different requests are shown on different rows simply for visualization purposes and are not representative of actual memory layouts - -![TRT-LLM Scheduler Visualization 2](../../media/TRTLLM_Scheduler_Vis_2.svg) - -Now the engine runs an iteration of execution, completing the context phases for both of the scheduled requests. After it is done, the kv-cache of the prompts for both requests have been created and the first token has been generated. Tokens that were generated are marked with "G(n)" - for example a token marked "G1" represents that it is the first token generated for its request. - -TRT-LLM prioritizes scheduling requests in generation phase first so the two generated tokens are queued to be processed in the next iteration. Now, since the two previously scheduled requests have entered generation phase and only take up two tokens out of the max num token budget of 12, the scheduler is able to schedule two additional requests, Request 3 and Request 4. It cannot schedule the last request, Request 5, even though there is space for it in the max num tokens budget because of the max batch size limit of 4. - -![TRT-LLM Scheduler Visualization 3](../../media/TRTLLM_Scheduler_Vis_3.svg) - -After the next iteration of execution, the second tokens for Requests 1 and 2 have been generated, and the first tokens for Request 3 and 4 have been generated. Lets say that G2 that was generated for Request 1 is the stop token, signifying that Request 1 is completed. In this case the scheduler would evict Request 1 before performing another execution iteration and prepare to return it to the user. This eviction puts the state of the engine below the max batch size limit and allows Request 5 to be scheduled. - -Another thing to note is that G1 that was generated for Request 2 has been added to the kv-cache for request 2, representing how kv-cache for a request grows as more and more tokens are generated. - -![TRT-LLM Scheduler Visualization 4](../../media/TRTLLM_Scheduler_Vis_4.svg) - -Overall, the max batch size and max num tokens limits play a key part in deciding when requests are actually executed, and tuning them can have significant impacts on throughput numbers as well as how the engine balances previously scheduled requests in generation phase with context phase on new requests - -> Note: This presents a simplified visualization of the scheduler to highlight how max batch size and max num tokens affect it. The scheduler also considers things like amount of free memory available to be used for kv-cache and has other configurable options that can affect its behavior. See the Runtime Flags of the Additional Options page for more on this. - - - ## Tuning Max Batch Size - - It is important to set the max batch size large enough that it doesn't bottleneck the scheduling of new requests. Due to this it's recommended to test additional values of max batch size for your workload. The default value is 2048. Powers of 2 are good initial values to sweep through. - -### How to change Max Batch Size - -You can specify the max batch size in the build config. - -```python -build_config = BuildConfig( - max_batch_size=512 -) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--max_batch_size ` to `trtllm-build` to tune max batch size. - -### Tuning Case Study - -Continuing with our case study of Llama-3.3-70B on 4 H100s, we pick up where left off in the previous section (building an engine with multiple profiles, gemm plugin, paged context attention, and reduce fusion enabled). Sweeping across max batch sizes of 64, 512, and the default 2048 produced the following results - -| Metric | Max Batch Size 64 | Max Batch Size 512 | Max Batch Size 2048 | -| -------------------------------- | ----------------- | ------------------ | ------------------- | -| Token Throughput (tokens/sec) | 1944.3031 | 2466.7933 | 2044.2628 | -| Request Throughput (req/sec) | 0.9494 | 1.2045 | 0.9982 | -| Average Time To First Token (ms) | 145.7607 | 147.7876 | 146.6628 | -| Average Inter-Token Latency (ms) | 14.6475 | 14.6554 | 14.4493 | - -From this its clear a max batch size of 64 results in bottlenecking whereas a max batch size of 512 is the sweet spot, boosting throughput by almost 20% with negligible effect on latency. - -## Tuning Max Num Tokens - -If max num tokens is too small it can bottleneck request scheduling. However, if it is too large it can result in prompt tokens taking up too much memory (especially in long-context workloads), not leaving enough for kv-cache and resulting in performance hits or even out-of-memory errors. The default value of max num tokens is 8192. It's recommended that you sweep across several values of max num tokens to identify the best number for your workload. Good values to try are powers of 2 >= 1024. Since max num tokens and max batch size both affect scheduling, grid searching across combinations of them if possible is ideal. - -### How to change Max Num Tokens - -Like Max Batch size, max num tokens is also specified in the build config. - -```python -build_config = BuildConfig( - max_batch_size=512 - max_num_tokens=2048 -) -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--max_num_tokens ` to `trtllm-build` to tune max_num_tokens. - -### Tuning Case Study - -Sweeping across max num tokens values of 2048, 8192, and 16384 yielded the following performance numbers with max batch size set to 512. - -| Metric | Max Num Tokens 2048 | Max Num Tokens 8192 | Max Num Tokens 16384 | -| -------------------------------- | ------------------- | ------------------- | -------------------- | -| Token Throughput (tokens/sec) | 2474.2581 | 2466.7933 | 2461.0165 | -| Request Throughput (req/sec) | 1.2081 | 1.2045 | 1.2017 | -| Average Time To First Token (ms) | 147.5742 | 147.7876 | 147.9623 | -| Average Inter-Token Latency (ms) | 14.6852 | 14.6554 | 14.6769 | - -For this particular workload max num tokens of 2048 provides the best performance, but not by an extremely large margin. This reflects the reality that for any given workload tuning various flags can have varied impact. However it is important to check different values to ensure that you are not giving up large gains due to scheduling imbalances. - -## Revisiting Paged Context Attention and Context Chunking - -[Previously](./useful-build-time-flags.md#paged-context-attention) we recommended enabling paged context attention even though in our case study it didn't affect performance significantly. Having now understood the TensorRT-LLM scheduler we can now explain why this is beneficial. In short, we recommend enabling it because it enables context chunking, which allows the context phase of a request to be broken up into pieces and processed over several execution iterations, allowing the engine to provide a more stable balance of context and generation phase execution. - -The [visualization](#understanding-the-trt-llm-scheduler) of the TensorRT-LLM scheduler showed that initially Request 3 couldn't be scheduled because it would put the scheduler over the max-num tokens limit. However with context chunking, this is no longer the case, and the first chunk of Request 3 would be able to be scheduled. - -![TRT-LLM Scheduler Visualization Chunked Context 1](../../media/TRTLLM_Scheduler_Vis_Chunked_Context_1.svg) - -This is extremely beneficial for several reasons. Firstly it eliminates the possibility of requests with large prompts relative to max num tokens being unable to be scheduled due to other requests that are already in-flight. In production workloads, this can help improve worst case TTFT numbers. Secondly it allows for setting smaller values of max num tokens since you no longer need max num tokens to be at least as large as the longest prompt you want to support. For long-context cases this is extremely important, because setting extremely large values of max-num tokens takes away from memory available to be used as kv-cache. Given that in the worst case scenario chunked context has minimal impact on performance but can significantly benefit it in many scenarios, it's recommended that you always enable it. - -## Conclusion - -The TensorRT-LLM Scheduler plays a large role in performance, and properly tuning it can provide significant performance uplifts. In the case-study example, tuning max batch size and max num tokens provided the following boosts to performance when compared to the [results from the previous page](./useful-build-time-flags.md#conclusion): - -| Metric | Build-Time Flags ON | Tuned Max Batch Size and Max Num Tokens | % Improvement | -| -------------------------------- | ------------------- | --------------------------------------- | ------------- | -| Token Throughput (tokens/sec) | 2044.2628 | 2474.2581 | 21.03 | -| Request Throughput (req/sec) | 0.9982 | 1.2081 | 21.03 | -| Average Time To First Token (ms) | 146.6628 | 147.5742 | -0.62 | -| Average Inter-Token Latency (ms) | 14.4493 | 14.6852 | -1.63 | - -Interpreting these results, tuning max batch size and max num tokens significantly improved throughput and remained at par on latency (the slight drops are within run to run variance). Compared to the [initial baseline](./benchmarking-default-performance.md#results) the tuned engine achieves the following uplifts: - -| Metric | Baseline | Build-Time Flags ON and Tuned Max Batch Size and Max Num Tokens | % Improvement | -| -------------------------------- | --------- | --------------------------------------- | ------------- | -| Token Throughput (tokens/sec) | 1564.3040 | 2474.2581 | 58.17 | -| Request Throughput (req/sec) | 0.7638 | 1.2081 | 58.17 | -| Average Time To First Token (ms) | 147.6976 | 147.5742 | 0.08 | -| Average Inter-Token Latency (ms) | 31.3276 | 14.6852 | 53.12 | diff --git a/docs/source/legacy/performance/performance-tuning-guide/useful-build-time-flags.md b/docs/source/legacy/performance/performance-tuning-guide/useful-build-time-flags.md deleted file mode 100644 index 54ecd97fa6bd..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/useful-build-time-flags.md +++ /dev/null @@ -1,171 +0,0 @@ -(useful-build-time-flags)= - -# Useful Build-Time Flags - -This page presents several build-time flags, set via the LLM-API's `BuildConfig` class that you can enable to improve upon the baseline performance. Build-time refers to the fact that these flags affect how the TensorRT-LLM engine is built and cannot be changed without rebuilding the engine. For each flag there is an explanation of what it does, a description of how to enable it, and then an example of running it through the benchmarking flow described in [Benchmarking Default Performance](./benchmarking-default-performance.md) to showcase its impact on performance. All options compatible with `trtllm-build` can be found in the Command Line Reference section of the docs. - -> Disclaimer: While performance numbers shown here are real, they are only for demonstration purposes. Differences in environment, SKU, interconnect, and workload can all significantly affect performance and lead to your results differing from what is shown here. - -## Multiple Profiles - -TensorRT-LLM is built on TensorRT, which handles engine building through "optimization profiles" defining min, optimal, and max input tensor shapes. TensorRT optimizes for the optimal shape while supporting the range between min and max. - -TensorRT-LLM abstracts away the need to create optimization profiles although flags like max_batch_size and max_num_tokens (covered later) influence how they are created. By default, only one profile is created. - -During inference serving, varying request loads can pose different tensor shapes to the engine. TensorRT addresses this by allowing multiple profiles, which TensorRT-LLM supports via the BuildConfig option in the LLM-API. Enabling multiple profiles increases build times but has no performance downsides, so it is recommended for production builds. - -The only thing to watch out for is that enabling this can lead to slightly different outputs when the same prompt is run multiple times as different profiles and consequently kernels might be used depending on the request load. However this variance should not affect output quality so it is safe to enable this flag as long as you don't need completely deterministic outputs. - -### Enabling building with multiple profiles - -Below is an example of how you can modify the baseline example to enable multiple profiles. - -```python -from tensorrt_llm import LLM, BuildConfig - -def main(): - build_config = BuildConfig() - build_config.plugin_config.multiple_profiles = True - - llm = LLM( - model="/scratch/Llama-3.3-70B-Instruct", - tensor_parallel_size=4, - build_config=build_config - ) - - llm.save("build_flags_multiple_profiles") - -if __name__ == '__main__': - main() -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--multiple_profiles` to `trtllm-build` to enable the feature. - - -### Performance with multiple profiles - -Baseline refers to the engine that was benchmarked in the previous Benchmarking Default Performance page. - -| Metric | Baseline | Multiple Profiles ON | -| -------------------------------- | --------- | -------------------- | -| Token Throughput (tokens/sec) | 1564.3040 | 1861.0881 | -| Request Throughput (req/sec) | 0.7638 | 0.9087 | -| Average Time To First Token (ms) | 147.6976 | 145.8958 | -| Average Inter-Token Latency (ms) | 31.3276 | 19.6452 | - -As you can see, enabling multiple profiles significantly improves the metrics across the board. - -## Paged Context Attention - -By default all the tokens of the prompt of a new request are processed in one iteration as the context phase. Enabling paged context attention allows TensorRT-LLM to break the context phase into chunks and handle the prompt over several iterations. This is particularly useful for workloads with large input length. In the worst case, this feature can provide a small performance hit in benchmarking runs (<2%) so it can be safely enabled. This feature is discussed further in the [next page](./tuning-max-batch-size-and-max-num-tokens.md#revisiting-paged-context-attention-and-context-chunking) of the guide. - - -### Enabling Paged Context Attention - -Add the following line to our multiple profiles example from above to enable paged context attention - -```python -build_config.plugin_config.use_paged_context_fmha=True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--use_paged_context_fmha` to `trtllm-build` to enable the feature. - -### Performance - -Paged Context OFF refers to the same engine shown as Multiple Profiles ON in the previous example. - -| Metric | Paged Context OFF | Paged Context ON | -| -------------------------------- | ----------------- | ---------------- | -| Token Throughput (tokens/sec) | 1861.0881 | 1866.6684 | -| Request Throughput (req/sec) | 0.9087 | 0.9115 | -| Average Time To First Token (ms) | 145.8958 | 145.4089 | -| Average Inter-Token Latency (ms) | 19.6452 | 19.6523 | - -In this case enabling paged context attention provides a small boost to performance, but a rerun of our tests found this to be within run to run variance of around 10 tok/s for token throughput and 2ms for average time to first token (ITL was stable with <1ms and request throughput corresponded directly to token throughput). In other cases naively enabling it might actually provide a small hit to performance. However, further guidance on how to reason about this flag and why we recommend enabling it is discussed in the [next page](./tuning-max-batch-size-and-max-num-tokens.md#revisiting-paged-context-attention-and-context-chunking) as it is closely intertwined with how TensorRT-LLM schedules requests as well as the max-num tokens flag. - -## GEMM Plugin - -TensorRT allows you to add "plugins" or custom kernels that can be used instead of the kernels that TensorRT selects for particular operations. TensorRT-LLM has a host of custom plugins that are specifically tailored to speed up supported modules. The GEMM plugin utilizes NVIDIA cuBLASLt and some custom kernels to perform GEMM operations. On FP16 and BF16, it’s recommended to be enabled for better performance and smaller GPU memory usage. On FP8, it’s recommended to be disabled. - -### Enabling GEMM Plugin - -Add the following line to the multiple profiles example from above to enable the GEMM plugin. - -```python -build_config.plugin_config.gemm_plugin = 'auto' -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--gemm_plugin auto` to `trtllm-build` to enable the feature. `'auto'` tells the GEMM plugin to have the same type as the model (fp16, bf16, etc). It is fine to leave it on auto unless you are trying to do mixed precision. - -### Performance with GEMM Plugin - -GEMM Plugin OFF refers to the same engine shown as Paged Context ON in the previous example. - -| Metric | GEMM Plugin OFF | GEMM Plugin ON | -| -------------------------------- | --------------- | -------------- | -| Token Throughput (tokens/sec) | 1866.6684 | 2033.2640 | -| Request Throughput (req/sec) | 0.9115 | 0.9928 | -| Average Time To First Token (ms) | 145.4089 | 147.8307 | -| Average Inter-Token Latency (ms) | 19.6523 | 15.4133 | - -In this case the GEMM plugin greatly improves throughput as well as ITL, with a slight hit to TTFT. - -## Reduce Norm Fusion Plugin for Llama models: - -TensorRT-LLM has custom kernels for AllReduce operations that are enabled by default. This feature extends this functionality by fusing the ResidualAdd and LayerNorm kernels that run after AllReduce into the AllReduce kernel, resulting in a single kernel that handles those operations and improves end-to-end performance. This feature is currently only available for Llama models. It is most beneficial in workloads that are generation-phase heavy. For extremely context-phase heavy workloads its worth checking performance with and without this. Additionally, since this is an optimization for AllReduce, it is only beneficial for cases with tensor-parallelism. For scenarios only using pipeline parallelism this should stay disabled since pipeline parallelism doesn't require any AllReduce operations. - -### Enabling Reduce Norm Fusion Plugin - -Add the following line to the multiple profiles example from above to enable reduce norm fusion. - -```python -build_config.plugin_config.reduce_fusion = True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) pass `--reduce_fusion enable` to `trtllm-build` to enable the feature. - -### Performance with Reduce Norm Fusion - -Reduce Fusion OFF refers to the same engine shown as GEMM Plugin ON in the previous example. - -| Metric | REDUCE FUSION OFF | REDUCE FUSION ON | -| -------------------------------- | ----------------- | ---------------- | -| Token Throughput (tokens/sec) | 2033.2640 | 2044.2628 | -| Request Throughput (req/sec) | 0.9928 | 0.9982 | -| Average Time To First Token (ms) | 147.8307 | 146.6628 | -| Average Inter-Token Latency (ms) | 15.4133 | 14.4493 | - -For the ISL/OSL pair of 2048/2048 enabling the reduce norm fusion plugin slightly improves performance all around. However, test reruns found that with run to run variance, in the worst case, they performed at par. Again this flag's effectiveness is dependent on the workload so users should check whether it provides meaningful performance boosts in their case. - - -## Pipeline Parallel Reduce Scatter Optimization - -This feature adds a pipeline parallelism optimization with ReduceScatter + AllGather targeting large mixture of experts models. -This can be enabled via the LLM-API as such -```python - build_config.plugin_config.pp_reduce_scatter = True -``` - -If you are using the [CLI flow for building engines](./benchmarking-default-performance.md#building-and-saving-engines-via-cli) flow you can enable this feature by adding `--pp_reduce_scatter` to `trtllm-build`. - -As the Llama model is not a MoE model this flag was not included as part of the case study. - -## Conclusion - -Overall, enabling these flags can greatly boost performance. However, the degree to which they are effective can vary from workload to workload, and it's recommended that you run sanity checks on your workloads to verify performance. - -The case-study example showed that enabling these flags provided the following performance uplifts from the baseline numbers. This included significant boosts in Token Throughput, Request Throughput, and Average Inter-Token Latency. TTFT remained largely unchanged. - -| Metric | Baseline | Build-Time Flags ON | % Improvement | -| -------------------------------- | --------- | ------------------- | ------------- | -| Token Throughput (tokens/sec) | 1564.3040 | 2044.2628 | 30.68 | -| Request Throughput (req/sec) | 0.7638 | 0.9982 | 30.69 | -| Average Time To First Token (ms) | 147.6976 | 146.6628 | 0.70 | -| Average Inter-Token Latency (ms) | 31.3276 | 14.4493 | 53.88 | - -### Summary of Configuration Option Recommendations: - -1. Multiple profiles: Always enable. It may increase build times a little but will only ever help performance. Enabling might cause engine to produce slightly different outputs when the same prompt is run multiple times depending on request load but it should not affect output quality, see [Multiple Profiles section](./useful-build-time-flags.md#multiple-profiles) for explanation. -2. Paged Context Attention: In the worst case it may hurt performance a little initially but typically helps with request scheduling and boosts performance after further tuning of max batch size and max num tokens. More on this topic is discussed in the next page. -3. GEMM Plugin: It's recommended to enable it for FP16 and BF16 models as it usually helps. However, it is a good idea to benchmark your workload and double check that it is helping. -4. Reduce Fusion: This feature is only supported on Llama and Mistral/Mixtral models. Effectiveness is workload dependent and it's recommend that you benchmark your workload with and without it and compare the results. diff --git a/docs/source/legacy/performance/performance-tuning-guide/useful-runtime-flags.md b/docs/source/legacy/performance/performance-tuning-guide/useful-runtime-flags.md deleted file mode 100644 index 8c5579431b96..000000000000 --- a/docs/source/legacy/performance/performance-tuning-guide/useful-runtime-flags.md +++ /dev/null @@ -1,224 +0,0 @@ -(useful-runtime-flags)= - -# Useful Runtime Options - -This part summarizes the runtime configuration knobs that can be tweaked to -enhance the performance of already built engines. As compared to previous examples where - the LLM-API was used to build and save an engine but not to process any requests, -runtime knobs would be specified when you are using the LLM-API to actually run inference -like in the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) - - -## Capacity Scheduler Policy - -TensorRT-LLM currently supports three batch scheduler policies: `GUARANTEED_NO_EVICT` (default), -`MAX_UTILIZATION` and `STATIC_BATCH`. - -The scheduling policy can be set to `MAX_UTILIZATION` to pack as many -requests as possible at each iteration of the forward loop, when in-flight -sequence batching is enabled. It maximizes the utilization of the GPUs by -aggressively scheduling requests at the risk of having to pause requests if the -KV cache size limit is reached. - -For a more conservative approach with respect to the KV cache limitations in -terms of memory allocation, `CapacitySchedulerPolicy` should be set to -`GUARANTEED_NO_EVICT` to guarantee that a started request is never paused. - -If the goal is to maximize the throughput, users should try `MAX_UTILIZATION`. -However, they need to keep in mind that it may have a negative impact on -latency if requests have to be paused. - -`STATIC_BATCH` is a legacy mode and is not recommended for production usage. - -To switch the capacity scheduler policy from the default of `GUARANTEED_NO_EVICT` to `MAX_UTILIZATION` -you would modify the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) to be: - -```python -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.bindings.executor import SchedulerConfig, CapacitySchedulerPolicy - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - scheduler_config = SchedulerConfig( - capacity_scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION - ) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", - tensor_parallel_size=4, - scheduler_config=scheduler_config - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` - -## Context Chunking Policy - -As discussed [previously](tuning-max-batch-size-and-max-num-tokens.md#revisiting-paged-context-attention-and-context-chunking) context chunking will increase the chance of batch processing between -the context and the generation phase, thereby balancing the calculation amount -of each iteration and typically increasing throughput. - -TensorRT-LLM currently supports two context chunking policies: `FIRST_COME_FIRST_SERVED` (default) which would prioritize scheduling all the context chunks of a request that comes in first, - and `EQUAL_PROGRESS` which schedules context chunks from all requests before scheduling the next chunk of any request. - -`FIRST_COME_FIRST_SERVED` should achieve overall better performance, while -`EQUAL_PROGRESS` can be helpful in theory to make sure time to first token (TTFT) -for most requests are relatively similar. - -To switch the context chunking policy from the default of `FIRST_COME_FIRST_SERVED` to `EQUAL_PROGRESS` -you would modify the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) to be: - -```python -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.bindings.executor import SchedulerConfig, ContextChunkingPolicy - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - scheduler_config = SchedulerConfig( - context_chunking_policy=ContextChunkingPolicy.EQUAL_PROGRESS - ) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", - tensor_parallel_size=4, - scheduler_config=scheduler_config - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` - -## Max Tokens in Paged KV Cache and KV Cache Free GPU Memory Fraction - -The `max_tokens_in_paged_kv_cache` and `kv_cache_free_gpu_mem_fraction` -parameters can be used to control the maximum number of tokens handled by the -KV cache manager. Setting them properly helps better control the amount of -available memory for the KV cache manager during inference. Keeping in mind -that increasing the amount of memory available to the KV cache manager tends to -translate to a higher achievable throughput. - -The `max_tokens_in_paged_kv_cache` flag directly sets the maximum number of -tokens in the KV cache manager. When left unset, that value will be computed -based on the `kv_cache_free_gpu_mem_fraction` setting. - -The `kv_cache_free_gpu_mem_fraction` is a floating-point number between `0.0` -and `1.0` that indicates the maximum fraction of GPU memory (after loading the -model) that will be used for the KV cache. The default value is `0.90` and -means that 90% of the free GPU memory will be used to save tokens in the KV -cache. Based on that value, TensorRT-LLM can determine the maximum number of -tokens in the KV cache manager. - -When both parameters are set, the maximum number of tokens in the KV cache -manager will be set to the smaller value between `max_tokens_in_paged_kv_cache` -and the value computed from the amount of memory available for the KV cache. - -Unless users clearly know the maximum number of tokens in the KV cache needed -by the model, it is recommended to leave `max_tokens_in_paged_kv_cache` unset. -For `kv_cache_free_gpu_mem_fraction`, if no other programs are executed on the -same GPU, it is recommended to test with a as high value as `0.95` to target a -high throughput. Note that the `kv_cache_free_gpu_mem_fraction` parameter -cannot be set to `1.0` because some amount of memory has to be reserved for -inputs and outputs. - -To set `kv_cache_free_gpu_mem_fraction` you would modify the [LLM-API end-to-end example](./benchmarking-default-performance.md#before-you-begin-tensorrt-llm-llm-api) to be: - -```python -from tensorrt_llm import LLM, SamplingParams -from tensorrt_llm.bindings.executor import KvCacheConfig - - -def main(): - prompts = [ - "Hello, I am", - "The president of the United States is", - "The capital of France is", - "The future of AI is", - ] - - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.95) - - llm = LLM( - model="meta-llama/Llama-3.3-70B-Instruct", - tensor_parallel_size=8, - kv_cache_config=kv_cache_config - ) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - -if __name__ == '__main__': - main() -``` -If you wanted to set `max_tokens_in_paged_kv_cache` instead, you would replace `free_gpu_memory_fraction` with `max_tokens` and specify the number. - -```python - kv_cache_config = KvCacheConfig(max_tokens=) -``` - - -## Maximum Attention Window Size - -The `max_attention_window_size` flag sets the maximum number of tokens that are -attended to in order to generate one token when using techniques like sliding window -attention. See this -[Document](../../advanced/gpt-attention.md#sliding-window-attention-cyclic-rolling-buffer-kv-cache) -for more details. It defaults to the maximum sequence length -(`max_seq_len` when building the engine), which means -that the feature is disabled by default. - -When set to a smaller value than `max_seq_len` (during -engine build), only the KV cache of the last `max_attention_window_size` tokens -will be stored. If the input sequence length at runtime exceeds the -`max_attention_window_size` value, the accuracy may start dropping, but the -runtime performance will be better (due to the reduction in terms of -computations and GPU memory allocation). Users can modify that value to -increase runtime performance at the expense of reduced accuracy. - -Just like [`kv_cache_free_gpu_mem_fraction`](./useful-runtime-flags.md#max-tokens-in-paged-kv-cache-and-kv-cache-free-gpu-memory-fraction), `max_attention_window_size` can be specified in the LLM-API -via `KVCacheConfig`. To specify `max_attention_window_size` you would instantiate `KVCacheConfig` like so - -```python - kv_cache_config = KvCacheConfig(max_attention_window=) -``` diff --git a/docs/source/legacy/python-api/tensorrt_llm.functional.rst b/docs/source/legacy/python-api/tensorrt_llm.functional.rst deleted file mode 100644 index bfa0368f1793..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.functional.rst +++ /dev/null @@ -1,11 +0,0 @@ -Functionals -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.functional - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.layers.rst b/docs/source/legacy/python-api/tensorrt_llm.layers.rst deleted file mode 100644 index 7baa17ab0d67..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.layers.rst +++ /dev/null @@ -1,69 +0,0 @@ -Layers -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -Activation ------------- -.. automodule:: tensorrt_llm.layers.activation - :members: - :undoc-members: - :show-inheritance: - -Attention ------------- -.. automodule:: tensorrt_llm.layers.attention - :members: - :undoc-members: - :show-inheritance: - -Cast ------------- -.. automodule:: tensorrt_llm.layers.cast - :members: - :undoc-members: - :show-inheritance: - -Conv ------------- -.. automodule:: tensorrt_llm.layers.conv - :members: - :undoc-members: - :show-inheritance: - -Embedding ------------- -.. automodule:: tensorrt_llm.layers.embedding - :members: - :undoc-members: - :show-inheritance: - -Linear ------------- -.. automodule:: tensorrt_llm.layers.linear - :members: - :undoc-members: - :show-inheritance: - -MLP ------------- -.. automodule:: tensorrt_llm.layers.mlp - :members: - :undoc-members: - :show-inheritance: - -Normalization ---------------- -.. automodule:: tensorrt_llm.layers.normalization - :members: - :undoc-members: - :show-inheritance: - -Pooling ------------- -.. automodule:: tensorrt_llm.layers.pooling - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.models.rst b/docs/source/legacy/python-api/tensorrt_llm.models.rst deleted file mode 100644 index ca0a5c06f50c..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.models.rst +++ /dev/null @@ -1,11 +0,0 @@ -Models -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.models - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.plugin.rst b/docs/source/legacy/python-api/tensorrt_llm.plugin.rst deleted file mode 100644 index b4459f98e2ce..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.plugin.rst +++ /dev/null @@ -1,10 +0,0 @@ -Plugin -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.plugin - :members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.quantization.rst b/docs/source/legacy/python-api/tensorrt_llm.quantization.rst deleted file mode 100644 index 1da9ec81fb6e..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.quantization.rst +++ /dev/null @@ -1,10 +0,0 @@ -Quantization -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.quantization - :members: - :show-inheritance: diff --git a/docs/source/legacy/python-api/tensorrt_llm.runtime.rst b/docs/source/legacy/python-api/tensorrt_llm.runtime.rst deleted file mode 100644 index 2bd02e3dca5c..000000000000 --- a/docs/source/legacy/python-api/tensorrt_llm.runtime.rst +++ /dev/null @@ -1,11 +0,0 @@ -Runtime -=========================== - -.. automodule:: tensorrt_llm - -.. currentmodule:: tensorrt_llm - -.. automodule:: tensorrt_llm.runtime - :members: - :undoc-members: - :show-inheritance: diff --git a/docs/source/legacy/reference/memory.md b/docs/source/legacy/reference/memory.md deleted file mode 100644 index b20429e061c2..000000000000 --- a/docs/source/legacy/reference/memory.md +++ /dev/null @@ -1,140 +0,0 @@ -(memory)= - -# Memory Usage of TensorRT-LLM - - -This document summarizes the memory usage of TensorRT-LLM, and addresses common issues and questions reported by users. - - -## Understand inference time GPU memory usage - - -At inference time, there are 3 major contributors to GPU memory usage for a given TRT engine generated from a TensorRT-LLM model: weights, internal activation tensors, and I/O tensors. For I/O tensors, the major memory footprint comes from the KV cache tensor. - - -### 1. Weights size - -Weights size is fixed depending on the model size, the chosen precision of the weights and the parallelization strategy. -Using lower precision like INT8 or FP8 can reduce the weights size. -When tensor parallelism or pipeline parallelism is used, each rank stores only some portion of the weights. -For example, each rank typically uses just 1/8 of the model weights when using 8-way tensor parallelism or 8-stages pipeline parallelism. - - -### 2. Activation size - - -TensorRT can optimize the memory usage by reusing memory for different tensors based on live analysis and tensor size. To avoid out of memory errors at runtime and to reduce the runtime cost of switching optimization profiles and changing shapes, **TensorRT pre-computes the activation tensors memory requirement at build time**. The memory requirement is computed based on an optimized TensorRT graph, one profile’s memory usage is computed by using the max tensor shape, and the memory requirement of one engine is computed by the maximum size between different profiles. There are external and internal factors that can affect the activation size returned by TensorRT, such as the network structure, kernel fusion, operation scheduling, etc. - -Once the TensorRT engine is built, the activation memory size of that engine **cannot be changed**, and can be queried by the API `trt.ICudaEngine.device_memory_size_v2`. - -Practically, for a given model, specified precision and parallelization strategy, one can tune the activation memory usage by adjusting the max batch size, max input length, max beam width, max number of tokens, padding removal on/off flag, context FMHA on/off flag. -Here some explanations on how these values affect the memory: - - -1. Reduce build time max number of input tokens (`max_num_tokens`) - - Most of the tensors inside a transformer network have a linear relationship with number of input tokens, so activation size will be close to `max number of input tokens * some constant factor`, the constant factor depends on the network structure and TRT internal optimization. The max number of input tokens is derived from build time arguments, one can change the parameters provided to the `prepare_inputs` function, like `PretrainedModel.prepare_inputs` to affect the memory usage, or one can change the command line options of the `trtllm-build` command used in the examples. - - When using the [packed tensors](../advanced/gpt-attention.md#padded-and-packed-tensors) format and `max_num_tokens` is specified, reducing its value will also reduce activation memory size. - - When using the [padded tensors](../advanced/gpt-attention.md#padded-and-packed-tensors) format, the max number of input tokens equals to `max_batch_size*max_input_len`, so reducing `max_batch_size` and `max_input_len` can almost linearly reduce the activation memory size. - - The packed tensors format is recommended, because it saves both memory and compute. - - The beam width will be folded into the batch size dimension when passing the tensors range into TensorRT, so reducing `max_beam_width` can also reduce the memory usage. - - -2. Turn on context FMHA - - When the GPT attention plugin is used, turning on the `context_fmha_type` of the plugin will reduce the memory footprint significantly. See the [Context Phase](../advanced/gpt-attention.md#context-phase) for details. When the `context_fmha_type` is set to disabled, a workspace size of the plugin will quadratically depend on the sequence length. - - -3. Tensor parallelism and pipeline parallelism - - TensorRT will reuse memory between layers as much as possible, for a typical example, given *N* decoder blocks in one transformer network, TRT will not allocate *N* copies of the activation memory for each block, since the memory of tensors in the 1st block can be released after the execution, memory can be reused for later blocks, only 1 block’s memory is needed. - - - When using tensor parallelism, some tensors are split into smaller chunks and each rank only holds one chunk of the tensor, the activation memory size of each rank will be smaller than when executing the network on a single GPU. When using pipeline parallelism, each rank executes several decoder blocks, and all the tensors are full-size tensors, so the activation memory size is equal to 1 block’s memory size. Thus tensor parallelism normally has higher memory efficiency than pipeline parallelism when all other parameters are the same. - - -### 3. I/O tensors - -#### 3.1 Runtime and decoder buffers except KV cache tensor - -##### C++ runtime - -Before KV cache blocks are allocated, some amount of GPU memory are pre-allocated by C++ runtime for storing I/O tensors of TensorRT engine and the decoupled dynamic decoder, it's allocated based on runtime max_batch_size and max_seq_len so that OOM can be avoided when there are indeed that amount of requests scheduled. - -#### 3.2 KV cache tensor - -##### C++ runtime - - TensorRT-LLM runtime pre-allocates paged KV cache pools during initialization for a configured number of blocks and distributes them at runtime. - - KV cache tensors are allocated based on the `KVCacheConfig` object when creating the `Executor`. If neither `maxTokens` nor `freeGpuMemoryFraction` is specified, KV cache will by default allocate 90% of the remaining free GPU memory. When either `maxTokens` or `freeGpuMemoryFraction` is specified, the specified value will be used to compute the KV cache memory size. And if both are specified, firstly the `freeGpuMemoryFraction` is used to compute the number of tokens in KV cache, and then the minimum between this computed number of tokens and `maxTokens` is used. - - In in-flight batching the scheduler can automatically schedule requests as long as enough KV cache space is available (exact behavior depends on the scheduler policy). - -##### Python runtime (Not recommended to be used) - -The Python runtime allocates KV cache tensors based on the parameters of the `GenerationSession.setup` function, the KV cache size is linearly dependent on the `batch_size` and `max_context_length+max_new_tokens`. **Note: This may change in the future, as the Python bindings of the C++ runtime may replace the current python runtime in the future. The Python bindings of C++ runtime behave like C++ runtime.** - -## Memory pool - -TensorRT-LLM C++ runtime is using stream-ordered memory allocator to allocate and free buffers, see [BufferManager::initMemoryPool](source:cpp/tensorrt_llm/runtime/bufferManager.cpp), which uses the default memory pool managed by the CUDA driver. When a `TrtGptModel` object is destroyed, memory is returned to the memory pool and can be reused by the next instance of a `TrtGptModel` object. Memory will be released from the pool if it is required for other memory allocations. - -However, `nvidia-smi` may still show high memory occupation after memory is returned to the CUDA driver's memory pool. This should not be a concern and is intended behavior. The amount of reserved and free memory in the pool can be inspected by [BufferManager::memoryPoolReserved())](source:cpp/tensorrt_llm/runtime/bufferManager.cpp) and [BufferManager::memoryPoolFree())](source:cpp/tensorrt_llm/runtime/bufferManager.cpp), respectively. - -## Known Issues - - -When FP8 GEMM is used, the activation memory might be larger than the theoretical optimized memory size, this will be enhanced in a future release. - -## FAQ - -1. How to debug the memory usage of TensorRT-LLM? - - When the `info` logging level is used, TensorRT and TensorRT-LLM will print messages about memory usage details. Here is part of a log example with `info` logging level at runtime: - ``` - [TensorRT-LLM][INFO] Loaded engine size: 6695 MiB - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 1134.01 MiB for execution context memory. - [TensorRT-LLM][INFO] [MS] Running engine with multi stream info - [TensorRT-LLM][INFO] [MS] Number of aux streams is 1 - [TensorRT-LLM][INFO] [MS] Number of total worker streams is 2 - [TensorRT-LLM][INFO] [MS] The main stream provided by execute/enqueue calls is the first worker stream - [TensorRT-LLM][INFO] [MemUsageChange] TensorRT-managed allocation in IExecutionContext creation: CPU +0, GPU +0, now: CPU 0, GPU 6678 (MiB) - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 43.29 MB GPU memory for runtime buffers. - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 180.30 MB GPU memory for decoder. - [TensorRT-LLM][INFO] Memory usage when calculating max tokens in paged kv cache: total: 79.10 GiB, available: 70.48 GiB - [TensorRT-LLM][INFO] Number of blocks in KV cache primary pool: 4060 - [TensorRT-LLM][INFO] Number of blocks in KV cache secondary pool: 0, onboard blocks to primary memory before reuse: true - [TensorRT-LLM][INFO] Max KV cache pages per sequence: 32 - [TensorRT-LLM][INFO] Number of tokens per block: 64. - [TensorRT-LLM][INFO] [MemUsageChange] Allocated 63.44 GiB for max tokens in paged KV cache (259840). - ``` - You can see that there are several GPU memory allocation started with `[MemUsageChange]` keyword happened at runtime. - - The line showing "Total Weights Memory" indicates the weights memory size, and the line "Total Activation Memory" indicates the activation memory size. - - Normally the weights memory size is close to the TensorRT engine size, since most of the content in the engine is from weights for LLM networks. - -2. Why is the memory size large even though a small batch size and sequence length are used in the runtime? - - As explained above, the activation memory size is computed based on the max tensor shapes at TensorRT engine building time, try to reduce the engine building time parameters like `max_num_token`, see [Activation size](#activation-size) for details. - - -3. Why can the engine be generated, but the inference will run out of memory (OOM) at runtime? - - At engine building time, TensorRT will tune the kernel selection layer by layer, it does not necessarily allocate all the memory required to run the entire engine. If the activation tensors required to run a single layer are small, while the I/O tensor (like KV cache) sizes required to run the engine are large, building will succeed since it may not need to allocate the large I/O tensors, runtime may fail with OOM errors on allocating large IO tensors. - - TensorRT-LLM has provided a `check_gpt_mem_usage` utility function to check the upper bound of the memory size given an engine, and the related batch size, I/O sequence length, etc., when the upper boundary check exceeded the GPU physical memory size, warning messages will be printed. - -4. For pipeline parallelism, is build time max batch size the limit of micro batch size? - - Yes, in pipeline parallel mode, TensorRT-LLM runtime will split the batch of requests into micro batches, and enqueue these micro batches into TRT engine sequentially. - - The `max_batch_size` at build time means that batch size of one engine enqueue call shall be smaller than it. The total batch size before splitting into micro batches can be larger than the build time `max_batch_size`. - - For example, if you have 4-stages pipeline parallelism, and intend to run the engine using micro batch size 2 and run 16 micro batches (total batch size 32) in one `generate` call. - - You could just set the `max_batch_size` at building time to 2, instead of 32. Setting build time `max_batch_size` 32 will occupy almost 16x more activation memory. diff --git a/docs/source/legacy/reference/multimodal-feature-support-matrix.md b/docs/source/legacy/reference/multimodal-feature-support-matrix.md deleted file mode 100644 index b6d99e24ca69..000000000000 --- a/docs/source/legacy/reference/multimodal-feature-support-matrix.md +++ /dev/null @@ -1,13 +0,0 @@ -# Multimodal Feature Support Matrix (PyTorch Backend) - -| Model | CUDA Graph | Encoder IFB | KV Cache Reuse | Chunked Prefill | -| :----------------- | :--------- | :------------------ | :------------- | :-------------- | -| Gemma 3 | Yes | Yes | N/A | N/A | -| HyperCLOVA | Yes | Yes | No | No | -| VILA | Yes | No | No | No | -| LLaVA-NeXT | Yes | Yes | Yes | Yes | -| Llama 4 | Yes | Yes | No | No | -| Mistral-Small-3.1 | Yes | Yes | Yes | Yes | -| Phi-4-multimodal | Yes | Yes | Yes | Yes | -| Qwen2-VL | Yes | Yes | Yes | Yes | -| Qwen2.5-VL | Yes | Yes | Yes | Yes | diff --git a/docs/source/legacy/reference/precision.md b/docs/source/legacy/reference/precision.md deleted file mode 100644 index b31eff6d6238..000000000000 --- a/docs/source/legacy/reference/precision.md +++ /dev/null @@ -1,184 +0,0 @@ -(precision)= - -# Numerical Precision - -This document describes the different quantization recipes implemented in TensorRT-LLM and contains a support matrix -for the different models. - -## FP32, FP16 and BF16 - -The different models implemented in TensorRT-LLM work with 32-bit IEEE -floating-point (FP32) numbers. When checkpoints are available, the models also -support 16-bit IEEE floating-point numbers (FP16) and 16-bit Bfloat16 (BF16) as -described [here](https://en.wikipedia.org/wiki/Bfloat16_floating-point_format). - -## Quantization and Dequantization (Q/DQ) - -Given a floating-point number `x` and a floating-point scaling factor `s`, -TensorRT-LLM implements INT8 quantization as: - -``` -q = int8.satfinite(x * s) -``` - -Given an INT8 number `q` and a floating-point scaling factor `s`, TensorRT-LLM -implements INT8 dequantization to the floating-point (FP) type as: - -``` -x = static_cast(q) * s -``` - -Given a matrix (2D tensor) of shape `M x N` (`M` rows and `N` columns) where -`M` is the number of tokens and `N` is the number of channels. TensorRT-LLM has -the three following modes to quantize and dequantize the elements of the -tensor: - - * Per-tensor: It uses a single scaling factor for all the elements, - * Per-token: It uses a different scaling factor for each token. There are `M` - scaling factors in that case, - * Per-channel: It uses a different scaling factor for each channel. There are - `N` scaling factors in that case. - -Note that per-token and per-channel scaling modes can be used together (i.e. -they are _not_ mutually exclusive). - -In pseudo-code, the quantization can be implemented as follows for the three -different modes: - -```python -# Per-tensor scaling. -for mi in range(M): - for ni in range(N): - q[mi][ni] = int8.satfinite(x[mi][ni] * s) - -# Per-token scaling. -for mi in range(M): - for ni in range(N): - q[mi][ni] = int8.satfinite(x[mi][ni] * s[mi]) - -# Per-channel scaling. -for mi in range(M): - for ni in range(N): - q[mi][ni] = int8.satfinite(x[mi][ni] * s[ni]) -``` - -## INT8 SmoothQuant (W8A8) - -The SmoothQuant technique was introduced in -[https://arxiv.org/abs/2211.10438](https://arxiv.org/abs/2211.10438). It is a -method to run inference using INT8 for both activations and weights while -maintaining the accuracy of the network (on downstream tasks). - -As explained in the research paper, preprocessing must be applied to the -weights of the model. TensorRT-LLM includes scripts to prepare the model to -run using the SmoothQuant method. - -Examples of how to enable SmoothQuant for GPT, GPT-J and LLaMA can be found in -the [examples/quantization](source:examples/quantization) folder of that release. - -## INT4 and INT8 Weight-Only (W4A16 and W8A16) - -The INT4 and INT8 Weight-Only techniques consist in quantizing the weights of -a model and dequantizing those weights on-the-fly in linear layers (Matmuls). -The activations are encoded using floating-point values (FP16 or BF16). - -To use INT4/INT8 Weight-Only methods, the user must determine the scaling -factors to use to quantize and dequantize the weights of the model. - -This release includes examples for [GPT](source:examples/models/core/gpt) and -[LLaMA](source:examples/models/core/llama). - -## GPTQ and AWQ (W4A16) - -The GPTQ and AWQ techniques are presented in -[https://arxiv.org/abs/2210.17323](https://arxiv.org/abs/2210.17323) -and -[https://arxiv.org/abs/2306.00978](https://arxiv.org/abs/2306.00978), -respectively. TensorRT-LLM supports per-group scaling factors and -zero-offsetting in linear layers to implement GPTQ and AWQ methods. See the -[WeightOnlyGroupwiseQuantMatmulPlugin](source:cpp/tensorrt_llm/plugins/weightOnlyGroupwiseQuantMatmulPlugin) -plugin and the corresponding -[`weight_only_groupwise_quant_matmul`](source:tensorrt_llm/quantization/functional.py) -Python function, for details. - -This release includes examples of applying GPTQ to [GPT-NeoX](source:examples/models/core/gpt) -and [LLaMA-v2](source:examples/models/core/llama), as well as an example of using AWQ with -[GPT-J](source:examples/models/contrib/gptj). - -## FP8 (Hopper) - -This release of TensorRT-LLM contains implementations of FP8 for GPT-NeMo, -GPT-J and LLaMA. Those examples can be found in -[examples/quantization](source:examples/quantization). - -## NVFP4 (Blackwell) - -LLama and Mixtral can run in NVFP4 datatype. Those examples can be found in Llama examples. - -## Support matrix - -This release of TensorRT-LLM contains the following examples: - -| Model | FP32 | FP16 | BF16 | FP8 | NVFP4 | W8A8 SQ | W8A16 | W4A16 | W4A16 AWQ | W4A16 GPTQ | -| :------------- | :---: | :---: | :---: | :---: | :---: | :-----: | :---: | :---: | :-------: | :--------: | -| Baichuan | Y | Y | Y | Y | . | Y | Y | Y | Y | Y | -| BERT | Y | Y | Y | . | . | . | . | . | . | . | -| BLIP-2 | Y | Y | Y | . | . | . | . | . | . | . | -| BLOOM | Y | Y | Y | Y | . | Y | Y | Y | . | . | -| ChatGLM | Y | Y | Y | . | . | . | . | . | . | . | -| ChatGLM-v2 | Y | Y | Y | . | . | . | . | . | . | . | -| ChatGLM-v3 | Y | Y | Y | . | . | . | . | . | . | . | -| DBRX | Y | Y | Y | . | . | . | Y | Y | . | . | -| Falcon | Y | Y | Y | Y | . | . | Y | Y | Y | . | -| Flan-T5 | Y | Y | Y | . | . | . | . | . | . | . | -| Gemma | Y | Y | Y | Y | . | Y | Y | Y | Y | . | -| GPT | Y | Y | Y | Y | . | Y | Y | Y | . | . | -| GPT-J | Y | Y | Y | Y | . | Y | Y | Y | Y | . | -| GPT-NeMo | Y | Y | Y | . | . | . | . | . | . | . | -| GPT-NeoX | Y | Y | Y | . | . | . | . | . | . | Y | -| InternLM | Y | Y | Y | . | . | Y | Y | Y | . | . | -| InternLM2 | Y | Y | Y | . | . | . | . | . | . | . | -| LLaMA | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | -| LLaMA-v2 | Y | Y | Y | Y | Y | Y | Y | Y | Y | Y | -| Mamba | Y | Y | Y | . | . | . | . | . | . | . | -| Mistral | Y | Y | Y | Y | . | Y | Y | Y | Y | . | -| Mixtral | Y | Y | Y | Y | Y | . | Y | Y | . | . | -| MPT | Y | Y | Y | Y | . | Y | Y | Y | Y | . | -| OPT | Y | Y | Y | . | . | . | . | . | . | . | -| Phi | Y | Y | Y | . | . | . | . | . | . | . | -| Qwen | Y | Y | Y | . | . | Y | Y | Y | Y | Y | -| RecurrentGemma | Y | Y | Y | Y | . | Y | . | . | Y | . | -| Replit Code | Y | Y | Y | . | . | . | . | . | . | . | -| SantaCoder | Y | Y | Y | . | . | . | Y | Y | . | . | -| Skywork | Y | Y | Y | . | . | . | . | . | . | . | -| StarCoder1 | Y | Y | Y | . | . | . | Y | Y | . | . | -| StarCoder2 | Y | Y | Y | Y | . | . | Y | Y | . | . | -| T5 | Y | Y | Y | . | . | . | . | . | . | . | -| Whisper | Y | Y | Y | . | . | . | Y | Y | . | . | -| BLIP2-OPT | Y | Y | Y | . | . | . | . | . | . | . | -| BLIP2-T5 | Y | Y | Y | . | . | . | . | . | . | . | -| LLaVA | Y | Y | Y | Y | . | Y | Y | Y | Y | Y | -| VILA | Y | Y | Y | Y | . | Y | Y | Y | Y | Y | -| Nougat | Y | Y | Y | . | . | . | . | . | . | . | - -Note: The vision component of multi-modal models(BLIP2-OPT/BLIP2-T5/LLaVA/VILA/Nougat) uses FP16 by default. -The language component decides which quantization methods are supported by a given multi-modal model. - -## Technical Detail: The `QuantMode` Flags - -The quantization method is controlled by the -[`QuantMode`](source:tensorrt_llm/quantization/mode.py) flags. The different fields -are: - - * `INT4_WEIGHTS`, the weights are quantized to 4 bits (W4A\*), - * `INT8_WEIGHTS`, the weights are quantized to 8 bits (W8A\*), - * `ACTIVATIONS`, the activations are quantized to 8 bits (W\*A8), - * `PER_CHANNEL`, the scaling factors are defined per channel, - * `PER_TOKEN`, the scaling factors are defined per token, - * `PER_GROUP`, the scaling factors are defined per group. - -There are three additional flags to control TensorRT-LLM: - - * `INT8_KV_CACHE`, the K/V cache stores K and V using 8-bit integers, - * `FP8_KV_CACHE`, the K/V cache stores K and V using 8-bit floating-point numbers, - * `FP8_QDQ`, TensorRT-LLM relies on automatic fusion of Q/DQ nodes in TensorRT. diff --git a/docs/source/legacy/reference/support-matrix.md b/docs/source/legacy/reference/support-matrix.md deleted file mode 100644 index 4c7629843299..000000000000 --- a/docs/source/legacy/reference/support-matrix.md +++ /dev/null @@ -1,185 +0,0 @@ -# Support Matrix - -```{deprecated} -This page is outdated and no longer maintained. Please refer to the up-to-date [Supported Models](../../models/supported-models.md) page instead. -``` - -TensorRT-LLM optimizes the performance of a range of well-known models on NVIDIA GPUs. The following sections provide a list of supported GPU architectures as well as important features implemented in TensorRT-LLM. - -## Models (PyTorch Backend) - -| Architecture | Model | HuggingFace Example | Modality | -|--------------|-------|---------------------|----------| -| `BertForSequenceClassification` | BERT-based | `textattack/bert-base-uncased-yelp-polarity` | L | -| `DeciLMForCausalLM` | Nemotron | `nvidia/Llama-3_1-Nemotron-51B-Instruct` | L | -| `DeepseekV3ForCausalLM` | DeepSeek-V3 | `deepseek-ai/DeepSeek-V3 `| L | -| `Exaone4ForCausalLM` | EXAONE 4.0 | `LGAI-EXAONE/EXAONE-4.0-32B` | L | -| `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it` | L | -| `Gemma3ForConditionalGeneration` | Gemma 3 | `google/gemma-3-27b-it` | L + I | -| `HCXVisionForCausalLM` | HyperCLOVAX-SEED-Vision | `naver-hyperclovax/HyperCLOVAX-SEED-Vision-Instruct-3B` | L + I | -| `LlavaLlamaModel` | VILA | `Efficient-Large-Model/NVILA-8B` | L + I + V | -| `LlavaNextForConditionalGeneration` | LLaVA-NeXT | `llava-hf/llava-v1.6-mistral-7b-hf` | L + I | -| `LlamaForCausalLM` | Llama 3.1, Llama 3, Llama 2, LLaMA | `meta-llama/Meta-Llama-3.1-70B` | L | -| `Llama4ForConditionalGeneration` | Llama 4 | `meta-llama/Llama-4-Scout-17B-16E-Instruct` | L + I | -| `MistralForCausalLM` | Bielik | `speakleash/Bielik-11B-v2.2-Instruct` | L | -| `MistralForCausalLM` | Mistral | `mistralai/Mistral-7B-v0.1` | L | -| `Mistral3ForConditionalGeneration` | Mistral3 | `mistralai/Mistral-Small-3.1-24B-Instruct-2503` | L + I | -| `MixtralForCausalLM` | Mixtral | `mistralai/Mixtral-8x7B-v0.1` | L | -| `MllamaForConditionalGeneration` | Llama 3.2 | `meta-llama/Llama-3.2-11B-Vision` | L | -| `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base` | L | -| `NemotronNASForCausalLM` | NemotronNAS | `nvidia/Llama-3_3-Nemotron-Super-49B-v1` | L | -| `Phi3ForCausalLM` | Phi-4 | `microsoft/Phi-4` | L | -| `Phi4MMForCausalLM` | Phi-4-multimodal | `microsoft/Phi-4-multimodal-instruct` | L + I + A | -| `Qwen2ForCausalLM` | QwQ, Qwen2 | `Qwen/Qwen2-7B-Instruct` | L | -| `Qwen2ForProcessRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-PRM-7B` | L | -| `Qwen2ForRewardModel` | Qwen2-based | `Qwen/Qwen2.5-Math-RM-72B` | L | -| `Qwen2VLForConditionalGeneration` | Qwen2-VL | `Qwen/Qwen2-VL-7B-Instruct` | L + I + V | -| `Qwen2_5_VLForConditionalGeneration` | Qwen2.5-VL | `Qwen/Qwen2.5-VL-7B-Instruct` | L + I + V | -| `Qwen3ForCausalLM` | Qwen3 | `Qwen/Qwen3-8B` | L | -| `Qwen3MoeForCausalLM` | Qwen3MoE | `Qwen/Qwen3-30B-A3B` | L | - -Note: -- L: Language -- I: Image -- V: Video -- A: Audio - -## Models (TensorRT Backend) - -### LLM Models - -- [Arctic](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/arctic) -- [Baichuan/Baichuan2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/baichuan) -- [BART](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [BERT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/bert) -- [BLOOM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/bloom) -- [ByT5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [ChatGLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/chatglm-6b) -- [ChatGLM2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/chatglm2-6b) -- [ChatGLM3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/chatglm3-6b-32k) -- [Code LLaMA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [DBRX](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/dbrx) -- [Exaone](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/exaone) -- [FairSeq NMT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [Falcon](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/falcon) -- [Flan-T5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) [^encdec] -- [Gemma/Gemma2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gemma) -- [GLM-4](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/glm-4-9b) -- [GPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [GPT-J](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/gptj) -- [GPT-Nemo](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [GPT-NeoX](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/gptneox) -- [Granite-3.0](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/granite) -- [Grok-1](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/grok) -- [InternLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples//models/contrib/internlm) -- [InternLM2](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/internlm2) -- [LLaMA/LLaMA 2/LLaMA 3/LLaMA 3.1](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [Mamba](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/mamba) -- [mBART](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [Minitron](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/nemotron) -- [Mistral](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [Mistral NeMo](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/llama) -- [Mixtral](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/mixtral) -- [MPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/mpt) -- [Nemotron](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/nemotron) -- [mT5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [OPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/opt) -- [Phi-1.5/Phi-2/Phi-3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/phi) -- [Qwen/Qwen1.5/Qwen2/Qwen3](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/qwen) -- [Qwen-VL](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/qwenvl) -- [RecurrentGemma](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/recurrentgemma) -- [Replit Code](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/mpt) [^replitcode] -- [RoBERTa](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/bert) -- [SantaCoder](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [Skywork](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/skywork) -- [Smaug](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/contrib/smaug) -- [StarCoder](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/gpt) -- [T5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/enc_dec) -- [Whisper](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/whisper) - - -### Multi-Modal Models [^multimod] - -- [BLIP2 w/ OPT](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [BLIP2 w/ T5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [CogVLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) [^bf16only] -- [Deplot](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Fuyu](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Kosmos](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLaVA-v1.5](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLaVa-Next](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLaVa-OneVision](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [NeVA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Nougat](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Phi-3-vision](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [Video NeVA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [VILA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [MLLaMA](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) -- [LLama 3.2 VLM](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/models/core/multimodal) - - -(support-matrix-hardware)= -## Hardware - -The following table shows the supported hardware for TensorRT-LLM. - -If a GPU architecture is not listed, the TensorRT-LLM team does not develop or test the software on the architecture and support is limited to community support. -In addition, older architectures can have limitations for newer software releases. - -```{list-table} -:header-rows: 1 -:widths: 20 80 - -* - - - Hardware Compatibility -* - Operating System - - TensorRT-LLM requires Linux x86_64 or Linux aarch64. -* - GPU Model Architectures - - - - [NVIDIA GB300 NVL72](https://www.nvidia.com/en-us/data-center/gb300-nvl72/) - - [NVIDIA GB200 NVL72](https://www.nvidia.com/en-us/data-center/gb200-nvl72/) - - [NVIDIA GB300 NVL72](https://www.nvidia.com/en-us/data-center/gb300-nvl72/) - - [NVIDIA Blackwell Architecture](https://www.nvidia.com/en-us/data-center/technologies/blackwell-architecture/) - - [NVIDIA Grace Hopper Superchip](https://www.nvidia.com/en-us/data-center/grace-hopper-superchip/) - - [NVIDIA Hopper Architecture](https://www.nvidia.com/en-us/data-center/technologies/hopper-architecture/) - - [NVIDIA Ada Lovelace Architecture](https://www.nvidia.com/en-us/technologies/ada-architecture/) - - [NVIDIA Ampere Architecture](https://www.nvidia.com/en-us/data-center/ampere-architecture/) -``` - -(support-matrix-software)= -## Software - -The following table shows the supported software for TensorRT-LLM. - -```{list-table} -:header-rows: 1 -:widths: 20 80 - -* - - - Software Compatibility -* - Container - - [26.02](https://docs.nvidia.com/deeplearning/frameworks/support-matrix/index.html) -* - TensorRT - - [10.14](https://docs.nvidia.com/deeplearning/tensorrt/release-notes/index.html) -* - Precision - - - - Blackwell (SM100/SM103/SM120) - FP32, FP16, BF16, FP8, FP4, INT8, INT4 - - Hopper (SM90) - FP32, FP16, BF16, FP8, INT8, INT4 - - Ada Lovelace (SM89) - FP32, FP16, BF16, FP8, INT8, INT4 - - Ampere (SM80, SM86) - FP32, FP16, BF16, INT8, INT4[^smgte89] -``` - -[^replitcode]: Replit Code is not supported with the transformers 4.45+. - -[^smgte89]: INT4 AWQ and GPTQ with FP8 activations require SM >= 89. - -[^encdec]: Encoder-Decoder provides general encoder-decoder functionality that supports many encoder-decoder models such as T5 family, BART family, Whisper family, NMT family, and so on. - -[^multimod]: Multi-modal provides general multi-modal functionality that supports many multi-modal architectures such as BLIP2 family, LLaVA family, and so on. - -[^bf16only]: Only supports bfloat16 precision. - - -```{note} -Support for FP8 and quantized data types (INT8 or INT4) is not implemented for all the models. Refer to {ref}`precision` and [examples](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples) folder for additional information. -``` diff --git a/docs/source/legacy/reference/troubleshooting.md b/docs/source/legacy/reference/troubleshooting.md deleted file mode 100644 index 1b29849bbe64..000000000000 --- a/docs/source/legacy/reference/troubleshooting.md +++ /dev/null @@ -1,328 +0,0 @@ -(troubleshooting)= - -# Troubleshooting - -This document describes some of the frequently asked questions and their solutions in TensorRT-LLM, including problems of installation, model-building, model-execution, and input / output size. - -## Installation Errors - -During compilation and installation of TensorRT-LLM, many build errors can be resolved by simply deleting the build tree and rebuilding again. - -In most occasions, these problems are caused by the workflow like: an old compilation -> some code change (update of the repo or users' writing) -> a later compilation. - -Solution: try running build script with `--clean`, or try running `rm -r build cpp/build` before running build script again. - -## Debug on Unit Tests - -Here is an example to print the values of the MLP output tensor in a unit test ([full example](../../../../tests/unittest/others/test_debugging_api.py)). - -1. Register the intermediate tensors as the network outputs with `register_network_output` API. - -```python -class MLP(Module): - - def __init__(self, ...): - super().__init__() - # Do not modify the definition in `__init__` method - self.fc = ... - self.proj = ... - - def forward(self, hidden_states): - inter = self.fc(hidden_states) - inter = tensorrt_llm.functional.relu(inter) - # Here register the tensor `inter` as our debug output tensor - self.register_network_output('inter', inter) - output = self.proj(inter) - return output -``` - -2. Mark the intermediate tensors as network outputs. - -```python -for k, v in gm.named_network_outputs(): - net._mark_output(v, k, dtype) -``` - -3. Print the tensors at runtime. - -```python -print(outputs.keys()) -print(outputs['inter']) -``` - -## Debug on E2E Models - -Here is an example to print the values of the MLP output tensor in the GPT model. - -1. Register the MLP output tensor in `tensorrt_llm/models/gpt/model.py`. - -```python - hidden_states = residual + attention_output.data - - residual = hidden_states - hidden_states = self.post_layernorm(hidden_states) - - hidden_states = self.mlp(hidden_states) - # Register as model output - # ------------------------------------------------------ - self.register_network_output('mlp_output', hidden_states) - # ------------------------------------------------------ - - hidden_states = residual + hidden_states -``` - -2. Build the TensorRT engine of the model. - -Enable the `--enable_debug_output` option when building engines with `trtllm-build` - -```bash -cd examples/models/core/gpt - -# Download hf gpt2 model -rm -rf gpt2 && git clone https://huggingface.co/gpt2-medium gpt2 -pushd gpt2 && rm pytorch_model.bin model.safetensors && wget -q https://huggingface.co/gpt2-medium/resolve/main/pytorch_model.bin && popd - -# Convert to TensorRT-LLM checkpoint -python3 convert_checkpoint.py \ - --model_dir gpt2 \ - --dtype float16 \ - --output_dir gpt2/trt_ckpt/fp16/1-gpu - -# Build TensorRT-LLM engines with --enable_debug_output -trtllm-build \ - --checkpoint_dir gpt2/trt_ckpt/fp16/1-gpu \ - --enable_debug_output \ - --output_dir gpt2/trt_engines/fp16/1-gpu -``` - -3. Print the intermediate output tensors. - -Add debug info in `tensorrt_llm/runtime/generation.py`. - -```python - stream = torch.cuda.current_stream().cuda_stream - instance_idx = step % 2 - if self.cuda_graph_mode and self.runtime.cuda_graph_instances[ - instance_idx] is not None: - # launch cuda graph - CUASSERT( - cudart.cudaGraphLaunch( - self.runtime.cuda_graph_instances[instance_idx], stream)) - ok = True - else: - ok = self.runtime._run(context, stream) - - if not ok: - raise RuntimeError(f"Executing TRT engine failed step={step}!") - if self.debug_mode: - torch.cuda.synchronize() - # ------------------------------------------- - if step == 0: - print(self.debug_buffer.keys()) - print(f"Step: {step}") - print(self.debug_buffer['transformer.layers.6.mlp_output']) - # ------------------------------------------- -``` - -4. Run `../run.py` with `--debug_mode` and `--use_py_session`. - -```bash -python3 ../run.py \ - --engine_dir gpt2/trt_engines/fp16/1-gpu \ - --tokenizer_dir gpt2 \ - --max_output_len 8 \ - --debug_mode \ - --use_py_session -``` - -5. See the value of the tensor. - -```txt -...... -dict_keys(['context_lengths', 'cache_indirection', 'position_ids', 'logits', 'last_token_ids', 'input_ids', 'kv_cache_block_pointers', 'host_kv_cache_block_pointers', 'sequence_length', 'host_past_key_value_lengths', 'host_sink_token_length', 'host_request_types', 'host_max_attention_window_sizes', 'host_context_lengths', 'transformer.layers.0.mlp_output', 'transformer.layers.1.mlp_output', 'transformer.layers.2.mlp_output', 'transformer.layers.3.mlp_output', 'transformer.layers.4.mlp_output', 'transformer.layers.5.mlp_output', 'transformer.layers.6.mlp_output', 'transformer.layers.7.mlp_output', 'transformer.layers.8.mlp_output', 'transformer.layers.9.mlp_output', 'transformer.layers.10.mlp_output', 'transformer.layers.11.mlp_output', 'transformer.layers.12.mlp_output', 'transformer.layers.13.mlp_output', 'transformer.layers.14.mlp_output', 'transformer.layers.15.mlp_output', 'transformer.layers.16.mlp_output', 'transformer.layers.17.mlp_output', 'transformer.layers.18.mlp_output', 'transformer.layers.19.mlp_output', 'transformer.layers.20.mlp_output', 'transformer.layers.21.mlp_output', 'transformer.layers.22.mlp_output', 'transformer.layers.23.mlp_output']) -Step: 0 -tensor([[ 0.0294, -0.0260, -0.0776, ..., -0.0560, -0.0235, 0.0273], - [-0.0071, 0.5879, 0.1993, ..., -1.0449, -0.6299, 0.5957], - [-0.8779, 0.1050, 0.7090, ..., 0.0910, 1.0713, -0.2939], - ..., - [ 0.1212, -0.0903, -0.5918, ..., -0.1045, -0.3445, 0.1082], - [-1.0723, -0.0732, 0.6157, ..., 0.3452, 0.2998, 0.2649], - [-0.7134, 0.9692, -0.1141, ..., -0.0096, 0.9521, 0.1437]], - device='cuda:0', dtype=torch.float16) -Step: 1 -tensor([[-0.2107, 0.5874, 0.8179, ..., 0.7900, -0.6890, 0.6064]], - device='cuda:0', dtype=torch.float16) -Step: 2 -tensor([[ 0.4192, -0.0047, 1.3887, ..., -0.9028, -0.0682, -0.2820]], - device='cuda:0', dtype=torch.float16) -Step: 3 -tensor([[-0.7949, -0.5073, -0.1721, ..., -0.5830, -0.1378, -0.0070]], - device='cuda:0', dtype=torch.float16) -Step: 4 -tensor([[-0.0804, 0.1272, -0.6255, ..., -0.1072, -0.0523, 0.7144]], - device='cuda:0', dtype=torch.float16) -Step: 5 -tensor([[-0.3328, -0.8828, 0.3442, ..., 0.8149, -0.0630, 1.2305]], - device='cuda:0', dtype=torch.float16) -Step: 6 -tensor([[-0.2225, -0.2079, -0.1459, ..., -0.3555, -0.1672, 0.1135]], - device='cuda:0', dtype=torch.float16) -Step: 7 -tensor([[ 0.1290, -0.1556, 0.3977, ..., -0.8218, -0.3291, -0.8672]], - device='cuda:0', dtype=torch.float16) -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef before moving to London in the early" -``` - -## Debug Execution Errors - -If problems come from plugins, try setting the environment variable `CUDA_LAUNCH_BLOCKING=1` to make kernels launch synchronously with their return status checked immediately. - -If problems come from runtime-shape of the input tensors, double-check the shape (rank and length of each rank) and location (CPU / GPU) of input tensors for the engine obey the build-time setting. - -For example, one possible reason of getting the error information like below is, we use mismatched configuration between engine building and running, including code change (update of repo or users' rewriting), too large or too small input shape, etc.. - -```txt -unexpected shape for input 'XXX' for model 'YYY'. Expected [-1,-1,-1], got [8,16]. NOTE: Setting a non-zero max_batch_size in the model config requires a batch dimension to be prepended to each input shape. If you want to specify the full shape including the batch dim in your input dims config, try setting max_batch_size to zero. See the model configuration docs for more info on max_batch_size. - -[TensorRT-LLM][ERROR] Assertion failed: Tensor 'input_ids' has invalid shape (8192), expected (-1) (/code/tensorrt_llm/cpp/tensorrt_llm/runtime/tllmRuntime.cpp:149) - -RuntimeError: Sizes of tensors must match except in dimension 0. Expected size 8192 but got size 1024 for tensor number 1 in the list. -``` - -By setting environment variable `export TLLM_LOG_LEVEL=TRACE`, we can get more information about the TensorRT engine and context at runtime. - -Before the first forward computation, the shapes of all input / output tensors and their corresponding allowed ranges are provided in a table like: - -```txt -[TensorRT-LLM][TRACE] Information of engine input / output. -[TensorRT-LLM][TRACE] ===================================================================== -[TensorRT-LLM][TRACE] Name |I/O|Location|DataType| Shape | -[TensorRT-LLM][TRACE] --------------------------------------------------------------------- -[TensorRT-LLM][TRACE] input_ids | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] position_ids | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] last_token_ids | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] kv_cache_block_offsets | I | GPU | INT32 |(1, -1, 2, -1)| -[TensorRT-LLM][TRACE] host_kv_cache_block_offsets | I | GPU | INT32 |(1, -1, 2, -1)| -[TensorRT-LLM][TRACE] host_kv_cache_pool_pointers | I | GPU | INT64 | (1, 2) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_mapping | I | GPU | INT32 | (28) | -[TensorRT-LLM][TRACE] sequence_length | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_request_types | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_past_key_value_lengths | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] context_lengths | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_runtime_perf_knobs | I | GPU | INT64 | (16) | -[TensorRT-LLM][TRACE] host_context_lengths | I | GPU | INT32 | (-1) | -[TensorRT-LLM][TRACE] host_max_attention_window_sizes| I | GPU | INT32 | (28) | -[TensorRT-LLM][TRACE] host_sink_token_length | I | GPU | INT32 | (1) | -[TensorRT-LLM][TRACE] cache_indirection | I | GPU | INT32 | (-1, 1, -1) | -[TensorRT-LLM][TRACE] logits | O | GPU | FP32 | (-1, 65024) | -[TensorRT-LLM][TRACE] ===================================================================== -[TensorRT-LLM][TRACE] Information of optimization profile. -[TensorRT-LLM][TRACE] Optimization Profile 0: -[TensorRT-LLM][TRACE] ============================================================================= -[TensorRT-LLM][TRACE] Name | Min | Opt | Max | -[TensorRT-LLM][TRACE] ----------------------------------------------------------------------------- -[TensorRT-LLM][TRACE] input_ids | (1) | (8) | (8192) | -[TensorRT-LLM][TRACE] position_ids | (1) | (8) | (8192) | -[TensorRT-LLM][TRACE] last_token_ids | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] kv_cache_block_offsets | (1, 1, 2, 1) |(1, 4, 2, 16) |(1, 8, 2, 32) | -[TensorRT-LLM][TRACE] host_kv_cache_block_offsets | (1, 1, 2, 1) |(1, 4, 2, 16) |(1, 8, 2, 32) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_pointers | (1, 2) | (1, 2) | (1, 2) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_mapping | (28) | (28) | (28) | -[TensorRT-LLM][TRACE] sequence_length | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_request_types | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_past_key_value_lengths | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] context_lengths | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_runtime_perf_knobs | (16) | (16) | (16) | -[TensorRT-LLM][TRACE] host_context_lengths | (1) | (4) | (8) | -[TensorRT-LLM][TRACE] host_max_attention_window_sizes| (28) | (28) | (28) | -[TensorRT-LLM][TRACE] host_sink_token_length | (1) | (1) | (1) | -[TensorRT-LLM][TRACE] cache_indirection | (1, 1, 1) | (4, 1, 1024) | (8, 1, 2048) | -[TensorRT-LLM][TRACE] logits | (1, 65024) | (4, 65024) | (8, 65024) | -[TensorRT-LLM][TRACE] ============================================================================= -``` - -Before each forward computation, the real shapes of all input / output tensors for TRT engine are provided by a table like: - -```txt -[TensorRT-LLM][TRACE] Information of context input / output. -[TensorRT-LLM][TRACE] Using Optimization Profile: 0 -[TensorRT-LLM][TRACE] ================================================= -[TensorRT-LLM][TRACE] Name |I/O| Shape | -[TensorRT-LLM][TRACE] ------------------------------------------------- -[TensorRT-LLM][TRACE] input_ids | I | (33) | -[TensorRT-LLM][TRACE] position_ids | I | (33) | -[TensorRT-LLM][TRACE] last_token_ids | I | (3) | -[TensorRT-LLM][TRACE] kv_cache_block_offsets | I |(1, 3, 2, 4)| -[TensorRT-LLM][TRACE] host_kv_cache_block_offsets | I |(1, 3, 2, 4)| -[TensorRT-LLM][TRACE] host_kv_cache_pool_pointers | I | (1, 2) | -[TensorRT-LLM][TRACE] host_kv_cache_pool_mapping | I | (28) | -[TensorRT-LLM][TRACE] sequence_length | I | (3) | -[TensorRT-LLM][TRACE] host_request_types | I | (3) | -[TensorRT-LLM][TRACE] host_past_key_value_lengths | I | (3) | -[TensorRT-LLM][TRACE] context_lengths | I | (3) | -[TensorRT-LLM][TRACE] host_runtime_perf_knobs | I | (16) | -[TensorRT-LLM][TRACE] host_context_progress | I | (1) | -[TensorRT-LLM][TRACE] host_context_lengths | I | (3) | -[TensorRT-LLM][TRACE] host_max_attention_window_sizes| I | (28) | -[TensorRT-LLM][TRACE] host_sink_token_length | I | (1) | -[TensorRT-LLM][TRACE] cache_indirection | I |(3, 2, 256) | -[TensorRT-LLM][TRACE] logits | O | (3, 65024) | -[TensorRT-LLM][TRACE] ================================================= -``` - -## Tips - -* It's recommended to add options `–shm-size=1g –ulimit memlock=-1` to the - docker or nvidia-docker run command. Otherwise you may see NCCL errors when - running multiple GPU inferences. See - https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/troubleshooting.html#errors - for details. - -* When building models, memory-related issues such as -``` -[09/23/2023-03:13:00] [TRT] [E] 9: GPTLMHeadModel/layers/0/attention/qkv/PLUGIN_V2_Gemm_0: could not find any supported formats consistent with input/output data types -[09/23/2023-03:13:00] [TRT] [E] 9: [pluginV2Builder.cpp::reportPluginError::24] Error Code 9: Internal Error (GPTLMHeadModel/layers/0/attention/qkv/PLUGIN_V2_Gemm_0: could not find any supported formats consistent with input/output data types) -``` -may happen. One possible solution is to reduce the amount of memory needed by -reducing the maximum batch size, input and output lengths. Another option is to -enable plugins, for example: `--gpt_attention_plugin`. - -* MPI + Slurm - -TensorRT-LLM is a -[MPI](https://en.wikipedia.org/wiki/Message_Passing_Interface)-aware package -that uses [`mpi4py`](https://mpi4py.readthedocs.io/en/stable/). If you are -running scripts in a [Slurm](https://slurm.schedmd.com/) environment, you might -encounter interferences: -``` --------------------------------------------------------------------------- -PMI2_Init failed to initialize. Return code: 14 --------------------------------------------------------------------------- --------------------------------------------------------------------------- -The application appears to have been direct launched using "srun", -but OMPI was not built with SLURM's PMI support and therefore cannot -execute. There are several options for building PMI support under -SLURM, depending upon the SLURM version you are using: - - version 16.05 or later: you can use SLURM's PMIx support. This - requires that you configure and build SLURM --with-pmix. - - Versions earlier than 16.05: you must use either SLURM's PMI-1 or - PMI-2 support. SLURM builds PMI-1 by default, or you can manually - install PMI-2. You must then build Open MPI using --with-pmi pointing - to the SLURM PMI library location. - -Please configure as appropriate and try again. --------------------------------------------------------------------------- -``` - -You may experience other problems like hanging on the program startup. - -As a rule of thumb, if you are running TensorRT-LLM interactively on a Slurm -node, prefix your commands with `mpirun -n 1` to run TensorRT-LLM in a -dedicated MPI environment, not the one provided by your Slurm allocation. - -For example: `mpirun -n 1 python3 examples/models/core/gpt/build.py ...` - -It's critical that it's always `-n 1` regardless of how many GPUs are being used. If you'd use `-n 2` for a 2 GPU program it will not work. `mpirun` here isn't being used to orchestrate multiple processes, but to invoke the right environment on SLURM. The internal MPI implementation deals with spawning the additional processes. diff --git a/docs/source/legacy/tensorrt-backend-removal.md b/docs/source/legacy/tensorrt-backend-removal.md new file mode 100644 index 000000000000..25ec1aae009f --- /dev/null +++ b/docs/source/legacy/tensorrt-backend-removal.md @@ -0,0 +1,93 @@ +# Migration Guide: TensorRT Backend Removed + +```{note} +**Breaking change.** The TensorRT engine backend has been **removed** in this +release. PyTorch is now the sole execution backend for TensorRT LLM. AutoDeploy +(built on the PyTorch backend) remains available. +``` + +This guide explains what was removed, the errors you will encounter if you rely +on the old TensorRT path, and how to migrate to the PyTorch backend. + +## What Was Removed + +The following public APIs, command-line tools, and packaging behaviors are no +longer available: + +- **`LLM(backend="tensorrt")`** — passing `backend="tensorrt"` to the `LLM` + constructor now raises a `ValueError`. PyTorch is the only supported backend. +- **`TrtLlmArgs`** — the TensorRT-specific arguments class has been removed. + Use `TorchLlmArgs` (the default) instead. +- **`tensorrt_llm._tensorrt_engine.LLM`** — the TensorRT engine `LLM` entry + point has been removed. +- **CLI build/engine tooling** — `trtllm-build`, `trtllm-refit`, and + `trtllm-prune` have been removed. The PyTorch backend loads HuggingFace + checkpoints directly and does not require a separate engine-build step. +- **`--backend tensorrt`** — the `tensorrt` choice has been removed from CLI + tools (for example `trtllm-serve`, `trtllm-bench`, `trtllm-eval`). Omit the + flag (PyTorch is the default) or pass `--backend pytorch`. +- **`tensorrt` pip dependency** — the `tensorrt` Python package is no longer a + dependency of TensorRT LLM and is not installed. +- **Checkpoint conversion scripts** — the per-model `convert_checkpoint.py` + scripts and the associated TensorRT example directories have been removed. + +## Breaking-Change Contract + +| Removed surface | Previous behavior | New behavior | +| --- | --- | --- | +| `LLM(backend="tensorrt")` | Built and ran a TensorRT engine | Raises `ValueError` | +| `TrtLlmArgs` | TensorRT argument schema | Removed; use `TorchLlmArgs` | +| `tensorrt_llm._tensorrt_engine.LLM` | TensorRT engine entry point | Removed | +| `trtllm-build` / `trtllm-refit` / `trtllm-prune` | Engine build / refit / prune CLIs | Removed | +| `--backend tensorrt` | Selected the TensorRT backend | Removed choice; use `pytorch` | +| `tensorrt` pip dependency | Installed with TensorRT LLM | Dropped | +| `convert_checkpoint.py` | Converted HF weights to TRT checkpoints | Removed | + +## How to Migrate + +### Python API + +Before: + +```python +from tensorrt_llm import LLM + +llm = LLM(model="", backend="tensorrt") +``` + +After (PyTorch is the default backend — no `backend` argument needed): + +```python +from tensorrt_llm import LLM + +llm = LLM(model="") +``` + +The PyTorch backend loads HuggingFace checkpoints directly. There is no separate +checkpoint-conversion or engine-build step, so `convert_checkpoint.py` and +`trtllm-build` are no longer part of the workflow. + +### Command Line + +Before: + +```bash +trtllm-build --checkpoint_dir --output_dir +trtllm-serve --backend tensorrt +``` + +After: + +```bash +trtllm-serve +``` + +PyTorch is the default backend, so `--backend` may be omitted. To be explicit, +pass `--backend pytorch`. + +## Where to Go Next + +- [Quick Start Guide](../quick-start-guide.md) +- [LLM API Reference](../llm-api/index.md) +- [Supported Models](../models/supported-models.md) +- [PyTorch Backend Architecture](../torch/arch_overview.md) diff --git a/docs/source/legacy/tensorrt_quickstart.md b/docs/source/legacy/tensorrt_quickstart.md deleted file mode 100644 index e74a0f5e9e28..000000000000 --- a/docs/source/legacy/tensorrt_quickstart.md +++ /dev/null @@ -1,9 +0,0 @@ -# LLM API with TensorRT Engine -A simple inference example with TinyLlama using the LLM API: - -```{literalinclude} ../../../examples/llm-api/_tensorrt_engine/quickstart_example.py - :language: python - :linenos: -``` - -For more advanced usage including distributed inference, multimodal, and speculative decoding, please refer to this [README](../../../examples/llm-api/README.md). diff --git a/docs/source/legacy/torch.md b/docs/source/legacy/torch.md deleted file mode 100644 index e3fcf1dd8835..000000000000 --- a/docs/source/legacy/torch.md +++ /dev/null @@ -1,44 +0,0 @@ -# PyTorch Backend - -```{note} -Note: -This feature is currently in beta, and the related API is subject to change in future versions. -``` -To enhance the usability of the system and improve developer efficiency, TensorRT LLM launches a new backend based on PyTorch. - -The PyTorch backend of TensorRT LLM is available in version 0.17 and later. You can try it via importing `tensorrt_llm._torch`. - -## Quick Start - -Here is a simple example to show how to use `tensorrt_llm.LLM` API with Llama model. - -```{literalinclude} ../../examples/llm-api/quickstart_example.py - :language: python - :linenos: -``` - -## Features - -- [Sampling](../torch/features/sampling.md) -- [Quantization](../torch/features/quantization.md) -- [Overlap Scheduler](../torch/features/overlap_scheduler.md) -- [Feature Combination Matrix](../features/feature-combination-matrix.md) - -## Developer Guide - -- [Architecture Overview](../torch/arch_overview.md) -- [Adding a New Model](../torch/adding_new_model.md) - -## Key Components - -- [Attention](../torch/attention.md) -- [KV Cache Manager](../torch/kv_cache_manager.md) -- [Scheduler](../torch/scheduler.md) - -## Known Issues - -- The PyTorch backend on SBSA is incompatible with bare metal environments like Ubuntu 24.04. Please use the [PyTorch NGC Container](https://catalog.ngc.nvidia.com/orgs/nvidia/containers/pytorch) for optimal support on SBSA platforms. - -## Prototype Features - -- [AutoDeploy: Seamless Model Deployment from PyTorch to TensorRT LLM](../features/auto_deploy/auto-deploy.md) diff --git a/docs/source/release-notes.md b/docs/source/release-notes.md index 9f0ed17d4f76..7b38dc720118 100644 --- a/docs/source/release-notes.md +++ b/docs/source/release-notes.md @@ -26,6 +26,7 @@ All published functionality in the Release Notes has been fully tested and verif ### API Changes +- **[BREAKING CHANGE] TensorRT backend removed.** PyTorch is now the sole execution backend. `LLM(backend="tensorrt")` now raises a `ValueError`; `TrtLlmArgs`, `tensorrt_llm._tensorrt_engine.LLM`, the `trtllm-build` / `trtllm-refit` / `trtllm-prune` CLIs, the `--backend tensorrt` CLI choice, and the per-model `convert_checkpoint.py` scripts have all been removed. The `tensorrt` pip dependency is no longer installed. See the [TensorRT Backend Removed migration guide](legacy/tensorrt-backend-removal.md) for details. - `trtllm-serve`, `trtllm-eval`, `trtllm-bench`: explicit CLI flags now take precedence over values in `--config` / `--extra_llm_api_options` YAML files (was: YAML overrode CLI). Un-set CLI flags continue to fall back to the YAML, then to model-specific and built-in defaults. ### Fixed Issues diff --git a/docs/source/torch/arch_overview.md b/docs/source/torch/arch_overview.md index d67a4668e259..5856dc92ee58 100644 --- a/docs/source/torch/arch_overview.md +++ b/docs/source/torch/arch_overview.md @@ -17,7 +17,7 @@ The `LLM` also manages the tokenization and detokenization processes of the inpu ## PyExecutor -Similar to the TensorRT backend, which uses [Executor API](../legacy/advanced/executor.md), the PyTorch backend employs a `PyExecutor` class. +The PyTorch backend employs a `PyExecutor` class. This class has a similar interface to Executor, allowing it to be integrated into LLM as an alternative backend. Key components of the `PyExecutor` include: diff --git a/docs/source/torch/kv_cache_manager.md b/docs/source/torch/kv_cache_manager.md index 57d02c8ab59f..25196a1b3164 100644 --- a/docs/source/torch/kv_cache_manager.md +++ b/docs/source/torch/kv_cache_manager.md @@ -4,7 +4,7 @@ In Transformer-based models, the KV (Key-Value) Cache is a mechanism used to opt Since KV Cache requires memory to store, it is also an important resource. In TensorRT LLM, KV Cache is managed by the `KVCacheManager`. -For details of the TensorRT LLM `KVCacheManager` implementation see [KV Cache Management](../legacy/advanced/kv-cache-management.md). +For details on KV cache behavior in TensorRT LLM, see [KV Cache Management](../features/kvcache.md). ## KV Cache Manager Introduction diff --git a/examples/draft_target_model/README.md b/examples/draft_target_model/README.md deleted file mode 100644 index 1703b3def7ec..000000000000 --- a/examples/draft_target_model/README.md +++ /dev/null @@ -1,364 +0,0 @@ -# Draft-Target-Model Speculative Decoding (DTM) - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using DTM speculative decoding (also known as `Speculative-Sampling`, [`Paper`](https://arxiv.org/abs/2302.01318)) in TensorRT LLM on single GPU, or single node multiple GPU. - -## Overview - -We provide two styles of running DTM now: using TensorRT-LLM-BLS in Triton Inference Server, or using TensorRT LLM directly. Here we introduce the detailed steps of running DTM in both workflows. - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16 / BF16 / FP8 (both draft and target model) - * Paged KV Cache - * Tensor Parallel - -## Usage - -### Build draft and target engines (the same in two workflows) - -+ We use open-source `llama-7B/13B` as draft and target models in this example, assuming the paths to the models' repository are `DRAFT_MODEL_PATH` and `TARGET_MODEL_PATH`. -+ `--use_paged_context_fmha=enable` must be specified since we need KV-Cache reuse in this approach. -+ `--speculative_decoding_mode=draft_tokens_external` and `--max_draft_len` must be specified for target model. -+ `--gather_generation_logits` is necessary if using generation logits for selecting tokens in target model. -+ `--tp_size` can be modified set if using TP mode for draft / target model. - -```bash -cd examples/models/core/llama -export DRAFT_CKPT_PATH=/workspace/ckpt-draft -export TARGET_CKPT_PATH=/workspace/ckpt-target -export DRAFT_ENGINE_PATH=/workspace/engine-draft -export TARGET_ENGINE_PATH=/workspace/engine-target -export MAX_BATCH_SIZE=4 -export MAX_DRAFT_LEN=10 -export MAX_INPUT_LEN=3200 -export MAX_SEQ_LEN=4800 - -python3 convert_checkpoint.py \ - --model_dir=${DRAFT_MODEL_PATH} \ - --output_dir=${DRAFT_CKPT_PATH} \ - --dtype=float16 - -python3 convert_checkpoint.py \ - --model_dir=${TARGET_MODEL_PATH} \ - --output_dir=${TARGET_CKPT_PATH} \ - --dtype=float16 - -trtllm-build \ - --checkpoint_dir=${DRAFT_CKPT_PATH} \ - --output_dir=${DRAFT_ENGINE_PATH} \ - --gemm_plugin=float16 \ - --use_paged_context_fmha=enable \ - --max_batch_size=${MAX_BATCH_SIZE} \ - --max_input_len=${MAX_INPUT_LEN} \ - --max_seq_len=${MAX_SEQ_LEN} - -trtllm-build \ - --checkpoint_dir=${TARGET_CKPT_PATH} \ - --output_dir=${TARGET_ENGINE_PATH} \ - --gemm_plugin=float16 \ - --use_paged_context_fmha=enable \ - --speculative_decoding_mode=draft_tokens_external \ - --max_batch_size=${MAX_BATCH_SIZE} \ - --max_draft_len=${MAX_DRAFT_LEN} \ - --max_input_len=${MAX_INPUT_LEN} \ - --max_seq_len=${MAX_SEQ_LEN} -``` - -### TensorRT LLM workflow - -+ `--draft_engine_dir` and `--engine_dir` must be specified for the draft and target engines respectively. -+ `--draft_target_model_config` is corresponding configuration of DTM, which has 4 hyperparameters that you need to specify to control the process of generation: - - `draft_len`: the number of tokens the draft model generated in one iteration, which the range is from 4 to 10 in common usage. Empirically, the larger the value is, the higher acceptance ratio but higher overhead is expected at the same time, so the right balance based on the models and application scenarios needs to be found. - - `draft_model_device_list`: the index list of device(s) to run the draft model. The length of it must be the same as the TP size of the draft model engine. For instances, `draft_model_device_list=[1]` means using tp_size=1 and GPU 1 for draft model, `draft_model_device_list=[4,5,6,7]` means using tp=4 and GPU from 4 to 7 for draft model. - - `target_model_device_list`: the index list of device(s) to run the target model. The length of it must be the same as the TP size of the target model engine. For instances, `draft_model_device_list=[0]` means using tp_size=1 and GPU 0 for target model, `draft_model_device_list=[2,3]` means using tp=2 and GPU from 2 to 3 for target model. - - `use_logits`: there are two methods to accept tokens proposed by draft model. When `use_logits=True`, the draft tokens are accepted based on the ratio of the logits from draft and target model (modified rejection sampling method in the original paper); When `use_logits=False`, the draft tokens are accepted based on per-token comparison with target predictions regardless of the logits. - - As an example, `[4,[0],[1],False]` means `draft_len=4`, device of draft model is `GPU0`, device of target model is `GPU1`, and use tokens rather than logits to accept. -+ `--kv_cache_enable_block_reuse` must be specified for this approach. -+ Only CPP session is supported, so `--use_py_session` must not be specified. -+ `--kv_cache_free_gpu_memory_fraction` should be specified if we want to place two models on one GPU, or one of the models would use out of the GPU memory. -+ `--num_beams` can not be specified as larger than 1 since beam search is not supported in this approach yet. -+ `--output_generation_logits` is optional. In original paper, we accept the tokens by comparing logits of draft and target models, so this parameter is needed. But for simplification, we can accept the tokens by comparing the output token directly, in this occasion, we can skip this parameter. - -```bash -python3 examples/run.py \ - --tokenizer_dir=${TARGET_MODEL_PATH} \ - --draft_engine_dir=/workspace/engine-draft \ - --engine_dir=/workspace/engine-target \ - --draft_target_model_config="[4,[0],[1],False]" \ - --max_output_len=256 \ - --kv_cache_enable_block_reuse \ - --kv_cache_free_gpu_memory_fraction=0.5 \ - --input_text="How does Draft-Sampling work?" -``` - -### Triton Inference Server workflow - -+ This example is based on TensorRT-LLM-0.18.0 and TRTLLM-backend-0.18.0 with docker image `nvcr.io/nvidia/tritonserver:25.03-trtllm-python-py3`. -+ DTM model approach is supported since TensorRT-LLM-0.7.0 (using two separate Tritonserver to maintain draft and target model respectively), but has significant optimization in TensorRT-LLM-0.10.0 (using one Tritonserver with [Business Logic Scripting](https://github.com/triton-inference-server/python_backend?tab=readme-ov-file#business-logic-scripting), BLS). - -#### Get related repository inside the container - -```bash -git clone https://github.com/triton-inference-server/tensorrtllm_backend.git -cd tensorrtllm_backend -git checkout rel -git lfs pull -git submodule update --init --recursive -pip install -r requirements.txt -pip install SentencePiece tritonclient - -export DRAFT_MODEL_NAME="tensorrt_llm_draft" -export TARGET_MODEL_NAME="tensorrt_llm" -export TRITON_MODEL_REPO=llama_dtm -``` - -#### Simple deploy - -+ Edit model configuration. - -```bash -export DRAFT_DEVICE_IDS="0" -export TARGET_DEVICE_IDS="1" - -rm -rf ${TRITON_MODEL_REPO} -cp -r all_models/inflight_batcher_llm/ ${TRITON_MODEL_REPO} -cp -r ${TRITON_MODEL_REPO}/tensorrt_llm ${TRITON_MODEL_REPO}/tensorrt_llm_draft -sed -i 's/name: "tensorrt_llm"/name: "tensorrt_llm_draft"/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/ensemble/config.pbtxt triton_max_batch_size:4,logits_datatype:TYPE_FP32 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/preprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},preprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/postprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},postprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,bls_instance_count:1,accumulate_tokens:False,tensorrt_llm_model_name:${TARGET_MODEL_NAME},tensorrt_llm_draft_model_name:${DRAFT_MODEL_NAME} - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${TARGET_ENGINE_PATH},gpu_device_ids:${TARGET_DEVICE_IDS} -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${DRAFT_ENGINE_PATH},gpu_device_ids:${DRAFT_DEVICE_IDS} -``` - -+ Start the triton inference server. - + Verbose log will be written in to file `triton_log.txt` if specifying `--log`. - -```bash -python3 scripts/launch_triton_server.py \ - --model_repo=${TRITON_MODEL_REPO} \ - --multi-model \ - --log -``` - -+ You can see the output below in the file if Triton server launches successfully: - -```txt -Started HTTPService at 0.0.0.0:8000 -Started GRPCInferenceService at 0.0.0.0:8001 -Started Metrics Service at 0.0.0.0:8002 -``` - -+ Send a request for inference. - -```bash -python3 inflight_batcher_llm/client/e2e_grpc_speculative_decoding_client.py \ - --url-target=localhost:8001 \ - --draft-tensorrt-llm-model-name=${DRAFT_MODEL_NAME} \ - --target-tensorrt-llm-model-name=${TARGET_MODEL_NAME} \ - --output-len=100 \ - --num-draft-tokens=4 \ - --end-id=2 \ - --pad-id=2 \ - --prompt "What is Ubuntu operation system?" -``` - -+ You can receive the following results if everything goes smoothly. - -```txt -Final text: - What is Ubuntu operation system? -Ubuntu is a free and open source operating system that runs from the desktop, to the cloud, to all your internet connected things. Ubuntu is used by millions of people around the world who want to explore new ideas and discover new opportunities. -Ubuntu is a community developed operating system that is perfect for laptops, desktops, servers, and cloud. It is used by millions of people around the world who want to explore new ideas and discover new opportunities. -``` - -+ Test DTM with a script. - + Prepare a JSON file `input_data.json` containing input data as below (more requests are acceptable). - -```json -[ - { - "input": "What is Ubuntu operation system?", - "instruction": "Answer the question shortly.", - "output": " " - } -] -``` - -+ Use command below to launch test. - -```bash -### Use BLS speculative decoding -python3 tools/inflight_batcher_llm/speculative_decoding_test.py \ - --max-input-len 2500 \ - --dataset input_data.json \ - --url-control=localhost:8001 \ - --url-target=localhost:8001 \ - --url-draft=localhost:8001 \ - --draft-tensorrt-llm-model-name="${DRAFT_MODEL_NAME}" \ - --target-tensorrt-llm-model-name="${TARGET_MODEL_NAME}" \ - --bls-speculative-tensorrt-llm-model-name="tensorrt_llm_bls" \ - --execute-bls-speculative-decoding \ - --disable-output-comparison \ - --num-draft-tokens=4 \ - --use-draft-logits - -### Use client-side speculative decoding -python3 tools/inflight_batcher_llm/speculative_decoding_test.py \ - --max-input-len 2500 \ - --dataset input_data.json \ - --url-control=localhost:8001 \ - --url-target=localhost:8001 \ - --url-draft=localhost:8001 \ - --draft-tensorrt-llm-model-name="${DRAFT_MODEL_NAME}" \ - --target-tensorrt-llm-model-name="${TARGET_MODEL_NAME}" \ - --bls-speculative-tensorrt-llm-model-name="tensorrt_llm_bls" \ - --disable-output-comparison \ - --num-draft-tokens=4 \ - --use-draft-logits -``` - -+ You can receive the following results if everything goes smoothly. - -```txt -Ubuntu is a free and open source operating system. It is a Linux based operating system. ... -``` - -+ Stop triton inference server after all work is done. - -```bash -pkill tritonserver -``` - -+ In addition, it appears better performance can be achieved with both draft and target engines deployed on a single GPU (llama-7B-FP8 + llama-30B-FP8, for a total of 40GiB on one H100-80GiB GPU for example). - -#### Usage of Tensor-Parallelization mode. - -+ In this example, we use draft engine with TP=1 and target engine with TP=2 (both symmetrical or asymmetrical TP size are acceptable), and want to place the draft engine on GPU0, target engine on GPU1 and GPU2. -+ Edit model configuration. - -```bash -export DRAFT_DEVICE_IDS="0" -export TARGET_DEVICE_IDS="1,2" - -rm -rf ${TRITON_MODEL_REPO} -cp -r all_models/inflight_batcher_llm/ ${TRITON_MODEL_REPO} -cp -r ${TRITON_MODEL_REPO}/tensorrt_llm ${TRITON_MODEL_REPO}/tensorrt_llm_draft -sed -i 's/name: "tensorrt_llm"/name: "tensorrt_llm_draft"/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/ensemble/config.pbtxt triton_max_batch_size:4,logits_datatype:TYPE_FP32 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/preprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},preprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/postprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},postprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,bls_instance_count:1,accumulate_tokens:False,tensorrt_llm_model_name:${TARGET_MODEL_NAME},tensorrt_llm_draft_model_name:${DRAFT_MODEL_NAME} - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${TARGET_ENGINE_PATH} -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,logits_datatype:TYPE_FP32,triton_backend:tensorrtllm,max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,engine_dir:${DRAFT_ENGINE_PATH} - -sed -i 's/\${gpu_device_ids}/'"${DRAFT_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt -sed -i 's/\${gpu_device_ids}/'"${TARGET_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt -``` - -+ As you see, the only difference is `gpu_device_ids`, which needs fix manually since comma is not supported in script `python3 tools/fill_template.py`. - -+ Start the triton inference server - + Use `--multi-model` to enable orchestrator mode in TP>1 scenario. See [model config](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/model_config.md) for more information. - -```bash -python3 scripts/launch_triton_server.py \ - --model_repo=${TRITON_MODEL_REPO} \ - --tensorrt_llm_model_name "tensorrt_llm,tensorrt_llm_draft" \ - --multi-model -``` - -+ All other operations are the same as `Simple deploy` part. - -#### Usage of Fast logits D2D transfer - -+ Fast logits boosts the performance (TPS) by hiding the latency of logits transfer from draft engine to target engine supported since TensorRT-LLM-0.15.0. -+ In this example, we use draft engine with TP=1 and target engine with TP=2 (both symmetrical or asymmetrical TP size are acceptable), and want to place the draft engine on GPU0, target engine on GPU1 and GPU2. -+ For `participant_ids`, rank 0 is reserved for the orchestrator; rank (`1` ~ `tp_size_draft`) are for draft engine; rank (`tp_size_draft+1` ~ `tp_size_draft+tp_size_target`) are for target engine. -+ Edit model configuration. - -```bash -cd examples/models/core/llama -export TARGET_CKPT_PATH=/workspace/ckpt-target -export TARGET_ENGINE_PATH=/workspace/engine-target -export MAX_BATCH_SIZE=4 -export MAX_DRAFT_LEN=10 -export MAX_INPUT_LEN=3200 -export MAX_SEQ_LEN=4800 - -python3 convert_checkpoint.py \ - --model_dir=${TARGET_MODEL_PATH} \ - --output_dir=${TARGET_CKPT_PATH} \ - --dtype=float16 \ - --tp_size=2 - -trtllm-build \ - --checkpoint_dir=${TARGET_CKPT_PATH} \ - --output_dir=${TARGET_ENGINE_PATH} \ - --gemm_plugin=float16 \ - --use_paged_context_fmha=enable \ - --speculative_decoding_mode=draft_tokens_external \ - --max_batch_size=${MAX_BATCH_SIZE} \ - --max_draft_len=${MAX_DRAFT_LEN} \ - --max_input_len=${MAX_INPUT_LEN} \ - --max_seq_len=${MAX_SEQ_LEN} -``` - -```bash -export DRAFT_DEVICE_IDS="0" -export TARGET_DEVICE_IDS="1,2" -export DRAFT_PARTICIPANT_IDS="1" -export TARGET_PARTICIPANT_IDS="2,3" - -cd /work/tensorrtllm_backend -rm -rf ${TRITON_MODEL_REPO} -cp -r all_models/inflight_batcher_llm/ ${TRITON_MODEL_REPO} -cp -r ${TRITON_MODEL_REPO}/tensorrt_llm ${TRITON_MODEL_REPO}/tensorrt_llm_draft -sed -i 's/name: "tensorrt_llm"/name: "tensorrt_llm_draft"/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/ensemble/config.pbtxt triton_max_batch_size:4,logits_datatype:TYPE_FP32 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/preprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},preprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/postprocessing/config.pbtxt triton_max_batch_size:4,tokenizer_dir:${HF_MODEL},postprocessing_instance_count:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_bls/config.pbtxt triton_max_batch_size:4,decoupled_mode:False,bls_instance_count:1,accumulate_tokens:False,tensorrt_llm_model_name:${TARGET_MODEL_NAME},logits_datatype:TYPE_FP32,tensorrt_llm_draft_model_name:${DRAFT_MODEL_NAME} - -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt triton_max_batch_size:4,triton_backend:tensorrtllm,decoupled_mode:False,max_beam_width:1,engine_dir:${TARGET_ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,logits_datatype:TYPE_FP32,gpu_device_ids:${TARGET_DEVICE_IDS},participant_ids:2,3,speculative_decoding_fast_logits:1 -python3 tools/fill_template.py -i ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt triton_max_batch_size:4,triton_backend:tensorrtllm,decoupled_mode:False,max_beam_width:1,engine_dir:${DRAFT_ENGINE_PATH},max_tokens_in_paged_kv_cache:2560,max_attention_window_size:2560,kv_cache_free_gpu_mem_fraction:0.5,exclude_input_in_output:False,enable_kv_cache_reuse:True,batching_strategy:inflight_fused_batching,max_queue_delay_microseconds:0,encoder_input_features_data_type:TYPE_FP16,logits_datatype:TYPE_FP32,gpu_device_ids:${DRAFT_DEVICE_IDS},participant_ids:1,speculative_decoding_fast_logits:1 - -sed -i 's/\${gpu_device_ids}/'"${DRAFT_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt -sed -i 's/\${participant_ids}/'"${DRAFT_PARTICIPANT_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm_draft/config.pbtxt -sed -i 's/\${gpu_device_ids}/'"${TARGET_DEVICE_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt -sed -i 's/\${participant_ids}/'"${TARGET_PARTICIPANT_IDS}"'/g' ${TRITON_MODEL_REPO}/tensorrt_llm/config.pbtxt -``` - -+ As you see, the differences are `participant_ids` and `speculative_decoding_fast_logits`. - -+ Start the triton inference server. - + Use `--disable-spawn-process` to enable pre-spawn variant in orchestrator mode. - + `--world_size` must be equal to `1 + tp_size_draft + tp_size_target`, which is 4 in this example. - -```bash -python3 scripts/launch_triton_server.py \ - --model_repo ${TRITON_MODEL_REPO} \ - --tensorrt_llm_model_name tensorrt_llm,tensorrt_llm_draft \ - --multi-model \ - --world_size 4 \ - --disable-spawn-processes -``` - -+ All other operations are the same as the `Simple deploy` part. - -### Additional information - -+ With the fast logits enabled and following optimization tips in [model configuration](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/model_config.md#some-tips-for-model-configuration), speculative decoding with draft logits achieves 2.x throughput in BS1, 1.x throughput in BS16 comparing to auto-regressive decoding using Llama 3.2 1B draft and Llama 3.1 70B target. -+ Streaming mode or batched-request mode are not supported in DTM yet. diff --git a/examples/draft_target_model/requirements.txt b/examples/draft_target_model/requirements.txt deleted file mode 100644 index 423ba2d4b0e0..000000000000 --- a/examples/draft_target_model/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -sentencepiece>=0.1.99 -evaluate diff --git a/examples/eagle/README.md b/examples/eagle/README.md deleted file mode 100644 index ac95e21516cd..000000000000 --- a/examples/eagle/README.md +++ /dev/null @@ -1,195 +0,0 @@ -# EAGLE speculative Decoding - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using EAGLE decoding ([`GitHub`](https://github.com/SafeAILab/EAGLE/tree/main), [`BLOG`](https://sites.google.com/view/eagle-llm)) in TensorRT LLM on a single node with one or multiple GPUs. - -## Overview -Different from other models, EAGLE decoding needs a base model and an EAGLE model. - -The TensorRT LLM EAGLE decoding implementation can be found in [tensorrt_llm/models/eagle/model.py](../../tensorrt_llm/models/eagle/model.py). -The implementation adds an EAGLE drafter network to a base model. - -For more info about EAGLE, refer to [speculative decoding documentation](https://nvidia.github.io/TensorRT-LLM/features/speculative-decoding.html). - -## Limitations - * EAGLE-2 is not supported. - * All EAGLE choices have to have exactly the same depth as `num_eagle_layers` of the engine. - * Pipeline parallelism is not supported. - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16/BF16 - * Paged KV cache - * Inflight-fused-batching - * C++ runtime - * Tensor Parallel - -This example is based on the Vicuna-7b v1.3 model, a fine-tuned Llama. With some modifications, you can add EAGLE to other base models as well. Some TensorRT LLM models might not work with EAGLE due to the missing head size in the speculative decoding XQA attention kernels. - -## Usage -The TensorRT LLM EAGLE example code is located in [`examples/eagle`](./). There is one [`convert_checkpoint.py`](./convert_checkpoint.py) file to convert and build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run models with EAGLE decoding support. -In this example, we use the model from HuggingFace [`yuhuili/EAGLE-Vicuna-7B-v1.3`](https://huggingface.co/yuhuili/EAGLE-Vicuna-7B-v1.3), which is a LLAMA-based model. - -### Build TensorRT engine(s) -Get the weights by downloading the base model [`vicuna-7b-v1.3`](https://huggingface.co/lmsys/vicuna-7b-v1.3) and the EAGLE draft model [`EAGLE-Vicuna-7B-v1.3`](https://huggingface.co/yuhuili/EAGLE-Vicuna-7B-v1.3) from HF. - -``` -pip install -r requirements.txt - -git lfs install -git clone https://huggingface.co/lmsys/vicuna-7b-v1.3 -https://huggingface.co/yuhuili/EAGLE-Vicuna-7B-v1.3 -``` - -Here is the example: -```bash -# Convert and Build EAGLE decoding support for vicuna-7b-v1.3 -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --eagle_model_dir EAGLE-Vicuna-7B-v1.3 \ - --output_dir ./tllm_checkpoint_1gpu_eagle \ - --dtype float16 \ - --max_draft_len 63 \ - --num_eagle_layers 4 \ - --max_non_leaves_per_layer 10 - -# Note: Increasing the batch size may have a negative impact on performance -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_eagle \ - --output_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --use_paged_context_fmha enable \ - --speculative_decoding_mode eagle \ - --max_batch_size 4 - -# Convert and Build EAGLE decoding support for vicuna-7b-v1.3 with 4-way tensor parallelism. -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --eagle_model_dir EAGLE-Vicuna-7B-v1.3 \ - --output_dir ./tllm_checkpoint_4gpu_eagle \ - --dtype float16 \ - --max_draft_len 63 \ - --num_eagle_layers 4 \ - --max_non_leaves_per_layer 10 \ - --tp_size 4 \ - --workers 4 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_eagle \ - --output_dir ./tmp/eagle/7B/trt_engines/fp16/4-gpu/ \ - --gemm_plugin float16 \ - --use_paged_context_fmha enable \ - --speculative_decoding_mode eagle \ - --max_batch_size 4 -``` - -### Run - -To run a TensorRT LLM model with EAGLE-1 decoding support, you can use `../run.py` script, with an additional argument -`--eagle_choices`. -The `--eagle_choices` argument is of type `list[list[int]]`. If you do not specify any choices, the -default, [mc_sim_7b_63](https://github.com/FasterDecoding/Medusa/blob/main/medusa/model/medusa_choices.py#L1) choices -are used. For more information regarding choices tree, refer -to [Medusa Tree](https://nvidia.github.io/TensorRT-LLM/legacy/advanced/speculative-decoding.html#medusa-tree). - -The number of non-leaf nodes at each level can not exceed `max_non_leaves_per_layer` set -to `convert_checkpoint`. For example, in the tree below (`mc_sim_7b_63`) the minimum number of -`max_non_leaves_per_layer` is 10. There are exactly 10 non leaf nodes at depth 0, check `[0, 0], [1, 0], ..., [9, 0]`. - -The maximum depth, meaning the maximum length of inner `list[int]` specified in the `--eagle_choices` argument, should -be equal to `num_eagle_layers`. - -To run non-greedy sampling and use typical acceptance, set `--eagle_posterior_threshold` to `run.py`. `eagle_posterior_threshold` corresponds to epsilon in typical acceptance criteria from [Medusa paper](https://arxiv.org/pdf/2401.10774). -`--temperature` can be specified as well. When no `--eagle_posterior_threshold` is specified or `--temperature=0.0` is set, greedy sampling is used. - -#### Run EAGLE-2 - -EAGLE-2 can be enabled with 2 runtime flags (`--eagle_use_dynamic_tree` and `--eagle_dynamic_tree_max_top_k=N`). The same engine can be used for EAGLE-1 and EAGLE-2. Eagle choices must not be set in case of EAGLE-2. EAGLE-2 will generate the tree corresponding to choices dynamically in the runtime. For more details, please refer to [EAGLE-2 paper](https://arxiv.org/pdf/2406.16858). - -When using EAGLE-2, please enable `--eagle_use_dynamic_tree`, which indicates whether to use a dynamic tree (default is `False`, i.e., use EAGLE-1 by default). Then set `--eagle_dynamic_tree_max_top_k=N`, which indicates how many new child nodes are expanded for the nodes in the dynamic tree. -- In EagleNet0, `N` draft tokens are generated. -- In EagleNet1, each draft token expands `N` new draft tokens. Therefore, this layer has `N * N` draft tokens. We select the top `N` as the output of this layer. -- In EagleNet2, the `N` output nodes of EagleNet1 are expanded, and each node expands `N` new draft tokens. Therefore, this layer also has a total of `N * N` draft tokens. And select the top `N` as the output of this layer. -- Etc. - -Finally, after `num_eagle_layer` EagleNets, `N + N * N * (num_eagle_layer - 1)` draft tokens are generated. We will rebuild the final tree based on all draft tokens and their scores. The final generated tree will have `min(N + N * N * (num_eagle_layer - 1), max_draft_len)` nodes. - - - -```bash -# Eagle greedy decoding using vicuna-7b-v1.3 model with 1 GPU -python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --input_text "Once upon" - -# Eagle typical acceptance decoding using vicuna-7b-v1.3 model with 1 GPU -python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --input_text "Once upon" \ - --temperature 0.7 \ - --eagle_posterior_threshold 0.09 - -# Eagle decoding using vicuna-7b-v1.3 model with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/4-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --input_text "Once upon" - -# Run EAGLE-2 -mpirun -np 1 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --eagle_use_dynamic_tree \ - --eagle_dynamic_tree_max_top_k 10 \ - --input_text "Once upon" -``` - -For greedy decoding, refer to the following example output: -```text -...... -Input [Text 0]: " Once upon" -Output [Text 0 Beam 0]: "a time, there was a young girl who loved to read. She would spend hours in the library, devouring books of all genres. She had a special love for fairy tales, and would often dream of living in a magical world where she could meet princes and princesses, and have adventures with talking animals. -One day, while she was reading a book, she came across a passage that spoke to her heart. It said, "You are the author of" -``` - -### Summarization using EAGLE decoding -```bash -# EAGLE decoding using vicuna-7b-v1.3 model with 1 GPU -python ../summarize.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --batch_size 1 - -# EAGLE decoding using vicuna-7b-v1.3 with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/4-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --eagle_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --batch_size 1 - -# Run EAGLE-2 -mpirun -np 1 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/eagle/7B/trt_engines/fp16/1-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --eagle_use_dynamic_tree \ - --eagle_dynamic_tree_max_top_k 10 \ - --batch_size 1 - -``` diff --git a/examples/eagle/convert_checkpoint.py b/examples/eagle/convert_checkpoint.py deleted file mode 100644 index 130faee4453c..000000000000 --- a/examples/eagle/convert_checkpoint.py +++ /dev/null @@ -1,493 +0,0 @@ -import argparse -import json -import os -import time -from pathlib import Path - -from tqdm import tqdm -from transformers import LlamaConfig - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.eagle.config import EagleConfig -from tensorrt_llm.models.eagle.model import EagleForCausalLM -from tensorrt_llm.models.model_weights_loader import ModelWeightsLoader -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--meta_ckpt_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4', 'int4_gptq'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument( - '--per_channel', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - - parser.add_argument( - '--per_group', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor to scale weights in the int4 range. ' - 'per_group chooses at run time, and for each group, a custom scaling factor. ' - 'The flag is built for GPTQ/AWQ quantization.') - - parser.add_argument('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - - parser.add_argument('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - - parser.add_argument("--storage-type", - "-t", - type=str, - default="fp32", - choices=["fp32", "fp16"]) - parser.add_argument("--dataset-cache-dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - parser.add_argument("--load-model-on-cpu", action="store_true") - parser.add_argument("--convert-model-on-cpu", action="store_true") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - - parser.add_argument('--eagle_model_dir', type=str, default=None) - parser.add_argument('--max_draft_len', type=int, default=63) - parser.add_argument( - '--num_eagle_layers', - type=int, - default=4, - help= - 'Maximum depth of the EAGLE choices tree, i.e. maximum number of accepted draft tokens.' - ) - parser.add_argument( - '--max_non_leaves_per_layer', - type=int, - default=10, - help='Maximum number of non-leaf nodes in the EAGLE choice tree.') - args = parser.parse_args() - return args - - -def convert_and_save_hf(config, args): - world_size = args.tp_size * args.pp_size - tllm_config = EagleConfig.from_dict(config) - for rank in range(world_size): - tllm_config.mapping = Mapping(world_size=world_size, - rank=rank, - cp_size=1, - tp_size=args.tp_size, - pp_size=args.pp_size) - - model = EagleForCausalLM(tllm_config) - - def check_and_update(module, dict): - if hasattr(module, 'tllm_to_externel_key_dict'): - module.tllm_to_externel_key_dict.update(dict) - else: - module.tllm_to_externel_key_dict = dict - - def copy(tensors): - if isinstance(tensors, list): - if None in tensors: - return tensors - else: - return [tensor.clone() for tensor in tensors] - elif tensors is None: - return tensors - else: - return tensors.clone() - - shared_weight_prefixs = [] - tllm_weights = {} - customized_dict = {"drafter": ""} - if args.eagle_model_dir is None: - # Single checkpoint for ModelOpt - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": "model"}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = 'eagle_module' - loader = ModelWeightsLoader(eagle_model_dir, customized_dict) - loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.named_parameters()): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update(loader.load(tllm_key, preprocess=copy)) - else: - tllm_weights.update(loader.load(tllm_key)) - loader.fill(tllm_weights) - else: - # Double checkpoint for HF - for idx, eagle_net in enumerate(model.eagle_nets): - check_and_update(eagle_net.drafter.fc, {"fc": "fc"}) - check_and_update(eagle_net.drafter.vocab_embedding, - {f"eagle_nets.{idx}": ""}) - check_and_update(eagle_net.lm_head, {f"eagle_nets.{idx}": ""}) - shared_weight_prefixs.append(f"eagle_nets.{idx}") - customized_dict[f'eagle_nets.{idx}'] = '' - - # Load base model - base_loader = ModelWeightsLoader(args.model_dir) - base_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.transformer.named_parameters()): - tllm_weights.update(base_loader.load("transformer." + tllm_key)) - tllm_weights.update(base_loader.load("lm_head.weight")) - for idx in range(args.num_eagle_layers): - tllm_weights.update( - base_loader.load(f"eagle_nets.{idx}.lm_head.weight", - preprocess=copy)) - - # Load eagle model - eagle_loader = ModelWeightsLoader(eagle_model_dir, customized_dict) - eagle_loader.update_key_mapping(model) - for tllm_key, _ in tqdm(model.eagle_nets.named_parameters()): - if not tllm_key.endswith("lm_head.weight"): - if any([ - tllm_key.startswith(prefix) - for prefix in shared_weight_prefixs - ]): - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key, - preprocess=copy)) - else: - tllm_weights.update( - eagle_loader.load("eagle_nets." + tllm_key)) - base_loader.fill(tllm_weights) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - - assert args.pp_size == 1, "Pipeline parallelism is not supported in EAGLE yet." - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - hf_config = None - eagle_model_dir = args.model_dir if args.eagle_model_dir is None else args.eagle_model_dir - if args.model_dir is not None: - hf_config = LlamaConfig.from_pretrained(args.model_dir) - - args.model_type = hf_config.model_type - args.n_head = hf_config.num_attention_heads - args.inter_size = hf_config.intermediate_size - args.n_layer = hf_config.num_hidden_layers - args.n_embd = hf_config.hidden_size - args.n_kv_head = hf_config.num_key_value_heads - args.rms_norm_eps = hf_config.rms_norm_eps - args.vocab_size = hf_config.vocab_size - args.rotary_scaling = hf_config.rope_scaling - args.rotary_base = get_hf_rope_theta(hf_config, 10000.0) - args.n_positions = hf_config.max_position_embeddings - args.dtype = str( - hf_config.torch_dtype)[6:] if args.dtype == 'auto' else args.dtype - if 'head_dim' in hf_config: - args.head_dim = hf_config.head_dim - else: - args.head_dim = args.n_embd // args.n_head - if 'head_size' in hf_config: - args.head_size = hf_config.head_size - else: - args.head_size = args.head_dim - - if args.eagle_model_dir is None: - hf_config_eagle = hf_config.eagle - args.n_head_eagle = hf_config_eagle['num_attention_heads'] - args.inter_size_eagle = hf_config_eagle['intermediate_size'] - args.n_layer_eagle = hf_config_eagle['num_hidden_layers'] - args.n_embd_eagle = hf_config_eagle['hidden_size'] - args.n_kv_head_eagle = hf_config_eagle['num_key_value_heads'] - args.rms_norm_eps_eagle = hf_config_eagle['rms_norm_eps'] - args.n_positions_eagle = hf_config_eagle['max_position_embeddings'] - if 'head_dim' in hf_config_eagle: - args.head_dim_eagle = hf_config_eagle['head_dim'] - else: - args.head_dim_eagle = args.n_embd_eagle // args.n_head_eagle - if 'head_size' in hf_config_eagle: - args.head_size_eagle = hf_config_eagle['head_size'] - else: - args.head_size_eagle = args.head_dim_eagle - else: - hf_config_eagle = LlamaConfig.from_pretrained(args.eagle_model_dir) - args.n_head_eagle = hf_config_eagle.num_attention_heads - args.inter_size_eagle = hf_config_eagle.intermediate_size - args.n_layer_eagle = hf_config_eagle.num_hidden_layers - args.n_embd_eagle = hf_config_eagle.hidden_size - args.n_kv_head_eagle = hf_config_eagle.num_key_value_heads - args.rms_norm_eps_eagle = hf_config_eagle.rms_norm_eps - args.n_positions_eagle = hf_config_eagle.max_position_embeddings - if 'head_dim' in hf_config_eagle: - args.head_dim_eagle = hf_config_eagle.head_dim - else: - args.head_dim_eagle = args.n_embd_eagle // args.n_head_eagle - if 'head_size' in hf_config_eagle: - args.head_size_eagle = hf_config_eagle.head_size - else: - args.head_size_eagle = args.head_dim_eagle - - elif args.meta_ckpt_dir is not None: - assert False, "meta ckpt is not supported yet" - - with open(Path(args.meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - args.n_embd = meta_config["dim"] - args.n_head = meta_config["n_heads"] - args.n_layer = meta_config["n_layers"] - args.n_kv_head = meta_config.get("n_kv_heads", args.n_head) - - if "hidden_dim" in meta_config: - args.inter_size = meta_config["hidden_dim"] - else: - args.multiple_of = meta_config.get("multiple_of", 1) - n_embd = int(4 * args.n_embd * 2 / 3) - args.ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - args.inter_size = args.multiple_of * ( - (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) - // args.multiple_of) - args.rms_norm_eps = meta_config["norm_eps"] - - if args.rotary_scaling is not None: - # assert args.use_gpt_attention_plugin, "RoPE scaling is only supported through GPT attention plugin." - rotary_scaling = { - "type": args.rotary_scaling["rope_type"], - } - args.rotary_scaling = rotary_scaling - - eagle_net_config = { - 'architecture': "LlamaForCausalLM", - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer_eagle, - 'num_attention_heads': args.n_head_eagle, - 'hidden_size': args.n_embd_eagle, - 'intermediate_size': args.inter_size_eagle, - 'num_key_value_heads': args.n_kv_head_eagle, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions_eagle, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps_eagle, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'head_dim': args.head_dim_eagle, - 'head_size': args.head_size_eagle - } - - config = { - 'architecture': 'EagleForCausalLM', - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': args.n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_draft_len': args.max_draft_len, - 'num_eagle_layers': args.num_eagle_layers, - 'max_non_leaves_per_layer': args.max_non_leaves_per_layer, - 'eagle_net_config': eagle_net_config, - 'head_dim': args.head_dim, - 'head_size': args.head_size - } - - assert args.max_draft_len <= 256, "args.max_draft_len > 256 is not supported" - - if args.use_weight_only: - if args.weight_only_precision == 'int8': - config['quantization']['quant_algo'] = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - config['quantization']['quant_algo'] = QuantAlgo.W4A16 - elif args.smoothquant: - if args.per_channel: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - config['quantization']['kv_cache_quant_algo'] = QuantAlgo.INT8 - - if args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - "group_size": args.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': QuantAlgo.W4A16_GPTQ - }) - - # Update quant config if hf_quant_config.json exists - quant_config = {} - try: - with open(eagle_model_dir + '/' + 'hf_quant_config.json') as f: - quant_config = json.load(f) - if "lm_head" in quant_config['quantization']['exclude_modules']: - quant_config['quantization']['exclude_modules'] += [ - f"eagle_nets.{i}.lm_head" - for i in range(args.num_eagle_layers) - ] - config['quantization'].update(quant_config['quantization']) - config['eagle_net_config']['quantization'].update( - quant_config['quantization']) - except IOError: - pass - - convert_and_save_hf(config, args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/eagle/requirements.txt b/examples/eagle/requirements.txt deleted file mode 100644 index 32bd8ce04018..000000000000 --- a/examples/eagle/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -SentencePiece~=0.1.99 -evaluate diff --git a/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py b/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py deleted file mode 100755 index 86b5ca28af4a..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_eagle2_decoding.py +++ /dev/null @@ -1,55 +0,0 @@ -### Generate Text Using Eagle2 Decoding - -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import (EagleDecodingConfig, KvCacheConfig, - SamplingParams) - - -def main(): - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - # The end user can customize the sampling configuration with the SamplingParams class - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - # The end user can customize the kv cache configuration with the KVCache class - kv_cache_config = KvCacheConfig(enable_block_reuse=True) - - llm_kwargs = {} - - model = "lmsys/vicuna-7b-v1.3" - - # The end user can customize the eagle decoding configuration by specifying the - # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices - # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK - # with the EagleDecodingConfig class - - speculative_config = EagleDecodingConfig( - speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3", - max_draft_len=63, - num_eagle_layers=4, - max_non_leaves_per_layer=10, - use_dynamic_tree=True, - dynamic_tree_max_topK=10) - - llm = LLM(model=model, - kv_cache_config=kv_cache_config, - speculative_config=speculative_config, - max_batch_size=1, - max_seq_len=1024, - **llm_kwargs) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py b/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py deleted file mode 100644 index e6e89a622eeb..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_eagle_decoding.py +++ /dev/null @@ -1,60 +0,0 @@ -### Generate Text Using Eagle Decoding - -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import EagleDecodingConfig, KvCacheConfig - - -def main(): - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - # The end user can customize the sampling configuration with the SamplingParams class - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - # The end user can customize the kv cache configuration with the KVCache class - kv_cache_config = KvCacheConfig(enable_block_reuse=True) - - llm_kwargs = {} - - model = "lmsys/vicuna-7b-v1.3" - - # The end user can customize the eagle decoding configuration by specifying the - # speculative_model, max_draft_len, num_eagle_layers, max_non_leaves_per_layer, eagle_choices - # greedy_sampling,posterior_threshold, use_dynamic_tree and dynamic_tree_max_topK - # with the EagleDecodingConfig class - - speculative_config = EagleDecodingConfig( - speculative_model="yuhuili/EAGLE-Vicuna-7B-v1.3", - max_draft_len=63, - num_eagle_layers=4, - max_non_leaves_per_layer=10, - eagle_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ - [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], \ - [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], \ - [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], \ - [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], \ - [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]] - ) - - llm = LLM(model=model, - kv_cache_config=kv_cache_config, - speculative_config=speculative_config, - max_batch_size=1, - max_seq_len=1024, - **llm_kwargs) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_inference_customize.py b/examples/llm-api/_tensorrt_engine/llm_inference_customize.py deleted file mode 100644 index 523501752459..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_inference_customize.py +++ /dev/null @@ -1,55 +0,0 @@ -### Generate text with customization -import tempfile - -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import BuildConfig, KvCacheConfig, SamplingParams - - -def main(): - # The end user can customize the build configuration with the build_config class and other arguments borrowed from the lower-level APIs - build_config = BuildConfig() - build_config.max_batch_size = 128 - build_config.max_num_tokens = 2048 - - build_config.max_beam_width = 4 - - # Model could accept HF model name or a path to local HF model. - - llm = LLM( - model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - build_config=build_config, - kv_cache_config=KvCacheConfig( - free_gpu_memory_fraction=0.8 - ), # Similar to `build_config`, you can also customize the runtime configuration with the `kv_cache_config`, `runtime_config`, `peft_cache_config` or \ - # other arguments borrowed from the lower-level APIs. - ) - - # You can save the engine to disk and load it back later, the LLM class can accept either a HF model or a TRT-LLM engine. - llm.save(tempfile.mkdtemp()) - - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - - # With SamplingParams, you can customize the sampling strategy, such as beam search, temperature, and so on. - sampling_params = SamplingParams(temperature=0.8, - top_p=0.95, - n=4, - use_beam_search=True) - - for output in llm.generate(prompts, sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - - # Got output like - # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming' - # Prompt: 'The capital of France is', Generated text: 'Paris.' - # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are' - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py b/examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py deleted file mode 100644 index 9a1d11782d58..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_inference_kv_events.py +++ /dev/null @@ -1,49 +0,0 @@ -### Get KV Cache Events - -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import KvCacheConfig - - -def main(): - - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - tensor_parallel_size=2, - enable_autotuner=False, - kv_cache_dtype='auto', - kv_cache_config=KvCacheConfig(enable_block_reuse=True, - event_buffer_max_size=1024), - backend="pytorch") - - # Sample prompts having a common prefix. - common_prefix = ( - "After the ghost's departure, Barnardo notes Horatio's pale appearance and asks if he's okay. " - "Horatio concedes that he's shaken and confesses that, without witnessing the ghost himself, he wouldn't have believed it existed. " - "He's also disturbed by the ghost's striking resemblance to the king. It even seems to be wearing the former king's armor. " - "Horatio thinks the ghost's presence foretells that something is about to go wrong in Denmark. " - "Marcellus concurs with Horatio, as he and the other guards have observed that their schedules have become more rigorous and have also noticed the preparations taking place within Elsinore, including the building of cannons, the storing of weapons, and the preparation of ships." - ) - prompts = [ - common_prefix, common_prefix + " Marcellus also notes that the king's" - ] - - # Create a sampling params. - sampling_params = SamplingParams(temperature=0.001, - top_p=0.001, - max_tokens=5) - - for output in llm.generate(prompts, sampling_params=sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - - kv_events = llm.get_kv_cache_events(10) - print(kv_events) - - # Got output like follows: - # [{'event_id': 0, 'data': {'type': 'created', 'num_blocks_per_cache_level': [101230, 0]}}, - # {'event_id': 1, 'data': {'type': 'stored', 'parent_hash': None, 'blocks': [{'type': 'stored_block', 'block_hash': 4203099703668305365, 'tokens': [{'type': 'unique_token', 'token_id': 1, 'token_extra_id': 0}, ... - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py b/examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py deleted file mode 100644 index ed2c94450dd4..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_lookahead_decoding.py +++ /dev/null @@ -1,38 +0,0 @@ -### Generate Text Using Lookahead Decoding -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import (BuildConfig, KvCacheConfig, - LookaheadDecodingConfig, SamplingParams) - - -def main(): - - # The end user can customize the build configuration with the build_config class - build_config = BuildConfig() - build_config.max_batch_size = 32 - - # The configuration for lookahead decoding - lookahead_config = LookaheadDecodingConfig(max_window_size=4, - max_ngram_size=4, - max_verification_set_size=4) - - kv_cache_config = KvCacheConfig(free_gpu_memory_fraction=0.4) - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - kv_cache_config=kv_cache_config, - build_config=build_config, - speculative_config=lookahead_config) - - prompt = "NVIDIA is a great company because" - print(f"Prompt: {prompt!r}") - - sampling_params = SamplingParams(lookahead_config=lookahead_config) - - output = llm.generate(prompt, sampling_params=sampling_params) - print(output) - - #Output should be similar to: - # Prompt: 'NVIDIA is a great company because' - #RequestOutput(request_id=2, prompt='NVIDIA is a great company because', prompt_token_ids=[1, 405, 13044, 10764, 338, 263, 2107, 5001, 1363], outputs=[CompletionOutput(index=0, text='they are always pushing the envelope. They are always trying to make the best graphics cards and the best processors. They are always trying to make the best', token_ids=[896, 526, 2337, 27556, 278, 427, 21367, 29889, 2688, 526, 2337, 1811, 304, 1207, 278, 1900, 18533, 15889, 322, 278, 1900, 1889, 943, 29889, 2688, 526, 2337, 1811, 304, 1207, 278, 1900], cumulative_logprob=None, logprobs=[], finish_reason='length', stop_reason=None, generation_logits=None)], finished=True) - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py b/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py deleted file mode 100644 index d371600d00fc..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_medusa_decoding.py +++ /dev/null @@ -1,93 +0,0 @@ -### Generate Text Using Medusa Decoding -import argparse -from pathlib import Path - -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import (BuildConfig, KvCacheConfig, - MedusaDecodingConfig, SamplingParams) - - -def run_medusa_decoding(use_modelopt_ckpt=False, model_dir=None): - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - # The end user can customize the sampling configuration with the SamplingParams class - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - # The end user can customize the build configuration with the BuildConfig class - build_config = BuildConfig( - max_batch_size=1, - max_seq_len=1024, - ) - - # The end user can customize the kv cache configuration with the KVCache class - kv_cache_config = KvCacheConfig(enable_block_reuse=True) - - llm_kwargs = {} - - if use_modelopt_ckpt: - # This is a Llama-3.1-8B combined with Medusa heads provided by Model Optimizer. - # Both the base model (except lm_head) and Medusa heads have been quantized in FP8. - model = model_dir or "nvidia/Llama-3.1-8B-Medusa-FP8" - - # ModelOpt ckpt uses 3 Medusa heads - speculative_config = MedusaDecodingConfig( - max_draft_len=63, - num_medusa_heads=3, - medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], \ - [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], \ - [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], \ - [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], \ - [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]] - ) - else: - # In this path, base model and Medusa heads are stored and loaded separately. - model = "lmsys/vicuna-7b-v1.3" - - # The end user can customize the medusa decoding configuration by specifying the - # speculative_model, max_draft_len, medusa heads num and medusa choices - # with the MedusaDecodingConfig class - speculative_config = MedusaDecodingConfig( - speculative_model="FasterDecoding/medusa-vicuna-7b-v1.3", - max_draft_len=63, - num_medusa_heads=4, - medusa_choices=[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], \ - [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], \ - [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], \ - [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], \ - [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], \ - [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]] - ) - - # Add 'tensor_parallel_size=2' if using ckpt for - # a larger model like nvidia/Llama-3.1-70B-Medusa. - llm = LLM(model=model, - build_config=build_config, - kv_cache_config=kv_cache_config, - speculative_config=speculative_config, - **llm_kwargs) - - outputs = llm.generate(prompts, sampling_params) - - # Print the outputs. - for output in outputs: - prompt = output.prompt - generated_text = output.outputs[0].text - print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}") - - -if __name__ == '__main__': - parser = argparse.ArgumentParser( - description="Generate text using Medusa decoding.") - parser.add_argument( - '--use_modelopt_ckpt', - action='store_true', - help="Use FP8-quantized checkpoint from Model Optimizer.") - # TODO: remove this arg after ModelOpt ckpt is public on HF - parser.add_argument('--model_dir', type=Path, default=None) - args = parser.parse_args() - - run_medusa_decoding(args.use_modelopt_ckpt, args.model_dir) diff --git a/examples/llm-api/_tensorrt_engine/llm_quantization.py b/examples/llm-api/_tensorrt_engine/llm_quantization.py deleted file mode 100644 index 9db421a7573a..000000000000 --- a/examples/llm-api/_tensorrt_engine/llm_quantization.py +++ /dev/null @@ -1,80 +0,0 @@ -### Generation with Quantization -import logging - -import torch - -from tensorrt_llm import SamplingParams -from tensorrt_llm._tensorrt_engine import LLM -from tensorrt_llm.llmapi import CalibConfig, QuantAlgo, QuantConfig - -major, minor = torch.cuda.get_device_capability() -enable_fp8 = major > 8 or (major == 8 and minor >= 9) -enable_nvfp4 = major >= 10 - -quant_and_calib_configs = [] - -if not enable_nvfp4: - # Example 1: Specify int4 AWQ quantization to QuantConfig. - # We can skip specifying CalibConfig or leave a None as the default value. - quant_and_calib_configs.append( - (QuantConfig(quant_algo=QuantAlgo.W4A16_AWQ), None)) - -if enable_fp8: - # Example 2: Specify FP8 quantization to QuantConfig. - # We can create a CalibConfig to specify the calibration dataset and other details. - # Note that the calibration dataset could be either HF dataset name or a path to local HF dataset. - quant_and_calib_configs.append( - (QuantConfig(quant_algo=QuantAlgo.FP8, - kv_cache_quant_algo=QuantAlgo.FP8), - CalibConfig(calib_dataset='cnn_dailymail', - calib_batches=256, - calib_max_seq_length=256))) -else: - logging.error( - "FP8 quantization only works on post-ada GPUs. Skipped in the example.") - -if enable_nvfp4: - # Example 3: Specify NVFP4 quantization to QuantConfig. - quant_and_calib_configs.append( - (QuantConfig(quant_algo=QuantAlgo.NVFP4, - kv_cache_quant_algo=QuantAlgo.FP8), - CalibConfig(calib_dataset='cnn_dailymail', - calib_batches=256, - calib_max_seq_length=256))) -else: - logging.error( - "NVFP4 quantization only works on Blackwell. Skipped in the example.") - - -def main(): - - for quant_config, calib_config in quant_and_calib_configs: - # The built-in end-to-end quantization is triggered according to the passed quant_config. - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - quant_config=quant_config, - calib_config=calib_config) - - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - - # Create a sampling params. - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - for output in llm.generate(prompts, sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - llm.shutdown() - - # Got output like - # Prompt: 'Hello, my name is', Generated text: 'Jane Smith. I am a resident of the city. Can you tell me more about the public services provided in the area?' - # Prompt: 'The capital of France is', Generated text: 'located in Paris, France. The population of Paris, France, is estimated to be 2 million. France is home to many famous artists, including Picasso' - # Prompt: 'The future of AI is', Generated text: 'an open and collaborative project. The project is an ongoing effort, and we invite participation from members of the community.\n\nOur community is' - - -if __name__ == '__main__': - main() diff --git a/examples/llm-api/_tensorrt_engine/quickstart_example.py b/examples/llm-api/_tensorrt_engine/quickstart_example.py deleted file mode 100644 index d02f55c46b35..000000000000 --- a/examples/llm-api/_tensorrt_engine/quickstart_example.py +++ /dev/null @@ -1,39 +0,0 @@ -from tensorrt_llm import BuildConfig, SamplingParams -from tensorrt_llm._tensorrt_engine import LLM # NOTE the change - - -def main(): - - build_config = BuildConfig() - build_config.max_batch_size = 256 - build_config.max_num_tokens = 1024 - - # Model could accept HF model name, a path to local HF model, - # or Model Optimizer's quantized checkpoints like nvidia/Llama-3.1-8B-Instruct-FP8 on HF. - llm = LLM(model="TinyLlama/TinyLlama-1.1B-Chat-v1.0", - build_config=build_config) - - # Sample prompts. - prompts = [ - "Hello, my name is", - "The capital of France is", - "The future of AI is", - ] - - # Create a sampling params. - sampling_params = SamplingParams(temperature=0.8, top_p=0.95) - - for output in llm.generate(prompts, sampling_params): - print( - f"Prompt: {output.prompt!r}, Generated text: {output.outputs[0].text!r}" - ) - - # Got output like - # Prompt: 'Hello, my name is', Generated text: '\n\nJane Smith. I am a student pursuing my degree in Computer Science at [university]. I enjoy learning new things, especially technology and programming' - # Prompt: 'The president of the United States is', Generated text: 'likely to nominate a new Supreme Court justice to fill the seat vacated by the death of Antonin Scalia. The Senate should vote to confirm the' - # Prompt: 'The capital of France is', Generated text: 'Paris.' - # Prompt: 'The future of AI is', Generated text: 'an exciting time for us. We are constantly researching, developing, and improving our platform to create the most advanced and efficient model available. We are' - - -if __name__ == '__main__': - main() diff --git a/examples/lookahead/README.md b/examples/lookahead/README.md deleted file mode 100644 index b7ac61fa3bcd..000000000000 --- a/examples/lookahead/README.md +++ /dev/null @@ -1,108 +0,0 @@ -# Lookahead Speculative Decoding - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using Lookahead speculative decoding ([Break the Sequential Dependency of LLM Inference Using Lookahead Decoidng](https://arxiv.org/pdf/2402.02057)) in TensorRT-LLM. - -## Overview - -Lookahead decoding algorithm operates through two parallel computation branches within the same LLM - a lookahead branch that generates n-grams using a fixed-sized 2D window, and a verification branch that validates promising n-gram candidates. - -Lookahead algorithm is configured with a tuple of `(windows_size, ngram_size, verification_set_size)` or, shortly, `(W, N, G)`. -+ `windows_size` is the Jacobi window size, meaning number of n-grams in lookahead branch that explores future draft tokens. -+ `ngram_size` is the n-gram size, meaning the maximum number of draft tokens accepted per iteration. -+ `verification_set_size` is the maximum number of n-grams considered for verification, meaning the number of draft token beam hypotheses. - -You can enable Lookahead decoding for any of decoder-only autoregressive LLM models without any fine-tuning. Some TensorRT LLM models might not work with Lookahead due to the missing head size in the speculative decoding XQA attention kernels. Lookahead performance greatly depends on the base model, hardware, batch size, sequence length, and the dataset. It is recommended to profile various configurations to find the best `(W, N, G)` configuration given the setup. - -Specify the Lookahead related flags in three places: - -1. *Build the engine* - -To build an engine with Lookahead support, specify the `--speculative_decoding_mode lookahead_decoding` and `--max_draft_len` arguments. -For Lookahead, the `max_draft_len` is defined as: -```python -def max_draft_len(windows_size, ngram_size, verification_set_size): - return (0 if (ngram_size == 1) else ngram_size - 2) - + (windows_size - 1 + verification_set_size) * (ngram_size - 1) -``` - -2. *Setup TensorRT LLM runtime* -When TensorRT LLM server starts, the server reserves resources according to the `executor_lookahead_config`. `executor_lookahead_config` is noted as `(W, N, G)`. Ensure the `max_draft_len` derived from `executor_lookahead_config` equals to the `max_draft_len` specified in the engine-building phase -- `--max_draft_len == max_draft_len(W, N, G)`. - -3. *Setup the request* -Each request can specify a Lookahead configuration, noted as `(w, n, g)`. If none are specified, the `executor_lookahead_config` is used. The minimum Lookahead config `(1, 1, 0)` forces non speculative, autoregressive mode. The meaningful minimum configuration is `(2, 2, 1)`. Ensure the Lookahead configuration for each request satisfies `w <= W, n <= N, g <= G`. - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16 / BF16 / FP8 - * Paged KV Cache - * Inflight-fused-batching - * C++ runtime - * Tensor Parallel - -## Usage -### Convert Checkpoint - -This example is based on the Vicuna-7b v1.3 model, a fine-tuned Llama model. -Checkpoint conversion is similar to any standard autoregressive model, such as the models located in the [examples/models/core/llama](../../examples/models/core/llama) directory. - -```bash -MODEL_DIR=/path/to/vicuna-7b-v1.3 -ENGINE_DIR=tmp/engine -CKPT_DIR=tmp/engine/ckpt - -python3 examples/models/core/llama/convert_checkpoint.py \ - --model_dir=$MODEL_DIR \ - --output_dir=$CKPT_DIR \ - --dtype=float16 \ - --tp_size=1 \ - --pp_size=1 -``` - -### Build engine - -```bash -trtllm-build \ - --checkpoint_dir=$CKPT_DIR \ - --output_dir=$ENGINE_DIR \ - --gpt_attention_plugin=float16 \ - --gemm_plugin=float16 \ - --max_batch_size=32 \ - --max_input_len=1024 \ - --max_seq_len=2048 \ - --max_beam_width=1 \ - --log_level=error \ - --max_draft_len=83 \ - --speculative_decoding_mode=lookahead_decoding -``` - -### Run decoding - -+ `--lookahead_config` is a server-level configuration of Lookahead decoding in the form of a triplet `[W, N, G]`. Note that `run.py` and `summarize.py` interfaces allow to set only per server but not per-request config. - -Run `examples/run.py` to generate sequences. -```bash -python examples/run.py \ - --tokenizer_dir=$MODEL_DIR \ - --engine_dir=$ENGINE_DIR \ - --max_output_len=32 \ - --lookahead_config=[7,7,7] \ - --log_level=verbose \ - --input_text 'Once upon' 'To be, or not' 'Be not afraid of greatness' -``` - -Run `examples/summarize.py` to summarize the CNN daily dataset. -```bash -python examples/summarize.py \ - --test_hf \ - --test_trt_llm \ - --hf_model_dir=$MODEL_DIR \ - --engine_dir=$ENGINE_DIR \ - --data_type=fp16 \ - --lookahead_config=[7,7,7] -``` diff --git a/examples/lookahead/requirements.txt b/examples/lookahead/requirements.txt deleted file mode 100644 index 423ba2d4b0e0..000000000000 --- a/examples/lookahead/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -sentencepiece>=0.1.99 -evaluate diff --git a/examples/medusa/README.md b/examples/medusa/README.md deleted file mode 100644 index a596608625ea..000000000000 --- a/examples/medusa/README.md +++ /dev/null @@ -1,232 +0,0 @@ -# Medusa Decoding - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a model using Medusa decoding([`Github`](https://github.com/FasterDecoding/Medusa), [`BLOG`](https://sites.google.com/view/medusa-llm)) in TensorRT LLM on single GPU, single node multiple GPU. - -## Overview -Different from other models, Medusa decoding needs a base model and Medusa heads. The TensorRT LLM Medusa Decoding implementation can be found in [tensorrt_llm/models/medusa/model.py](../../tensorrt_llm/models/medusa/model.py). The implementation adds Medusa heads to a base model. - -For more info about Medusa visit [speculative decoding documentation](https://nvidia.github.io/TensorRT-LLM/features/speculative-decoding.html). - -## Support Matrix - * GPU Compute Capability >= 8.0 (Ampere or newer) - * FP16 - * BF16 - * FP8 (base model) - * PAGED_KV_CACHE - * Tensor Parallel - -## Usage -The TensorRT LLM Medusa example code is located in [`examples/medusa`](./). There is one [`convert_checkpoint.py`](./convert_checkpoint.py) file to convert and build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run models with Medusa decoding support. -In this example, we demonstrate the usage of two models: -1. The Vucuna 7B model from Hugging Face [`FasterDecoding/medusa-vicuna-7b-v1.3`](https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3) with its Medusa heads [`medusa-vicuna-7b-v1.3`](https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3). -2. The quantized checkpoint [`nvidia/Llama-3.1-8B-Medusa-FP8`](https://huggingface.co/nvidia/Llama-3.1-8B-Medusa-FP8) on Hugging Face by [Model Optimizer](https://github.com/NVIDIA/Model-Optimizer) (ModelOpt). This model is based on [Llama-3.1 8B](https://huggingface.co/meta-llama/Llama-3.1-8B) and enhanced with Medusa heads, with both the base model (except lm_head) and Medusa heads already quantized in FP8. - -### Build TensorRT engine(s) -Get the weights by downloading base model [`vicuna-7b-v1.3`](https://huggingface.co/lmsys/vicuna-7b-v1.3) and Medusa Heads [`medusa-vicuna-7b-v1.3`](https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3) from HF. - -``` -pip install -r requirements.txt - -git lfs install -git clone https://huggingface.co/lmsys/vicuna-7b-v1.3 -https://huggingface.co/FasterDecoding/medusa-vicuna-7b-v1.3 -``` - -We use `convert_checkpoint.py` script to convert the model for Medusa decoding into TensorRT LLM checkpoint format. -We could use `--num_medusa_heads` to set the number of medusa heads that we want to use. If not, `num_medusa_heads` will be set according to the `medusa_num_heads` from medusa weights' `config.json`. - -Here is the example: -```bash -# Convert and Build Medusa decoding support for vicuna-7b-v1.3 -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --medusa_model_dir medusa-vicuna-7b-v1.3 \ - --output_dir ./tllm_checkpoint_1gpu_medusa \ - --dtype float16 \ - --num_medusa_heads 4 - -# Note: Increasing the batch size may have a negative impact on performance -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_medusa \ - --output_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - -# Convert and Build Medusa decoding support for vicuna-13b-v1.3 with 4-way tensor parallelism. -python convert_checkpoint.py --model_dir ./vicuna-7b-v1.3 \ - --medusa_model_dir medusa-vicuna-7b-v1.3 \ - --output_dir ./tllm_checkpoint_1gpu_medusa \ - --dtype float16 \ - --num_medusa_heads 4 \ - --tp_size 4 \ - --workers 4 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_medusa \ - --output_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - -# Convert and Build Llama-3.1-8B-Medusa by ModelOpt -python convert_checkpoint.py --model_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --output_dir ./tllm_checkpoint_1gpu_modelopt_llama_medusa \ - --dtype float16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_modelopt_llama_medusa \ - --output_dir ./tmp/modelopt/llama-8B-medusa/trt_engines/1-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - - -# Convert and Build Llama-3.1-70B-Medusa by ModelOpt with 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --output_dir ./tllm_checkpoint_2gpu_modelopt_llama_medusa_70b \ - --dtype float16 - --tp_size 2 - --workers 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_modelopt_llama_medusa_70b \ - --output_dir ./tmp/modelopt/llama-70B-medusa/trt_engines/2-gpu/ \ - --gemm_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 - - -``` - -### FP8 Post-Training Quantization for Base Model -The example below quantizes the base model to FP8, while keeping the weight of the medusa head non-quantize. -```bash -# Quantize base model into FP8 and export trtllm checkpoint -python ../quantization/quantize.py --model_dir /path/to/base-model-hf/ \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_base_model_fp8_medusa_fp16 \ - --calib_size 512 \ - --tp_size 1 \ - --medusa_model_dir /path/to/medusa_head/ \ - --num_medusa_heads 4 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_base_model_fp8_medusa_fp16 \ - --output_dir ./trt_engine_1gpu_base_model_fp8_medusa_fp16 \ - --gemm_plugin float16 \ - --gpt_attention_plugin float16 \ - --speculative_decoding_mode medusa \ - --max_batch_size 4 -``` - -### Run -To run a TensorRT LLM model with Medusa decoding support, we can use `../run.py` script, with an additional argument `--medusa_choices`. -The `--medusa_choices` is of type `list[list[int]]`. - -Medusa decoding is supported by Python runtime and C++ runtime with inflight-batching. C++ runtime is recommended for performance. -For Python runtime use `--use_py_session` flag to `run.py`. - -Medusa decoding only supporting greedy decoding, indicated by `temperature=1.0` argument. The output is equivalent to the base model inference with `--temperature 0.0` (equivalent to `--temperature 1.0 --top-k 1`). - -```bash -# Medusa decoding using vicuna-7b-v1.3 model with 1 GPU -python ../run.py --engine_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -# Medusa decoding using vicuna-13b-v1.3 with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/medusa/13B/trt_engines/fp16/4-gpu/ \ - --tokenizer_dir ./vicuna-13b-v1.3/ \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -# Medusa decoding using Llama-3.1-8B-Medusa by ModelOpt with 1 GPU -python ../run.py --engine_dir ./tmp/modelopt/llama-8B-medusa/trt_engines/1-gpu/ \ - --tokenizer_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -# Medusa decoding using Llama-3.1-70B-Medusa by ModelOpt with 2 GPUs -mpirun -np 2 --allow-run-as-root --oversubscribe \ - python ../run.py --engine_dir ./tmp/modelopt/llama-70B-medusa/trt_engines/2-gpu/ \ - --tokenizer_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --max_output_len=100 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --temperature 1.0 \ - --input_text "Once upon" - -``` - -And you will see output like this if run successfully: -```text -...... -Input [Text 0]: " Once upon" -Output [Text 0 Beam 0]: "a time, there was a young girl who loved to read. She would spend hours in the library, devouring books of all genres. She had a special love for fairy tales, and would often dream of living in a magical world where she could meet princes and princesses, and have adventures with talking animals. -One day, while she was reading a book, she came across a passage that spoke to her heart. It said, "You are the author of" -``` - -### Summarization using Medusa decoding - -```bash -# Medusa decoding using vicuna-7b-v1.3 model with 1 GPU -python ../summarize.py --engine_dir ./tmp/medusa/7B/trt_engines/fp16/1-gpu/ \ - --hf_model_dir ./vicuna-7b-v1.3/ \ - --tokenizer_dir ./vicuna-7b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 - -# Medusa decoding using vicuna-13b-v1.3 with 4 GPUs -mpirun -np 4 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/medusa/13B/trt_engines/fp16/4-gpu/ \ - --hf_model_dir ./vicuna-13b-v1.3/ \ - --tokenizer_dir ./vicuna-13b-v1.3/ \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 - -# Medusa decoding using Llama-3.1-8B-Medusa by ModelOpt with 1 GPU -python ../summarize.py --engine_dir ./tmp/modelopt/llama-8B-medusa/trt_engines/1-gpu/ \ - --hf_model_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --tokenizer_dir ./llama3.1-medusa-8b-hf_v0.1 \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 - -# Medusa decoding using Llama-3.1-70B-Medusa by ModelOpt with 2 GPUs -mpirun -np 2 --allow-run-as-root --oversubscribe \ - python ../summarize.py --engine_dir ./tmp/modelopt/llama-70B-medusa/trt_engines/2-gpu/ \ - --hf_model_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --tokenizer_dir ./llama-3.1-70b-medusa_vfp8-fp8-fp8 \ - --test_trt_llm \ - --data_type fp16 \ - --medusa_choices="[[0], [0, 0], [1], [0, 1], [2], [0, 0, 0], [1, 0], [0, 2], [3], [0, 3], [4], [0, 4], [2, 0], [0, 5], [0, 0, 1], [5], [0, 6], [6], [0, 7], [0, 1, 0], [1, 1], [7], [0, 8], [0, 0, 2], [3, 0], [0, 9], [8], [9], [1, 0, 0], [0, 2, 0], [1, 2], [0, 0, 3], [4, 0], [2, 1], [0, 0, 4], [0, 0, 5], [0, 0, 0, 0], [0, 1, 1], [0, 0, 6], [0, 3, 0], [5, 0], [1, 3], [0, 0, 7], [0, 0, 8], [0, 0, 9], [6, 0], [0, 4, 0], [1, 4], [7, 0], [0, 1, 2], [2, 0, 0], [3, 1], [2, 2], [8, 0], [0, 5, 0], [1, 5], [1, 0, 1], [0, 2, 1], [9, 0], [0, 6, 0], [0, 0, 0, 1], [1, 6], [0, 7, 0]]" \ - --use_py_session \ - --temperature 1.0 \ - --batch_size 1 -``` - -### Medusa with Qwen2 - -To use Medusa with Qwen2 models, specify `--model_type qwen2` to `convert_checkpoint.py`. You have to provide a Qwen2 model checkpoint and the medusa heads. After TRT-LLM checkpoint is generated, trllm-build and `../run.py` use the same arguments as for LLaMA models. diff --git a/examples/medusa/convert_checkpoint.py b/examples/medusa/convert_checkpoint.py deleted file mode 100644 index 09eb55b46105..000000000000 --- a/examples/medusa/convert_checkpoint.py +++ /dev/null @@ -1,487 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -import safetensors -import torch -from transformers import (LlamaConfig, LlamaForCausalLM, LlamaTokenizer, - Qwen2Config) - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta, numpy_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import (LLaMAForCausalLM, PretrainedConfig, - QWenForCausalLM) -from tensorrt_llm.models.convert_utils import load_calib_dataset -from tensorrt_llm.models.llama.convert import load_weights_from_hf_by_shard -from tensorrt_llm.models.medusa.weight import (capture_activation_range, - convert_hf_llama, load_medusa_hf) -from tensorrt_llm.quantization import QuantAlgo - -try: - from transformers import MixtralForCausalLM -except ImportError: - MixtralForCausalLM = None - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument( - '--model_dir', - type=str, - help="base model or Medusa-enhanced quantized model from ModelOpt") - parser.add_argument('--meta_ckpt_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4', 'int4_gptq'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument( - '--per_channel', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - - parser.add_argument( - '--per_group', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor to scale weights in the int4 range. ' - 'per_group chooses at run time, and for each group, a custom scaling factor. ' - 'The flag is built for GPTQ/AWQ quantization.') - - parser.add_argument('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - - parser.add_argument('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - - parser.add_argument("--storage-type", - "-t", - type=str, - default="fp32", - choices=["fp32", "fp16"]) - parser.add_argument("--dataset-cache-dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - parser.add_argument("--load-model-on-cpu", action="store_true") - parser.add_argument("--convert-model-on-cpu", action="store_true") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - - parser.add_argument('--num_medusa_heads', type=int, default=4) - parser.add_argument('--num_medusa_layers', type=int, default=1) - parser.add_argument('--max_medusa_token_len', type=int, default=63) - parser.add_argument('--medusa_hidden_act', type=str, default="silu") - parser.add_argument('--medusa_model_dir', type=str, default=None) - parser.add_argument('--model_type', type=str, default="llama") - args = parser.parse_args() - return args - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - assert args.model_type in ["llama", "mixtral", - "qwen2"], "Invalid model type" - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - if args.model_dir is not None: - config_cls = Qwen2Config if args.model_type == "qwen2" else LlamaConfig - hf_config = config_cls.from_pretrained(args.model_dir) - - args.model_type = hf_config.model_type - args.n_head = hf_config.num_attention_heads - args.inter_size = hf_config.intermediate_size - args.n_layer = hf_config.num_hidden_layers - args.n_embd = hf_config.hidden_size - args.n_kv_head = hf_config.num_key_value_heads - args.rms_norm_eps = hf_config.rms_norm_eps - args.vocab_size = hf_config.vocab_size - args.n_positions = hf_config.max_position_embeddings - args.rotary_base = get_hf_rope_theta(hf_config, 10000.0) - args.rotary_scaling = hf_config.rope_scaling - - elif args.meta_ckpt_dir is not None: - - with open(Path(args.meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - args.n_embd = meta_config["dim"] - args.n_head = meta_config["n_heads"] - args.n_layer = meta_config["n_layers"] - args.n_kv_head = meta_config.get("n_kv_heads", args.n_head) - - if "hidden_dim" in meta_config: - args.inter_size = meta_config["hidden_dim"] - else: - args.multiple_of = meta_config.get("multiple_of", 1) - n_embd = int(4 * args.n_embd * 2 / 3) - args.ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - args.inter_size = args.multiple_of * ( - (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) - // args.multiple_of) - args.rms_norm_eps = meta_config["norm_eps"] - - if args.rotary_scaling is not None: - # assert args.use_gpt_attention_plugin, "RoPE scaling is only supported through GPT attention plugin." - rotary_scaling = { - "type": args.rotary_scaling["rope_type"], - } - args.rotary_scaling = rotary_scaling - - # ModelOpt quantized ckpt - quant_config_file = Path(args.model_dir) / "hf_quant_config.json" - - if quant_config_file.exists(): - with open(quant_config_file, 'r') as f: - quant_config = json.load(f) - - is_modelopt_ckpt = quant_config.get("producer", - {}).get("name") == "modelopt" - if is_modelopt_ckpt: # quantized ckpt - args.medusa_model_dir = args.model_dir - args.num_medusa_heads = hf_config.medusa["num_medusa_heads"] - args.num_medusa_layers = hf_config.medusa["num_medusa_layers"] - if args.smoothquant or args.calib_dataset: - logger.warning( - "Checkpoint is already quantized. All quantization-related flags will be ignored." - ) - else: - is_modelopt_ckpt = False - - config = { - 'architecture': 'MedusaForCausalLM', - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': args.n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_draft_len': args.max_medusa_token_len, - 'num_medusa_heads': args.num_medusa_heads, - 'num_medusa_layers': args.num_medusa_layers, - 'model_type': args.model_type, - } - if args.model_type == "qwen2": - config['qwen_type'] = args.model_type - - if is_modelopt_ckpt: - with open(quant_config_file, "r") as f: - hf_quant_config = json.load(f) - for key, value in hf_quant_config.get("quantization", {}).items(): - config['quantization'][key] = value - elif args.use_weight_only: - if args.weight_only_precision == 'int8': - config['quantization']['quant_algo'] = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - config['quantization']['quant_algo'] = QuantAlgo.W4A16 - elif args.smoothquant: - if args.per_channel: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - config['quantization'][ - 'quant_algo'] = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - config['quantization']['kv_cache_quant_algo'] = QuantAlgo.INT8 - - if args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - "group_size": args.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': QuantAlgo.W4A16_GPTQ - }) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - if args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - elif args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - - act_range = {} - llama_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - llama_smoother = {} - model = None - if args.model_dir is not None: - if args.model_type == "qwen2": - model = QWenForCausalLM.from_hugging_face(args.model_dir, - args.dtype) - elif not is_modelopt_ckpt: - hf_model = LlamaForCausalLM if args.model_type != "mixtral" else MixtralForCausalLM - model = hf_model.from_pretrained( - args.model_dir, - dtype='auto', - device_map='auto' if not args.load_model_on_cpu else 'cpu', - trust_remote_code=True) - - if args.smoothquant is not None or args.int8_kv_cache: - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - if args.load_model_on_cpu: - logger.warning( - "Note that running capture_activation_range on cpu would be very slow." - ) - tokenizer = LlamaTokenizer.from_pretrained(args.model_dir, - padding_side='left') - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - - act_range = capture_activation_range(model, tokenizer, dataset) - if args.smoothquant is not None: - smooth_llama_model(model, act_range, args.smoothquant, - llama_qkv_para, llama_smoother) - convert_args = { - 'hf_model': model, - 'act_range': act_range, - 'llama_qkv_para': llama_qkv_para, - 'llama_smoother': llama_smoother, - 'config': config, - 'is_modelopt_ckpt': is_modelopt_ckpt - } - - def covert_and_save(rank, convert_args): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if convert_args["is_modelopt_ckpt"]: - convert_args["hf_model"] = LLaMAForCausalLM.from_hugging_face( - args.model_dir, - args.dtype, - mapping=mapping, - quant_config=convert_args["config"]["quantization"], - load_by_shard=args.load_by_shard) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - assert False, "Never supported" - else: - if args.load_by_shard: - weights = load_weights_from_hf_by_shard( - args.model_dir, PretrainedConfig.from_dict(config)) - - else: - if args.model_type == "qwen2" or convert_args[ - "is_modelopt_ckpt"]: - weights = { - name: numpy_to_torch(param.raw_value) - for name, param in - convert_args['hf_model'].named_parameters() - } - else: - weights = convert_hf_llama( - convert_args['hf_model'], - mapping, - rank, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type= - plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_smooth_quant=args.smoothquant, - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.int8_kv_cache, - act_range=convert_args['act_range'], - qkv_para=convert_args['llama_qkv_para'], - smoother=convert_args['llama_smoother']) - - if args.medusa_model_dir is not None: - config_file = Path(args.medusa_model_dir) / "config.json" - with open(config_file) as fp: - config = json.load(fp) - if not convert_args["is_modelopt_ckpt"]: - num_medusa_heads_from_config = config.get( - 'medusa_num_heads', args.num_medusa_heads) - args.num_medusa_layers = config.get( - 'medusa_num_layers', args.num_medusa_layers) - if args.num_medusa_heads is None: - args.num_medusa_heads = num_medusa_heads_from_config - - assert args.max_medusa_token_len > 0, "should have max_medusa_token_len > 0" - - medusa_weights = load_medusa_hf( - medusa_path=args.medusa_model_dir, - num_medusa_heads=args.num_medusa_heads, - num_medusa_layers=args.num_medusa_layers, - mapping=mapping, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type= - plugin_weight_only_quant_type, - is_modelopt_ckpt=convert_args["is_modelopt_ckpt"]) - weights.update(medusa_weights) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank, convert_args) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank, convert_args) - for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/medusa/requirements.txt b/examples/medusa/requirements.txt deleted file mode 100644 index 423ba2d4b0e0..000000000000 --- a/examples/medusa/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -sentencepiece>=0.1.99 -evaluate diff --git a/examples/models/contrib/baichuan/README.md b/examples/models/contrib/baichuan/README.md deleted file mode 100644 index b1da0a6b0319..000000000000 --- a/examples/models/contrib/baichuan/README.md +++ /dev/null @@ -1,314 +0,0 @@ -# Baichuan - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a Baichuan models (including `v1_7b`/`v1_13b`/`v2_7b`/`v2_13b`) in TensorRT LLM on both single GPU and single node multi-GPU. - -## Table of Contents - -- [Baichuan](#baichuan) - - [Table of Contents](#table-of-contents) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [SmoothQuant](#smoothquant) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Groupwise quantization (AWQ/GPTQ)](#groupwise-quantization-awqgptq) - - [AWQ](#awq) - - [GPTQ](#gptq) - - [INT8 KV cache](#int8-kv-cache) - - [Run](#run) - - [Summarization using the Baichuan model](#summarization-using-the-baichuan-model) - -## Overview - -The TensorRT LLM Baichuan implementation can be found in [tensorrt_llm/models/baichuan/model.py](../../tensorrt_llm/models/baichuan/model.py). The TensorRT LLM Baichuan example code is located in [`examples/models/contrib/baichuan`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert supported checkpoints into TensorRT LLM format. - -The script accepts an argument named model_version, whose value should be `v1_7b`/`v1_13b`/`v2_7b`/`v2_13b` and the default value is `v1_13b`. - -In addition, there are two shared files in the folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * FP8 - * BF16 - * INT4 & INT8 Weight-Only - * INT8 SmoothQuant - * Groupwise quantization (AWQ/GPTQ) - * INT8 KV CACHE (+ AWQ/per-channel weight-only/SmoothQuant) - -## Usage - -The TensorRT LLM Baichuan example code locates at [examples/models/contrib/baichuan](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -Need to specify the HF Baichuan checkpoint path. For `v1_13b`, you should use whether [baichuan-inc/Baichuan-13B-Chat](https://huggingface.co/baichuan-inc/Baichuan-13B-Chat) or [baichuan-inc/Baichuan-13B-Base](https://huggingface.co/baichuan-inc/Baichuan-13B-Base). For `v2_13b`, you should use whether [baichuan-inc/Baichuan2-13B-Chat](https://huggingface.co/baichuan-inc/Baichuan2-13B-Chat) or [baichuan-inc/Baichuan2-13B-Base](https://huggingface.co/baichuan-inc/Baichuan2-13B-Base). More Baichuan models could be found on [baichuan-inc](https://huggingface.co/baichuan-inc). - -TensorRT LLM Baichuan builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -***For all kinds of checkpoints, they share the same trtllm-build command like:*** - -```bash -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. - -# The TensorRT LLM GPT Attention plugin (--gpt_attention_plugin) is -# enabled by default to increase runtime performance. -# 7B models should always enable `gpt_attention_plugin`` since RoPE is only -# supported with GPTAttention plugin now. -# Try gemm_plugin to prevent accuracy issue. -trtllm-build --checkpoint_dir ./tmp/baichuan_v1_13b/trt_ckpts/fp16/1-gpu/ \ - --output_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size=32 \ - --max_input_len=1024 \ - --max_seq_len=1536 -``` - - -Here're some examples for checkpoint conversion that take `v1_13b` as example: - -```bash -# Convert the Baichuan V1 13B model using a single GPU and FP16. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/fp16/1-gpu/ - -# Convert the Baichuan V1 13B model using a single GPU and BF16. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype bfloat16 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/bf16/1-gpu/ - -# Convert the Baichuan V1 13B model using a single GPU and apply INT8 weight-only quantization. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --use_weight_only \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/int8_weight_only/1-gpu/ - -# Convert the Baichuan V1 13B model using a single GPU and apply INT4 weight-only quantization. -# Note that Baichuan V1 7B performs not well when using INT4 weight-only quantization. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/int4_weight_only/1-gpu/ - -# Convert Baichuan V1 13B using 2-way tensor parallelism. -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --output_dir ./tmp/baichuan_v1_13b/trt_ckpts/fp16/1-gpu/ \ - --tp_size 2 -``` - -#### SmoothQuant - -The SmoothQuant supports all Baichuan model variants. Unlike the FP16 build where the HF weights are processed and loaded into the TensorRT LLM directly, the SmoothQuant needs to load INT8 weights which should be pre-processed before building an engine. - -`--smoothquant` is the starting point of INT8 inference. By default, it -will run the model in the _per-tensor_ mode. - -Then, you can add any combination of `--per-token` and `--per-channel` to get the corresponding behaviors. - -Examples of build invocations: -```bash -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --smoothquant 0.8 \ - --per_channel \ - --per_token \ - --output_dir ./tmp/baichuan_v1_13b/sq0.8/1-gpu/ -``` - -#### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -```bash -# Quantize HF Baichuan v2 13B into FP8 and export a single-rank checkpoint -python ../../../quantization/quantize.py --model_dir /code/model/Baichuan2-13B-Chat/ \ - --dtype float16 \ - --qformat fp8 \ - --output_dir ./quantized_fp8 \ - --calib_size 256 -``` - -The quantized model checkpoint is saved to `./quantized_fp8/` for future TensorRT LLM engine build directly with the `trtllm-build` command mentioned above. -Note that you can enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` when building the engines. - -#### Groupwise quantization (AWQ/GPTQ) -##### AWQ -NVIDIA Modelopt toolkit is used for AWQ weight quantization. Please see [examples/quantization/README.md](/examples/quantization/README.md#preparation) for Modelopt installation instructions. -```bash -# Quantize HF Baichuan 13B checkpoint into INT4 AWQ format -python ../../../quantization/quantize.py --model_dir /code/model/Baichuan2-13B-Chat/ \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir ./quantized_int4-awq_gs128 \ - --calib_size 32 -``` -The quantized model checkpoint is saved to `./quantized_int4-awq_gs128/` for future TensorRT LLM engine build directly with the `trtllm-build` command mentioned above. - -##### GPTQ -To run the GPTQ Baichuan example, the following steps are required: - -1. Weight quantization: - - Quantized weights for GPTQ can be generated using an open source project such as [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa.git). - - Let us build the TensorRT LLM engine with the saved `./baichuan-2-13b-4bit-gs64.safetensors`. - -2. Checkpoint conversion: - - ```bash - # Build the Baichuan2 13B model using 2-way tensor parallelism and apply INT4 GPTQ quantization. - # Compressed checkpoint safetensors are generated separately from GPTQ. - python convert_checkpoint.py --model_version v2_13b \ - --quant_ckpt_path ./baichuan-2-13b-4bit-gs64.safetensors \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --group_size 64 \ - --tp_size 2 \ - --output_dir ./tmp/baichuan_v2_13b/trt_ckpts/int4_gptq_gs64/2-gpu/ - ``` - The quantized model checkpoint is saved for future TensorRT LLM engine build directly with the `trtllm-build` command mentioned above. - -#### INT8 KV cache -INT8 KV cache could be enabled to reduce memory footprint. It will bring more performance gains when batch size gets larger. - -You can get the INT8 scale of KV cache through NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit, which features a -`--kv_cache_dtype` option. - -Example: - -```bash -python ../../../quantization/quantize.py --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/baichuan_int8kv_tp1 \ - --calib_size 512 -``` - -**INT8 KV cache + per-channel weight-only quantization** - -INT8 KV cache could be combined with per-channel weight-only quantization, as follows: -```bash -python ../../../quantization/quantize.py --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --qformat int4_wo \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/baichuan_int4wo_int8kv_tp1 \ - --calib_size 512 \ -``` - -**INT8 KV cache + AWQ** - -In addition, you can enable INT8 KV cache together with AWQ (per-group INT4 weight-only quantization), as follows: - -```bash -python ../../../quantization/quantize.py --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --qformat int4_awq \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/baichuan_int4awq_int8kv_tp1 \ - --calib_size 512 -``` - -**INT8 KV cache + INT8 SmoothQuant** - -In addition, you can enable INT8 KV cache together with INT8 SmoothQuant, as follows: - -```bash -python convert_checkpoint.py --model_version v1_13b \ - --model_dir baichuan-inc/Baichuan-13B-Chat \ - --dtype float16 \ - --smoothquant 0.8 \ - --per_channel \ - --per_token \ - --int8_kv_cache \ - --output_dir ./tmp/baichuan_v1_13b/sq0.8/1-gpu/ -``` - -### Run - -To run a TensorRT LLM Baichuan model using the engines generated by `trtllm-build` - -```bash -# With fp16 inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/fp16/1-gpu/ - -# With bf16 inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/bf16/1-gpu/ - -# With INT8 weight-only quantization inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir=baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/int8_weight_only/1-gpu/ - -# With INT4 weight-only quantization inference -python ../../../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir=baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/int8_weight_only/1-gpu/ - -# With 2-way tensor parallelism inference -mpirun -n 2 --allow-run-as-root \ - python ../run.py --input_text "世界上第二高的山峰是哪座？" \ - --max_output_len=50 \ - --tokenizer_dir=baichuan-inc/Baichuan-13B-Chat \ - --engine_dir=./tmp/baichuan_v1_13b/trt_engines/fp16/2-gpu/ -``` - -### Summarization using the Baichuan model - -```bash -# Run summarization using the Baichuan V1 13B model in FP16. -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir baichuan-inc/Baichuan-13B-Chat \ - --data_type fp16 \ - --engine_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/1-gpu/ - -# Run summarization using the Baichuan V1 13B model quantized to INT8. -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir baichuan-inc/Baichuan-13B-Chat \ - --data_type fp16 \ - --engine_dir ./tmp/baichuan_v1_13b/trt_engines/int8_weight_only/1-gpu/ - -# Run summarization using the Baichuan V1 13B model in FP16 using two GPUs. -mpirun -n 2 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir baichuan-inc/Baichuan-13B-Chat \ - --data_type fp16 \ - --engine_dir ./tmp/baichuan_v1_13b/trt_engines/fp16/2-gpu/ -``` diff --git a/examples/models/contrib/baichuan/convert_checkpoint.py b/examples/models/contrib/baichuan/convert_checkpoint.py deleted file mode 100644 index f5a296deee87..000000000000 --- a/examples/models/contrib/baichuan/convert_checkpoint.py +++ /dev/null @@ -1,263 +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. -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.baichuan.config import BaichuanConfig -from tensorrt_llm.models.baichuan.convert import load_weights_from_gptq -from tensorrt_llm.models.baichuan.model import BaichuanForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--quant_ckpt_path', type=str, default=None) - parser.add_argument('--tp_size', type=int, default=1) - parser.add_argument('--pp_size', type=int, default=1) - parser.add_argument('--model_version', - type=str, - default='v1_13b', - choices=['v1_7b', 'v1_13b', 'v2_7b', 'v2_13b']) - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - '--per_channel', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4', 'int4_gptq'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - args = parser.parse_args() - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - config = QuantConfig(group_size=args.group_size) - - if args.smoothquant: - config.smoothquant_val = args.smoothquant - if args.per_token and args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif not args.per_token and not args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - elif not args.per_token and args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif args.per_token and not args.per_channel: - config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - if args.use_weight_only and args.weight_only_precision == 'int8': - config.quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - config.quant_algo = QuantAlgo.W4A16 - elif args.use_weight_only and args.weight_only_precision == 'int4_gptq': - config.quant_algo = QuantAlgo.W4A16_GPTQ - - if args.int8_kv_cache: - config.kv_cache_quant_algo = QuantAlgo.INT8 - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - config.has_zero_point = True - - return config - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - quantization_config = args_to_quant_config(args) - - if args.smoothquant is not None or args.int8_kv_cache: - mapping = Mapping( - world_size=world_size, - tp_size=args.tp_size, - pp_size=args.pp_size, - ) - BaichuanForCausalLM.quantize( - args.model_dir, - args.output_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quantization_config, - calib_dataset=args.calib_dataset, - model_version=args.model_version, - ) - else: - - def convert_and_save_rank(args, rank): - mapping = Mapping( - world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - ) - - model = BaichuanForCausalLM.from_hugging_face( - args.model_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quantization_config, - model_version=args.model_version, - logits_dtype=args.logits_dtype, - ) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - execute(args.workers, [convert_and_save_rank] * world_size, args) - - -def convert_and_save_gptq(args, rank): - mapping = Mapping(world_size=args.tp_size * args.pp_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - config = BaichuanConfig.from_hugging_face( - args.model_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=args_to_quant_config(args), - model_version=args.model_version, - logits_dtype=args.logits_dtype) - - config.vocab_size = int((config.vocab_size + 63) / 64) * 64 - model = BaichuanForCausalLM(config) - weights = load_weights_from_gptq(config, args.quant_ckpt_path) - model.load(weights) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - assert args.model_dir is not None - assert args.quant_ckpt_path is not None - execute(args.workers, [convert_and_save_gptq] * world_size, args) - else: - assert args.model_dir is not None - assert args.quant_ckpt_path is None, "only gptq weights only needs this option" - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/baichuan/requirements.txt b/examples/models/contrib/baichuan/requirements.txt deleted file mode 100644 index d234cb9d00fa..000000000000 --- a/examples/models/contrib/baichuan/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 -cpm-kernels~=1.0.11 -transformers_stream_generator~=0.0.4 diff --git a/examples/models/contrib/bloom/.gitignore b/examples/models/contrib/bloom/.gitignore deleted file mode 100644 index 49fd9fe3a8fa..000000000000 --- a/examples/models/contrib/bloom/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -__pycache__/ -bloom/ diff --git a/examples/models/contrib/bloom/README.md b/examples/models/contrib/bloom/README.md deleted file mode 100644 index a086294d5165..000000000000 --- a/examples/models/contrib/bloom/README.md +++ /dev/null @@ -1,242 +0,0 @@ -# BLOOM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a BLOOM model in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -## Table of Contents - -- [BLOOM](#bloom) - - [Table of Contents](#table-of-contents) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [INT8 weight only + INT8 KV cache](#int8-weight-only--int8-kv-cache) - - [SmoothQuant](#smoothquant) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Run](#run) - -## Overview - -The TensorRT LLM BLOOM implementation can be found in [tensorrt_llm/models/bloom/model.py](../../tensorrt_llm/models/bloom/model.py). The TensorRT LLM BLOOM example code is located in [`examples/models/contrib/bloom`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 & INT4 Weight-Only - * INT8 KV CACHE - * Smooth Quant - * Tensor Parallel - * FP8 and FP8 KV cache - -## Usage - -The TensorRT LLM BLOOM example code locates at [examples/models/contrib/bloom](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -Need to prepare the HF BLOOM checkpoint by following the guides here https://huggingface.co/docs/transformers/main/en/model_doc/bloom. - -e.g. To install BLOOM-560M - -```bash -# Setup git-lfs -git lfs install -rm -rf ./bloom/560M -mkdir -p ./bloom/560M && git clone https://huggingface.co/bigscience/bloom-560m ./bloom/560M -``` - -TensorRT LLM BLOOM builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `workers` feature only supports single node. - -Here're some examples: - -```bash -# Build a single-GPU float16 engine from HF weights. -# Try gemm_plugin to prevent accuracy issue. TODO check this holds for BLOOM - -# Single GPU on BLOOM 560M -python convert_checkpoint.py --model_dir ./bloom/560M/ \ - --dtype float16 \ - --output_dir ./bloom/560M/trt_ckpt/fp16/1-gpu/ -trtllm-build --checkpoint_dir ./bloom/560M/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560M/trt_engines/fp16/1-gpu/ - -# Build the BLOOM 560M using a single GPU and apply INT8 weight-only quantization. -python convert_checkpoint.py --model_dir ./bloom/560M/ \ - --dtype float16 \ - --use_weight_only \ - --output_dir ./bloom/560M/trt_ckpt/int8_weight_only/1-gpu/ -trtllm-build --checkpoint_dir ./bloom/560M/trt_ckpt/int8_weight_only/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560M/trt_engines/int8_weight_only/1-gpu/ - -# Use 2-way tensor parallelism on BLOOM 560M -python convert_checkpoint.py --model_dir ./bloom/560M/ \ - --dtype float16 \ - --output_dir ./bloom/560M/trt_ckpt/fp16/2-gpu/ \ - --tp_size 2 -trtllm-build --checkpoint_dir ./bloom/560M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560M/trt_engines/fp16/2-gpu/ - -# Use 8-way tensor parallelism on BLOOM 176B -# Currently, TensorRT does not support tensors with more than 2^31-1 elements, -# so we have to shard the embedding table to multi-GPUs. - -# sharding embedding table in the vocab dimension (the lookup plugin is optional) -python convert_checkpoint.py --model_dir ./bloom/176B/ \ - --dtype float16 \ - --output_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --tp_size 8 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 -trtllm-build --checkpoint_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/176B/trt_engines/fp16/8-gpu/ \ - --workers 2 - -# sharding embedding table in the hidden dimension -python convert_checkpoint.py --model_dir ./bloom/176B/ \ - --dtype float16 \ - --output_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --tp_size 8 \ - --use_parallel_embedding \ - --embedding_sharding_dim 1 -trtllm-build --checkpoint_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/176B/trt_engines/fp16/8-gpu/ \ - --workers 2 - -# share embedding table between embedding() and lm_head() layers -# To reduce the generated engine size, we can turn off gemm plugin and shard the embedding table in the vocab dimension. -python convert_checkpoint.py --model_dir ./bloom/176B/ \ - --dtype float16 \ - --output_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --tp_size 8 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 -trtllm-build --checkpoint_dir ./bloom/176B/trt_ckpt/fp16/8-gpu/ \ - --output_dir ./bloom/176B/trt_engines/fp16/8-gpu/ \ - --workers 2 -``` - -#### INT8 weight only + INT8 KV cache - - -For INT8 KV cache, [`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 KV cache. - -`--int8_kv_cache` is the command-line option to enable INT8 KV cache. - -In addition, it could be combined with INT8 weight-only quantization, as follows: - -Examples of INT8 weight-only quantization + INT8 KV cache - -```bash -# Build model with both INT8 weight-only and INT8 KV cache enabled -python convert_checkpoint.py --model_dir ./bloom/560m/ \ - --dtype float16 \ - --int8_kv_cache \ - --use_weight_only --output_dir ./bloom/560m/trt_ckpt/int8/1-gpu/ -trtllm-build --checkpoint_dir ./bloom/560m/trt_ckpt/int8/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./bloom/560m/trt_engines/int8/1-gpu/ \ -``` - - -#### SmoothQuant - -Unlike the FP16 build where the HF weights are processed and loaded into the TensorRT LLM directly, the SmoothQuant needs to load INT8 weights which should be pre-processed before building an engine. - -Example: -```bash -python convert_checkpoint.py --model_dir bloom/560M/ --output_dir tllm_checkpoint_1gpu --smoothquant 0.5 --per_token --per_channel -``` - -[`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 inference of SmoothQuant models. - -`--smoothquant` is the starting point of INT8 inference. By default, it -will run the model in the _per-tensor_ mode. - -Then, you can add any combination of `--per-token` and `--per-channel` to get the corresponding behaviors. - - - -```bash -# Build model for SmoothQuant with below command. - -trtllm-build --checkpoint_dir tllm_checkpoint_1gpu --output_dir ./engine_outputs -``` -Note that GPT attention plugin is required to be enabled for SmoothQuant for now. - - -Note we use `--bin_model_dir` instead of `--model_dir` since SmoothQuant model needs INT8 weights and various scales from the binary files. - -#### FP8 Post-Training Quantization - -``` -# Quantize HF Bloom 3B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir /home/scratch.trt_llm_data_ci/llm-models/bloom-3b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir /tmp/bloom/3b/trt_ckpts/fp8/1-gpu/ \ - --calib_size 512 \ - --tp_size 1 - -trtllm-build --checkpoint_dir /tmp/bloom/3b/trt_ckpts/fp8/1-gpu/ \ - --output_dir /tmp/bloom/3b/trt_engines/fp8/1-gpu/ \ - --gemm_plugin float16 \ - --use_fp8_context_fmha enable \ - --workers 1 -``` - -### Run - -```bash -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./bloom/560M/ \ - --data_type fp16 \ - --engine_dir ./bloom/560M/trt_engines/fp16/1-gpu/ - -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./bloom/560M/ \ - --data_type fp16 \ - --engine_dir ./bloom/560M/trt_engines/int8_weight_only/1-gpu/ - -mpirun -n 2 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./bloom/560M/ \ - --data_type fp16 \ - --engine_dir ./bloom/560M/trt_engines/fp16/2-gpu/ - -mpirun -n 8 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./bloom/176B/ \ - --data_type fp16 \ - --engine_dir ./bloom/176B/trt_engines/fp16/8-gpu/ - -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir /home/scratch.trt_llm_data_ci/llm-models/bloom-3b \ - --data_type fp16 \ - --engine_dir /tmp/bloom/3b/trt_engines/fp8/1-gpu/ -``` diff --git a/examples/models/contrib/bloom/convert_checkpoint.py b/examples/models/contrib/bloom/convert_checkpoint.py deleted file mode 100644 index 8ab16f2b23bf..000000000000 --- a/examples/models/contrib/bloom/convert_checkpoint.py +++ /dev/null @@ -1,962 +0,0 @@ -import argparse -import functools -import json -import os -import time -from collections import defaultdict -from pathlib import Path -from typing import Any, Dict, Iterable, Optional, Union - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import BloomConfig, BloomForCausalLM, BloomTokenizerFast -from transformers.models.bloom.modeling_bloom import BloomBlock -from transformers.pytorch_utils import Conv1D - -# isort: off -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm import logger -from tensorrt_llm.quantization import QuantAlgo, QuantMode -from tensorrt_llm.models.convert_utils import iterate_shard_files, load_state_dict, \ - load_calib_dataset, split_matrix_tp, get_weight_and_bias, split, smooth_gemm, \ - generate_int8,get_weight -# isort: on - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=512, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip(1e-8, - None).max(dim=1)[0] - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - input_ids = tokenizer(dataset[i], - return_tensors="pt", - max_length=seq_len, - truncation=True).input_ids.to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -def reorder_torch_qkv_weight_or_bias(v, model, is_bias=False): - """ Reorder the qkv weight. - - Note that the shape of the fused QKV weights in HF is different from the - shape that TRT-LLM requires. - HF: (num_heads x 3 x head_dim, hidden_size) - TRT-LLM: (3 x num_heads x head_dim, hidden_size) - This is unlike to the other models in HF e.g. GPT where they have the - same shape with TRT-LLM, i.e., (3 x num_heads x head_dim, hidden_size). We reshape the qkv - weight: (3 x num_heads x head_dim, hidden). - bias : (3 x num_heads x head_dim). - """ - - n_head = model.transformer.num_heads - hidden_size = model.transformer.embed_dim - head_dim = hidden_size // n_head - - # (3 x hidden, ...) view as (num_heads, 3, head_dim, ...) - v = v.reshape(n_head, 3, head_dim, -1) - # permute to (3, num_heads, head_dim, ...) - v = v.permute((1, 0, 2, 3)) - # final shape: weight=(3 x hidden, hidden) or bias=(3 x hidden) - if is_bias: - return v.reshape(3 * hidden_size) - return v.reshape(3 * hidden_size, hidden_size) - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=Path, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=Path, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--calib_dataset', - type=str, - default='lambada', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument( - '--per_channel', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers to convert checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - return args - - -def reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=False): - """ Reorder the qkv weight. - - Note that the shape of the fused QKV weights in HF is different from the - shape that TRT-LLM requires. - HF: (num_heads x 3 x head_dim, hidden_size) - TRT-LLM: (3 x num_heads x head_dim, hidden_size) - This is unlike to the other models in HF e.g. GPT where they have the - same shape with TRT-LLM, i.e., (3 x num_heads x head_dim, hidden_size). Also, - to split across attention heads in tensor parallel, we reshape the qkv - weight: (3, num_heads x head_dim, hidden). - bias : (3, num_heads x head_dim). - """ - - head_dim = n_hidden // n_head - - # (3 x hidden, ...) view as (num_heads, 3, head_dim, ...) - v = v.reshape(n_head, 3, head_dim, -1) - # permute to (3, num_heads, head_dim, ...) - v = v.transpose(0, 1) - # final shape: weight=(3, hidden, hidden) or bias=(3, hidden) - if is_bias: - return v.reshape(3, n_hidden) - return v.reshape(3, n_hidden, n_hidden) - - -def split_qkv_tp(v, n_head, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - v = reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=False) - split_v = split(v, tensor_parallel, rank, dim=1) - split_v = split_v.reshape(3 * (n_hidden // tensor_parallel), n_hidden) - return split_v.contiguous() - - -def split_qkv_bias_tp(v, n_head, n_hidden, tensor_parallel, rank): - """ - Splits the QKV bias according to tensor parallelism - """ - v = reorder_qkv_weight_or_bias(v, n_head, n_hidden, is_bias=True) - split_v = split(v, tensor_parallel, rank, dim=1) - split_v = split_v.reshape(3 * (n_hidden // tensor_parallel)) - return split_v.contiguous() - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.cpu().t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + 'weight'] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + 'weight'] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def add_tllm_weight( - weights: Dict[str, torch.Tensor], - name: str, - param: torch.Tensor, - quant_mode: QuantMode = QuantMode(0), -): - assert name not in weights, f'{name} is already added.' - - if name.endswith('.weight') and quant_mode.is_weight_only(): - if quant_mode.is_int8_weight_only(): - quant_dtype = torch.int8 - elif quant_mode.is_int4_weight_only(): - quant_dtype = torch.quint4x2 - else: - raise ValueError( - f'Invalid configuration, got quant_mode={quant_mode}') - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - param.t().contiguous(), quant_dtype) - weights[name] = processed_torch_weights - scale_name = name.replace('.weight', '.per_channel_scale') - weights[scale_name] = torch_weight_scales - else: - weights[name] = param.contiguous() - - -@torch.no_grad() -def smooth_bloom_model(model, scales, alpha, bloom_qkv_param, bloom_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, BloomBlock): - continue - - # reorder qkv weight/bias and scales - param = module.self_attention.query_key_value.weight - param = reorder_torch_qkv_weight_or_bias(param, model, is_bias=False) - - layer_name = name + ".self_attention.query_key_value" - act_range_qkv = scales.get(layer_name) - # (n_head x 3 x head_dim) -> (3 x n_head x head_dim) - act_range_qkv['w'] = reorder_torch_qkv_weight_or_bias( - act_range_qkv['w'], model, is_bias=True) - act_range_qkv['y'] = reorder_torch_qkv_weight_or_bias( - act_range_qkv['y'], model, is_bias=True) - scales[layer_name] = act_range_qkv - - # qkv_proj - smoother = smooth_gemm(param, scales[layer_name]["x"], - module.input_layernorm.weight, - module.input_layernorm.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = param.abs().max(dim=1)[0] - bloom_qkv_param[layer_name] = param - - # dense - # enabled for better accuracy with perf overhead of quantization - layer_name = name + ".self_attention.dense" - smoother = smooth_gemm(module.self_attention.dense.weight, - scales[layer_name]["x"], None, None, alpha) - bloom_smoother[layer_name] = smoother - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.self_attention.dense.weight.abs().max( - dim=1)[0] - - # fc1 - layer_name = name + ".mlp.dense_h_to_4h" - smoother = smooth_gemm(module.mlp.dense_h_to_4h.weight, - scales[layer_name]["x"], - module.post_attention_layernorm.weight, - module.post_attention_layernorm.bias, alpha) - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.dense_h_to_4h.weight.abs().max( - dim=1)[0] - - # fc2 - # enabled for better accuracy with perf overhead of quantization - layer_name = name + ".mlp.dense_4h_to_h" - smoother = smooth_gemm(module.mlp.dense_4h_to_h.weight, - scales[layer_name]["x"], None, None, alpha) - bloom_smoother[layer_name] = smoother - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.mlp.dense_4h_to_h.weight.abs().max( - dim=1)[0] - - -def get_tllm_linear_sq_weight( - vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, -): - results = {} - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - - original_weights = vals["weight.int8.col"] - cur_weights = np.split(original_weights, tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if smoother_value is None: - cur_per_channel_value = np.split(vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - else: - original_weights = vals["weight.int8"] - cur_weights = np.split(original_weights, tensor_parallel, - axis=cat_dim)[rank] - - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + 'weight'] = cur_weights.t().contiguous() - - cur_per_channel_value = vals["scale_y_accum_quant"] - - results[prefix + 'per_channel_scale'] = cur_per_channel_value.reshape( - col_shape).contiguous() - - results[last_prefix] = vals['scale_x_orig_quant'].contiguous() - - results[prefix + 'act_scale'] = vals["scale_y_quant_orig"].contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_bloom(hf_bloom, - rank=0, - tensor_parallel=1, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8, - use_smooth_quant=False, - bloom_qkv_param={}, - act_range=None, - smoother=None, - per_channel=False, - per_token=False, - int8_kv_cache=False): - - weights = {} - tik = time.time() - - model_params = dict(hf_bloom.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_bloom.config.n_head - hidden_size = hf_bloom.config.hidden_size - - for l in range(hf_bloom.config.num_hidden_layers): - prefix = f'transformer.h.{l}.' - tllm_prex = f'transformer.layers.{l}.' - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, prefix + 'self_attention.query_key_value', dtype) - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, rank) - bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size, - tensor_parallel, rank) - - if use_smooth_quant: - qkv_weight = bloom_qkv_param[prefix + - 'self_attention.query_key_value'].t() - - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8( - qkv_weight, - act_range.get( - (tllm_prex + 'self_attention.query_key_value').replace( - ".layers.", ".h.")), - is_qkv=True) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.qkv.', - [1, 3 * hidden_size // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'input_layernorm.scale_to_int', - bias=bias, - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - else: - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, rank) - bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size, - tensor_parallel, rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', - bias, use_weight_only, - plugin_weight_only_quant_type)) - if int8_kv_cache: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - - int8_weights = generate_int8( - qkv_weight, - act_range.get( - (tllm_prex + 'self_attention.query_key_value').replace( - ".layers.", ".h.")), - is_qkv=True) - - kv_cache_weights = {} - - kv_cache_weights[ - tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.from_numpy( - np.array([int8_weights['scale_y_quant_orig']], - dtype=np.float32)).contiguous() - weights.update(kv_cache_weights) - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, prefix + 'self_attention.dense', dtype) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t() - int8_weights = generate_int8( - attn_dense_weight, - act_range.get((tllm_prex + 'self_attention.dense').replace( - ".layers.", ".h."))) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - bias=attn_dense_bias, - smoother_value=smoother[(tllm_prex + - 'self_attention.dense').replace( - ".layers.", ".h.")], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_h_to_4h', dtype) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t() - int8_weights = generate_int8( - mlp_fc_weight, - act_range.get((tllm_prex + 'mlp.dense_h_to_4h').replace( - ".layers.", ".h."))) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, 4 * hidden_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - bias=bias, - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - else: - split_v = split_matrix_tp(mlp_fc_weight, - tensor_parallel, - rank, - dim=0) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_4h_to_h', dtype) - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t() - int8_weights = generate_int8( - mlp_proj_weight, - act_range.get((tllm_prex + 'mlp.dense_4h_to_h').replace( - ".layers.", ".h."))) - - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - bias=mlp_proj_bias, - smoother_value=smoother[(tllm_prex + - 'mlp.dense_4h_to_h').replace( - ".layers.", ".h.")], - smoother_shape=[1, 4 * hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - split_v = split_matrix_tp(mlp_proj_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, prefix + 'input_layernorm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - weights[tllm_prex + 'input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - weights[tllm_prex + 'post_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'transformer.word_embeddings', dtype) - weights['lm_head.weight'] = split_matrix_tp(embed_w.clone(), - tensor_parallel, - rank, - dim=0) - - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - assert hf_bloom.config.vocab_size % tensor_parallel == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix_tp( - embed_w, tensor_parallel, rank, dim=sharding_dim) - - embed_f_w, embed_f_b = get_weight_and_bias( - model_params, 'transformer.word_embeddings_layernorm', dtype) - weights['transformer.ln_embed.weight'] = embed_f_w - weights['transformer.ln_embed.bias'] = embed_f_b - - ln_f_w, ln_f_b = get_weight_and_bias(model_params, 'transformer.ln_f', - dtype) - weights['transformer.ln_f.weight'] = ln_f_w - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def rename_hf_to_tllm(name: str): - """ Rename a HF parameter name by the corresponding TRT-LLM style name. """ - if 'word_embeddings_layernorm.' in name: - name = name.replace('word_embeddings_layernorm', 'ln_embed') - if not name.startswith('transformer.'): - name = f'transformer.{name}' - elif 'word_embeddings.' in name: - name = name.replace('word_embeddings', 'vocab_embedding') - if name.startswith(('ln_embed.', 'vocab_embedding.', 'ln_f.')): - name = f'transformer.{name}' - - # Parameter names in layers - if name.startswith(('transformer.h.', 'h.')): - import re - name = re.sub(r'^(transformer.h.|h.)', 'transformer.layers.', name, 1) - if 'post_attention_layernorm' in name: - name = name.replace('post_attention_layernorm', 'post_layernorm') - elif 'self_attention.query_key_value' in name: - name = name.replace('self_attention.query_key_value', 'attention.qkv') - elif 'self_attention.dense' in name: - name = name.replace('self_attention.dense', 'attention.dense') - elif 'mlp.dense_h_to_4h' in name: - name = name.replace('mlp.dense_h_to_4h', 'mlp.fc') - elif 'mlp.dense_4h_to_h' in name: - name = name.replace('mlp.dense_4h_to_h', 'mlp.proj') - return name - - -def contain_any(name: str, words: Iterable[str]): - for word in words: - if word in name: - return True - return False - - -def convert_from_hf_checkpoint( - model_dir: Union[str, Path], - rank=0, - tensor_parallel=1, - dtype: Union[str, torch.dtype] = torch.float32, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8, - use_smooth_quant: bool = False, - bloom_qkv_param: Optional[Dict] = None, - act_range: Optional[Any] = None, - smoother: Optional[Any] = None, - per_channel: bool = False, - per_token: bool = False, - int8_kv_cache: bool = False, -): - logger.info('Loading weights from HF BLOOM...') - tik = time.time() - - weights = {} - hf_config = BloomConfig.from_pretrained(model_dir) - num_heads = hf_config.n_head - hidden_size = hf_config.hidden_size - if isinstance(dtype, str): - dtype = tensorrt_llm.str_dtype_to_torch(dtype) - tp_rank = rank - tp_size = tensor_parallel - - if use_smooth_quant: - quant_mode = QuantMode.use_smooth_quant(per_token, per_channel) - elif use_weight_only: - quant_mode = QuantMode.from_description( - quantize_weights=True, - quantize_activations=False, - per_token=False, - per_channel=False, - use_int8_kv_cache=int8_kv_cache, - use_int4_weights=plugin_weight_only_quant_type == torch.quint4x2) - else: - quant_mode = QuantMode(0) - - def is_bias(_name): - return 'bias' in _name - - for model_file in iterate_shard_files(model_dir, tp_rank): - logger.debug(f'Loading file {str(model_file)}...') - model_params = load_state_dict(model_file, dtype=dtype) - for name, param in model_params.items(): - logger.debug(f'Converting weight {name}...') - tllm_name = rename_hf_to_tllm(name) - param = param.detach().cpu() - - # TODO: Support SmmothQuant. - - if 'self_attention.query_key_value' in name: - if not is_bias(name): - param = split_qkv_tp(param, num_heads, hidden_size, tp_size, - tp_rank) - # TODO: Add KV scalers when quantizing KV cache. - else: - param = split_qkv_bias_tp(param, num_heads, hidden_size, - tp_size, tp_rank) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'self_attention.dense' in name: - if not is_bias(name): - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'mlp.dense_h_to_4h' in name: - if not is_bias(name): - param = split_matrix_tp(param, tp_size, tp_rank, dim=0) - else: - param = split_matrix_tp(param, tp_size, tp_rank, dim=0) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'mlp.dense_4h_to_h' in name: - if not is_bias(name): - param = split_matrix_tp(param, tp_size, tp_rank, dim=1) - add_tllm_weight(weights, tllm_name, param, quant_mode) - elif 'word_embeddings.' in name: - # TODO: safetensor doesn't allow to save a shared tensor. - # Currently, we clone the weight but to save the disk, it - # would be better to skip saving lm_head weights and - # handle it at the loading phase. - lm_head = split_matrix_tp(param, tp_size, tp_rank, dim=0) - weights['lm_head.weight'] = lm_head.clone() - - if not use_parallel_embedding: - weights[tllm_name] = param - else: - assert hf_config.vocab_size % tp_size == 0 - weights[tllm_name] = split_matrix_tp(param, - tp_size, - tp_rank, - dim=sharding_dim) - elif contain_any(name, - ('input_layernorm', 'post_attention_layernorm', - 'word_embeddings_layernorm.', 'ln_f.')): - weights[tllm_name] = param - del model_params - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') - return weights - - -def do_convert_from_ckpt(args): - return (args.model_dir.exists() and args.smoothquant is None - and not args.use_weight_only and not args.int8_kv_cache) - - -def convert(worker_rank, world_size, args, convert_args): - convert_from_ckpt = do_convert_from_ckpt(args) - for rank in range(worker_rank, world_size, args.workers): - if convert_from_ckpt: - weights = convert_from_hf_checkpoint(rank=rank, **convert_args) - else: - weights = convert_hf_bloom(rank=rank, **convert_args) - safetensors.torch.save_file(weights, - args.output_dir / f'rank{rank}.safetensors') - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - - args = parse_arguments() - world_size = args.tp_size * args.pp_size - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - logger.set_level(args.log_level) - tik = time.time() - - args.output_dir.mkdir(exist_ok=True, parents=True) - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - elif args.smoothquant: - if args.per_channel and args.per_token: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif args.per_channel and not args.per_token: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif not args.per_channel and args.per_token: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - kv_cache_quant_algo = None - if args.int8_kv_cache: - kv_cache_quant_algo = QuantAlgo.INT8 - - hf_config = BloomConfig.from_pretrained(args.model_dir) - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'hidden_size': hf_config.hidden_size, - 'norm_epsilon': hf_config.layer_norm_epsilon, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'alibi', - 'hidden_act': 'gelu', - 'intermediate_size': hf_config.hidden_size * 4, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - } - - with (args.output_dir / 'config.json').open('w') as f: - json.dump(config, f, indent=4) - - # TODO: convert_from_hf_checkpoint is memory efficient but has not - # supported quantization yet. Will enable once implemented. - convert_from_ckpt = do_convert_from_ckpt(args) - if not convert_from_ckpt: - logger.info(f'Convert by using model') - hf_bloom = BloomForCausalLM.from_pretrained(args.model_dir, - dtype="auto", - device_map="auto", - trust_remote_code=True) - else: - logger.info(f'Convert by using checkpoint') - hf_bloom = None - - act_range = {} - bloom_qkv_param = {} - bloom_smoother = {} - - if args.smoothquant is not None or args.int8_kv_cache: - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - tokenizer = BloomTokenizerFast.from_pretrained(args.model_dir) - dataset = load_calib_dataset(args.calib_dataset) - - act_range = capture_activation_range(hf_bloom, tokenizer, dataset) - if args.smoothquant is not None: - smooth_bloom_model(hf_bloom, act_range, args.smoothquant, - bloom_qkv_param, bloom_smoother) - - convert_args = dict( - tensor_parallel=args.tp_size, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_smooth_quant=args.smoothquant, - act_range=act_range, - bloom_qkv_param=bloom_qkv_param, - smoother=bloom_smoother, - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.int8_kv_cache, - ) - if convert_from_ckpt: - convert_args['model_dir'] = args.model_dir - else: - convert_args['hf_bloom'] = hf_bloom - - if args.workers == 1: - convert(0, world_size, args, convert_args) - else: - if args.workers > world_size: - args.workers = world_size - logger.info(f'Convert checkpoint using {args.workers} workers.') - import torch.multiprocessing as mp - mp.spawn(convert, - nprocs=args.workers, - args=(world_size, args, convert_args)) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/bloom/requirements.txt b/examples/models/contrib/bloom/requirements.txt deleted file mode 100644 index 88232baef811..000000000000 --- a/examples/models/contrib/bloom/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 diff --git a/examples/models/contrib/cogvlm/convert_checkpoint.py b/examples/models/contrib/cogvlm/convert_checkpoint.py deleted file mode 100644 index 26a56de58438..000000000000 --- a/examples/models/contrib/cogvlm/convert_checkpoint.py +++ /dev/null @@ -1,501 +0,0 @@ -import argparse -import copy -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -import safetensors -import torch -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import PretrainedConfig -from tensorrt_llm.models.cogvlm.convert import convert_hf_cogvlm -from tensorrt_llm.models.convert_utils import load_calib_dataset -from tensorrt_llm.models.llama.convert import (capture_activation_range, - load_weights_from_gptq, - load_weights_from_hf_by_shard, - load_weights_from_meta_ckpt, - smooth_llama_model) - -try: - from transformers import LlavaConfig, LlavaForConditionalGeneration -except ImportError: - pass - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--meta_ckpt_dir', type=str, default=None) - - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - parser.add_argument('--n_head', type=int, default=32) - parser.add_argument('--n_kv_head', type=int, default=None) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--inter_size', type=int, default=11008) - parser.add_argument('--rms_norm_eps', type=float, default=1e-06) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4', 'int4_gptq'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument( - '--per_channel', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - parser.add_argument( - '--quant_ckpt_path', - type=str, - default=None, - help='Path of a quantized model checkpoint in .npz format') - - parser.add_argument( - '--per_group', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor to scale weights in the int4 range. ' - 'per_group chooses at run time, and for each group, a custom scaling factor. ' - 'The flag is built for GPTQ/AWQ quantization.') - - parser.add_argument( - '--enable_fp8', - default=False, - action='store_true', - help='Use FP8 Linear layer for Attention QKV/Dense and MLP.') - parser.add_argument( - '--fp8_kv_cache', - default=False, - action="store_true", - help='By default, we use dtype for KV cache. fp8_kv_cache chooses int8 ' - 'quantization for KV') - parser.add_argument('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - - parser.add_argument('--group_size', - type=int, - default=128, - help='Group size used in GPTQ/AWQ quantization.') - - parser.add_argument("--storage-type", - "-t", - type=str, - default="fp32", - choices=["fp32", "fp16"]) - parser.add_argument("--dataset-cache-dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - parser.add_argument("--load_model_on_cpu", action="store_true") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--use_prompt_tuning', - action="store_true", - default=False) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--enable_pos_shift', - default=False, - action='store_true', - help='Enable position shift for streamingllm method') - parser.add_argument( - '--dense_context_fmha', - default=False, - action='store_true', - help= - 'Enable dense fmha in context phase, otherwise sliding window attention.' - 'If dense_context_fmha=False, the sliding window size is the max attention window size.' - ) - args = parser.parse_args() - return args - - -def update_quantization_from_args(config: dict, args: argparse.Namespace): - '''update the given config dict in-place based on the command line args - ''' - if args.use_weight_only: - if args.weight_only_precision == 'int8': - config['quantization']['quant_algo'] = 'W8A16' - elif args.weight_only_precision == 'int4': - config['quantization']['quant_algo'] = 'W4A16' - elif args.smoothquant: - config['quantization']['sq_use_plugin'] = True - if args.per_channel: - if args.per_token: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN' - else: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN' - else: - if args.per_token: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN' - else: - config['quantization'][ - 'quant_algo'] = 'W8A8_SQ_PER_TENSOR_PLUGIN' - - if args.int8_kv_cache: - config['quantization']['kv_cache_quant_algo'] = 'INT8' - - if args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - "group_size": args.group_size, - "has_zero_point": True, - "pre_quant_scale": False, - 'quant_algo': 'W4A16_GPTQ' - }) - - -def create_config_from_args(args: argparse.Namespace): - config = { - 'architecture': args.architecture, - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': args.n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'learned_absolute', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'norm_epsilon': args.rms_norm_eps, - 'quantization': { - 'quant_algo': None, - 'kv_cache_quant_algo': None, - }, - 'mapping': { - 'world_size': args.tp_size * args.pp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'use_prompt_tuning': args.use_prompt_tuning, - 'enable_pos_shift': args.enable_pos_shift, - 'dense_context_fmha': args.dense_context_fmha, - } - update_quantization_from_args(config, args) - return config - - -def smooth_quant(model, args): - assert model is not None - act_range = {} - llama_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - llama_smoother = {} - - os.environ["TOKENIZERS_PARALLELISM"] = os.environ.get( - "TOKENIZERS_PARALLELISM", "false") - if args.load_model_on_cpu: - logger.warning( - "Note that running capture_activation_range on cpu would be very slow." - ) - tokenizer = AutoTokenizer.from_pretrained(args.model_dir, - trust_remote_code=True, - use_fast=False, - padding_side='left') - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - - act_range = capture_activation_range(model, tokenizer, dataset) - if args.smoothquant is not None: - smooth_llama_model(model, act_range, args.smoothquant, llama_qkv_para, - llama_smoother) - return act_range, llama_qkv_para, llama_smoother - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - logger.info(tensorrt_llm.__version__) - args = parse_arguments() - if args.model_dir is None and args.meta_ckpt_dir is None: - raise AssertionError( - "One of the model_dir or meta_ckpt_dir must be specified to generate the checkpoint" - ) - - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - hf_config = None - if args.model_dir is not None: - hf_config = AutoConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - if hf_config.model_type == "llava": - # LLaVA = Vision model + Llama LLM - # We load a llava config and use its' text config as llama config - hf_config = LlavaConfig.from_pretrained(args.model_dir).text_config - hf_config.model_type = "llava" # Replace llama with llava - - if hf_config.architectures[0] == "CogVLMForCausalLM": - hf_config.model_type = 'cogvlm' - args.model_type = hf_config.model_type - args.n_head = hf_config.num_attention_heads - args.inter_size = hf_config.intermediate_size - args.n_layer = hf_config.num_hidden_layers - args.n_embd = hf_config.hidden_size - if hasattr(hf_config, "num_key_value_heads"): - args.n_kv_head = hf_config.num_key_value_heads - if args.n_kv_head is None: - args.n_kv_head = args.n_head - args.rms_norm_eps = hf_config.rms_norm_eps - args.vocab_size = hf_config.vocab_size - args.n_positions = hf_config.max_position_embeddings - - args.architecture = hf_config.architectures[0] - - elif args.meta_ckpt_dir is not None: - with open(Path(args.meta_ckpt_dir, "params.json")) as fp: - meta_config: dict = json.load(fp) - args.n_embd = meta_config["dim"] - args.n_head = meta_config["n_heads"] - args.n_layer = meta_config["n_layers"] - args.n_kv_head = meta_config.get("n_kv_heads", args.n_head) - - if "hidden_dim" in meta_config: - args.inter_size = meta_config["hidden_dim"] - else: - args.multiple_of = meta_config.get("multiple_of", 1) - n_embd = int(4 * args.n_embd * 2 / 3) - args.ffn_dim_multiplier = meta_config.get("ffn_dim_multiplier", 1) - args.inter_size = args.multiple_of * ( - (int(n_embd * args.ffn_dim_multiplier) + args.multiple_of - 1) - // args.multiple_of) - args.rms_norm_eps = meta_config["norm_eps"] - args.architecture = "LlamaForCausalLM" - else: - args.n_kv_head = args.n_kv_head or args.n_head - args.architecture = "LlamaForCausalLM" - - if args.rotary_scaling is not None: - # assert args.use_gpt_attention_plugin, "RoPE scaling is only supported through GPT attention plugin." - rotary_scaling = { - "type": args.rotary_scaling[0], - "factor": float(args.rotary_scaling[1]) - } - assert rotary_scaling["type"] in ["linear", "dynamic"] - assert rotary_scaling["factor"] > 1.0 - args.rotary_scaling = rotary_scaling - - config = create_config_from_args(args) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - act_range = {} - llama_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - llama_smoother = {} - model = None - if args.model_dir is not None: - - if args.model_type == "llava": - hf_llava = LlavaForConditionalGeneration.from_pretrained( - args.model_dir, dtype="auto") - model = hf_llava.language_model - else: - model = AutoModelForCausalLM.from_pretrained( - args.model_dir, - device_map='auto' if not args.load_model_on_cpu else 'cpu', - dtype='auto' if not args.smoothquant else torch.float16, - trust_remote_code=True, - ) - if args.smoothquant is not None or args.int8_kv_cache: - act_range, llama_qkv_para, llama_smoother = smooth_quant( - model, args) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - weights = load_weights_from_gptq( - args.quant_ckpt_path, - PretrainedConfig.from_dict(copy.deepcopy(config)), - ) - - elif args.meta_ckpt_dir is not None: - weights = load_weights_from_meta_ckpt( - args.meta_ckpt_dir, - PretrainedConfig.from_dict(copy.deepcopy(config)), - ) - - else: - if args.load_by_shard: - weights = load_weights_from_hf_by_shard( - args.model_dir, - PretrainedConfig.from_dict(copy.deepcopy(config)), - ) - - else: - if args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - elif args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - weights = convert_hf_cogvlm( - model, - mapping, - vocab_size=args.vocab_size, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - use_gemm_woq_plugin=not args. - disable_weight_only_quant_plugin, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_smooth_quant=args.smoothquant, - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.int8_kv_cache, - act_range=act_range, - qkv_para=llama_qkv_para, - smoother=llama_smoother) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - logger.info(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/dbrx/README.md b/examples/models/contrib/dbrx/README.md deleted file mode 100644 index 85002397a83c..000000000000 --- a/examples/models/contrib/dbrx/README.md +++ /dev/null @@ -1,238 +0,0 @@ -# DBRX - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a DBRX model in TensorRT-LLM. DBRX is a leading large language model trained by Databricks. Read more details about the model: [DBRX Technical Blog](https://www.databricks.com/blog/introducing-dbrx-new-state-art-open-llm). - -- [DBRX](#dbrx) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Download weights from HuggingFace Transformers](#download-weights-from-huggingface-transformers) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [Weight Only Quantization](#weight-only-quantization) - - [INT8 KV Cache](#int8-kv-cache) - - [Run inference](#run-inference) - -## Overview - -The TensorRT LLM DBRX implementation can be found in [tensorrt_llm/models/dbrx/model.py](../../../../tensorrt_llm/models/dbrx/model.py). - -## Support Matrix - * BF16 - * FP16 - * INT8 Weight-Only - * INT4 Weight-Only - * INT8 KV CACHE - * Tensor Parallel - * Pipeline Parallel - * Expert Parallel - -## Usage - -### Download weights from HuggingFace Transformers - -Install the dependencies and setup `git-lfs`. - -```bash -# Install dependencies -# DBRX uses tiktoken as the tokenizer; make sure it is installed -pip install -r requirements.txt - -# Setup git-lfs -git lfs install -``` - -Download one or more DBRX models that you would like to build to TensorRT LLM engines. You can download from the [HuggingFace](https://huggingface.co) hub: - -```bash -# Download dbrx-base -git clone https://huggingface.co/databricks/dbrx-base - -# Download dbrx-instruct -git clone https://huggingface.co/databricks/dbrx-instruct -``` - -### Build TensorRT engine(s) - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. A DBRX model has 132B parameters, so you need at least 4 x 80GB GPUs to load the model in 16-bit precision for weight conversion. - -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is same to the number of GPUs used to run inference. Normally, `trtllm-build` uses one GPU by default, but if you have already more GPUs available at build time, you may enable parallel builds to make the engine building process faster by adding the `--workers` argument. - -Here are some examples: - -```bash -# 8-way tensor parallelism, dtype bfloat16 -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype bfloat16 \ - --tp_size 8 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/bf16/tp8 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/bf16/tp8 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/bf16/tp8 -``` - -```bash -# 8-way tensor parallelism, dtype float16 -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --tp_size 8 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/fp16/tp8 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/fp16/tp8 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/fp16/tp8 -``` - -```bash -# 4-way tensor parallelism and 2-way pipeline parallelism, dtype bfloat16 -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype bfloat16 \ - --tp_size 4 \ - --pp_size 2 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/bf16/tp4pp2 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/bf16/tp4pp2 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/bf16/tp4pp2 -``` - - -```bash -# Build DBRX with expert parallelism for DbrxExperts layer and tensor parallelism for rest -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype bfloat16 \ - --tp_size 8 \ - --moe_tp_size 1 \ - --moe_ep_size 8 \ - --workers 8 \ - --output_dir dbrx/trt_ckpt/bf16/ep8 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/bf16/ep8 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --workers 8 \ - --output_dir dbrx/trt_engines/bf16/ep8 -``` - -#### Weight Only Quantization - -[`convert_checkpoint.py`](./convert_checkpoint.py) features a `--use_weight_only` option that can enable weight-only quantization. You can further set the weight-only precision by passing `int8` or `int4` to the `--weight_only_precision` flag. - -```bash -# 4-way tensor parallelism, int8 weight-only -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --tp_size 4 \ - --workers 4 \ - --output_dir dbrx/trt_ckpt/int8-wo/tp4 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/int8-wo/tp4 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 4 \ - --output_dir dbrx/trt_engines/int8-wo/tp4 -``` - -```bash -# 4-way tensor parallelism, int4 weight-only -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --tp_size 4 \ - --workers 4 \ - --output_dir dbrx/trt_ckpt/int4-wo/tp4 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/int4-wo/tp4 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 4 \ - --output_dir dbrx/trt_engines/int4-wo/tp4 -``` - -#### INT8 KV Cache -INT8 KV cache can be enabled to reduce the memory footprint. It will bring performance gains at large batch sizes and long sequence lengths. - -For INT8 KV cache, [`convert_checkpoint.py`](./convert_checkpoint.py) features a `--int8_kv_cache` option. Setting `--int8_kv_cache` will calibrate the model, and then export the scaling factors needed for INT8 KV cache inference. - -```bash -# 4-way tensor parallelism, int8 kv-cache -python convert_checkpoint.py --model_dir dbrx-base \ - --dtype float16 \ - --int8_kv_cache \ - --tp_size 4 \ - --workers 4 \ - --output_dir dbrx/trt_ckpt/int8kv/tp4 - -trtllm-build --checkpoint_dir dbrx/trt_ckpt/int8kv/tp4 \ - --gpt_attention_plugin float16 \ - --gemm_plugin float16 \ - --moe_plugin float16 \ - --workers 4 \ - --output_dir dbrx/trt_engines/int8kv/tp4 -``` - -### Run inference - -You can test your engines with the [run.py](../../../run.py) script: - -```bash -mpirun -n 8 \ - python3 ../run.py --engine_dir dbrx/trt_engines/bf16/tp8 \ - --tokenizer_dir dbrx-base \ - --max_output_len 10 \ - --input_text "What is AGI?" -``` - -If the engines are run successfully, you will see output like: -``` -...... -Input [Text 0]: "What is AGI?" -Output [Text 0 Beam 0]: " How do I find it? -AGI stands for" -``` - - -You can also evaluate with the [summarize.py](../../../summarize.py) script: -```bash -mpirun -n 8 \ - python ../summarize.py --engine_dir dbrx/trt_engines/bf16/tp8 \ - --hf_model_dir dbrx-base \ - --test_trt_llm -``` - -If the engines are run successfully, you will see output like: -``` -...... -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM (total latency: 9.962657451629639 sec) -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1189) -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM (tokens per second: 119.34566713477734) -[04/02/2024-11:16:37] [TRT-LLM] [I] TensorRT LLM beam 0 result -[04/02/2024-11:16:37] [TRT-LLM] [I] rouge1 : 26.842471264679535 -[04/02/2024-11:16:37] [TRT-LLM] [I] rouge2 : 9.979512100961314 -[04/02/2024-11:16:37] [TRT-LLM] [I] rougeL : 19.50336050538688 -[04/02/2024-11:16:37] [TRT-LLM] [I] rougeLsum : 22.00400189383231 -``` diff --git a/examples/models/contrib/dbrx/convert_checkpoint.py b/examples/models/contrib/dbrx/convert_checkpoint.py deleted file mode 100644 index 1ca287f2588a..000000000000 --- a/examples/models/contrib/dbrx/convert_checkpoint.py +++ /dev/null @@ -1,654 +0,0 @@ -import argparse -import copy -import functools -import json -import os -import time -import traceback -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Dict, Optional, Tuple - -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer -from transformers.pytorch_utils import Conv1D - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta, release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (generate_int8, - load_calib_dataset, split) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - parser.add_argument( - '--per_channel', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument("--dataset_cache_dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--rotary_base', type=float, default=10000.0) - parser.add_argument('--rotary_scaling', nargs=2, type=str, default=None) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - parser.add_argument('--n_head', type=int, default=32) - parser.add_argument('--n_kv_head', type=int, default=None) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--inter_size', type=int, default=11008) - parser.add_argument('--max_seq_len', type=int, default=4096) - parser.add_argument('--clip_qkv', type=int, default=None) - parser.add_argument('--hidden_act', - type=str, - default='gelu', - help='Set to swiglu to use GLU in MoEs') - parser.add_argument( - '--moe_num_experts', - default=0, - type=int, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - default=0, - type=int, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MOE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MOE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_renorm_mode', - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - type=int, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values', - ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--dense_context_fmha', - default=False, - action='store_true', - help= - 'Enable dense fmha in context phase, otherwise sliding window attention.' - 'If dense_context_fmha=False, the sliding window size is the max attention window size.' - ) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--use_prompt_tuning', - action="store_true", - default=False) - args = parser.parse_args() - - return args - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin - } - - -def get_weight(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> torch.Tensor: - if f'{prefix}' in params: - return params[f'{prefix}'].to(dtype).detach().cpu() - elif f'{prefix}.weight' not in params: - return None - return params[f'{prefix}.weight'].to(dtype).detach().cpu() - - -def get_bias(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> torch.Tensor: - if f'{prefix}.bias' not in params: - return None - return params[f'{prefix}.bias'].to(dtype).detach().cpu() - - -def get_weight_and_bias(params: Dict[str, torch.Tensor], prefix: str, - dtype: torch.dtype) -> Tuple[torch.Tensor]: - return get_weight(params, prefix, dtype), get_bias(params, prefix, dtype) - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=1, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip( - 1e-8, None).max(dim=1)[0].float() - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - input_ids = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -def split_qkv_tp(qkv, n_head, n_kv_heads, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - kv_head_size = n_kv_heads * (n_hidden // n_head) - q, k, v = torch.split(qkv, [n_hidden, kv_head_size, kv_head_size], dim=0) - q = split(q, tensor_parallel, rank, dim=0) - k = split(k, tensor_parallel, rank, dim=0) - v = split(v, tensor_parallel, rank, dim=0) - return torch.concatenate([q, k, v], dim=0).contiguous() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8, - postfix='weight', - quant_scale_name=None) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - if weight.dim() > 2: - v = weight.transpose(1, 2).contiguous().clone() - else: - v = weight.t().contiguous().clone() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v.cpu(), plugin_weight_only_quant_type) - results[prefix + postfix] = processed_torch_weights - if quant_scale_name is not None: - results[quant_scale_name] = torch_weight_scales - else: - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + postfix] = weight.contiguous() - - if bias is not None: - results[f'{prefix}bias'] = bias - - return results - - -def convert_hf_dbrx(model_params: dict, - hf_config: AutoConfig, - mapping: Mapping, - dtype: str = 'float32', - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8, - moe_config: MoeConfig = None, - int8_kv_cache=False, - act_range=[]): - - weights = {} - tik = time.time() - - dtype = getattr(torch, dtype) - num_hidden_layers = hf_config.n_layers - num_head = hf_config.n_heads - num_kv_heads = hf_config.attn_config.kv_n_heads - num_hidden = hf_config.d_model - mlp_hidden_size = hf_config.ffn_config.ffn_hidden_size - layers_range = mapping.pp_layers(num_hidden_layers) - multi_query_mode = (num_kv_heads != num_head) - - for l in layers_range: - prefix = f'transformer.blocks.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # Attention QKV (no bias) - qkv_w = get_weight(model_params, f'{prefix}.norm_attn_norm.attn.Wqkv', - dtype) - qkv_w = split_qkv_tp(qkv_w, num_head, num_kv_heads, num_hidden, - mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv.', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Attention dense (no bias) - attn_dense_weight = get_weight( - model_params, f'{prefix}.norm_attn_norm.attn.out_proj', dtype) - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, - f'{tllm_prex}.attention.dense.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_weight = get_weight(model_params, - f'{prefix}.norm_attn_norm.attn.Wqkv', dtype) - qkv_weight = qkv_weight.t() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(num_hidden, 3, num_hidden) - int8_weights = generate_int8( - qkv_weight, - act_range.get(f'{prefix}.norm_attn_norm.attn.Wqkv'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[ - f'{tllm_prex}.attention.kv_cache_scaling_factor'] = int8_weights[ - 'scale_y_quant_orig'].contiguous() - - # input layer_norm - input_ln_weight = get_weight(model_params, - f'{prefix}.norm_attn_norm.norm_1', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, - f'{prefix}.norm_attn_norm.norm_2', dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - - if moe_config and moe_config.has_moe(): - # experts mlp w1 -> mlp gate - mlp_gate_weight = get_weight(model_params, - f'{prefix}.ffn.experts.mlp.w1', dtype) - mlp_gate_weight = mlp_gate_weight.reshape(-1, mlp_hidden_size, - num_hidden) - # moe expert parallel - mlp_gate_weight = split_matrix(mlp_gate_weight, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - # moe tensor parallel - mlp_gate_w = split_matrix(mlp_gate_weight, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - - # experts mlp v1 -> mlp fc - mlp_fc_weight = get_weight(model_params, - f'{prefix}.ffn.experts.mlp.v1', dtype) - mlp_fc_weight = mlp_fc_weight.reshape(-1, mlp_hidden_size, - num_hidden) - # moe expert parallel - mlp_fc_weight = split_matrix(mlp_fc_weight, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - # moe tensor parallel - mlp_fc_w = split_matrix(mlp_fc_weight, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=1) - mlp_fc_w = torch.concat([mlp_fc_w, mlp_gate_w], dim=-2) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc.', None, - use_weight_only, - plugin_weight_only_quant_type)) - - # experts mlp w2 -> mlp proj - mlp_proj_weight = get_weight(model_params, - f'{prefix}.ffn.experts.mlp.w2', dtype) - mlp_proj_weight = mlp_proj_weight.reshape(-1, mlp_hidden_size, - num_hidden).transpose( - 1, 2) - # moe expert parallel - mlp_proj_weight = split_matrix(mlp_proj_weight, - mapping.moe_ep_size, - mapping.moe_ep_rank, - dim=0) - # moe tensor parallel - mlp_proj_w = split_matrix(mlp_proj_weight, - mapping.moe_tp_size, - mapping.moe_tp_rank, - dim=2) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj.', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # router mlp - router_weights = get_weight(model_params, - f'{prefix}.ffn.router.layer', - torch.float32) - weights[f'{tllm_prex}.mlp.router.weight'] = router_weights - - embed_w = get_weight(model_params, 'transformer.wte', dtype) - lm_head = get_weight(model_params, 'lm_head', dtype) - if mapping.is_first_pp_rank(): - # Embedding - weights['transformer.vocab_embedding.weight'] = embed_w - if mapping.is_last_pp_rank(): - if lm_head is None: - lm_head = embed_w.clone() - ln_f_w = get_weight(model_params, 'transformer.norm_f', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - weights['lm_head.weight'] = split_matrix(lm_head, - mapping.tp_size, - mapping.tp_rank, - dim=0) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def execute(workers, func, hf_model): - if workers == 1: - for rank, f in enumerate(func): - f(hf_model, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [ - p.submit(f, hf_model, rank) for rank, f in enumerate(func) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - quant_algo = None - kv_cache_quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only: - if args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - - if args.int8_kv_cache: - kv_cache_quant_algo = QuantAlgo.INT8 - - hf_config = None - - if args.model_dir is not None: - hf_config = AutoConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - args.n_kv_head = hf_config.attn_config.kv_n_heads - args.n_layer = hf_config.n_layers - args.n_head = hf_config.n_heads - args.vocab_size = hf_config.vocab_size - args.n_embd = hf_config.d_model - args.inter_size = hf_config.ffn_config.ffn_hidden_size - args.max_seq_len = hf_config.max_seq_len - args.moe_num_experts = getattr(hf_config.ffn_config, "moe_num_experts", - 0) - args.moe_top_k = getattr(hf_config.ffn_config, "moe_top_k", 0) - if args.moe_num_experts and args.moe_top_k == 0: - args.moe_top_k = 1 - args.clip_qkv = hf_config.attn_config.clip_qkv - args.hidden_act = 'swiglu' - args.rotary_base = get_hf_rope_theta(hf_config.attn_config, 10000.0) - args.moe_config = MoeConfig( - num_experts=args.moe_num_experts, - top_k=args.moe_top_k, - normalization_mode=args.moe_renorm_mode).validate() - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': args.logits_dtype, - 'vocab_size': args.vocab_size, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'num_key_value_heads': args.n_kv_head, - 'max_position_embeddings': args.max_seq_len, - 'norm_epsilon': 1e-5, - 'position_embedding_type': 'rope_gpt_neox', - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'rotary_scaling': args.rotary_scaling, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'moe': { - "num_experts": args.moe_num_experts, - "top_k": args.moe_top_k, - "normalization_mode": args.moe_renorm_mode - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - 'moe_tp_size': args.moe_tp_size, - 'moe_ep_size': args.moe_ep_size, - }, - 'clip_qkv': args.clip_qkv, - 'dense_context_fmha': args.dense_context_fmha, - } - - config.update(args_to_build_options(args)) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def load_from_hf(model_dir): - hf_model = AutoModelForCausalLM.from_pretrained(model_dir, - trust_remote_code=True, - device_map="auto", - dtype=getattr( - torch, args.dtype), - config=hf_config) - return hf_model - - def convert_and_save(hf_model, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - act_range = {} - if args.int8_kv_cache: - tokenizer = AutoTokenizer.from_pretrained(args.model_dir, - padding_side='left', - trust_remote_code=True) - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - act_range = capture_activation_range(hf_model, tokenizer, dataset) - - hf_model = dict(hf_model.named_parameters()) - weights = convert_hf_dbrx( - hf_model, - hf_config, - mapping, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - moe_config=args.moe_config, - int8_kv_cache=args.int8_kv_cache, - act_range=act_range) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - del weights - release_gc() - - if args.model_dir: - hf_model = load_from_hf(args.model_dir) - execute(args.workers, [convert_and_save] * world_size, hf_model) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/dbrx/requirements.txt b/examples/models/contrib/dbrx/requirements.txt deleted file mode 100644 index 0b8f8e5e0f4f..000000000000 --- a/examples/models/contrib/dbrx/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -tiktoken==0.6.0 diff --git a/examples/models/contrib/deepseek_v1/README.md b/examples/models/contrib/deepseek_v1/README.md deleted file mode 100755 index 4a49c8266705..000000000000 --- a/examples/models/contrib/deepseek_v1/README.md +++ /dev/null @@ -1,103 +0,0 @@ -# Deepseek-v1 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run [deepseek-v1](https://arxiv.org/pdf/2401.06066) model in TensorRT-LLM. - -- [Deepseek-v1](#deepseek-v1) - - [Prerequisite](#prerequistie) - - [Hardware](#hardware) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - -## Prerequisite - -First, please download Deepseek-v1 weights from HF https://huggingface.co/deepseek-ai/deepseek-moe-16b-base. - -```bash -git lfs install -git clone https://huggingface.co/deepseek-ai/deepseek-moe-16b-base -``` - -## Hardware - -The Deepseek-v1 model requires 1x80G GPU memory. - -## Overview - -The TensorRT LLM Deepseek-v1 implementation can be found in [tensorrt_llm/models/deepseek_v1/model.py](../../tensorrt_llm/models/deepseek_v1/model.py). The TensorRT LLM Deepseek-v1 example code is located in [`examples/models/contrib/deepseek_v1`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the Deepseek-v1 model into TensorRT LLM checkpoint format. - -In addition, there are three shared files in the parent folder [`examples`](../../../) can be used for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the model inference output by given an input text. -* [`../../../summarize.py`](../../../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset, it can running the summarize from HF model and TensorRT LLM model. -* [`../../../mmlu.py`](../../../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT LLM model on the MMLU dataset. - -## Support Matrix - -- [x] FP16 -- [x] TENSOR PARALLEL -- [x] FP8 - -## Usage - -The TensorRT LLM Deepseek-v1 example code locates at [examples/models/contrib/deepseek_v1](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Below is the step-by-step to run Deepseek-v1 with TensorRT-LLM. - -First the checkpoint will be converted to the TensorRT LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be build with TensorRT LLM checkpoint. - -```bash -# Build the bfloat16 engine from Deepseek-v1 HF weights. -python convert_checkpoint.py --model_dir ./deepseek_moe_16b/ \ - --output_dir ./trtllm_checkpoint_deepseek_v1_1gpu_bf16 \ - --dtype bfloat16 \ - --tp_size 1 -trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v1_1gpu_bf16 \ - --output_dir ./trtllm_engines/deepseek_v1/bf16/tp1 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ -``` - -Then, test the engine with [run.py](../../../run.py) script: - -```bash -python ../../../run.py --engine_dir ./trtllm_engines/deepseek_v1/bf16/tp1 \ - --tokenizer_dir ./deepseek_moe_16b/ \ - --max_output_len 32 \ - --top_p 0 \ - --input_text "The president of the United States is person who" -``` - -### FP8 Quantization - -The [`../../../quantization/quantize.py`](../../../quantization/quantize.py) script can be used to quantize the models and export TensorRT LLM checkpoints. - -```bash -# Deepseek-v1: single gpu, fp8 quantization -python ../../../quantization/quantize.py --model_dir deepseek_moe_16b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir trt_ckpt/deepseek_moe_16b/fp8/1-gpu \ - --calib_size 512 - -# Deepseek-v1: single-gpu engine with fp8 quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir ./trt_ckpt/deepseek_moe_16b/fp8/1-gpu \ - --gemm_plugin float16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./trt_engines/fp8/1-gpu/ -``` -## Credits -This Deepseek-v1 model example exists thanks to @akhoroshev(https://github.com/akhoroshev) community contribution! diff --git a/examples/models/contrib/deepseek_v1/__init__.py b/examples/models/contrib/deepseek_v1/__init__.py deleted file mode 100644 index 71bf6d298d42..000000000000 --- a/examples/models/contrib/deepseek_v1/__init__.py +++ /dev/null @@ -1,14 +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. diff --git a/examples/models/contrib/deepseek_v1/convert_checkpoint.py b/examples/models/contrib/deepseek_v1/convert_checkpoint.py deleted file mode 100644 index 3e7b1aaf2825..000000000000 --- a/examples/models/contrib/deepseek_v1/convert_checkpoint.py +++ /dev/null @@ -1,210 +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. -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import DeepseekForCausalLM - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None, required=True) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MoE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MoE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') - parser.add_argument('--output_dir', - type=str, - default='trtllm_checkpoint', - required=True, - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--moe_num_experts', - type=int, - default=0, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - type=int, - default=0, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_renorm_mode', - type=int, - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values' - ) - parser.add_argument( - '--save_config_only', - action="store_true", - default=False, - help= - 'Only save the model config w/o read and converting weights, be careful, this is for debug only' - ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact' - ) - # Add quantization related feature later - args = parser.parse_args() - - return args - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin - } - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - deepseekv1 = DeepseekForCausalLM.from_hugging_face( - args.model_dir, args.dtype, mapping, **override_fields) - deepseekv1.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del deepseekv1 - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - - tik = time.time() - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/deepseek_v1/requirements.txt b/examples/models/contrib/deepseek_v1/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/deepseek_v1/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/deepseek_v2/README.md b/examples/models/contrib/deepseek_v2/README.md deleted file mode 100644 index aefbeec4120a..000000000000 --- a/examples/models/contrib/deepseek_v2/README.md +++ /dev/null @@ -1,267 +0,0 @@ -# Deepseek-v2 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run [deepseek-v2](https://arxiv.org/pdf/2405.04434) model in TensorRT-LLM. - -- [Deepseek-v2](#deepseek-v2) - - [Prerequisite](#prerequisite) - - [Hardware](#hardware) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - -## Prerequisite - -First, please download Deepseek-v2 weights from HF https://huggingface.co/deepseek-ai/DeepSeek-V2. - -```bash -git lfs install -git clone https://huggingface.co/deepseek-ai/DeepSeek-V2 -``` - -## Hardware - -The Deepseek-v2 model requires least 8x80G GPU memory, model contains 236B parameters roughly 472GB memory (with BF16 precision). - -***Caution: Current TRT-LLM MLA kernel only support Hopper architecture (SM90). Ampere architecture (SM80 & SM86) will be supported in next release.*** - -## Overview - -The TensorRT LLM Deepseek-v2 implementation can be found in [tensorrt_llm/models/deepseek_v2/model.py](../../../../tensorrt_llm/models/deepseek_v2/model.py). The TensorRT LLM Deepseek-v2 example code is located in [`examples/models/contrib/deepseek_v2`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the Deepseek-v2 model into TensorRT LLM checkpoint format. - -In addition, there are three shared files in the parent folder [`examples`](../../../) can be used for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the model inference output by given an input text. -* [`../../../summarize.py`](../../../summarize.py) to summarize the article from [cnn_dailmail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset, it can running the summarize from HF model and TensorRT LLM model. -* [`../../../mmlu.py`](../../../mmlu.py) to running score script from https://github.com/declare-lab/instruct-eval to compare HF model and TensorRT LLM model on the MMLU dataset. - -## Support Matrix - -- [x] BF16 -- [ ] FP8 - -***Caution: prefer using BF16 over FP16 for Deepseek-v2 since model original training precision is BF16 and we found direct convert BF16 -> FP16 will suffer accuracy drop.*** - -## Usage - -The TensorRT LLM Deepseek-v2 example code locates at [examples/models/contrib/deepseek_v2](./). It takes PyTorch weights as input, and builds corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Below is the step-by-step to run Deepseek-v2 with TensorRT-LLM. - -First the checkpoint will be converted to the TensorRT LLM checkpoint format by apply [`convert_checkpoint.py`](./convert_checkpoint.py). After that, the TensorRT engine(s) can be build with TensorRT LLM checkpoint. - -```bash -# Convert Deepseek-v2 HF weights to TensorRT LLM checkpoint format. -python convert_checkpoint.py --model_dir ./DeepSeek-V2 \ - --output_dir ./trtllm_checkpoint_deepseek_v2_8gpu_bf16 \ - --dtype bfloat16 \ - --tp_size 8 - -# With '--load_model_on_cpu' option if total GPU memory is insufficient -python convert_checkpoint.py --model_dir ./DeepSeek-V2 \ - --output_dir ./trtllm_checkpoint_deepseek_v2_cpu_bf16 \ - --dtype bfloat16 \ - --tp_size 8 \ - --load_model_on_cpu -``` - - -We observe use GPUs(8xH200) the checkpoint conversion time took ~ 34 mints, while use CPUs took ~ 21 mints and CPU memory required >= 770GB. - -After the checkpoint conversion, the TensorRT engine(s) can be built with the TensorRT LLM checkpoint. - -```bash -# Build engine -trtllm-build --checkpoint_dir ./trtllm_checkpoint_deepseek_v2_8gpu_bf16 \ - --output_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --max_batch_size 4 \ - --max_seq_len 4096 \ - --max_input_len 2048 \ - --use_paged_context_fmha enable -``` - -***Caution: `--max_batch_size` and `--max_seq_len` are the main factors to determine how many GPU memory will be used during runtime, so later when try to run e.g., `summarize.py` or `mmlu.py` or `gptManagerBenchmark.cpp`may need adjust `--max_batch_size` and `--max_seq_len` accordingly to avoid OOM.(meaning rebuild TensorRT engine with smaller `--max_batch_size` and `--max_seq_len` if needed based on GPU memory size), there is beautiful technical log perf-best-practices.md (https://github.com/NVIDIA/TensorRT-LLM/blob/main/docs/source/performance/perf-best-practices.md) explained the mechanism.*** - -Test the engine with [run.py](../../../run.py) script: - -```bash -mpirun --allow-run-as-root -n 8 python ../../../run.py --engine_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --tokenizer_dir ./DeepSeek-V2 \ - --max_output_len 40 \ - --input_text "The president of the United States is person who" -``` - -and the output will be like: - -``` -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31490468978882 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31163835525513 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31164216995239 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31491041183472 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3116364479065 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118085861206 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.3118691444397 sec -[10/28/2024-15:03:14] [TRT-LLM] [I] Load engine takes: 78.31516337394714 sec -Input [Text 0]: "<｜begin▁of▁sentence｜>The president of the United States is person who" -Output [Text 0 Beam 0]: " is elected by the people of the United States to lead the country. The president is the head of the executive branch of the government. The president is also the commander in chief of the armed forces." -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -[TensorRT-LLM][INFO] Refreshed the MPI local session -``` - -If we want to evaluate the model summarization ability, we can use [summarize.py](../../../summarize.py) script: - -```bash -mpirun --allow-run-as-root -n 8 python ../../../summarize.py --engine_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --hf_model_dir ./DeepSeek-V2 \ - --data_type bfloat16 \ - --batch_size 1 \ - --test_trt_llm \ - --test_hf -``` - -and the output will be like: - - -``` -[10/28/2024-16:46:22] [TRT-LLM] [I] HF Generated : -[10/28/2024-16:46:22] [TRT-LLM] [I] Input : ['(CNN)James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88. Best died in hospice in Hickory, North Carolina, of complications from pneumonia, said Steve Latshaw, a longtime friend and Hollywood colleague. Although he\'d been a busy actor for decades in theater and in Hollywood, Best didn\'t become famous until 1979, when "The Dukes of Hazzard\'s" cornpone charms began beaming into millions of American homes almost every Friday night. For seven seasons, Best\'s Rosco P. Coltrane chased the moonshine-running Duke boys back and forth across the back roads of fictitious Hazzard County, Georgia, although his "hot pursuit" usually ended with him crashing his patrol car. Although Rosco was slow-witted and corrupt, Best gave him a childlike enthusiasm that got laughs and made him endearing. His character became known for his distinctive "kew-kew-kew" chuckle and for goofy catchphrases such as "cuff \'em and stuff \'em!" upon making an arrest. Among the most popular shows on TV in the early \'80s, "The Dukes of Hazzard" ran until 1985 and spawned TV movies, an animated series and video games. Several of Best\'s "Hazzard" co-stars paid tribute to the late actor on social media. "I laughed and learned more from Jimmie in one hour than from anyone else in a whole year," co-star John Schneider, who played Bo Duke, said on Twitter. "Give Uncle Jesse my love when you see him dear friend." "Jimmy Best was the most constantly creative person I have ever known," said Ben Jones, who played mechanic Cooter on the show, in a Facebook post. "Every minute of his long life was spent acting, writing, producing, painting, teaching, fishing, or involved in another of his life\'s many passions." Born Jewel Guy on July 26, 1926, in Powderly, Kentucky, Best was orphaned at 3 and adopted by Armen and Essa Best, who renamed him James and raised him in rural Indiana. Best served in the Army during World War II before launching his acting career. In the 1950s and 1960s, he accumulated scores of credits, playing a range of colorful supporting characters in such TV shows as "The Twilight Zone," "Bonanza," "The Andy Griffith Show" and "Gunsmoke." He later appeared in a handful of Burt Reynolds\' movies, including "Hooper" and "The End." But Best will always be best known for his "Hazzard" role, which lives on in reruns. "Jimmie was my teacher, mentor, close friend and collaborator for 26 years," Latshaw said. "I directed two of his feature films, including the recent \'Return of the Killer Shrews,\' a sequel he co-wrote and was quite proud of as he had made the first one more than 50 years earlier." People we\'ve lost in 2015 . CNN\'s Stella Chan contributed to this story.'] -[10/28/2024-16:46:22] [TRT-LLM] [I] - Reference : ['James Best, who played the sheriff on "The Dukes of Hazzard," died Monday at 88 .\n"Hazzard" ran from 1979 to 1985 and was among the most popular shows on TV .'] -[10/28/2024-16:46:22] [TRT-LLM] [I] - Output : [[' James Best, best known for his portrayal of bumbling sheriff Rosco P. Coltrane on TV\'s "The Dukes of Hazzard," died Monday after a brief illness. He was 88.']] -[10/28/2024-16:46:22] [TRT-LLM] [I] --------------------------------------------------------- -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM (total latency: 32.02327513694763 sec) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1394) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM (tokens per second: 43.53083793080361) -[10/28/2024-16:49:33] [TRT-LLM] [I] TensorRT LLM beam 0 result -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 17.85755990133811 -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.273032755727469 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 14.768323033457317 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 15.700915348496391 -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total latency: 189.76398921012878 sec) -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (total output tokens: 1376) -[10/28/2024-16:49:33] [TRT-LLM] [I] Hugging Face (tokens per second: 7.2511123197159) -[10/28/2024-16:49:33] [TRT-LLM] [I] HF beam 0 result -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge1 : 18.542590123197257 -[10/28/2024-16:49:33] [TRT-LLM] [I] rouge2 : 6.345777100488389 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeL : 15.235695878419156 -[10/28/2024-16:49:33] [TRT-LLM] [I] rougeLsum : 16.64935135356226 -``` - -At last, we can evaluate the model with [mmlu.py](../../../mmlu.py) script: - -```bash -# Download MMLU dataset -mkdir mmlu_data && cd mmlu_data -wget https://people.eecs.berkeley.edu/~hendrycks/data.tar && tar -xf data.tar -# Run MMLU evaluation -mpirun --allow-run-as-root -n 8 python ../../../mmlu.py --engine_dir ./trtllm_engines/deepseek_v2/bf16/tp8-sel4096-isl2048-bs4 \ - --hf_model_dir ./DeepSeek-V2 \ - --data_type bfloat16 \ - --batch_size 1 \ - --test_trt_llm \ - --test_hf \ - --data_dir ./mmlu_data/data/ -``` - -and the output will be like: - -``` -Average accuracy 0.480 - abstract_algebra -Average accuracy 0.741 - anatomy -Average accuracy 0.888 - astronomy -Average accuracy 0.790 - business_ethics -Average accuracy 0.845 - clinical_knowledge -Average accuracy 0.924 - college_biology -Average accuracy 0.600 - college_chemistry -Average accuracy 0.720 - college_computer_science -Average accuracy 0.510 - college_mathematics -Average accuracy 0.751 - college_medicine -Average accuracy 0.618 - college_physics -Average accuracy 0.860 - computer_security -Average accuracy 0.796 - conceptual_physics -Average accuracy 0.675 - econometrics -Average accuracy 0.800 - electrical_engineering -Average accuracy 0.741 - elementary_mathematics -Average accuracy 0.643 - formal_logic -Average accuracy 0.570 - global_facts -Average accuracy 0.910 - high_school_biology -Average accuracy 0.714 - high_school_chemistry -Average accuracy 0.890 - high_school_computer_science -Average accuracy 0.891 - high_school_european_history -Average accuracy 0.909 - high_school_geography -Average accuracy 0.953 - high_school_government_and_politics -Average accuracy 0.826 - high_school_macroeconomics -Average accuracy 0.522 - high_school_mathematics -Average accuracy 0.916 - high_school_microeconomics -Average accuracy 0.576 - high_school_physics -Average accuracy 0.923 - high_school_psychology -Average accuracy 0.704 - high_school_statistics -Average accuracy 0.907 - high_school_us_history -Average accuracy 0.932 - high_school_world_history -Average accuracy 0.834 - human_aging -Average accuracy 0.893 - human_sexuality -Average accuracy 0.909 - international_law -Average accuracy 0.880 - jurisprudence -Average accuracy 0.853 - logical_fallacies -Average accuracy 0.598 - machine_learning -Average accuracy 0.874 - management -Average accuracy 0.953 - marketing -Average accuracy 0.880 - medical_genetics -Average accuracy 0.920 - miscellaneous -Average accuracy 0.850 - moral_disputes -Average accuracy 0.613 - moral_scenarios -Average accuracy 0.830 - nutrition -Average accuracy 0.859 - philosophy -Average accuracy 0.883 - prehistory -Average accuracy 0.635 - professional_accounting -Average accuracy 0.625 - professional_law -Average accuracy 0.857 - professional_medicine -Average accuracy 0.833 - professional_psychology -Average accuracy 0.745 - public_relations -Average accuracy 0.869 - security_studies -Average accuracy 0.935 - sociology -Average accuracy 0.930 - us_foreign_policy -Average accuracy 0.596 - virology -Average accuracy 0.877 - world_religions -Average accuracy 0.632 - math -Average accuracy 0.801 - health -Average accuracy 0.738 - physics -Average accuracy 0.897 - business -Average accuracy 0.914 - biology -Average accuracy 0.677 - chemistry -Average accuracy 0.762 - computer science -Average accuracy 0.832 - economics -Average accuracy 0.800 - engineering -Average accuracy 0.736 - philosophy -Average accuracy 0.821 - other -Average accuracy 0.902 - history -Average accuracy 0.909 - geography -Average accuracy 0.883 - politics -Average accuracy 0.876 - psychology -Average accuracy 0.919 - culture -Average accuracy 0.660 - law -Average accuracy 0.727 - STEM -Average accuracy 0.740 - humanities -Average accuracy 0.873 - social sciences -Average accuracy 0.821 - other (business, health, misc.) -Average accuracy: 0.785 -``` diff --git a/examples/models/contrib/deepseek_v2/convert_checkpoint.py b/examples/models/contrib/deepseek_v2/convert_checkpoint.py deleted file mode 100755 index 5a0c6a7df330..000000000000 --- a/examples/models/contrib/deepseek_v2/convert_checkpoint.py +++ /dev/null @@ -1,231 +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. -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import DeepseekV2ForCausalLM -from tensorrt_llm.models.deepseek_v2.convert import load_hf_deepseek - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None, required=True) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MoE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MoE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--load_model_on_cpu', - default=False, - action="store_true", - help='Choose to load HF cpkt into CPU') - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0)' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim=0') - parser.add_argument('--output_dir', - type=str, - default='trtllm_checkpoint', - required=True, - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--moe_num_experts', - type=int, - default=0, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - type=int, - default=0, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_renorm_mode', - type=int, - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values' - ) - parser.add_argument('--use_preloading', - action="store_true", - default=False, - help='Loading weights from HF model') - parser.add_argument('--use_safetensors_loading', - action="store_true", - default=False, - help='Loading weights from HF safetensors') - parser.add_argument( - '--save_config_only', - action="store_true", - default=False, - help= - 'Only save the model config w/o read and converting weights, be careful, this is for debug only' - ) - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact' - ) - # Add quantization related feature later - args = parser.parse_args() - - return args - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin, - 'load_model_on_cpu': args.load_model_on_cpu - } - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - use_preloading = args.use_preloading - use_safetensors_loading = args.use_safetensors_loading - - args.load_model_on_cpu - - if os.environ.get("TRTLLM_DISABLE_UNIFIED_CONVERTER" - ) is not None and not use_safetensors_loading: - hf_model = load_hf_deepseek(args.model_dir, args.load_model_on_cpu) - else: - hf_model = None - - def convert_and_save_rank(args, rank): - - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - deepseekv2 = DeepseekV2ForCausalLM.from_hugging_face( - args.model_dir, args.dtype, hf_model, use_preloading, - use_safetensors_loading, mapping, **override_fields) - deepseekv2.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del deepseekv2 - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - - tik = time.time() - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/deepseek_v2/requirements.txt b/examples/models/contrib/deepseek_v2/requirements.txt deleted file mode 100644 index 6be396146581..000000000000 --- a/examples/models/contrib/deepseek_v2/requirements.txt +++ /dev/null @@ -1,3 +0,0 @@ -datasets==3.1.0 -evaluate~=0.4.1 -rouge_score~=0.1.2 diff --git a/examples/models/contrib/dit/README.md b/examples/models/contrib/dit/README.md deleted file mode 100644 index 868ecbf8adef..000000000000 --- a/examples/models/contrib/dit/README.md +++ /dev/null @@ -1,133 +0,0 @@ -# DiT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a [DiT](https://arxiv.org/abs/2212.09748) model with TensorRT-LLM. - -## Overview - -The TensorRT LLM DiT implementation can be found in [tensorrt_llm/models/dit/model.py](../../../../tensorrt_llm/models/dit/model.py). The TensorRT LLM DiT example code is located in [`examples/dit`](./). There are main files to build and run DiT with TensorRT-LLM: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the DiT model into TensorRT LLM checkpoint format. -* [`sample.py`](./sample.py) to generate images with TensorRT engine(s). - -## Support Matrix - -- [x] FP16 -- [x] TP -- [x] FP8 -- [x] CP - -## Usage - -The TensorRT LLM DiT example code locates at [examples/dit](./). It takes PyTorch weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build DiT TensorRT engine(s) - -First, download the pretrained DiT-XL/2 PyTorch checkpoint from the official pytorch implementation repo [here](https://github.com/facebookresearch/DiT/tree/main?tab=readme-ov-file#sampling--), please review its license items before use. - -This checkpoint will be converted to the TensorRT LLM checkpoint format by [`convert_checkpoint.py`](./convert_checkpoint.py). After that, we can build TensorRT engine(s) with the TensorRT LLM checkpoint. - -As for run inference with FP8 quantization, currently only linear layers are supported to be quantized. Make sure that scaling factors for weights are also stored in the quantized checkpoint. - -``` -# Convert to TRT-LLM with float16(by default) -python convert_checkpoint.py -trtllm-build --checkpoint_dir ./tllm_checkpoint/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable - -# Convert to TRT-LLM with float8 -python convert_checkpoint.py --fp8_linear --timm_ckpt= --output_dir=tllm_checkpoint_fp8 -trtllm-build --checkpoint_dir ./tllm_checkpoint_fp8/ \ - --output_dir ./engine_outputs_fp8/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable -``` - -The default quantized checkpoint is based on official DiT by Facebook. If the checkpoint is obtained with APIs of `diffusers` (provided by HuggingFace), please enable the option `--diffusers_dit`: - -``` -python convert_checkpoint.py --fp8_linear --diffusers_dit --timm_ckpt= --output_dir=tllm_checkpoint_fp8 -trtllm-build --checkpoint_dir ./tllm_checkpoint_fp8/ \ - --output_dir ./engine_outputs_fp8/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable -``` - -Set `--max_batch_size` to tell how many images at most you would like to generate. We disable `--remove_input_padding` since we don't need to padding DiT's patches. Besides, we disable `--bert_attention_plugin` for better performance, since the plugin's fmha is not supported for DiT's hidden size (72 for DiT-XL). - -After build, we can find a `./engine_output` directory, it is ready for running DiT model with TensorRT LLM now. - -### Build VAE TensorRT engine -We can further accelerate VAE decoder by TensorRT. -``` -python vae_decoder_trt.py --max_batch_size 8 -``` - -### Generate images - -A [`sample.py`](./sample.py) is provided to generated images with the optimized TensorRT engines. - -We can simply run `python sample.py`. After that, an image named `sample.png` will be generated: -![sample.png](./figs/sample.png). - -Image generated with quantized FP8 DiT: - -![sample.fp8.png](./figs/sample.fp8.png). - -By default, we set batch size to 2, you can set batch size up to `--max-batch-size`, e.g., `python sample.py --batch-size 8` - -We can also run `python sample.py --tllm_model_dir=` to apply specific engines. - -### Tensor Parallel - -We can levaerage tensor parallel to further reduce latency and memory consumption on each GPU. We can also run TP for FP8 DiT. - -``` -# build dit engine -python convert_checkpoint.py --tp_size 4 -trtllm-build --checkpoint_dir ./tllm_checkpoint/ \ - --max_batch_size 8 \ - --remove_input_padding disable \ - --bert_attention_plugin disable -# build vae engine -python vae_decoder_trt.py --max_batch_size 8 -# run -mpirun -n 4 --allow-run-as-root python sample.py -``` - -### Context Parallel - -Context parallel can also be used to reduce latency and memory consumption on each GPU. We can also run CP for FP8 DiT. - -``` -# build dit engine -python convert_checkpoint.py --cp_size 4 -trtllm-build --checkpoint_dir ./tllm_checkpoint/ --max_batch_size 8 --remove_input_padding disable --bert_attention_plugin disable -# build vae engine -python vae_decoder_trt.py --max_batch_size 8 -# run -mpirun -n 4 --allow-run-as-root python sample.py -``` - -### Combine Tensor Parallel and Context Parallel - -Tensor Parallel and Context Parallel can be used together to better balance latency and memory consumption. - -``` -# build dit engine -python convert_checkpoint.py --cp_size 2 --tp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint/ --max_batch_size 8 --remove_input_padding disable --bert_attention_plugin disable -# build vae engine -python vae_decoder_trt.py --max_batch_size 8 -# run -mpirun -n 4 --allow-run-as-root python sample.py -``` diff --git a/examples/models/contrib/dit/convert_checkpoint.py b/examples/models/contrib/dit/convert_checkpoint.py deleted file mode 100644 index e36c9f2b9c2c..000000000000 --- a/examples/models/contrib/dit/convert_checkpoint.py +++ /dev/null @@ -1,334 +0,0 @@ -import argparse -import json -import os -import re -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import safetensors.torch -import torch - -import tensorrt_llm -from tensorrt_llm import str_dtype_to_torch -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) - -FACEBOOK_DIT_NAME_MAPPING = { - '^t_embedder.mlp.0.weight$': - 't_embedder.mlp1.weight', - '^t_embedder.mlp.0.bias$': - 't_embedder.mlp1.bias', - '^t_embedder.mlp.2.weight$': - 't_embedder.mlp2.weight', - '^t_embedder.mlp.2.bias$': - 't_embedder.mlp2.bias', - '^t_embedder.mlp.0.weights_scaling_factor$': - 't_embedder.mlp1.weights_scaling_factor', - '^t_embedder.mlp.0.activation_scaling_factor$': - 't_embedder.mlp1.activation_scaling_factor', - '^t_embedder.mlp.2.weights_scaling_factor$': - 't_embedder.mlp2.weights_scaling_factor', - '^t_embedder.mlp.2.activation_scaling_factor$': - 't_embedder.mlp2.activation_scaling_factor', - - # Add negative lookhead for scaling matching. - r'^blocks.(\d+).mlp.fc1.weight$': - 'blocks.*.mlp.fc.weight', - r'^blocks.(\d+).mlp.fc1.bias$': - 'blocks.*.mlp.fc.bias', - r'^blocks.(\d+).mlp.fc2.weight$': - 'blocks.*.mlp.proj.weight', - r'^blocks.(\d+).mlp.fc2.bias$': - 'blocks.*.mlp.proj.bias', - r'^blocks.(\d+).mlp.fc1.weights_scaling_factor$': - 'blocks.*.mlp.fc.weights_scaling_factor', - r'^blocks.(\d+).mlp.fc1.activation_scaling_factor$': - 'blocks.*.mlp.fc.activation_scaling_factor', - r'^blocks.(\d+).mlp.fc2.weights_scaling_factor$': - 'blocks.*.mlp.proj.weights_scaling_factor', - r'^blocks.(\d+).mlp.fc2.activation_scaling_factor$': - 'blocks.*.mlp.proj.activation_scaling_factor', - - # Add negative lookhead for scaling matching. - r'^blocks.(\d+).attn.proj.weight$': - 'blocks.*.attn.dense.weight', - r'^blocks.(\d+).attn.proj.bias$': - 'blocks.*.attn.dense.bias', - r'^blocks.(\d+).attn.proj.weights_scaling_factor$': - 'blocks.*.attn.dense.weights_scaling_factor', - r'^blocks.(\d+).attn.proj.activation_scaling_factor$': - 'blocks.*.attn.dense.activation_scaling_factor', - r'^blocks.(\d+).adaLN_modulation.1.weight$': - 'blocks.*.adaLN_modulation.weight', - r'^blocks.(\d+).adaLN_modulation.1.bias$': - 'blocks.*.adaLN_modulation.bias', - r'^blocks.(\d+).adaLN_modulation.1.weights_scaling_factor$': - 'blocks.*.adaLN_modulation.weights_scaling_factor', - r'^blocks.(\d+).adaLN_modulation.1.activation_scaling_factor$': - 'blocks.*.adaLN_modulation.activation_scaling_factor', - '^final_layer.adaLN_modulation.1.weight$': - 'final_layer.adaLN_modulation.weight', - '^final_layer.adaLN_modulation.1.bias$': - 'final_layer.adaLN_modulation.bias', - '^final_layer.adaLN_modulation.1.weights_scaling_factor$': - 'final_layer.adaLN_modulation.weights_scaling_factor', - '^final_layer.adaLN_modulation.1.activation_scaling_factor$': - 'final_layer.adaLN_modulation.activation_scaling_factor', -} - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--timm_ckpt', - type=str, - default="./DiT-XL-2-512x512.pt") - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument('--input_size', - type=int, - default=64, - help='The input latent size') - parser.add_argument('--patch_size', - type=int, - default=2, - help='The patch size for patchify') - parser.add_argument('--in_channels', - type=int, - default=4, - help='The channels of input latent') - parser.add_argument('--hidden_size', - type=int, - default=1152, - help='The hidden size of DiT') - parser.add_argument('--depth', - type=int, - default=28, - help='The number of DiTBlock layers') - parser.add_argument('--num_heads', - type=int, - default=16, - help='The number of heads of attention module') - parser.add_argument( - '--mlp_ratio', - type=float, - default=4.0, - help= - 'The ratio of hidden size compared to input hidden size in MLP layer') - parser.add_argument( - '--class_dropout_prob', - type=float, - default=0.1, - help='The probability to drop class token when training') - parser.add_argument('--num_classes', - type=int, - default=1000, - help='The number of classes for conditional control') - parser.add_argument('--learn_sigma', - type=bool, - default=True, - help='Whether the model learn sigma') - parser.add_argument('--cfg_scale', type=float, default=4.0) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--cp_size', - type=int, - default=1, - help='Context parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--fp8_linear', - action='store_true', - help='Whether use FP8 for linear layers') - parser.add_argument( - '--diffusers_dit', - action='store_true', - help='Convert checkpoint provided by `HuggingFace/diffusers`') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - return args - - -def convert_timm_dit(args, mapping, dtype='float32'): - - weights = {} - tik = time.time() - torch_dtype = str_dtype_to_torch(dtype) - tensor_parallel = mapping.tp_size - if args.diffusers_dit and args.fp8_linear: - from utils_modelopt import remap_model - converted_ckpt = "transformer.fp8.converted.state_dict.pt" - remap_model(quantized_ckpt=args.timm_ckpt, output_ckpt=converted_ckpt) - model_params = dict(torch.load(converted_ckpt)) - else: - model_params = dict(torch.load(args.timm_ckpt)) - timm_to_trtllm_name = FACEBOOK_DIT_NAME_MAPPING - - def get_trtllm_name(timm_name): - for k, v in timm_to_trtllm_name.items(): - m = re.match(k, timm_name) - if m is not None: - if "*" in v: - v = v.replace("*", m.groups()[0]) - return v - return timm_name - - weights = dict() - for name, param in model_params.items(): - if param.dtype in [torch.int8, torch.float8_e4m3fn - ] or 'scaling_factor' in name: - if 'scaling_factor' in name: - assert param.dtype == torch.float32 - weights[get_trtllm_name(name)] = param.contiguous() - else: - weights[get_trtllm_name(name)] = param.contiguous().to(torch_dtype) - - assert len(weights) == len(model_params) - - for k, v in weights.items(): - if re.match('^blocks.*.attn.qkv.weight$', k): - weights[k] = split_qkv_tp(v, args.num_heads, args.hidden_size, - tensor_parallel, mapping.tp_rank) - elif re.match('^blocks.*.attn.qkv.bias$', k): - weights[k] = split_qkv_bias_tp(v, args.num_heads, args.hidden_size, - tensor_parallel, mapping.tp_rank) - elif re.match('^blocks.*.attn.dense.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=1) - elif re.match('^blocks.*.mlp.fc.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=0) - elif re.match('^blocks.*.mlp.fc.bias$', k): - weights[k] = split(v, tensor_parallel, mapping.tp_rank) - elif re.match('^blocks.*.mlp.proj.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=1) - elif re.match(r'.*adaLN_modulation.weight$', k): - weights[k] = split_matrix_tp(v, - tensor_parallel, - mapping.tp_rank, - dim=0) - elif re.match(r'.*adaLN_modulation.bias$', k): - weights[k] = split(v, tensor_parallel, mapping.tp_rank) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def save_config(args): - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - config = { - 'architecture': "DiT", - 'dtype': args.dtype, - 'input_size': args.input_size, - 'patch_size': args.patch_size, - 'in_channels': args.in_channels, - 'hidden_size': args.hidden_size, - 'num_hidden_layers': args.depth, - 'num_attention_heads': args.num_heads, - 'mlp_ratio': args.mlp_ratio, - 'class_dropout_prob': args.class_dropout_prob, - 'num_classes': args.num_classes, - 'learn_sigma': args.learn_sigma, - 'cfg_scale': args.cfg_scale, - 'mapping': { - 'world_size': args.cp_size * args.tp_size * args.pp_size, - 'cp_size': args.cp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - } - } - if args.fp8_linear: - config['quantization'] = { - 'quant_algo': "FP8", - # TODO: add support for exclude modules. - # 'exclude_modules': ["*final_layer*"], - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - -def covert_and_save(args, rank): - if rank == 0: - save_config(args) - - mapping = Mapping(world_size=args.cp_size * args.tp_size * args.pp_size, - rank=rank, - cp_size=args.cp_size, - tp_size=args.tp_size, - pp_size=args.pp_size) - - weights = convert_timm_dit(args, mapping, dtype=args.dtype) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.cp_size * args.tp_size * args.pp_size - - assert args.pp_size == 1, "PP is not supported yet." - - tik = time.time() - - if args.timm_ckpt is None: - return - - execute(args.workers, [covert_and_save] * world_size, args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/dit/diffusion.py b/examples/models/contrib/dit/diffusion.py deleted file mode 100644 index 510001b2ede2..000000000000 --- a/examples/models/contrib/dit/diffusion.py +++ /dev/null @@ -1,116 +0,0 @@ -import numpy as np -import torch -from tqdm.auto import tqdm - - -def space_timesteps(num_timesteps, timestep_respacing): - if num_timesteps < timestep_respacing: - raise ValueError( - f"cannot divide section of {num_timesteps} steps into {timestep_respacing}" - ) - frac_stride = (num_timesteps - 1) / (timestep_respacing - 1) - cur_idx = 0.0 - taken_steps = [] - for _ in range(timestep_respacing): - taken_steps.append(round(cur_idx)) - cur_idx += frac_stride - return set(taken_steps) - - -def _extract_into_tensor(arr, timesteps, broadcast_shape): - res = torch.from_numpy(arr).to(device=timesteps.device)[timesteps].float() - while len(res.shape) < len(broadcast_shape): - res = res[..., None] - return res + torch.zeros(broadcast_shape, device=timesteps.device) - - -class DiTDiffusionPipeline: - - def __init__(self, dit_model, timestep_respacing, diffusion_steps=1000): - self.dit_model = dit_model - self.timesteps = space_timesteps(diffusion_steps, timestep_respacing) - self.timestep_map = [] - self.original_num_steps = diffusion_steps - betas = self._setup_betas() - self.betas = betas - self.num_timesteps = int(betas.shape[0]) - - alphas = 1.0 - betas - self.alphas_cumprod = np.cumprod(alphas, axis=0) - self.alphas_cumprod_prev = np.append(1.0, self.alphas_cumprod[:-1]) - self.alphas_cumprod_next = np.append(self.alphas_cumprod[1:], 0.0) - assert self.alphas_cumprod_prev.shape == (self.num_timesteps, ) - - self.sqrt_recip_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod) - self.sqrt_recipm1_alphas_cumprod = np.sqrt(1.0 / self.alphas_cumprod - - 1) - - self.posterior_variance = (betas * (1.0 - self.alphas_cumprod_prev) / - (1.0 - self.alphas_cumprod)) - - self.posterior_log_variance_clipped = np.log( - np.append(self.posterior_variance[1], - self.posterior_variance[1:])) if len( - self.posterior_variance) > 1 else np.array([]) - - self.posterior_mean_coef1 = (betas * np.sqrt(self.alphas_cumprod_prev) / - (1.0 - self.alphas_cumprod)) - self.posterior_mean_coef2 = ((1.0 - self.alphas_cumprod_prev) * - np.sqrt(alphas) / - (1.0 - self.alphas_cumprod)) - - def _setup_betas(self): - scale = 1000 / self.original_num_steps - betas = np.linspace(scale * 0.0001, - scale * 0.02, - self.original_num_steps, - dtype=np.float64) - last_alpha_cumprod = 1.0 - new_betas = [] - alphas = 1.0 - betas - alphas_cumprod = np.cumprod(alphas, axis=0) - for i, alpha_cumprod in enumerate(alphas_cumprod): - if i in self.timesteps: - new_betas.append(1 - alpha_cumprod / last_alpha_cumprod) - last_alpha_cumprod = alpha_cumprod - self.timestep_map.append(i) - betas = np.array(new_betas) - return betas - - def run(self, sample, labels): - indices = list(range(self.num_timesteps))[::-1] - - for i in tqdm(indices): - t = torch.tensor([i] * sample.shape[0], device=sample.device) - B, C = sample.shape[:2] - assert t.shape == (B, ) - map_tensor = torch.tensor(self.timestep_map, - device=t.device, - dtype=t.dtype) - model_output = self.dit_model(sample, map_tensor[t], labels) - assert model_output.shape == (B, C * 2, *sample.shape[2:]) - model_output, model_var_values = torch.split(model_output, C, dim=1) - min_log = _extract_into_tensor(self.posterior_log_variance_clipped, - t, sample.shape) - max_log = _extract_into_tensor(np.log(self.betas), t, sample.shape) - - frac = (model_var_values + 1) / 2 - model_log_variance = frac * max_log + (1 - frac) * min_log - pred_xstart = ( - _extract_into_tensor(self.sqrt_recip_alphas_cumprod, t, - sample.shape) * sample - - _extract_into_tensor(self.sqrt_recipm1_alphas_cumprod, t, - sample.shape) * model_output) - - model_mean = (_extract_into_tensor(self.posterior_mean_coef1, t, - sample.shape) * pred_xstart + - _extract_into_tensor(self.posterior_mean_coef2, t, - sample.shape) * sample) - - assert model_mean.shape == model_log_variance.shape == pred_xstart.shape == sample.shape - noise = torch.randn_like(sample) - nonzero_mask = ((t != 0).float().view( - -1, *([1] * (len(sample.shape) - 1)))) - sample = model_mean + nonzero_mask * torch.exp( - 0.5 * model_log_variance) * noise - return sample diff --git a/examples/models/contrib/dit/figs/sample.fp8.png b/examples/models/contrib/dit/figs/sample.fp8.png deleted file mode 100644 index 333318aaa4de..000000000000 Binary files a/examples/models/contrib/dit/figs/sample.fp8.png and /dev/null differ diff --git a/examples/models/contrib/dit/figs/sample.png b/examples/models/contrib/dit/figs/sample.png deleted file mode 100644 index ba3293298600..000000000000 Binary files a/examples/models/contrib/dit/figs/sample.png and /dev/null differ diff --git a/examples/models/contrib/dit/sample.py b/examples/models/contrib/dit/sample.py deleted file mode 100644 index 2b4b7aebb2bc..000000000000 --- a/examples/models/contrib/dit/sample.py +++ /dev/null @@ -1,268 +0,0 @@ -import argparse -import json -import os -from functools import wraps - -import tensorrt as trt -import torch -from cuda import cudart -from diffusion import DiTDiffusionPipeline -from torchvision.utils import save_image - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch, trt_dtype_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper -from tensorrt_llm.runtime.session import Session, TensorInfo - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -class TllmDiT(object): - - def __init__(self, - config, - debug_mode=True, - stream: torch.cuda.Stream = None): - self.dtype = config['pretrained_config']['dtype'] - - rank = tensorrt_llm.mpi_rank() - world_size = config['pretrained_config']['mapping']['world_size'] - cp_size = config['pretrained_config']['mapping']['cp_size'] - tp_size = config['pretrained_config']['mapping']['tp_size'] - pp_size = config['pretrained_config']['mapping']['pp_size'] - assert pp_size == 1 - self.mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=rank, - cp_size=cp_size, - tp_size=tp_size, - pp_size=1, - gpus_per_node=args.gpus_per_node) - - local_rank = rank % self.mapping.gpus_per_node - self.device = torch.device(f'cuda:{local_rank}') - torch.cuda.set_device(self.device) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_file = os.path.join(args.tllm_model_dir, f"rank{rank}.engine") - logger.info(f'Loading engine from {engine_file}') - with open(engine_file, "rb") as f: - engine_buffer = f.read() - - assert engine_buffer is not None - - self.session = Session.from_serialized_engine(engine_buffer) - - self.debug_mode = debug_mode - - self.inputs = {} - self.outputs = {} - self.buffer_allocated = False - - expected_tensor_names = ['latent', 'timestep', 'label', 'output'] - - if self.mapping.tp_size > 1: - self.buffer, self.all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.mapping.tp_size)) - self.inputs['all_reduce_workspace'] = self.all_reduce_workspace - expected_tensor_names += ['all_reduce_workspace'] - - found_tensor_names = [ - self.session.engine.get_tensor_name(i) - for i in range(self.session.engine.num_io_tensors) - ] - if not self.debug_mode and set(expected_tensor_names) != set( - found_tensor_names): - logger.error( - f"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}" - ) - logger.error( - f"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}" - ) - logger.error(f"Expected tensor names: {expected_tensor_names}") - logger.error(f"Found tensor names: {found_tensor_names}") - raise RuntimeError( - "Tensor names in engine are not the same as expected.") - if self.debug_mode: - self.debug_tensors = list( - set(found_tensor_names) - set(expected_tensor_names)) - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name)) - return dtype - - def _setup(self, batch_size): - for i in range(self.session.engine.num_io_tensors): - name = self.session.engine.get_tensor_name(i) - if self.session.engine.get_tensor_mode( - name) == trt.TensorIOMode.OUTPUT: - shape = list(self.session.engine.get_tensor_shape(name)) - shape[0] = batch_size // 2 if name in [ - 'cond_eps', 'uncond_eps' - ] else batch_size - self.outputs[name] = torch.empty(shape, - dtype=self._tensor_dtype(name), - device=self.device) - - self.buffer_allocated = True - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - @cuda_stream_guard - def forward(self, latent: torch.Tensor, timestep: torch.Tensor, - label: torch.Tensor): - """ - Forward pass of DiT. - latent: (N, C, H, W) - timestep: (N,) - label: (N,) - """ - self._setup(latent.shape[0]) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - inputs = { - 'latent': latent.to(str_dtype_to_torch(self.dtype)), - "timestep": timestep.int(), - "label": label.int() - } - self.inputs.update(**inputs) - self.session.set_shapes(self.inputs) - ok = self.session.run(self.inputs, self.outputs, - self.stream.cuda_stream) - - if not ok: - raise RuntimeError('Executing TRT engine failed!') - if self.debug_mode: - torch.cuda.synchronize() - for k, v in self.inputs.items(): - print(k, v.sum()) - for k, v in self.outputs.items(): - print(k, v.sum()) - return self.outputs['output'] - - -def vae_decode(samples, engine_path): - # Load standard plugins - TRT_LOGGER = trt.Logger(trt.Logger.WARNING) - trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="") - - logger.info(f'Loading vae engine from {engine_path}') - with open(engine_path, 'rb') as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine {engine_path}') - session_vae = Session.from_serialized_engine(engine_buffer) - inputs = {'input': samples} - output_info = session_vae.infer_shapes( - [TensorInfo('input', trt.DataType.FLOAT, samples.shape)]) - outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - stream = torch.cuda.current_stream().cuda_stream - ok = session_vae.run(inputs, outputs, stream) - - assert ok, "Runtime execution failed for vae session" - - samples = outputs['output'] - return samples - - -def main(args): - tensorrt_llm.logger.set_level(args.log_level) - - torch.manual_seed(args.seed) - assert torch.cuda.is_available() - device = "cuda" - - # Load model: - config_file = os.path.join(args.tllm_model_dir, 'config.json') - with open(config_file) as f: - config = json.load(f) - model = TllmDiT(config, debug_mode=args.debug_mode) - - diffusion = DiTDiffusionPipeline(model.forward, - timestep_respacing=args.num_sampling_steps) - - latent_size = args.image_size // 8 - latent = torch.randn(args.batch_size, - 4, - latent_size, - latent_size, - device=device) - labels = torch.randint(args.num_classes, [args.batch_size], device=device) - - latent = torch.cat([latent, latent], 0) - labels_null = torch.tensor([1000] * args.batch_size, device=device) - labels = torch.cat([labels, labels_null], 0) - - samples = diffusion.run(latent, labels) - samples, _ = samples.chunk(2, dim=0) - - samples = vae_decode(samples / 0.18215, args.vae_decoder_engine) - - save_image(samples, - "sample.png", - nrow=4, - normalize=True, - value_range=(-1, 1)) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument('--vae_decoder_engine', - type=str, - default='vae_decoder/plan/visual_encoder_fp16.plan', - help='') - parser.add_argument("--batch-size", type=int, default=2) - parser.add_argument("--image-size", - type=int, - choices=[256, 512], - default=512) - parser.add_argument("--num-classes", type=int, default=1000) - parser.add_argument("--num-sampling-steps", type=int, default=250) - parser.add_argument("--seed", type=int, default=0) - parser.add_argument("--tllm_model_dir", - type=str, - default='./engine_outputs/') - parser.add_argument("--gpus_per_node", type=int, default=8) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument("--debug_mode", type=bool, default=False) - args = parser.parse_args() - main(args) diff --git a/examples/models/contrib/dit/utils_modelopt.py b/examples/models/contrib/dit/utils_modelopt.py deleted file mode 100644 index e5f620e05325..000000000000 --- a/examples/models/contrib/dit/utils_modelopt.py +++ /dev/null @@ -1,186 +0,0 @@ -import os -import re -from collections import OrderedDict - -import modelopt.torch.opt as mto -import torch -from diffusers import DiTPipeline -from modelopt.torch.export.layer_utils import (get_activation_scaling_factor, - get_weight_scaling_factor) -from modelopt.torch.export.model_config_utils import to_quantized_weight -from torchvision.datasets.utils import download_url - -HUGGINGFACE_TO_FACEBOOK_DIT_NAME_MAPPING = { - r"^transformer_blocks.(\d+).norm1.emb.class_embedder.embedding_table.weight$": - "y_embedder.embedding_table.weight", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.weight$": - "t_embedder.mlp.0.weight", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.bias$": - "t_embedder.mlp.0.bias", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.weight$": - "t_embedder.mlp.2.weight", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.bias$": - "t_embedder.mlp.2.bias", - "^pos_embed.proj.weight$": - "x_embedder.proj.weight", - "^pos_embed.proj.bias$": - "x_embedder.proj.bias", - r"^transformer_blocks.(\d+).attn1.to_qkv.weight$": - "blocks.*.attn.qkv.weight", - r"^transformer_blocks.(\d+).attn1.to_qkv.bias$": - "blocks.*.attn.qkv.bias", - r"^transformer_blocks.(\d+).attn1.to_out.0.weight$": - "blocks.*.attn.proj.weight", - r"^transformer_blocks.(\d+).attn1.to_out.0.bias$": - "blocks.*.attn.proj.bias", - r"^transformer_blocks.(\d+).ff.net.0.proj.weight$": - "blocks.*.mlp.fc1.weight", - r"^transformer_blocks.(\d+).ff.net.0.proj.bias$": - "blocks.*.mlp.fc1.bias", - r"^transformer_blocks.(\d+).ff.net.2.weight$": - "blocks.*.mlp.fc2.weight", - r"^transformer_blocks.(\d+).ff.net.2.bias$": - "blocks.*.mlp.fc2.bias", - r"^transformer_blocks.(\d+).norm1.linear.weight$": - "blocks.*.adaLN_modulation.1.weight", - r"^transformer_blocks.(\d+).norm1.linear.bias$": - "blocks.*.adaLN_modulation.1.bias", - "^proj_out_2.weight$": - "final_layer.linear.weight", - "^proj_out_2.bias$": - "final_layer.linear.bias", - "^proj_out_1.weight$": - "final_layer.adaLN_modulation.1.weight", - "^proj_out_1.bias$": - "final_layer.adaLN_modulation.1.bias", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.weights_scaling_factor$": - "t_embedder.mlp.0.weights_scaling_factor", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_1.activation_scaling_factor": - "t_embedder.mlp.0.activation_scaling_factor", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.weights_scaling_factor": - "t_embedder.mlp.2.weights_scaling_factor", - r"^transformer_blocks.(\d+).norm1.emb.timestep_embedder.linear_2.activation_scaling_factor": - "t_embedder.mlp.2.activation_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_qkv.weights_scaling_factor$": - "blocks.*.attn.qkv.weights_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_qkv.activation_scaling_factor$": - "blocks.*.attn.qkv.activation_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_out.0.weights_scaling_factor$": - "blocks.*.attn.proj.weights_scaling_factor", - r"^transformer_blocks.(\d+).attn1.to_out.0.activation_scaling_factor$": - "blocks.*.attn.proj.activation_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.0.proj.weights_scaling_factor$": - "blocks.*.mlp.fc1.weights_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.0.proj.activation_scaling_factor$": - "blocks.*.mlp.fc1.activation_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.2.weights_scaling_factor$": - "blocks.*.mlp.fc2.weights_scaling_factor", - r"^transformer_blocks.(\d+).ff.net.2.activation_scaling_factor$": - "blocks.*.mlp.fc2.activation_scaling_factor", - r"^transformer_blocks.(\d+).norm1.linear.weights_scaling_factor$": - "blocks.*.adaLN_modulation.1.weights_scaling_factor", - r"^transformer_blocks.(\d+).norm1.linear.activation_scaling_factor$": - "blocks.*.adaLN_modulation.1.activation_scaling_factor", - "^proj_out_2.weights_scaling_factor$": - "final_layer.linear.weights_scaling_factor", - "^proj_out_2.activation_scaling_factor$": - "final_layer.linear.activation_scaling_factor", - "^proj_out_1.weights_scaling_factor$": - "final_layer.adaLN_modulation.1.weights_scaling_factor", - "^proj_out_1.activation_scaling_factor$": - "final_layer.adaLN_modulation.1.activation_scaling_factor", -} - - -def convert_amax_to_scaling_factor(model, state_dict): - ret_dict = state_dict.copy() - for name, module in model.named_modules(): - if isinstance(module, torch.nn.Linear): - activation_scaling_factor = get_activation_scaling_factor(module) - weight_scaling_factor = get_weight_scaling_factor(module) - if activation_scaling_factor: - ret_dict[ - f'{name}.activation_scaling_factor'] = activation_scaling_factor - if weight_scaling_factor: - ret_dict[ - f'{name}.weights_scaling_factor'] = weight_scaling_factor - - weight = module.weight.detach().cpu() - if weight_scaling_factor: - # only module with valid weight scaling factor - weight = to_quantized_weight( - weight=weight, - weights_scaling_factor=weight_scaling_factor, - quantization="fp8", - ) - # replace the quantized weight - ret_dict[f'{name}.weight'] = weight - return ret_dict - - -def get_weights_map(state_dict): - - def _get_fb_dit_name(dit_name): - for k, v in HUGGINGFACE_TO_FACEBOOK_DIT_NAME_MAPPING.items(): - m = re.match(k, dit_name) - if m is not None: - if "*" in v: - v = v.replace("*", m.groups()[0]) - return v - return dit_name - - weights_map = OrderedDict() - for key, value in state_dict.items(): - if ("to_q." in key) or ("to_k." in key) or ("to_v." in key): - continue - if _get_fb_dit_name(key) in weights_map: - continue - else: - weights_map[_get_fb_dit_name(key)] = value - return weights_map - - -def download_model(ckpt="DiT-XL-2-512x512.pt"): - """ - Downloads a pre-trained DiT model from the web. - """ - if not os.path.isfile(ckpt): - os.makedirs('pretrained_models', exist_ok=True) - web_path = f"https://dl.fbaipublicfiles.com/DiT/models/{ckpt}" - download_url(web_path, '.') - model = torch.load(ckpt, map_location=lambda storage, loc: storage) - return model - - -def remap_model(model_name="facebook/DiT-XL-2-512", - quantized_ckpt="dit.quantized.pt", - output_ckpt="dit.converted.pt", - fuse_qkv=True, - dtype=torch.float32): - pipe = DiTPipeline.from_pretrained(model_name, torch_dtype=dtype) - transformer = pipe.transformer - - # fuse qkv gemm - if fuse_qkv: - from diffusers.models.attention_processor import (Attention, - AttnProcessor2_0) - for module in transformer.modules(): - if isinstance(module, Attention): - assert (isinstance(module.processor, AttnProcessor2_0)) - module.fuse_projections(fuse=True) - pretrained_state_dict = transformer.state_dict() - - mto.restore(transformer, quantized_ckpt) - quantized_dict = convert_amax_to_scaling_factor(transformer, - pretrained_state_dict) - assert set(pretrained_state_dict.keys()).issubset(set( - quantized_dict.keys())) - - remapped_dict = get_weights_map(quantized_dict) - # TODO: currently we use weights of pos_embed and x_embedder from official DiT because - # in the implementation by HuggingFace, there are some modifications for these modules. - stored_params = download_model(ckpt="DiT-XL-2-512x512.pt") - for name in ["pos_embed", "x_embedder.proj.weight", "x_embedder.proj.bias"]: - remapped_dict[name] = stored_params[name] - - torch.save(remapped_dict, output_ckpt) diff --git a/examples/models/contrib/dit/vae_decoder_trt.py b/examples/models/contrib/dit/vae_decoder_trt.py deleted file mode 100644 index 31803a6690d7..000000000000 --- a/examples/models/contrib/dit/vae_decoder_trt.py +++ /dev/null @@ -1,139 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2022-2023 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. -import argparse -import os - -import torch -from diffusers.models import AutoencoderKL - - -class TRT_Exporter(object): - - def __init__(self, pytorch_model, max_batch_size, latent_channel, - latent_shape): - self.pytorch_model = pytorch_model - self.max_batch_size = max_batch_size - self.latent_channel = latent_channel - self.latent_shape = latent_shape - - def export_onnx(self, onnxFile): - print(f"Start exporting ONNX model to {onnxFile}!") - latent = torch.randn(self.max_batch_size, self.latent_channel, - *self.latent_shape).cuda() - self.pytorch_model.cuda().eval() - with torch.inference_mode(): - torch.onnx.export( - self.pytorch_model, - latent, - onnxFile, - opset_version=17, - input_names=['input'], - output_names=['output'], - dynamic_axes={'input': { - 0: 'batch' - }}, - # Required for pytorch>=2.9.0 as dynamo becomes the default and introduces bugs as it does not support opset_version=17 natively - dynamo=False) - - def generate_trt_engine(self, onnxFile, planFile): - print(f"Start exporting TRT model to {planFile}!") - from time import time - - import tensorrt as trt - logger = trt.Logger(trt.Logger.VERBOSE) - builder = trt.Builder(logger) - network = builder.create_network( - 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)) - profile = builder.create_optimization_profile() - config = builder.create_builder_config() - config.set_flag(trt.BuilderFlag.FP16) - parser = trt.OnnxParser(network, logger) - - with open(onnxFile, 'rb') as model: - if not parser.parse(model.read(), "/".join(onnxFile.split("/"))): - print("Failed parsing %s" % onnxFile) - for error in range(parser.num_errors): - print(parser.get_error(error)) - print("Succeeded parsing %s" % onnxFile) - - nBS = -1 - nMinBS = 1 - nMaxBS = self.max_batch_size - nOptBS = (nMaxBS + nMinBS) // 2 - inputT = network.get_input(0) - inputT.shape = [nBS, self.latent_channel, *self.latent_shape] - profile.set_shape(inputT.name, - [nMinBS, self.latent_channel, *self.latent_shape], - [nOptBS, self.latent_channel, *self.latent_shape], - [nMaxBS, self.latent_channel, *self.latent_shape]) - - config.add_optimization_profile(profile) - - t0 = time() - engineString = builder.build_serialized_network(network, config) - t1 = time() - if engineString is None: - print("Failed building %s" % planFile) - else: - print("Succeeded building %s in %d s" % (planFile, t1 - t0)) - print("plan file is", planFile) - with open(planFile, 'wb') as f: - f.write(engineString) - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--vae", - type=str, - choices=["ema", "mse"], - default="mse") - parser.add_argument('--max_batch_size', type=int, default=8) - parser.add_argument("--image-size", - type=int, - choices=[256, 512], - default=512) - parser.add_argument('--onnxFile', - type=str, - default='vae_decoder/onnx/visual_encoder.onnx', - help='') - parser.add_argument('--planFile', - type=str, - default='vae_decoder/plan/visual_encoder_fp16.plan', - help='') - parser.add_argument('--only_trt', - action='store_true', - help='Run only convert the onnx to TRT engine.') - args = parser.parse_args() - return args - - -if __name__ == '__main__': - args = parse_arguments() - onnx_file_dir = os.path.dirname(args.onnxFile) - if not onnx_file_dir == '' and not os.path.exists(onnx_file_dir): - os.makedirs(onnx_file_dir) - plan_file_dir = os.path.dirname(args.planFile) - if not os.path.exists(plan_file_dir): - os.makedirs(plan_file_dir) - - vae = AutoencoderKL.from_pretrained(f"stabilityai/sd-vae-ft-{args.vae}") - vae.forward = vae.decode - latant_shape = [args.image_size // 8] * 2 - onnx_trt_obj = TRT_Exporter(vae, args.max_batch_size, 4, latant_shape) - if args.only_trt: - onnx_trt_obj.generate_trt_engine(args.onnxFile, args.planFile) - else: - onnx_trt_obj.export_onnx(args.onnxFile) - onnx_trt_obj.generate_trt_engine(args.onnxFile, args.planFile) diff --git a/examples/models/contrib/falcon/README.md b/examples/models/contrib/falcon/README.md deleted file mode 100644 index 3d77aa079cff..000000000000 --- a/examples/models/contrib/falcon/README.md +++ /dev/null @@ -1,364 +0,0 @@ -# Falcon - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a Falcon model in TensorRT LLM on single GPU, single node multi-GPU, and multi-node multi-GPU. - -- [Falcon](#falcon) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace Transformers](#1-download-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Run summarization task with the TensorRT engine(s)](#4-run-summarization-task-with-the-tensorrt-engines) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Groupwise quantization (AWQ)](#groupwise-quantization-awq) - - [W4A16 AWQ with FP8 GEMM (W4A8 AWQ)](#w4a16-awq-with-fp8-gemm-w4a8-awq) - - [Troubleshooting](#troubleshooting) - - [1. The HuggingFace Falcon may raise an error when using the `accelerate` package.](#1-the-huggingface-falcon-may-raise-an-error-when-using--the-accelerate-package) - -## Overview - -The TensorRT LLM Falcon implementation can be found in [tensorrt_llm/models/falcon/model.py](../../tensorrt_llm/models/falcon/model.py). The TensorRT LLM Falcon example code is located in [`examples/models/contrib/falcon`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * BF16 - * FP8 - * FP8 KV CACHE - * Groupwise quantization (AWQ) - * Tensor Parallel - * STRONGLY TYPED - -## Usage - -The next two sections describe how to convert the weights from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) -format to the TensorRT LLM format. - -### 1. Download weights from HuggingFace Transformers - -Install the dependency packages and setup `git-lfs`. - -```bash -# Install dependencies -pip install -r requirements.txt - -# Setup git-lfs -git lfs install -``` - -There are four HF checkpoints available. Use one of the following commands to fetch the checkpoint you are interested in. Follow the guides here https://huggingface.co/docs/transformers/main/en/model_doc/falcon. - -```bash -# falcon-rw-1b -git clone https://huggingface.co/tiiuae/falcon-rw-1b falcon/rw-1b - -# falcon-7b-instruct -git clone https://huggingface.co/tiiuae/falcon-7b-instruct falcon/7b-instruct - -# falcon-40b-instruct -git clone https://huggingface.co/tiiuae/falcon-40b-instruct falcon/40b-instruct - -# falcon-180b -git clone https://huggingface.co/tiiuae/falcon-180B falcon/180b - -# falcon-11b (Falcon 2) -git clone https://huggingface.co/tiiuae/falcon-11B falcon/11b -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# falcon-rw-1b: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir ./falcon/rw-1b \ - --dtype float16 \ - --output_dir ./falcon/rw-1b/trt_ckpt/fp16/1-gpu/ - -# falcon-7b-instruct: single gpu, dtype bfloat16 -python3 convert_checkpoint.py --model_dir ./falcon/7b-instruct \ - --dtype bfloat16 \ - --output_dir ./falcon/7b-instruct/trt_ckpt/bf16/1-gpu/ - -# falcon-40b-instruct: 2-way tensor parallelism -python3 convert_checkpoint.py --model_dir ./falcon/40b-instruct \ - --dtype bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp1/ \ - --tp_size 2 - -# falcon-40b-instruct: 2-way tensor parallelism and 2-way pipeline parallelism -python3 convert_checkpoint.py --model_dir ./falcon/40b-instruct \ - --dtype bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp2/ \ - --tp_size 2 \ - --pp_size 2 - -# falcon-180b: 8-way tensor parallelism, loading weights shard-by-shard -python3 convert_checkpoint.py --model_dir ./falcon/180b \ - --dtype bfloat16 \ - --output_dir ./falcon/180b/trt_ckpt/bf16/tp8-pp1/ \ - --tp_size 8 \ - --load_by_shard \ - --workers 8 - -# falcon-180b: 4-way tensor parallelism and 2-way pipeline parallelism, loading weights shard-by-shard -python3 convert_checkpoint.py --model_dir ./falcon/180b \ - --dtype bfloat16 \ - --output_dir ./falcon/180b/trt_ckpt/bf16/tp4-pp2/ \ - --tp_size 4 \ - --pp_size 2 \ - --load_by_shard \ - --workers 8 - -# falcon-11b (Falcon 2): single gpu, dtype bfloat16 -python3 convert_checkpoint.py --model_dir ./falcon/11b \ - --dtype bfloat16 \ - --output_dir ./falcon/11b/trt_ckpt/bf16/1-gpu/ -``` - -Note that in order to use N-way tensor parallelism, the number of attention heads must be a multiple of N. -For example, you can't configure 2-way tensor parallelism for [falcon-7b](https://huggingface.co/tiiuae/falcon-7b) or [falcon-7b-instruct](https://huggingface.co/tiiuae/falcon-7b-instruct), because the number of attention heads is 71 (not divisible by 2). - - -### 3. Build TensorRT engine(s) -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is also same to the number of GPUs used to run inference. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# falcon-rw-1b -trtllm-build --checkpoint_dir ./falcon/rw-1b/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --output_dir ./falcon/rw-1b/trt_engines/fp16/1-gpu/ - -# falcon-7b-instruct -# Enabling --gpt_attention_plugin is necessary for rotary positional embedding (RoPE) -trtllm-build --checkpoint_dir ./falcon/7b-instruct/trt_ckpt/bf16/1-gpu/ \ - --gemm_plugin bfloat16 \ - --remove_input_padding enable \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/7b-instruct/trt_engines/bf16/1-gpu/ - -# falcon-40b-instruct: 2-way tensor parallelism -trtllm-build --checkpoint_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp1/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp1/ - -# falcon-40b-instruct: 2-way tensor parallelism and 2-way pipeline parallelism -trtllm-build --checkpoint_dir ./falcon/40b-instruct/trt_ckpt/bf16/tp2-pp2/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp2/ - -# falcon-180b: 8-way tensor parallelism -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/bf16/tp8-pp1/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/180b/trt_engines/bf16/tp8-pp1/ \ - --workers 8 - -# falcon-180b: 4-way tensor parallelism and 2-way pipeline parallelism -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/bf16/tp4-pp2/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/180b/trt_engines/bf16/tp4-pp2/ \ - --workers 8 - -# falcon-11b (Falcon 2) -trtllm-build --checkpoint_dir ./falcon/11b/trt_ckpt/bf16/1-gpu/ \ - --gemm_plugin bfloat16 \ - --gpt_attention_plugin bfloat16 \ - --output_dir ./falcon/11b/trt_engines/bf16/1-gpu/ -``` - -If the engines are built successfully, you will see output like (falcon-rw-1b as the example): -``` -...... -[12/27/2023-03:46:29] [TRT] [I] Engine generation completed in 35.0677 seconds. -[12/27/2023-03:46:29] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 393 MiB, GPU 2699 MiB -[12/27/2023-03:46:29] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +2699, now: CPU 0, GPU 2699 (MiB) -[12/27/2023-03:46:29] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 10624 MiB -[12/27/2023-03:46:29] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:00:36 -[12/27/2023-03:46:31] [TRT-LLM] [I] Serializing engine to ./falcon/rw-1b/trt_engines/fp16/1-gpu/rank0.engine... -[12/27/2023-03:46:59] [TRT-LLM] [I] Engine serialized. Total time: 00:00:28 -[12/27/2023-03:46:59] [TRT-LLM] [I] Total time of building all engines: 00:01:59 -``` - -### 4. Run summarization task with the TensorRT engine(s) -The `../../../summarize.py` script can run the built engines to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -```bash -# falcon-rw-1b -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/rw-1b \ - --engine_dir ./falcon/rw-1b/trt_engines/fp16/1-gpu/ - -# falcon-7b-instruct -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/7b-instruct \ - --engine_dir ./falcon/7b-instruct/trt_engines/bf16/1-gpu/ - -# falcon-40b-instruct: 2-way tensor parallelism -mpirun -n 2 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/40b-instruct \ - --engine_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp1/ - -# falcon-40b-instruct: 2-way tensor parallelism and 2-way pipeline parallelism -mpirun -n 4 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/40b-instruct \ - --engine_dir ./falcon/40b-instruct/trt_engines/bf16/tp2-pp2/ - -# falcon-180b: 8-way tensor parallelism -mpirun -n 8 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/bf16/tp8-pp1/ - -# falcon-180b: 4-way tensor parallelism and 2-way pipeline parallelism -mpirun -n 8 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/bf16/tp4-pp2/ - -# falcon-11b (Falcon 2) -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/11b \ - --engine_dir ./falcon/11b/trt_engines/bf16/1-gpu/ -``` - -If the engines are run successfully, you will see output like (falcon-rw-1b as the example): -``` -...... -[12/27/2023-03:57:02] [TRT-LLM] [I] TensorRT LLM (total latency: 5.816917419433594 sec) -[12/27/2023-03:57:02] [TRT-LLM] [I] TensorRT LLM beam 0 result -[12/27/2023-03:57:02] [TRT-LLM] [I] rouge1 : 15.061493342516243 -[12/27/2023-03:57:02] [TRT-LLM] [I] rouge2 : 4.495335888974063 -[12/27/2023-03:57:02] [TRT-LLM] [I] rougeL : 11.800002670828547 -[12/27/2023-03:57:02] [TRT-LLM] [I] rougeLsum : 13.458777656925877 -``` - -### FP8 Post-Training Quantization - -The examples below use the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -Now quantize HF Falcon weights and export trtllm checkpoint. - -```bash -# Quantize HF Falcon 180B checkpoint into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./falcon/180b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./falcon/180b/trt_ckpt/fp8/tp8-pp1 \ - --tp_size 8 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/fp8/tp8-pp1 \ - --gemm_plugin float16 \ - --output_dir ./falcon/180b/trt_engines/fp8/tp8-pp1 \ - --workers 8 - -# Run the summarization task -mpirun -n 8 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/fp8/tp8-pp1 -``` - -Note that you can enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` when building the engines. - -### Groupwise quantization (AWQ) - -The examples below use the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -Now quantize HF Falcon weights and export trtllm checkpoint. - -```bash -# Quantize HF Falcon 180B checkpoint into INT4-AWQ and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./falcon/180b \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir ./falcon/180b/trt_ckpt/int4_awq/tp2 \ - --tp_size 2 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/int4_awq/tp2 \ - --gemm_plugin float16 \ - --output_dir ./falcon/180b/trt_engines/int4_awq/tp2 \ - --workers 2 - -# Run the summarization task -mpirun -n 2 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/int4_awq/tp2 -``` - -#### W4A16 AWQ with FP8 GEMM (W4A8 AWQ) -For Hopper GPUs, TRT-LLM also supports employing FP8 GEMM for accelerating linear layers. This mode is noted with `w4a8_awq` for Modelopt and TRT-LLM, in which both weights and activations are converted from W4A16 to FP8 for GEMM calculation. - -Please make sure your system contains a Hopper GPU before trying the commands below. - -```bash -# Quantize HF Falcon 180B checkpoint into W4A8-AWQ and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./falcon/180b \ - --dtype float16 \ - --qformat w4a8_awq \ - --output_dir ./falcon/180b/trt_ckpt/w4a8_awq/tp2 \ - --tp_size 2 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./falcon/180b/trt_ckpt/w4a8_awq/tp2 \ - --gemm_plugin float16 \ - --output_dir ./falcon/180b/trt_engines/w4a8_awq/tp2 \ - --workers 2 - -# Run the summarization task -mpirun -n 2 --allow-run-as-root --oversubscribe \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./falcon/180b \ - --engine_dir ./falcon/180b/trt_engines/w4a8_awq/tp2 -``` - -## Troubleshooting - -### 1. The HuggingFace Falcon may raise an error when using the `accelerate` package. - -One may find the following message. -``` -Traceback (most recent call last): - File "build.py", line 10, in - from transformers import FalconConfig, FalconForCausalLM - File "", line 1039, in _handle_fromlist - File "/usr/local/lib/python3.8/dist-packages/transformers/utils/import_utils.py", line 1090, in __getattr__ - value = getattr(module, name) - File "/usr/local/lib/python3.8/dist-packages/transformers/utils/import_utils.py", line 1089, in __getattr__ - module = self._get_module(self._class_to_module[name]) - File "/usr/local/lib/python3.8/dist-packages/transformers/utils/import_utils.py", line 1101, in _get_module - raise RuntimeError( -RuntimeError: Failed to import transformers.models.falcon.modeling_falcon because of the following error (look up to see its traceback): -``` -It may be resolved by pinning the version of `typing-extensions` package by `4.5.0`. -```bash -pip install typing-extensions==4.5.0 -``` diff --git a/examples/models/contrib/falcon/convert_checkpoint.py b/examples/models/contrib/falcon/convert_checkpoint.py deleted file mode 100644 index ac59316519eb..000000000000 --- a/examples/models/contrib/falcon/convert_checkpoint.py +++ /dev/null @@ -1,202 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.falcon.model import FalconForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - - parser.add_argument('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - tensorrt_llm.logger.set_level(args.log_level) - - # WAR for modelopt multithreading issue. - import importlib.metadata - if (importlib.metadata.version('nvidia-modelopt') < '0.27' - and args.workers > 1): - args.workers = 1 - tensorrt_llm.logger.warning( - 'Reducing workers=1 when converting checkpoint because ' - 'modelopt has an issue in multi-threading, which will be ' - 'fixed in 0.27.') - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - config = QuantConfig() - if args.use_weight_only and args.weight_only_precision == 'int8': - config.quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - config.quant_algo = QuantAlgo.W4A16 - return config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - } - - -def convert_and_save_hf(args: argparse.Namespace): - model_dir = args.model_dir - load_by_shard = args.load_by_shard - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - - quant_config = args_to_quant_config(args) - hf_model = None - import transformers - if not args.load_by_shard and quant_config.quant_mode.has_any_quant(): - hf_model = transformers.FalconForCausalLM.from_pretrained( - model_dir, trust_remote_code=True, dtype='auto', device_map='auto') - else: - # Initialize huggingface local cache. - # Huggingface copies the external configuration source (`configuration_falcon.py` here) into its local cache at - # `~/.cache/huggingface/modules/transformers_modules/`, - # and if multiple threads attempt to do this concurrently, weird issue can happen: - # Some threads may see an empty configuration_falcon.py file and fail. - # Preload the config once to initialize local cache, so subsequent multithread loading won't fail. - _ = transformers.FalconConfig.from_pretrained(model_dir, - trust_remote_code=True, - dtype='auto', - device_map='auto') - - def convert_and_save_rank(args, rank: int): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - falcon = FalconForCausalLM.from_hugging_face( - model_dir if hf_model is None else hf_model, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - load_by_shard=load_by_shard, - **override_fields, - ) - falcon.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del falcon - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/falcon/requirements.txt b/examples/models/contrib/falcon/requirements.txt deleted file mode 100644 index e69190e50443..000000000000 --- a/examples/models/contrib/falcon/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -transformers>=4.56.0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece>=0.1.99 -tqdm diff --git a/examples/models/contrib/gptj/.gitignore b/examples/models/contrib/gptj/.gitignore deleted file mode 100644 index 18404db569c0..000000000000 --- a/examples/models/contrib/gptj/.gitignore +++ /dev/null @@ -1,4 +0,0 @@ -__pycache__/ -gptj_model/ -*.log -*.json diff --git a/examples/models/contrib/gptj/README.md b/examples/models/contrib/gptj/README.md deleted file mode 100644 index c470109a95f7..000000000000 --- a/examples/models/contrib/gptj/README.md +++ /dev/null @@ -1,269 +0,0 @@ -# GPT-J - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [GPT-J](https://huggingface.co/EleutherAI/gpt-j-6b) model using TensorRT LLM and run on a single GPU. - -- [GPT-J](#gpt-j) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace (HF) Transformers](#1-download-weights-from-huggingface-hf-transformers) - - [2. Build TensorRT engine(s)](#2-build-tensorrt-engines) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [AWQ INT4 weight only quantization](#awq-int4-weight-only-quantization) - - [SmoothQuant (W8A8) quantization](#smoothquant-w8a8-quantization) - - [Fused MultiHead Attention (FMHA)](#fused-multihead-attention-fmha) - - [INT8 KV cache](#int8-kv-cache) - - [3. Run](#3-run) - - [Summarization using the GPT-J model](#summarization-using-the-gpt-j-model) - -## Overview - -The TensorRT LLM GPT-J implementation can be found in [`tensorrt_llm/models/gptj/model.py`](../../tensorrt_llm/models/gptj/model.py). The TensorRT LLM GPT-J example -code is located in [`examples/models/contrib/gptj`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 & INT4 per-channel weight-only - * FP8 (with FP8 kv cache) - * Groupwise quantization (AWQ) - * INT8 KV CACHE (+ AWQ/per-channel weight-only) - * INT8 SmoothQuant - -## Usage - -### 1. Download weights from HuggingFace (HF) Transformers - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -```bash -# 1. Weights & config -git clone https://huggingface.co/EleutherAI/gpt-j-6b gptj_model -pushd gptj_model && \ - rm -f pytorch_model.bin && \ - wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/pytorch_model.bin && \ -popd - -# 2. Vocab and merge table -wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/vocab.json -wget https://huggingface.co/EleutherAI/gpt-j-6b/resolve/main/merges.txt -``` - -### 2. Build TensorRT engine(s) - -TensorRT LLM builds TensorRT engine(s) using a HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) using -dummy weights. - -Examples of build invocations: - -```bash -# Build a float16 engine using HF weights. -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --output_dir ./trt_ckpt/gptj_fp16_tp1/ -``` - -***For all kinds of checkpoints, they share the same trtllm-build command like:*** - -```bash -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. -trtllm-build --checkpoint_dir ./trt_ckpt/gptj_fp16_tp1/ \ - --output_dir ./trt_engines/gptj_fp16_tp1/ \ - --gemm_plugin float16 \ - --max_batch_size=32 \ - --max_input_len=1919 \ - --max_seq_len=2047 -``` - -INT8 weight-only - -```bash -# Build an int8 weight-only engine using HF weights with TP=2 -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir ./trt_ckpt/gptj_int8_tp2/ \ - --tp_size 2 -``` -Building command is identical to the common one above. - -INT4 weight-only - -```bash -# Build an int4 weight only quantization engine using int4 weight only quantized weights. -python convert_checkpoint.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --output_dir ./trt_ckpt/gptj_int4wo_tp1/ -``` -Building command is identical to the common one above. - -#### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -One can quantize HF GPT-J weights in FP8 as follows. - -```bash -# Quantize HF GPT-J 6B checkpoint into FP8 format -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./trt_ckpt/gptj_fp8_tp1 \ - --calib_size 512 -``` -Building command is identical to the common one above. -Note that you can enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` when building the engines. - -#### AWQ INT4 weight only quantization - -One can enable AWQ INT4 weight only quantization like the following command. - -```bash -# Enable AWQ int4 group-wise weight-only quantization. -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir ./trt_ckpt/gptj_int4_awq_tp1 \ - --calib_size 512 -``` - -```bash -# Enable AWQ int4 group-wise weight-only quantization with tp2. -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_awq \ - --tp_size 2 \ - --output_dir ./trt_ckpt/gptj_int4_awq_tp2 \ - --calib_size 512 -``` - -And build command is identical to the common one above. - -#### SmoothQuant (W8A8) quantization - -One can enable smoothquant W8A8 (weight per-channel, activation per-tensor) quantization like the following command. - -```bash -# Enable smoothquant (W8A8 kernel). -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int8_sq \ - --output_dir ./trt_ckpt/gptj_sq_tp1 \ - --calib_size 512 -``` -Building command is identical to the common one above. - -#### Fused MultiHead Attention (FMHA) - -You can enable the FMHA kernels for GPT by adding `--context_fmha` to the invocation of `trtllm-build`. Note that it is enabled by default. - -If you find that the default fp16 accumulation (`--context_fmha`) cannot meet the requirement, you can try to enable fp32 accumulation by adding `--enable_context_fmha_fp32_acc` to the inference command (`run.py` or `summarize.py`). However, it is expected to see performance drop. - -Note `--context_fmha` has to be used together with `--gpt_attention_plugin float16`. - -#### INT8 KV cache -INT8 KV cache could be enabled to reduce memory footprint. It will bring more performance gains when batch size gets larger. - -You can get the INT8 scale of KV cache through Modelopt: - -```bash -# INT8 calibration -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/gptj_fp16_int8kv_tp1 \ - --calib_size 512 -``` - -And build command is identical to the common one above. - -**INT8 KV cache + per-channel weight-only quantization** - -For example, you can enable INT8 KV cache together with per-channel INT8/INT4 weight-only quantization like the following command. - -```bash -# Enable INT8 KV cache together with per-channel INT8 weight-only quantization -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_wo \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/gptj_int4wo_int8kv_tp1 \ - --calib_size 512 -``` - -Building command is identical to the common one above. - -**INT8 KV cache + AWQ** - -In addition, you can enable INT8 KV cache together with AWQ (per-group INT4 weight-only quantization)like the following command. - -```bash -# Enable INT8 KV cache together with group-wise 4bit AWQ quantization -python ../../../quantization/quantize.py --model_dir ./gpt-j-6b \ - --dtype float16 \ - --qformat int4_awq \ - --kv_cache_dtype int8 \ - --output_dir ./trt_ckpt/gptj_int4awq_int8kv_tp1 \ - --calib_size 512 -``` - -Building command is identical to the common one above. - - -### 3. Run - - -To run a TensorRT LLM GPT-J model: - -```bash -python3 ../../../run.py --max_output_len=50 --engine_dir=gptj_engine --tokenizer_dir=gptj_model -``` - -## Summarization using the GPT-J model - -The following section describes how to run a TensorRT LLM GPT-J model to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -The script can also perform the same summarization using the HF GPT-J model. - -As previously explained, the first step is to build the TensorRT engine as described above using HF weights. You also have to install the requirements: - -```bash -pip install -r requirements.txt -``` - -The summarization can be done using the [`../../../summarize.py`](../../../summarize.py) script as follows: - -```bash -# Run the summarization task. -python3 ../../../summarize.py --engine_dir ./trt_engines/gptj_fp16_tp1 \ - --hf_model_dir ./gpt-j-6b \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --tensorrt_llm_rouge1_threshold 14 \ - --data_type fp16 \ - --check_accuracy -``` diff --git a/examples/models/contrib/gptj/convert_checkpoint.py b/examples/models/contrib/gptj/convert_checkpoint.py deleted file mode 100644 index 885b74c6a926..000000000000 --- a/examples/models/contrib/gptj/convert_checkpoint.py +++ /dev/null @@ -1,166 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -from transformers import AutoModelForCausalLM - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.llmapi import QuantConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import GPTJConfig, GPTJForCausalLM -from tensorrt_llm.models.convert_utils import infer_dtype -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--vocab_size', type=int, default=50400) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=28) - parser.add_argument('--n_head', type=int, default=16) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--norm_eps', type=float, default=1e-05) - parser.add_argument('--rotary_dim', type=int, default=64) - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -def args_to_quant_config(args): - quant_algo = None - if args.use_weight_only and args.weight_only_precision == 'int8': - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - quant_algo = QuantAlgo.W4A16 - return QuantConfig(quant_algo=quant_algo) - - -def convert_and_save_hf(args): - model_dir = args.model_dir - world_size = args.tp_size * args.pp_size - quant_config = args_to_quant_config(args) - - hf_model = AutoModelForCausalLM.from_pretrained(model_dir, - dtype='auto', - trust_remote_code=True) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - model = GPTJForCausalLM.from_hugging_face(hf_model, - args.dtype, - mapping=mapping, - quant_config=quant_config) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - if args.workers == 1: - for rank in range(world_size): - convert_and_save_rank(args, rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(convert_and_save_rank, args, rank) - for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - del hf_model - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_dir is None: - config = GPTJConfig(architecture='GPTJForCausalLM', - dtype=infer_dtype(args.dtype), - num_hidden_layers=args.n_layer, - num_attention_heads=args.n_head, - hidden_size=args.n_embd, - norm_epsilon=args.norm_eps, - vocab_size=args.vocab_size, - position_embedding_type='rope_gptj', - max_position_embeddings=args.n_positions, - hidden_act='gelu', - rotary_dim=args.rotary_dim, - mapping=Mapping(world_size=args.tp_size * - args.pp_size, - tp_size=args.tp_size, - pp_size=args.pp_size), - quantization=args_to_quant_config(args)) - config.to_json_file(os.path.join(args.output_dir, 'config.json')) - else: - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/gptj/requirements.txt b/examples/models/contrib/gptj/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/gptj/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/gptneox/.gitignore b/examples/models/contrib/gptneox/.gitignore deleted file mode 100644 index d51d01117f3a..000000000000 --- a/examples/models/contrib/gptneox/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -__pycache__/ -gptneox_model/ -*.log diff --git a/examples/models/contrib/gptneox/README.md b/examples/models/contrib/gptneox/README.md deleted file mode 100644 index d1928bfdd8d9..000000000000 --- a/examples/models/contrib/gptneox/README.md +++ /dev/null @@ -1,235 +0,0 @@ -# GPT-NeoX - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [GPT-NeoX](https://huggingface.co/EleutherAI/gpt-neox-20b) model using TensorRT LLM and run on a single GPU and a single node with -multiple GPUs. - -- [GPT-NeoX](#gpt-neox) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace (HF) Transformers](#1-download-weights-from-huggingface-hf-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Summarization using the GPT-NeoX model](#4-summarization-using-the-gpt-neox-model) - - [Apply groupwise quantization GPTQ](#apply-groupwise-quantization-gptq) - - [1. Download weights from HuggingFace (HF)](#1-download-weights-from-huggingface-hf) - - [2. Generating quantized weights](#2-generating-quantized-weights) - - [3. Convert weights from HF Transformers to TensorRT LLM format](#3-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [4. Build TensorRT engine(s)](#4-build-tensorrt-engines) - - [5. Summarization using the GPT-NeoX model](#5-summarization-using-the-gpt-neox-model) - -## Overview - -The TensorRT LLM GPT-NeoX implementation can be found in [`tensorrt_llm/models/gptneox/model.py`](../../tensorrt_llm/models/gptneox/model.py). The TensorRT LLM GPT-NeoX example code is located in [`examples/models/contrib/gptneox`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 Weight-Only - * INT4 GPTQ - * Tensor Parallel - -## Usage - -The TensorRT LLM GPT-NeoX example code locates at [examples/models/contrib/gptneox](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### 1. Download weights from HuggingFace (HF) Transformers - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -```bash -# Weights & config -git clone https://huggingface.co/EleutherAI/gpt-neox-20b gptneox_model -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -If you want to use Int8 weight only quantization, just need to add `--use_weight_only` flag. - -```bash -# Single GPU -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --output_dir ./gptneox/20B/trt_ckpt/fp16/1-gpu/ -# With 2-way Tensor Parallel -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --tp_size 2 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_ckpt/fp16/2-gpu/ -# Single GPU with int8 weight only -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --output_dir ./gptneox/20B/trt_ckpt/int8_wo/1-gpu/ -# With 2-way Tensor Parallel with int8 weight only -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --tp_size 2 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_ckpt/int8_wo/2-gpu/ -``` - -### 3. Build TensorRT engine(s) -```bash -# Single GPU -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./gptneox/20B/trt_engines/fp16/1-gpu/ -# With 2-way Tensor Parallel -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_engines/fp16/2-gpu/ -# Single GPU with int8 weight only -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int8_wo/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./gptneox/20B/trt_engines/int8_wo/1-gpu/ -# With 2-way Tensor Parallel with int8 weight only -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int8_wo/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_engines/int8_wo/2-gpu/ -``` - -### 4. Summarization using the GPT-NeoX model - -The following section describes how to run a TensorRT LLM GPT-NeoX model to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -The script can also perform the same summarization using the HF GPT-NeoX model. - -```bash -# Single GPU -python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/fp16/1-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# With 2-way Tensor Parallel -mpirun -np 2 --oversubscribe --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/fp16/2-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# Single GPU with int8 weight only -python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int8_wo/1-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# With 2-way Tensor Parallel with int8 weight only -mpirun -np 2 --oversubscribe --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int8_wo/2-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -``` - -## Apply groupwise quantization GPTQ - -### 1. Download weights from HuggingFace (HF) - -```bash -# Weights & config -sh get_weights.sh -``` - -### 2. Generating quantized weights - -In this example, the weights are quantized using [GPTQ-for-LLaMa](https://github.com/qwopqwop200/GPTQ-for-LLaMa). Note that the parameter `--act-order` referring to whether to apply the activation order GPTQ heuristic is **not supported** by TRT-LLM. - -```bash -sh gptq_convert.sh -``` - -### 3. Convert weights from HF Transformers to TensorRT LLM format - -To apply groupwise quantization GPTQ, addition command-line flags need to be passed to `convert_checkpoint.py`: -Here `--quant_ckpt_path` flag specifies the output safetensors of `gptq_convert.sh` script. - -```bash -# Single GPU -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --quant_ckpt_path ./gptneox_model/gptneox-20b-4bit-gs128.safetensors \ - --output_dir ./gptneox/20B/trt_ckpt/int4_gptq/1-gpu/ -# With 2-way Tensor Parallel -python3 convert_checkpoint.py --model_dir ./gptneox_model \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --tp_size 2 \ - --workers 2 \ - --quant_ckpt_path ./gptneox_model/gptneox-20b-4bit-gs128.safetensors \ - --output_dir ./gptneox/20B/trt_ckpt/int4_gptq/2-gpu/ -``` - -### 4. Build TensorRT engine(s) - -The command to build TensorRT engines to apply GPTQ does not change: - -```bash -# Single GPU -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int4_gptq/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./gptneox/20B/trt_engines/int4_gptq/1-gpu/ -# With 2-way Tensor Parallel -trtllm-build --checkpoint_dir ./gptneox/20B/trt_ckpt/int4_gptq/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --workers 2 \ - --output_dir ./gptneox/20B/trt_engines/int4_gptq/2-gpu/ -``` - -### 5. Summarization using the GPT-NeoX model - -The command to run summarization with GPTQ quantized model also does not change: - -```bash -# Single GPU -python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int4_gptq/1-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -# With 2-way Tensor Parallel -mpirun -np 2 --oversubscribe --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./gptneox/20B/trt_engines/int4_gptq/2-gpu/ \ - --test_trt_llm \ - --hf_model_dir gptneox_model \ - --data_type fp16 -``` diff --git a/examples/models/contrib/gptneox/convert_checkpoint.py b/examples/models/contrib/gptneox/convert_checkpoint.py deleted file mode 100644 index 7c4c76413c42..000000000000 --- a/examples/models/contrib/gptneox/convert_checkpoint.py +++ /dev/null @@ -1,732 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import List, Optional - -import safetensors -import safetensors.torch -import torch -from safetensors import safe_open -from transformers import AutoConfig, AutoModelForCausalLM - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import str_dtype_to_torch -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (get_weight, get_weight_and_bias, - split_matrix_tp) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4', 'int4_gptq'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument('--quant_ckpt_path', - type=str, - default=None, - help='Path of a quantized model checkpoint') - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -# TODO: Seems all convert checkpoints may use following utility functions. -# Maybe in one common version. -def reorder_qkv_weight_or_bias(weight: torch.Tensor, - head_dim: int, - num_heads: int, - num_kv_heads: Optional[int] = None, - tp_size: int = 1, - is_bias: bool = False) -> torch.Tensor: - """ Reorder the qkv weight for TRT-LLM use. - - The shape of the fused QKV weights in HF is different from the shape that - TRT-LLM requires. In particular, the weight of HF consists of interleaved - q, k, v head weights, while that of TRT-LLM is contiguous. - HF : [q1, k1, v1, ..., qh, kh, vh] - TRT-LLM: [q1, ..., qh, k1, ..., kh, v1, vh] - where qi, vi, ki are weight vectors corresponding to attention head i. - It's similar to multi/grouped query attention cases. - - We reorder and split the weight of an attention layer to fit into TRT-LLM. - The reordered weight and bias will be - weight: (T, Qh * D + 2 * KVh * D, H) - bias : (T, Qh * D + 2 * KVh * D) - where T=tp_size, Qh=local_num_q_heads, KVh=local_num_kv_heads, D=head_dim, - H=hidden_dim. In the multi/grouped query attention, the number of K/V - attention heads are less than that of Q attention, so that K/V attention - heads may be shared across different ranks if necessary. - - For tensor parallelism, we use the first dimension to select the - corresponding weights. - """ - - # Query types and expected kv heads. - # - Conventional MHA: num_heads = num_kv_heads - # - Multi-Query Attention: num_kv_heads = 1 - # - Grouped-Query Attention: num_heads % num_kv_heads = 0 - num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads - assert num_heads % num_kv_heads == 0, \ - f'num_heads({num_heads}) must be divisible by ' \ - f'num_kv_heads({num_kv_heads})).' - - # The number of attention heads per group: N q head + 1 k head + 1 v head. - num_group_heads = num_heads // num_kv_heads + 2 - assert weight.shape[0] == num_kv_heads * num_group_heads * head_dim, \ - f'{weight.shape[0]} != {num_kv_heads} * {num_group_heads} * {head_dim}' - - qkv_in = num_heads * head_dim if not is_bias else 1 - - # Split Q/K/V weights - weight = weight.reshape(num_kv_heads, num_heads // num_kv_heads + 2, - head_dim, qkv_in) - q_w = weight[:, :-2, ...] # (nKV, num_heads // nKV, head_dim, qkv_in) - k_w = weight[:, -2:-1, ...] # (nKV, 1, head_dim, qkv_in) - v_w = weight[:, -1:, ...] # (nKV, 1, head_dim, qkv_in) - - if num_kv_heads < num_heads and num_kv_heads < tp_size: - # Duplicate K/V heads to make sure that each rank has at least one - # K/V heads. For instance, num_heads=8, num_kv_heads=2, tp_size=4, - # we will make the qkv weight as below. - # Orig: [q0 q1 q2 q3 k0 v0 q4 q5 q6 q7 k1 v0 v1] - # >>>> [[q0 q1 k0 v0], [q2 q3 k0 v0], [q4 q5 k1 v1], [q6 q7 k1 v1]] - assert tp_size % num_kv_heads == 0 - num_dups = tp_size // num_kv_heads - - # k_w and v_w have the same shape. - new_shape = (num_kv_heads, num_dups) + k_w.shape[2:] - k_w = torch.broadcast_to(k_w, size=new_shape) - v_w = torch.broadcast_to(v_w, size=new_shape) - - # Update the number of kv heads. - num_kv_heads = tp_size - - reordered = torch.concat( - [ - q_w.reshape(tp_size, num_heads // tp_size, head_dim, qkv_in), - k_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - v_w.reshape(tp_size, num_kv_heads // tp_size, head_dim, qkv_in), - ], - dim=1, - ) - - qkv_out = (num_heads + 2 * num_kv_heads) // tp_size * head_dim - return reordered.reshape((tp_size, qkv_out, -1)) - - -def get_gptq_gptneox_group_size(quant_ckpt_path, hf_config): - gptq_model = safe_open(quant_ckpt_path, framework="pt", device=0) - gptq_prefix = "gpt_neox." - split_sym = "." - - def load(key, no_prefix=0): - if no_prefix: - return gptq_model.get_tensor(key).cpu() - else: - return gptq_model.get_tensor(gptq_prefix + key).cpu() - - hidden_size = hf_config.hidden_size - prefix = "layers" + split_sym + "0" + split_sym - - scales_fp16 = load(prefix + 'attention.query_key_value.scales') - return hidden_size // scales_fp16.shape[0] - - -def load_from_gptq_gptneox(quant_ckpt_path, - hf_config=None, - use_parallel_embedding=False, - sharding_dim=0, - mapping=Mapping(), - dtype='float16'): - tensorrt_llm.logger.info( - 'Loading weights from groupwise GPTQ LLaMA safetensors...') - weights = {} - tik = time.time() - - gptq_model = safe_open(quant_ckpt_path, framework="pt", device=0) - gptq_prefix = "gpt_neox." - gptq_suffix_list = [".qweight", ".qzeros", ".scales"] - split_sym = "." - - packer = torch.ops.trtllm.pack_int8_tensor_to_packed_int4 - preprocessor = torch.ops.trtllm.preprocess_weights_for_mixed_gemm - torch_dtype = str_dtype_to_torch(dtype) - - def load(key, no_prefix=0): - if no_prefix: - return gptq_model.get_tensor(key).cpu() - else: - return gptq_model.get_tensor(gptq_prefix + key).cpu() - - def torch_split(v, dim): - if v.shape[dim] % mapping.tp_size != 0: - tensorrt_llm.logger.error( - "Current weight shape is invalid for mapping.tp_size=" + - str(mapping.tp_size)) - assert False, "Invalid TP size" - return v.split(v.shape[dim] // mapping.tp_size, - dim=dim)[mapping.tp_rank].contiguous() - - def unpack_int32_into_int8(w_packed): - # Unpack inputs packed in int32/float32 into uint4 and store them in int8 format - w_packed_int4x2 = w_packed.contiguous().view(torch.uint8) - w_unpacked = torch.zeros(w_packed_int4x2.shape[0], - w_packed_int4x2.shape[1] * 2, - dtype=torch.int8, - device=w_packed.device) - w_unpacked[:, ::2] = w_packed_int4x2 % 16 - w_unpacked[:, 1::2] = w_packed_int4x2 // 16 - return w_unpacked.contiguous() - - def process_and_assign_weight(v: List[torch.Tensor], - tllm_prex: str, - tp_dim: int = -1): - if tp_dim == -1: - qweight_int32, qzeros_int32, scales_fp16 = [ - item.cpu() for item in v - ] - else: - qweight_int32, qzeros_int32, scales_fp16 = [ - torch_split(item, tp_dim).cpu() for item in v - ] - - USE_UINT4_INPUT = 1 # Set to true if checkpoint store UINT4 weights - USE_GPTQ_FOR_LLAMA = 1 # GPTQ-for-LLaMA added 1 to zeros - - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).T.contiguous() - 8 - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - # zeros = zeros * scales - qzeros_unpacked_int32 = unpack_int32_into_int8(qzeros_int32) - if not USE_UINT4_INPUT: - # Correcting UINT4 values back to INT4 order - mask_negative = qzeros_unpacked_int32[qzeros_unpacked_int32 < 0] - mask_positive = qzeros_unpacked_int32[qzeros_unpacked_int32 >= 0] - qzeros_unpacked_int32 = qzeros_unpacked_int32 + 16 * mask_negative - 16 * mask_positive - zeros_x_scales_fp16 = (-qzeros_unpacked_int32 + 8 * USE_UINT4_INPUT - - USE_GPTQ_FOR_LLAMA) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - results = { - f'{tllm_prex}.weight': qweight_interleaved, - f'{tllm_prex}.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.zero': zeros_x_scales_fp16, - } - return results - - def preprocess_groupwise_weight_params(qweight_unpacked_int8, scales_fp16, - qzeros_unpacked_int8): - UINT4_TO_INT4_FLAG = 1 - GPTQ_FLAG = 1 - - qweight_interleaved = preprocessor(packer(qweight_unpacked_int8), - torch.quint4x2, - torch.float16).view(torch.float16) - - # zeros = zeros * scales - zeros_x_scales_fp16 = (-qzeros_unpacked_int8 + 8 * UINT4_TO_INT4_FLAG - - GPTQ_FLAG) * scales_fp16 - zeros_x_scales_fp16 = zeros_x_scales_fp16.half() - - # return processed interleaved weight, original scales and zeros * scales - return qweight_interleaved.contiguous(), scales_fp16.contiguous( - ), zeros_x_scales_fp16.contiguous() - - # Load weights from GPTQ checkpoint into TRT-LLM module - # 1. vocab_embedding - v = load('embed_in.weight') - if mapping.is_first_pp_rank(): - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = v.to(torch_dtype) - else: - assert hf_config.vocab_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = torch_split( - v, sharding_dim).to(torch_dtype) - # 2. lm_head - if mapping.is_last_pp_rank(): - v = load('embed_out.weight', no_prefix=1) - weights['lm_head.weight'] = torch_split(v, 0).to(torch_dtype) - - # 3. ln_f - v = load('final_layer_norm.weight') - b = load('final_layer_norm.bias') - if mapping.is_last_pp_rank(): - weights['transformer.ln_f.weight'] = v.to(torch_dtype) - weights['transformer.ln_f.bias'] = b.to(torch_dtype) - # 4. Weights inside each layer - num_hidden_layers = hf_config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - layer_idx = l - layers_range[0] - prefix = "layers" + split_sym + str(l) + split_sym - tensorrt_llm.logger.info(f'Process weights in layer: {layer_idx}') - # layer = tensorrt_llm_llama.layers[layer_idx] - tllm_prex = f'transformer.layers.{l - layers_range[0]}' - # 4.1 attention.qkv - num_heads = hf_config.num_attention_heads - hidden_size = hf_config.hidden_size - head_size = hidden_size // num_heads - - qweight_int32 = load(prefix + 'attention.query_key_value.qweight') - scales_fp16 = load(prefix + 'attention.query_key_value.scales') - qzeros_int32 = load(prefix + 'attention.query_key_value.qzeros') - biases_fp16 = load(prefix + 'attention.query_key_value.bias') - GROUP_SIZE = hidden_size // scales_fp16.shape[0] - - # [hidden_size // 8, hidden_size * 3] -> [hidden_size * 3, hidden_size] - qweight_unpacked_int8 = unpack_int32_into_int8( - qweight_int32.T).contiguous() - 8 - # [hidden_size // GROUP_SIZE, hidden_size * 3 // 8] -> - # [hidden_size // GROUP_SIZE, hidden_size * 3] - qzeros_unpacked_int8 = unpack_int32_into_int8(qzeros_int32) - - # qkv_weights [num_heads x (q|k|v), hidden_size] -> - # [(num_heads x q)|(num_heads x k)|(num_heads x v), hidden_size] - new_qkv_weight_shape = torch.Size( - [num_heads, 3, head_size * qweight_unpacked_int8.size()[-1]]) - # [hidden_size * 3, hidden_size] - qweight_unpacked_int8 = qweight_unpacked_int8.view( - new_qkv_weight_shape).permute(1, 0, 2).reshape( - [hidden_size * 3, hidden_size]).contiguous() - - new_qkv_scale_shape = torch.Size( - [num_heads, 3, head_size * (hidden_size // GROUP_SIZE)]) - # [hidden_size * 3, hidden_size // GROUP_SIZE] - scales_fp16 = scales_fp16.T.contiguous().view( - new_qkv_scale_shape).permute(1, 0, 2).reshape( - [hidden_size * 3, hidden_size // GROUP_SIZE]).contiguous() - - new_qkv_zero_shape = torch.Size( - [num_heads, 3, head_size * (hidden_size // GROUP_SIZE)]) - # [hidden_size * 3, hidden_size // GROUP_SIZE] - qzeros_unpacked_int8 = qzeros_unpacked_int8.T.contiguous().view( - new_qkv_zero_shape).permute(1, 0, 2).reshape( - [hidden_size * 3, hidden_size // GROUP_SIZE]).contiguous() - - new_qkv_bias_shape = torch.Size([num_heads, 3, head_size]) - biases_fp16 = biases_fp16.view(new_qkv_bias_shape).permute( - 1, 0, 2).reshape([hidden_size * 3]) - - tp_size = mapping.tp_size - - if tp_size > 1: - qweight_unpacked_int8 = qweight_unpacked_int8.reshape( - [3, hidden_size, hidden_size]) - qweight_unpacked_int8 = torch_split(qweight_unpacked_int8, dim=1) - qweight_unpacked_int8 = qweight_unpacked_int8.reshape( - [3 * hidden_size // tp_size, hidden_size]) - - scales_fp16 = scales_fp16.reshape( - [3, hidden_size, hidden_size // GROUP_SIZE]) - scales_fp16 = torch_split(scales_fp16, dim=1) - scales_fp16 = scales_fp16.reshape( - [3 * hidden_size // tp_size, hidden_size // GROUP_SIZE]) - - qzeros_unpacked_int8 = qzeros_unpacked_int8.reshape( - [3, hidden_size, hidden_size // GROUP_SIZE]) - qzeros_unpacked_int8 = torch_split(qzeros_unpacked_int8, dim=1) - qzeros_unpacked_int8 = qzeros_unpacked_int8.reshape( - [3 * hidden_size // tp_size, hidden_size // GROUP_SIZE]) - - biases_fp16 = biases_fp16.reshape([3, hidden_size]) - biases_fp16 = torch_split(biases_fp16, dim=1) - biases_fp16 = biases_fp16.reshape([3 * hidden_size // tp_size]) - - qweight_fp32, scales_fp16, zeros_fp16 = preprocess_groupwise_weight_params( - qweight_unpacked_int8.T.contiguous(), scales_fp16.T.contiguous(), - qzeros_unpacked_int8.T.contiguous()) - weights.update({ - f'{tllm_prex}.attention.qkv.weight': qweight_fp32, - f'{tllm_prex}.attention.qkv.weights_scaling_factor': scales_fp16, - f'{tllm_prex}.attention.qkv.zero': zeros_fp16, - f'{tllm_prex}.attention.qkv.bias': biases_fp16, - }) - - # 4.2 attention.dense - v = [load(prefix + 'attention.dense' + suf) for suf in gptq_suffix_list] - # pre scaling down for duplicated bias add between different tp ranks - b = load(prefix + 'attention.dense.bias') / mapping.tp_size - - weights.update( - process_and_assign_weight(v, - f'{tllm_prex}.attention.dense', - tp_dim=0)) - weights.update({f'{tllm_prex}.attention.dense.bias': b.to(torch_dtype)}) - # 4.3 mlp.fc - v = [ - load(prefix + 'mlp.dense_h_to_4h' + suf) for suf in gptq_suffix_list - ] - b = load(prefix + 'mlp.dense_h_to_4h.bias') - weights.update( - process_and_assign_weight(v, f'{tllm_prex}.mlp.fc', tp_dim=1)) - weights.update( - {f'{tllm_prex}.mlp.fc.bias': torch_split(b, dim=0).to(torch_dtype)}) - # 4.4 mlp.proj - v = [ - load(prefix + 'mlp.dense_4h_to_h' + suf) for suf in gptq_suffix_list - ] - # pre scaling down for duplicated bias add between different tp ranks - b = load(prefix + 'mlp.dense_4h_to_h.bias') / mapping.tp_size - - weights.update( - process_and_assign_weight(v, f'{tllm_prex}.mlp.proj', tp_dim=0)) - weights.update({f'{tllm_prex}.mlp.proj.bias': b.to(torch_dtype)}) - # 4.5 input_layernorm - v = load(prefix + 'input_layernorm.weight') - b = load(prefix + 'input_layernorm.bias') - weights[f'{tllm_prex}.input_layernorm.weight'] = v.to(torch_dtype) - weights[f'{tllm_prex}.input_layernorm.bias'] = b.to(torch_dtype) - - # 4.6 post_layernorm - v = load(prefix + 'post_attention_layernorm.weight') - b = load(prefix + 'post_attention_layernorm.bias') - weights[f'{tllm_prex}.post_attention_layernorm.weight'] = v.to( - torch_dtype) - weights[f'{tllm_prex}.post_attention_layernorm.bias'] = b.to( - torch_dtype) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - tensorrt_llm.logger.info(f'Weights loaded. Total time: {t}') - - return weights - - -def split_qkv_weight(weight: torch.Tensor, - hidden_size: int, - num_heads: int, - tp_size: int, - rank: int, - is_bias: bool, - num_kv_heads: Optional[int] = None) -> torch.Tensor: - """ Splits the QKV matrix according to tensor parallelism """ - head_dim = hidden_size // num_heads - weight = reorder_qkv_weight_or_bias(weight, - head_dim=head_dim, - num_heads=num_heads, - num_kv_heads=num_kv_heads, - tp_size=tp_size, - is_bias=is_bias) - - # Copy a sliced tensor to prevent memory leak. A sliced tensor shares the - # memory buffer of the original tensor. So, returning without copying makes - # the buffer of a loaded "qkv" be referenced, resulting GC can't release - # those weights until the whole process ends. - if not is_bias: - return weight[rank, ...].clone().contiguous() - else: - return weight[rank, ...].ravel().clone().contiguous() - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + 'weight'] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + 'weight'] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_gptneox(hf_model, - mapping: Mapping, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - tensor_parallel = mapping.tp_size - rank = mapping.rank - - for l in range(hf_model.config.num_hidden_layers): - prefix = f'gpt_neox.layers.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - qkv_weight, qkv_bias = get_weight_and_bias( - model_params, prefix + 'attention.query_key_value', dtype) - qkv_w = split_qkv_weight(qkv_weight, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_attention_heads) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = split_qkv_weight(qkv_bias, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_attention_heads) - weights.update( - get_tllm_linear_weight(qkv_w, tllm_prex + 'attention.qkv.', qkv_b, - use_weight_only, - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, prefix + 'attention.dense', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_h_to_4h', dtype) - split_v = split_matrix_tp(mlp_fc_weight, tensor_parallel, rank, dim=0) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, prefix + 'mlp.dense_4h_to_h', dtype) - split_v = split_matrix_tp(mlp_proj_weight, tensor_parallel, rank, dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, prefix + 'input_layernorm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - weights[tllm_prex + 'input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, prefix + 'post_attention_layernorm', dtype) - weights[tllm_prex + 'post_attention_layernorm.weight'] = post_ln_weight - weights[tllm_prex + 'post_attention_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'gpt_neox.embed_in', dtype) - lm_head_w = get_weight(model_params, 'embed_out', dtype) - - weights['lm_head.weight'] = split_matrix_tp(lm_head_w, - tensor_parallel, - rank, - dim=0) - - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - assert hf_model.config.vocab_size % tensor_parallel == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix_tp( - embed_w, tensor_parallel, rank, dim=sharding_dim) - - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'gpt_neox.final_layer_norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - tensorrt_llm.logger.set_level('info') - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - elif args.use_weight_only and args.weight_only_precision == 'int4_gptq': - quant_algo = QuantAlgo.W4A16_GPTQ - - hf_config = AutoConfig.from_pretrained(args.model_dir) - hf_model = AutoModelForCausalLM.from_pretrained(args.model_dir, - dtype="auto") - - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'hidden_size': hf_config.hidden_size, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': hf_config.max_position_embeddings, - 'rotary_emb_base': hf_config.rotary_emb_base, - 'rotary_pct': hf_config.rotary_pct, - 'hidden_act': hf_config.hidden_act, - 'quantization': { - 'quant_algo': quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - } - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - config['quantization'].update({ - 'has_zero_point': - True, - 'group_size': - get_gptq_gptneox_group_size(args.quant_ckpt_path, hf_config) - }) - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if args.use_weight_only and args.weight_only_precision == 'int4_gptq': - weights = load_from_gptq_gptneox( - args.quant_ckpt_path, - hf_config, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - mapping=mapping, - dtype=args.dtype) - else: - weights = convert_hf_gptneox( - hf_model, - mapping, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim) - safe_save_path = os.path.join(args.output_dir, - f'rank{rank}.safetensors') - tensorrt_llm.logger.info(f'Saving safetensors to: {safe_save_path}') - safetensors.torch.save_file(weights, safe_save_path) - tensorrt_llm.logger.info(f'Saved safetensors to: {safe_save_path}') - return True - - if args.workers == 1: - for rank in range(world_size): - passed = covert_and_save(rank) - assert passed, "Convert checkpoint failed" - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/gptneox/gptq_convert.sh b/examples/models/contrib/gptneox/gptq_convert.sh deleted file mode 100644 index b0a89462cd1a..000000000000 --- a/examples/models/contrib/gptneox/gptq_convert.sh +++ /dev/null @@ -1,9 +0,0 @@ -git clone https://github.com/qwopqwop200/GPTQ-for-LLaMa.git GPTQ-for-LLaMa - -pip install -r ./GPTQ-for-LLaMa/requirements.txt - -CUDA_VISIBLE_DEVICES=0 python3 GPTQ-for-LLaMa/neox.py ./gptneox_model \ -wikitext2 \ ---wbits 4 \ ---groupsize 128 \ ---save_safetensors ./gptneox_model/gptneox-20b-4bit-gs128.safetensors diff --git a/examples/models/contrib/gptneox/requirements.txt b/examples/models/contrib/gptneox/requirements.txt deleted file mode 100644 index cc77e27f78ee..000000000000 --- a/examples/models/contrib/gptneox/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -rouge_score -evaluate diff --git a/examples/models/contrib/grok/README.md b/examples/models/contrib/grok/README.md deleted file mode 100644 index 41bcd8505b62..000000000000 --- a/examples/models/contrib/grok/README.md +++ /dev/null @@ -1,92 +0,0 @@ -# Grok-1 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run grok-1 model in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -- [Grok1](#Grok-1) - - [Prerequisite](#prerequisite) - - [Hardware](#hardware) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - -## Prerequisite -First of all, please clone the official grok-1 code repo with below commands and install the dependencies. -```bash -git clone https://github.com/xai-org/grok-1.git /path/to/folder -``` -And then downloading the weights per [instructions](https://github.com/xai-org/grok-1?tab=readme-ov-file#downloading-the-weights). - -## Hardware -The grok-1 model requires a node with 8x80GB GPU memory(at least). - -## Overview - -The TensorRT LLM Grok-1 implementation can be found in [tensorrt_llm/models/grok/model.py](../../../../tensorrt_llm/models/grok/model.py). The TensorRT LLM Grok-1 example code is located in [`examples/models/contrib/grok`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the Grok-1 model into TensorRT LLM checkpoint format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * INT8 Weight-Only - * Tensor Parallel - * STRONGLY TYPED - -## Usage - -The TensorRT LLM Grok-1 example code locates at [examples/models/contrib/grok](./). It takes xai weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first to make sure the example uses matched `tensorrt_llm` version: - -```bash -pip install -r requirements.txt -``` - -Need to prepare the Grok-1 checkpoint by following the guides here https://github.com/xai-org/grok-1. - -TensorRT LLM Grok-1 builds TensorRT engine(s) from Xai's checkpoints. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `workers` feature only supports single node. - - -Below is the step-by-step to run Grok-1 with TensorRT LLM. - -```bash -# Build the bfloat16 engine from xai official weights. -python convert_checkpoint.py --model_dir ./tmp/grok-1/ \ - --output_dir ./tllm_checkpoint_8gpus_bf16 \ - --dtype bfloat16 \ - --use_weight_only \ - --tp_size 8 \ - --workers 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpus_bf16 \ - --output_dir ./tmp/grok-1/trt_engines/bf16/8-gpus \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 \ - --moe_plugin bfloat16 \ - --paged_kv_cache enable \ - --remove_input_padding enable \ - --workers 8 - - -# Run Grok-1 with 8 GPUs -mpirun -n 8 --allow-run-as-root \ - python ../../../run.py \ - --input_text "The answer to life the universe and everything is of course" \ - --engine_dir ./tmp/grok-1/trt_engines/bf16/8-gpus \ - --max_output_len 50 --top_p 1 --top_k 8 --temperature 0.3 \ - --vocab_file ./tmp/grok-1/tokenizer.model -``` diff --git a/examples/models/contrib/grok/convert_checkpoint.py b/examples/models/contrib/grok/convert_checkpoint.py deleted file mode 100644 index 18fb429eb04b..000000000000 --- a/examples/models/contrib/grok/convert_checkpoint.py +++ /dev/null @@ -1,366 +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. -import argparse -import json -import os -import sys -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import numpy as np - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.layers import MoeConfig -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import GrokForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--weights_dir', type=str, default=None) - - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--vocab_size', type=int, default=32000) - parser.add_argument('--n_positions', type=int, default=2048) - parser.add_argument('--n_layer', type=int, default=32) - parser.add_argument('--n_head', type=int, default=32) - parser.add_argument('--n_kv_head', type=int, default=None) - parser.add_argument('--n_embd', type=int, default=4096) - parser.add_argument('--inter_size', type=int, default=11008) - parser.add_argument('--rms_norm_eps', type=float, default=1e-06) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - - parser.add_argument('--load_by_shard', - action='store_true', - help='Load a pretrained model shard-by-shard.') - parser.add_argument('--hidden_act', type=str, default='silu') - - parser.add_argument('--rotary_base', type=float, default=10000.0) - - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument( - '--moe_num_experts', - default=0, - type=int, - help='Specify the number of experts to use for MOE layers') - parser.add_argument( - '--moe_top_k', - default=0, - type=int, - help= - 'Specify the top_k value to use for MOE layers. Default to 1 if --moe_num_experts is set' - ) - parser.add_argument( - '--moe_tp_size', - type=int, - default=-1, - help= - 'N-way tensor parallelism size for MOE, default is tp_size, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_ep_size', - type=int, - default=-1, - help= - 'N-way expert parallelism size for MOE, default is 1, which will do tp-only for MoE' - ) - parser.add_argument( - '--moe_renorm_mode', - default=MoeConfig.ExpertScaleNormalizationMode.RENORMALIZE, - type=int, - help= - 'Controls renormalization after gate logits. Check layers/moe.py for accepted values', - ) - parser.add_argument( - '--save_config_only', - action="store_true", - default=False, - help= - 'Only save the model config w/o read and converting weights, be careful, this is for debug only' - ) - - args = parser.parse_args() - # changing the default to be consistent as the cli help said. - if args.moe_num_experts and args.moe_top_k == 0: - args.moe_top_k = 1 - return args - - -def args_to_quantization(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - - return quant_config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin - } - - -def from_cli_args(args): - n_kv_head = args.n_kv_head if args.n_kv_head is not None else args.n_head - config = { - 'architecture': "LlamaForCausalLM", - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': args.n_layer, - 'num_attention_heads': args.n_head, - 'hidden_size': args.n_embd, - 'intermediate_size': args.inter_size, - 'num_key_value_heads': n_kv_head, - 'vocab_size': args.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'max_position_embeddings': args.n_positions, - 'hidden_act': args.hidden_act, - 'rotary_base': args.rotary_base, - 'norm_epsilon': args.rms_norm_eps, - 'moe_num_experts': args.moe_num_experts, - 'moe_top_k': args.moe_top_k, - 'moe_normalization_mode': args.moe_renorm_mode, - 'mapping': { - 'world_size': args.tp_size * args.pp_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - 'moe_tp_size': args.moe_tp_size, - 'moe_ep_size': args.moe_ep_size, - }, - 'quantization': args_to_quantization(args).asdict() - } - config.update(args_to_build_options(args)) - return config - - -def preload_model(model_dir, weights_dir=None): - sys.path.append(model_dir) - from model import LanguageModelConfig, TransformerConfig - from runners import ModelRunner - if weights_dir and os.path.exists(weights_dir): - CKPT_PATH = weights_dir - else: - CKPT_PATH = os.path.join(model_dir, "checkpoints") - - grok_1_model = LanguageModelConfig( - vocab_size=128 * 1024, - pad_token=0, - eos_token=2, - sequence_len=8192, - embedding_init_scale=1.0, - output_multiplier_scale=0.5773502691896257, - embedding_multiplier_scale=78.38367176906169, - model=TransformerConfig( - emb_size=48 * 128, - widening_factor=8, - key_size=128, - num_q_heads=48, - num_kv_heads=8, - num_layers=64, - attn_output_multiplier=0.08838834764831845, - shard_activations=True, - # MoE. - num_experts=8, - num_selected_experts=2, - # Activation sharding. - data_axis="data", - model_axis="model", - ), - ) - - runner = ModelRunner( - model=grok_1_model, - bs_per_device=0.125, - checkpoint_path=CKPT_PATH, - ) - dummy_data = dict( - inputs=np.zeros((1, 256), dtype=np.int32), - targets=np.zeros((1, 256), dtype=np.int32), - ) - runner.transform_forward = True - runner.initialize(dummy_data, (1, 8), (1, 1)) - - params = runner.load_or_init(dummy_data) - - return params - - -def convert_and_save_xai(args): - model_dir = args.model_dir - load_by_shard = args.load_by_shard - world_size = args.tp_size * args.pp_size - if (args.moe_tp_size == -1 and args.moe_ep_size == -1): - # moe default to tp-only - args.moe_tp_size = args.tp_size - args.moe_ep_size = 1 - elif (args.moe_tp_size == -1): - args.moe_tp_size = args.tp_size // args.moe_ep_size - elif (args.moe_ep_size == -1): - args.moe_ep_size = args.tp_size // args.moe_tp_size - assert (args.moe_tp_size * args.moe_ep_size == args.tp_size - ), "moe_tp_size * moe_ep_size must equal to tp_size" - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - quantization = args_to_quantization(args) - override_fields.update(args_to_build_options(args)) - - # When not loading by shard, preload one complete model and then slice per rank weights from this - # this saves the disk reloading time - xai_model = preload_model( - model_dir, args.weights_dir) if not args.load_by_shard else None - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - moe_tp_size=args.moe_tp_size, - moe_ep_size=args.moe_ep_size) - grok = GrokForCausalLM.from_hugging_face( - model_dir, - args.dtype, - mapping=mapping, - quantization=quantization, - load_by_shard=load_by_shard, - override_fields=override_fields, - preloaded_model=xai_model, - ) - grok.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del grok - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - - args.tp_size * args.pp_size - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_dir is None: # generate fake config.json - config = from_cli_args(args) - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - else: # all other non-gptq paths from hf model - assert args.model_dir is not None - convert_and_save_xai(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/grok/requirements.txt b/examples/models/contrib/grok/requirements.txt deleted file mode 100644 index 3ab319c0efc5..000000000000 --- a/examples/models/contrib/grok/requirements.txt +++ /dev/null @@ -1,10 +0,0 @@ --f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -sentencepiece==0.2.0 -jax[cuda12-pip]==0.4.29 -jaxlib[cuda12-pip]==0.4.29 -dm_haiku==0.0.12 diff --git a/examples/models/contrib/mmdit/README.md b/examples/models/contrib/mmdit/README.md deleted file mode 100644 index 4a3d80986566..000000000000 --- a/examples/models/contrib/mmdit/README.md +++ /dev/null @@ -1,95 +0,0 @@ -# MMDiT in SD 3 & SD 3.5 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a [MMDiT](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_sd3.py) in Stable Diffusion 3/3.5 with TensorRT-LLM. - -## Overview - -The TensorRT LLM implementation of MMDiT can be found in [tensorrt_llm/models/sd3/model.py](../../../../tensorrt_llm/models/mmdit_sd3/model.py). The TensorRT LLM MMDiT (SD 3/3.5) example code is located in [`examples/models/contrib/mmdit`](./). There are main files to build and run MMDiT with TensorRT-LLM: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the MMDiT model into TensorRT LLM checkpoint format. -* [`sample.py`](./sample.py) to run the [diffusers](https://huggingface.co/docs/diffusers/index) pipeline with TensorRT engine(s) to generate images. - -## Support Matrix - -- [x] TP -- [x] CP -- [ ] FP8 - -## Usage - -The TensorRT LLM MMDiT example code locates at [examples/models/contrib/mmdit](./). It takes HuggingFace checkpoint as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build MMDiT TensorRT engine(s) - -This checkpoint will be converted to the TensorRT LLM checkpoint format by [`convert_checkpoint.py`](./convert_checkpoint.py). After that, we can build TensorRT engine(s) with the TensorRT LLM checkpoint. - -``` -# Convert to TRT-LLM -python convert_checkpoint.py --model_path='stabilityai/stable-diffusion-3.5-medium' -trtllm-build --checkpoint_dir=./tllm_checkpoint/ \ - --max_batch_size=2 \ - --remove_input_padding=disable \ - --bert_attention_plugin=auto -``` - -Set `--max_batch_size` to tell how many images at most you would like to generate. We disable `--remove_input_padding` since we don't need to padding MMDiT's patches. - -After build, we can find a `./engine_output` directory, it is ready for running MMDiT model with TensorRT LLM now. - -### Generate images - -A [`sample.py`](./sample.py) is provided to generated images with the optimized TensorRT engines. - -If using `float16` for inference, `FusedRMSNorm` from `Apex` used by T5-encoder should be disabled in the [huggingface/transformers](https://github.com/huggingface/transformers/blob/v4.48.3/src/transformers/models/t5/modeling_t5.py#L259) or just uninstall the `apex`: -```python -try: - from apex.normalization import FusedRMSNorm - - # [NOTE] Avoid using `FusedRMSNorm` for T5 encoder. - # T5LayerNorm = FusedRMSNorm # noqa - - logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of T5LayerNorm") -except ImportError: - # using the normal T5LayerNorm - pass -except Exception: - logger.warning("discovered apex but it failed to load, falling back to T5LayerNorm") - pass - -ALL_LAYERNORM_LAYERS.append(T5LayerNorm) -``` - -Just run `python sample.py` and we can see an image named `sd3.5-mmdit.png` will be generated: -![sd3.5-mmdit.png](./assets/sd3.5-mmdit.png). - -### Tensor Parallel - -``` -# Convert to TRT-LLM -python convert_checkpoint.py --tp_size=2 --model_path='stabilityai/stable-diffusion-3.5-medium' -trtllm-build --checkpoint_dir=./tllm_checkpoint/ \ - --max_batch_size=2 \ - --remove_input_padding=disable \ - --bert_attention_plugin=auto -mpirun -n 2 --allow-run-as-root python sample.py "A capybara holding a sign that reads 'Hello World' in the forrest." -``` - -### Context Parallel - -Pipeline with CP is similar to that with TP, but it doesn't support `BertAttention` plugin. And make sure `tensorrt>=10.8.0.43`. - -``` -# Convert to TRT-LLM -python convert_checkpoint.py --tp_size=2 --model_path='stabilityai/stable-diffusion-3.5-medium' -trtllm-build --checkpoint_dir=./tllm_checkpoint/ \ - --max_batch_size=2 \ - --remove_input_padding=disable \ - --bert_attention_plugin=disable -mpirun -n 2 --allow-run-as-root python sample.py "A capybara holding a sign that reads 'Hello World' in the forrest." -``` diff --git a/examples/models/contrib/mmdit/assets/sd3.5-mmdit.png b/examples/models/contrib/mmdit/assets/sd3.5-mmdit.png deleted file mode 100644 index 9ccee43d85f5..000000000000 Binary files a/examples/models/contrib/mmdit/assets/sd3.5-mmdit.png and /dev/null differ diff --git a/examples/models/contrib/mmdit/convert_checkpoint.py b/examples/models/contrib/mmdit/convert_checkpoint.py deleted file mode 100644 index 3584ab309288..000000000000 --- a/examples/models/contrib/mmdit/convert_checkpoint.py +++ /dev/null @@ -1,123 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import SD3Transformer2DModel - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_path', - type=str, - default="stabilityai/stable-diffusion-3.5-medium") - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--cp_size', - type=int, - default=1, - help='N-way context parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='float16', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - - args = parser.parse_args() - return args - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.cp_size - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - cp_size=args.cp_size) - tik = time.time() - mmdit = SD3Transformer2DModel.from_pretrained(args.model_path, - args.dtype, - mapping=mapping) - print(f'Total time of reading and converting: {time.time()-tik:.3f} s') - tik = time.time() - mmdit.save_checkpoint(args.output_dir, save_config=(rank == 0)) - # post-process checkpoint config - with open(f"{args.output_dir}/config.json", 'r') as f: - config = json.load(f) - config["model_path"] = args.model_path - config["use_pretrained_pos_emb"] = True if "pos_embed.pos_embed" in [ - name for name, _ in mmdit.named_parameters() - ] else False - with open(f"{args.output_dir}/config.json", 'w') as f: - json.dump(config, f, indent=4) - del mmdit - print(f'Total time of saving checkpoint: {time.time()-tik:.3f} s') - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - assert args.model_path is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/mmdit/requirements.txt b/examples/models/contrib/mmdit/requirements.txt deleted file mode 100644 index 876bf841255d..000000000000 --- a/examples/models/contrib/mmdit/requirements.txt +++ /dev/null @@ -1,2 +0,0 @@ -transformers>=4.56.0 -diffusers>=0.32.2 diff --git a/examples/models/contrib/mmdit/sample.py b/examples/models/contrib/mmdit/sample.py deleted file mode 100644 index 31da097bd3fd..000000000000 --- a/examples/models/contrib/mmdit/sample.py +++ /dev/null @@ -1,266 +0,0 @@ -import argparse -import json -import os -from functools import wraps -from typing import List, Optional - -import numpy as np -import tensorrt as trt -import torch -from cuda import cudart -from diffusers import StableDiffusion3Pipeline -from diffusers.models.modeling_outputs import Transformer2DModelOutput - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch, trt_dtype_to_torch -from tensorrt_llm.logger import logger -from tensorrt_llm.models.mmdit_sd3.config import SD3Transformer2DModelConfig -from tensorrt_llm.runtime.session import Session - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -class TllmMMDiT(object): - - def __init__(self, - config, - skip_layers: Optional[List[int]] = None, - debug_mode: bool = True, - stream: torch.cuda.Stream = None, - pos_embedder: torch.nn.Module = None): - self.dtype = config['pretrained_config']['dtype'] - # Update config with `skip_layers` - config['pretrained_config']['skip_layers'] = skip_layers - self.config = SD3Transformer2DModelConfig.from_dict( - config['pretrained_config']) - - rank = tensorrt_llm.mpi_rank() - world_size = config['pretrained_config']['mapping']['world_size'] - cp_size = config['pretrained_config']['mapping']['cp_size'] - tp_size = config['pretrained_config']['mapping']['tp_size'] - pp_size = config['pretrained_config']['mapping']['pp_size'] - gpus_per_node = config['pretrained_config']['mapping']['gpus_per_node'] - assert pp_size == 1 - self.mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=rank, - cp_size=cp_size, - tp_size=tp_size, - pp_size=1, - gpus_per_node=gpus_per_node) - - local_rank = rank % self.mapping.gpus_per_node - self.device = torch.device(f'cuda:{local_rank}') - torch.cuda.set_device(self.device) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_file = os.path.join(args.tllm_model_dir, f"rank{rank}.engine") - logger.info(f'Loading engine from {engine_file}') - with open(engine_file, "rb") as f: - engine_buffer = f.read() - - assert engine_buffer is not None - - self.session = Session.from_serialized_engine(engine_buffer) - - self.debug_mode = debug_mode - - self.inputs = {} - self.outputs = {} - self.buffer_allocated = False - - expected_tensor_names = [ - 'hidden_states', 'encoder_hidden_states', 'pooled_projections', - 'timestep', 'output' - ] - - found_tensor_names = [ - self.session.engine.get_tensor_name(i) - for i in range(self.session.engine.num_io_tensors) - ] - if not self.debug_mode and set(expected_tensor_names) != set( - found_tensor_names): - logger.error( - f"The following expected tensors are not found: {set(expected_tensor_names).difference(set(found_tensor_names))}" - ) - logger.error( - f"Those tensors in engine are not expected: {set(found_tensor_names).difference(set(expected_tensor_names))}" - ) - logger.error(f"Expected tensor names: {expected_tensor_names}") - logger.error(f"Found tensor names: {found_tensor_names}") - raise RuntimeError( - "Tensor names in engine are not the same as expected.") - if self.debug_mode: - self.debug_tensors = list( - set(found_tensor_names) - set(expected_tensor_names)) - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name)) - return dtype - - def _setup(self, outputs_shape): - for i in range(self.session.engine.num_io_tensors): - name = self.session.engine.get_tensor_name(i) - if self.session.engine.get_tensor_mode( - name) == trt.TensorIOMode.OUTPUT: - shape = list(self.session.engine.get_tensor_shape(name)) - if self.debug_mode: - shape = list(self.session.engine.get_tensor_shape(name)) - if shape[0] == -1: - shape[0] = outputs_shape['output'][0] - else: - shape = outputs_shape[name] - self.outputs[name] = torch.empty(shape, - dtype=self._tensor_dtype(name), - device=self.device) - - self.buffer_allocated = True - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - def __call__(self, *args, **kwargs): - return self.forward(*args, **kwargs) - - @cuda_stream_guard - def forward( - self, - hidden_states: torch.Tensor, - encoder_hidden_states: torch.Tensor = None, - pooled_projections: torch.Tensor = None, - timestep: torch.LongTensor = None, - block_controlnet_hidden_states: List = None, - joint_attention_kwargs=None, - return_dict: bool = True, - skip_layers: Optional[List[int]] = None, - ): - if block_controlnet_hidden_states is not None: - raise NotImplementedError("ControlNet is not supported yet.") - if joint_attention_kwargs is not None: - if joint_attention_kwargs.get("scale", None) is not None: - raise NotImplementedError("LoRA is not supported yet.") - if "ip_adapter_image_embeds" in joint_attention_kwargs: - raise NotImplementedError("IP-Adapter is not supported yet.") - if skip_layers is not None: - raise RuntimeWarning( - "Please initialize model with `skip_layers` instead of passing via `forward`." - ) - - self._setup(outputs_shape={'output': hidden_states.shape}) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - inputs = { - 'hidden_states': - hidden_states.to(str_dtype_to_torch(self.dtype)), - 'encoder_hidden_states': - encoder_hidden_states.to(str_dtype_to_torch(self.dtype)), - 'pooled_projections': - pooled_projections.to(str_dtype_to_torch(self.dtype)), - 'timestep': - timestep.to(str_dtype_to_torch(self.dtype)), - } - for k, v in inputs.items(): - inputs[k] = v.cuda().contiguous() - self.inputs.update(**inputs) - self.session.set_shapes(self.inputs) - ok = self.session.run(self.inputs, self.outputs, - self.stream.cuda_stream) - - if not ok: - raise RuntimeError('Executing TRT engine failed!') - output = self.outputs['output'].to(hidden_states.device) - - if self.debug_mode: - torch.cuda.synchronize() - output_np = { - k: v.cpu().float().numpy() - for k, v in self.outputs.items() - } - np.savez("tllm_output.npz", **output_np) - - if not return_dict: - return (output, ) - else: - return Transformer2DModelOutput(sample=output) - - -def main(args): - tensorrt_llm.logger.set_level(args.log_level) - assert torch.cuda.is_available() - - config_file = os.path.join(args.tllm_model_dir, 'config.json') - with open(config_file) as f: - config = json.load(f) - - pipe = StableDiffusion3Pipeline.from_pretrained( - config['pretrained_config']['model_path'], torch_dtype=torch.float16) - pipe.to("cuda") - - # replace sd3.5 transformer with TRTLLM model - torch_pos_embed = pipe.transformer.pos_embed - del pipe.transformer - torch.cuda.empty_cache() - - # Load model: - model = TllmMMDiT(config, - debug_mode=args.debug_mode, - pos_embedder=torch_pos_embed) - - pipe.transformer = model - - image = pipe(args.prompt, - height=1024, - width=1024, - guidance_scale=4.5, - num_inference_steps=40, - generator=torch.Generator("cpu").manual_seed(0)).images[0] - if tensorrt_llm.mpi_rank() == 0: - image.save("sd3.5-mmdit.png") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - 'prompt', - nargs='*', - default= - "A capybara holding a sign that reads 'Hello World' in the forrest.", - help="Text prompt(s) to guide image generation") - parser.add_argument("--tllm_model_dir", - type=str, - default='./engine_outputs/') - parser.add_argument("--gpus_per_node", type=int, default=8) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument("--debug_mode", action='store_true') - args = parser.parse_args() - main(args) diff --git a/examples/models/contrib/mpt/.gitignore b/examples/models/contrib/mpt/.gitignore deleted file mode 100644 index a6c46e9436fc..000000000000 --- a/examples/models/contrib/mpt/.gitignore +++ /dev/null @@ -1,3 +0,0 @@ -mpt* -*.log -c-model diff --git a/examples/models/contrib/mpt/README.md b/examples/models/contrib/mpt/README.md deleted file mode 100644 index 60aacbc5d1c2..000000000000 --- a/examples/models/contrib/mpt/README.md +++ /dev/null @@ -1,182 +0,0 @@ -# MPT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [MPT](https://huggingface.co/mosaicml/mpt-7b) model using TensorRT LLM and run on a single GPU and a single node with multiple GPUs. - -- [MPT](#mpt) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [MPT 7B](#mpt-7b) - - [1.1 Convert from HF Transformers in FP](#11-convert-from-hf-transformers-in-fp) - - [1.2 Convert from HF Transformers with weight-only quantization](#12-convert-from-hf-transformers-with-weight-only-quantization) - - [1.3 Convert from HF Transformers with SmoothQuant quantization](#13-convert-from-hf-transformers-with-smoothquant-quantization) - - [1.4 Convert from HF Transformers with INT8 KV cache quantization](#14-convert-from-hf-transformers-with-int8-kv-cache-quantization) - - [1.5 AWQ weight-only quantization with Modelopt](#15-awq-weight-only-quantization-with-modelopt) - - [1.6 FP8 Post-Training Quantization with Modelopt](#16-fp8-post-training-quantization-with-modelopt) - - [1.6 Weight-only quantization with Modelopt](#16-weight-only-quantization-with-modelopt) - - [1.7 SmoothQuant and INT8 KV cache with Modelopt](#17-smoothquant-and-int8-kv-cache-with-modelopt) - - [2.1 Build TensorRT engine(s)](#21-build-tensorrt-engines) - - [MPT 30B](#mpt-30b) - - [1. Convert weights from HF Transformers to TRTLLM format](#1-convert-weights-from-hf-transformers-to-trtllm-format) - - [2. Build TensorRT engine(s)](#2-build-tensorrt-engines) - - [3. Run TRT engine to check if the build was correct](#3-run-trt-engine-to-check-if-the-build-was-correct) - -## Overview - -The TensorRT LLM MPT implementation can be found in [`tensorrt_llm/models/mpt/model.py`](../../tensorrt_llm/models/mpt/model.py). The TensorRT LLM MPT example code is located in [`examples/models/contrib/mpt`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * FP8 (with FP8 KV Cache) - * INT8 & INT4 Weight-Only - * INT8 Smooth Quant - * INT4 AWQ - * Tensor Parallel - * MHA, MQA & GQA - * STRONGLY TYPED - -### MPT 7B - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script allows you to convert weights from HF Transformers format to TRTLLM checkpoints. - -#### 1.1 Convert from HF Transformers in FP - -```bash -# Generate FP16 checkpoints. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp16/ --dtype float16 - -# Generate FP32 checkpoints with TP=4. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp32_tp4/ --dtype float32 --tp_size 4 -``` - -#### 1.2 Convert from HF Transformers with weight-only quantization - -```bash -# Use int8 weight-only quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int8_wo/ --use_weight_only - -# Use int4 weight-only quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int4_wo/ --use_weight_only --weight_only_precision int4 -``` - -#### 1.3 Convert from HF Transformers with SmoothQuant quantization - -```bash -# Use int8 smoothquant (weight and activation) quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int8_sq/ --smoothquant 0.5 -``` - -#### 1.4 Convert from HF Transformers with INT8 KV cache quantization - -```bash -# Use int8 kv cache quantization. -python convert_checkpoint.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp16_int8kv/ --dtype float16 --calibrate_kv_cache -``` -***INT8-KV-cache can be used with SQ and Weight-only at the same time*** - - -***We now introduce Modelopt to do all quantization*** -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -#### 1.5 AWQ weight-only quantization with Modelopt - -```bash -# INT4 AWQ quantization using Modelopt. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int4_awq/ --qformat int4_awq -``` - -#### 1.6 FP8 Post-Training Quantization with Modelopt - -```bash -# FP8 quantization using Modelopt. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/fp8/ --qformat fp8 --kv_cache_dtype fp8 -``` - -#### 1.6 Weight-only quantization with Modelopt - -```bash -# INT8 Weight-only quantization using Modelopt with TP=2. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int8_wo/ --qformat int8_wo --tp_size 2 - -# INT4 Weight-only quantization using Modelopt. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/int4_wo/ --qformat int4_wo -``` - -#### 1.7 SmoothQuant and INT8 KV cache with Modelopt - -```bash -# Use int4 awq quantization. -python ../../../quantization/quantize.py --model_dir mosaicml/mpt-7b --output_dir ./ckpts/mpt-7b/sq_int8kv/ --qformat int8_sq --kv_cache_dtype int8 -``` -***INT8-KV-cache can also be used with Weight-only at the same time*** - - -### 2.1 Build TensorRT engine(s) - -All of the checkpoint generated by `convert_checkpoint.py` or `quantize.py` (Modelopt) can share the same building commands. - -```bash -# Build a single-GPU float16 engine using TRTLLM checkpoints. -trtllm-build --checkpoint_dir=./ckpts/mpt-7b/fp16 \ - --max_batch_size 32 \ - --max_input_len 1024 \ - --max_seq_len 1536 \ - --gemm_plugin float16 \ - --workers 1 \ - --output_dir ./trt_engines/mpt-7b/fp16 -``` - -### MPT 30B - -Same commands can be changed to convert MPT 30B to TRT LLM format. Below is an example to build MPT30B fp16 4-way tensor parallelized TRT engine - -#### 1. Convert weights from HF Transformers to TRTLLM format - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script allows you to convert weights from HF Transformers format to TRTLLM format. - -```bash -python convert_checkpoint.py --model_dir mosaicml/mpt-30b --output_dir ./ckpts/mpt-30b/fp16_tp4/ --tp_szie 4 --dtype float16 -``` - -#### 2. Build TensorRT engine(s) - -Examples of build invocations: - -```bash -# Build 4-GPU MPT-30B float16 engines -trtllm-build --checkpoint_dir ./ckpts/mpt-30b/fp16_tp4 \ - --max_batch_size 32 \ - --max_input_len 1024 \ - --max_seq_len 1536 \ - --gemm_plugin float16 \ - --workers 4 \ - --output_dir ./trt_engines/mpt-30b/fp16_tp4 -``` - -#### 3. Run TRT engine to check if the build was correct - -```bash -# Run 4-GPU MPT-30B TRT engine on a sample input prompt -mpirun -n 4 --allow-run-as-root \ - python ../../../run.py --max_output_len 10 \ - --engine_dir ./trt_engines/mpt-30b/fp16/4-gpu/ \ - --tokenizer_dir mosaicml/mpt-30b -``` diff --git a/examples/models/contrib/mpt/convert_checkpoint.py b/examples/models/contrib/mpt/convert_checkpoint.py deleted file mode 100644 index d8af794a93d0..000000000000 --- a/examples/models/contrib/mpt/convert_checkpoint.py +++ /dev/null @@ -1,923 +0,0 @@ -import argparse -import copy -import functools -import json -import os -import time -import traceback -from collections import defaultdict -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Dict, Optional - -import numpy as np -import safetensors -import torch -import torch.nn as nn -from tqdm import tqdm -from transformers import AutoTokenizer, MptConfig, MptForCausalLM -from transformers.pytorch_utils import Conv1D - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.convert_utils import (generate_int8, get_weight, - load_calib_dataset, smooth_gemm, - split) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - "--calibrate_kv_cache", - "-kv", - action="store_true", - help= - "Generate scaling factors for KV cache. Used for storing KV cache in int8." - ) - parser.add_argument( - '--per_channel', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument("--dataset_cache_dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -@torch.no_grad() -def capture_activation_range(model, - tokenizer, - dataset, - num_samples=1, - seq_len=512): - model.eval() - device = next(model.parameters()).device - act_scales = defaultdict(lambda: {"x": None, "y": None, "w": None}) - - tokenizer.pad_token = tokenizer.eos_token - - def stat_tensor(name, tensor, act_scales, key): - hidden_dim = tensor.shape[-1] - tensor = tensor.view(-1, hidden_dim).abs().detach() - comming_max = torch.max(tensor, dim=0)[0].float() - - if act_scales[name][key] is None: - act_scales[name][key] = comming_max - else: - act_scales[name][key] = torch.max(act_scales[name][key], - comming_max) - - def stat_input_hook(m, x, y, name): - if isinstance(x, tuple): - x = x[0] - stat_tensor(name, x, act_scales, "x") - stat_tensor(name, y, act_scales, "y") - - if act_scales[name]["w"] is None: - act_scales[name]["w"] = m.weight.abs().clip( - 1e-8, None).max(dim=1)[0].float() - - hooks = [] - for name, m in model.named_modules(): - if isinstance(m, nn.Linear) or isinstance(m, Conv1D): - hooks.append( - m.register_forward_hook( - functools.partial(stat_input_hook, name=name))) - - for i in tqdm(range(num_samples), desc="calibrating model"): - datapoint = dataset[i:i + 1] - line = copy.copy(datapoint) - line[0] = line[0] + ' TL;DR: ' - line[0] = line[0].strip() - line[0] = line[0].replace(" n't", "n't") - input_ids = tokenizer(line, - return_tensors="pt", - max_length=seq_len, - padding=True, - truncation=True).input_ids.to(device) - model(input_ids) - - for h in hooks: - h.remove() - - return act_scales - - -@torch.no_grad() -def smooth_mpt_model(model, scales, alpha, mpt_qkv_para, mpt_smoother): - # Smooth the activation and weights with smoother = $\diag{s}$ - for name, module in model.named_modules(): - if not isinstance(module, type(model.transformer.blocks[0])): - continue - # qkv_proj - layer_name_qkv = name + ".attn.Wqkv" - weight = module.attn.Wqkv.weight - smoother = smooth_gemm(weight, scales[layer_name_qkv]["x"], - module.norm_1.weight, module.norm_1.bias, alpha) - scales[layer_name_qkv]["x"] = scales[layer_name_qkv]["x"] / smoother - scales[layer_name_qkv]["w"] = weight.abs().max(dim=1)[0] - # see transpose_weights function - mpt_qkv_para[layer_name_qkv] = weight.transpose(0, 1) - - # ================================================================= - layer_name = name + ".attn.out_proj" - smoother = smooth_gemm(module.attn.out_proj.weight, - scales[layer_name]["x"], None, None, alpha) - mpt_smoother[layer_name] = smoother.float() - - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.attn.out_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - fc1_layer_name = name + ".ffn.up_proj" - - smoother = smooth_gemm(module.ffn.up_proj.weight, - scales[fc1_layer_name]["x"], - module.norm_2.weight, module.norm_2.bias, alpha) - - scales[fc1_layer_name]["x"] = scales[fc1_layer_name]["x"] / smoother - scales[fc1_layer_name]["w"] = module.ffn.up_proj.weight.abs().max( - dim=1)[0] - - # ================================================================== - layer_name = name + ".ffn.down_proj" - smoother = smooth_gemm(module.ffn.down_proj.weight, - scales[layer_name]["x"], None, None, alpha) - mpt_smoother[layer_name] = smoother.float() - scales[layer_name]["x"] = scales[layer_name]["x"] / smoother - scales[layer_name]["w"] = module.ffn.down_proj.weight.abs().max( - dim=1)[0] - - -def get_tllm_linear_sq_weight(vals, - prefix, - shape, - tensor_parallel, - is_qkv=False, - per_token=False, - per_channel=False, - last_prefix=None, - bias=None, - smoother_value=None, - smoother_shape=None, - rank=0, - cat_dim=0, - multi_query_mode=False): - results = {} - - def multi_query_split(data, local_dim, head_size, tp_size, cur_rank): - q, k, v = np.split(data, [local_dim, local_dim + head_size], axis=-1) - q_split = np.split(q, tp_size, axis=-1) - k_split = np.split(k, tp_size, axis=-1) - v_split = np.split(v, tp_size, axis=-1) - return [ - np.concatenate((q_split[ii], k_split[ii], v_split[ii]), axis=-1) - for ii in range(tp_size) - ][cur_rank] - - col_shape = shape if (is_qkv or per_channel) else [1, 1] - - if per_token: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - if smoother_value is None: - results[last_prefix] = torch.from_numpy( - np.array([1.0], dtype=np.float32)) - - if per_channel: - cur_per_channel_value = vals["scale_w_quant_orig.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_w_quant_orig.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_w_quant_orig"] - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_w_quant_orig"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split(vals["scale_w_quant_orig"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array(cur_per_channel_value, - dtype=np.float32).reshape(col_shape)).contiguous() - else: - if per_channel: - original_weights = np.array(vals["weight.int8.col"]) - else: - original_weights = np.array(vals["weight.int8"]) - local_dim = original_weights.shape[0] - head_size = (original_weights.shape[1] - local_dim) // 2 - - if multi_query_mode: - cur_weights = multi_query_split(original_weights, local_dim, - head_size, tensor_parallel, rank) - else: - cur_weights = np.split(original_weights, - tensor_parallel, - axis=cat_dim)[rank] - if is_qkv: - hidden_dim = cur_weights.shape[0] - cur_weights = cur_weights.reshape(hidden_dim, -1) - results[prefix + - 'weight'] = torch.from_numpy(cur_weights).t().contiguous() - - if per_channel: - cur_per_channel_value = vals["scale_y_accum_quant.col"] - if smoother_value is None: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant.col"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant.col"], - tensor_parallel, - axis=cat_dim)[rank] - else: - cur_per_channel_value = vals["scale_y_accum_quant"] - # QKV is always per_channel - if is_qkv: - if multi_query_mode: - cur_per_channel_value = multi_query_split( - vals["scale_y_accum_quant"], local_dim, head_size, - tensor_parallel, rank) - else: - cur_per_channel_value = np.split( - vals["scale_y_accum_quant"], - tensor_parallel, - axis=cat_dim)[rank] - - results[prefix + 'per_channel_scale'] = torch.from_numpy( - np.array([cur_per_channel_value], - dtype=np.float32).reshape(col_shape)).contiguous() - - results[last_prefix] = torch.from_numpy( - np.array([vals['scale_x_orig_quant']], - dtype=np.float32)).contiguous() - - results[prefix + 'act_scale'] = torch.from_numpy( - np.array([[vals["scale_y_quant_orig"]]], - dtype=np.float32)).contiguous() - - if smoother_value is not None: - cur_smoother_value = np.split(smoother_value, - tensor_parallel, - axis=cat_dim)[rank] - results[prefix + 'smoother'] = cur_smoother_value.reshape( - smoother_shape).contiguous().to(torch.float32) - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def split_qkv_tp(qkv, n_head, n_kv_heads, n_hidden, tensor_parallel, rank): - """ - Splits the QKV matrix according to tensor parallelism - """ - kv_head_size = n_kv_heads * (n_hidden // n_head) - q, k, v = torch.split(qkv, [n_hidden, kv_head_size, kv_head_size], dim=0) - q = split(q, tensor_parallel, rank, dim=0) - k = split(k, tensor_parallel, rank, dim=0) - v = split(v, tensor_parallel, rank, dim=0) - return torch.concatenate([q, k, v], dim=0).contiguous() - - -def split_matrix(weight: torch.Tensor, tp_size: int, rank: int, - dim: int) -> torch.Tensor: - return split(weight, tp_size, rank, dim=dim) - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def get_tllm_param( - param: torch.Tensor, - name: str, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if name.endswith('.weight') and use_weight_only: - v = param.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[name] = processed_torch_weights - results[name.replace('weight', - 'per_channel_scale')] = torch_weight_scales - else: - results[name] = param - - return results - - -def convert_hf_mpt_legacy(hf_model, - hf_config, - mapping, - rank=0, - dtype='float32', - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only=False, - plugin_weight_only_quant_type='int8', - use_smooth_quant=False, - per_channel=False, - per_token=False, - int8_kv_cache=False, - act_range=[], - qkv_para=[], - smoother=[]): - weights = {} - tik = time.time() - tensor_parallel = mapping.tp_size - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_attention_heads = hf_model.config.n_heads - hidden_size = hf_model.config.d_model - vocab_size = hf_model.config.vocab_size - num_key_value_heads = hf_config.attn_config['kv_n_heads'] if 'kv_n_heads' in hf_config.attn_config \ - else hf_config.n_heads - multi_query_mode = (num_key_value_heads != num_attention_heads) - - for l in range(hf_model.config.n_layers): - prefix = f'transformer.blocks.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - # attn.Wqkv -> attention.qkv - qkv_weight = get_weight(model_params, prefix + 'attn.Wqkv', dtype) - - if use_smooth_quant: - qkv_out_dim = qkv_weight.shape[0] - qkv_weight = qkv_weight.t().numpy() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + 'attn.Wqkv'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights.update( - get_tllm_linear_sq_weight(int8_weights, - tllm_prex + 'attention.qkv.', - [1, qkv_out_dim // tensor_parallel], - tensor_parallel, - is_qkv=True, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'input_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1, - multi_query_mode=multi_query_mode)) - else: - qkv_weight = split_qkv_tp(qkv_weight, num_attention_heads, - num_key_value_heads, hidden_size, - mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(qkv_weight, tllm_prex + 'attention.qkv', - None, use_weight_only, - plugin_weight_only_quant_type)) - - if int8_kv_cache: - qkv_weight = get_weight(model_params, prefix + 'attn.Wqkv', dtype) - qkv_weight = qkv_weight.t().numpy() - if not multi_query_mode: - qkv_weight = qkv_weight.reshape(hidden_size, 3, hidden_size) - int8_weights = generate_int8(qkv_weight, - act_range.get(prefix + 'attn.Wqkv'), - is_qkv=True, - multi_query_mode=multi_query_mode) - weights[tllm_prex + - 'attention.kv_cache_scaling_factor'] = torch.from_numpy( - np.array([int8_weights['scale_y_quant_orig']], - dtype=np.float32)).contiguous() - - # attn.out_proj -> attention.dense - attn_dense_weight = get_weight(model_params, prefix + 'attn.out_proj', - dtype) - if use_smooth_quant: - attn_dense_weight = attn_dense_weight.t().numpy() - int8_weights = generate_int8( - attn_dense_weight, act_range.get(prefix + 'attn.out_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'attention.dense.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + - 'attention.quantization_scaling_factor', - smoother_value=smoother[(prefix + 'attn.out_proj')], - smoother_shape=[1, hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, - tllm_prex + 'attention.dense', None, - use_weight_only, - plugin_weight_only_quant_type)) - - # ffn.up_proj -> mlp.fc - mlp_fc_weight = get_weight(model_params, prefix + 'ffn.up_proj', dtype) - if use_smooth_quant: - mlp_fc_weight = mlp_fc_weight.t().numpy() - int8_weights = generate_int8(mlp_fc_weight, - act_range.get(prefix + 'ffn.up_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.fc.', - [1, 4 * hidden_size // tensor_parallel], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'post_layernorm.scale_to_int', - smoother_value=None, - smoother_shape=None, - rank=rank, - cat_dim=-1)) - else: - mlp_fc_weight = split_matrix(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_weight, tllm_prex + 'mlp.fc', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # ffn.down_proj -> mlp.proj - mlp_proj_weight = get_weight(model_params, prefix + 'ffn.down_proj', - dtype) - if use_smooth_quant: - mlp_proj_weight = mlp_proj_weight.t().numpy() - int8_weights = generate_int8( - mlp_proj_weight, act_range.get(prefix + 'ffn.down_proj')) - weights.update( - get_tllm_linear_sq_weight( - int8_weights, - tllm_prex + 'mlp.proj.', [1, hidden_size], - tensor_parallel, - is_qkv=False, - per_token=per_token, - per_channel=per_channel, - last_prefix=tllm_prex + 'mlp.quantization_scaling_factor', - smoother_value=smoother[prefix + 'ffn.down_proj'], - smoother_shape=[1, 4 * hidden_size // tensor_parallel], - rank=rank, - cat_dim=0)) - else: - mlp_proj_weight = split_matrix(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_weight, tllm_prex + 'mlp.proj', - None, use_weight_only, - plugin_weight_only_quant_type)) - - # input layer_norm - input_ln_weight = get_weight(model_params, prefix + 'norm_1', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, prefix + 'norm_2', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - - embed_w = get_weight(model_params, 'transformer.wte', dtype) - if mapping.is_first_pp_rank(): - # Embedding - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - if sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix( - embed_w, mapping.tp_size, mapping.tp_rank, sharding_dim) - if mapping.is_last_pp_rank(): - # lm_head weight and bias - weights['lm_head.weight'] = split_matrix(embed_w.clone(), - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'transformer.norm_f', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -def convert_hf_mpt(hf_model: MptForCausalLM, - hf_config: MptConfig, - mapping: Mapping, - dtype: str = 'float32', - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8): - - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - num_hidden_layers = hf_config.n_layers - num_head = hf_config.n_heads - num_kv_heads = getattr(hf_config.attn_config, 'kv_n_heads', - hf_config.n_heads) - hidden_size = hf_config.d_model - vocab_size = hf_config.vocab_size - - layers_range = mapping.pp_layers(num_hidden_layers) - for l in layers_range: - prefix = f'transformer.blocks.{l}' - tllm_prex = f'transformer.layers.{l-layers_range[0]}' - # Attention QKV (no bias) - qkv_w = get_weight(model_params, f'{prefix}.attn.Wqkv', dtype) - qkv_w = split_qkv_tp(qkv_w, num_head, num_kv_heads, hidden_size, - mapping.tp_size, mapping.tp_rank) - weights.update( - get_tllm_linear_weight(qkv_w, f'{tllm_prex}.attention.qkv', None, - use_weight_only, - plugin_weight_only_quant_type)) - # Attention dense (no bias) - attn_dense_weight = get_weight(model_params, f'{prefix}.attn.out_proj', - dtype) - attn_dense_w = split_matrix(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(attn_dense_w, f'{tllm_prex}.attention.dense', - None, use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_in (no bias) - mlp_fc_weight = get_weight(model_params, f'{prefix}.ffn.up_proj', dtype) - mlp_fc_w = split_matrix(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', None, - use_weight_only, - plugin_weight_only_quant_type)) - # MLP fc_out (no bias) - mlp_proj_weight = get_weight(model_params, f'{prefix}.ffn.down_proj', - dtype) - mlp_proj_w = split_matrix(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', None, - use_weight_only, - plugin_weight_only_quant_type)) - # input layer_norm - input_ln_weight = get_weight(model_params, f'{prefix}.norm_1', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - - # post layer_norm - post_ln_weight = get_weight(model_params, f'{prefix}.norm_2', dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - - embed_w = get_weight(model_params, 'transformer.wte', dtype) - if mapping.is_first_pp_rank(): - # Embedding - if not use_parallel_embedding: - weights['transformer.vocab_embedding.weight'] = embed_w - else: - if sharding_dim == 0: - assert vocab_size % mapping.tp_size == 0 - else: - assert hidden_size % mapping.tp_size == 0 - weights['transformer.vocab_embedding.weight'] = split_matrix( - embed_w, mapping.tp_size, mapping.tp_rank, sharding_dim) - if mapping.is_last_pp_rank(): - # lm_head weight and bias - weights['lm_head.weight'] = split_matrix(embed_w.clone(), - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = get_weight(model_params, 'transformer.norm_f', dtype) - # ln_f weight and bias - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.fastertransformer.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - world_size = args.tp_size * args.pp_size - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - - if args.smoothquant: - if args.per_token and args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif not args.per_token and not args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - elif not args.per_token and args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif args.per_token and not args.per_channel: - quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - - if args.calibrate_kv_cache: - kv_cache_quant_algo = QuantAlgo.INT8 - else: - kv_cache_quant_algo = None - - hf_config = MptConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - num_kv_heads = getattr(hf_config.attn_config, 'kv_n_heads', - hf_config.n_heads) - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': args.logits_dtype, - 'vocab_size': hf_config.vocab_size, - 'hidden_size': hf_config.d_model, - 'intermediate_size': hf_config.d_model * 4, - 'num_hidden_layers': hf_config.n_layers, - 'num_attention_heads': hf_config.n_heads, - 'num_key_value_heads': num_kv_heads, - 'position_embedding_type': 'alibi', - 'hidden_act': 'gelu', - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'bias': (not hf_config.no_bias), - 'clip_qkv': hf_config.attn_config.clip_qkv, - 'alibi_bias_max': hf_config.attn_config.alibi_bias_max - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - hf_model = MptForCausalLM.from_pretrained(args.model_dir, - device_map="auto", - dtype=getattr(torch, args.dtype)) - - act_range = {} - mpt_qkv_para = {} - # smoother for inputs of self_attn.o_proj and mlp.down_proj - mpt_smoother = {} - if args.smoothquant is not None or args.calibrate_kv_cache: - tokenizer = AutoTokenizer.from_pretrained(args.model_dir, - padding_side='left') - dataset = load_calib_dataset(args.calib_dataset, - cache_dir=args.dataset_cache_dir) - - act_range = capture_activation_range(hf_model, tokenizer, dataset) - if args.smoothquant is not None: - smooth_mpt_model(hf_model, act_range, args.smoothquant, - mpt_qkv_para, mpt_smoother) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - if args.smoothquant is not None or args.calibrate_kv_cache: - weights = convert_hf_mpt_legacy( - hf_model, - hf_config, - mapping, - rank, - dtype=args.dtype, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_smooth_quant=(args.smoothquant is not None), - per_channel=args.per_channel, - per_token=args.per_token, - int8_kv_cache=args.calibrate_kv_cache, - act_range=act_range, - qkv_para=mpt_qkv_para, - smoother=mpt_smoother) - else: - weights = convert_hf_mpt( - hf_model, - hf_config, - mapping, - dtype=args.dtype, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - del hf_model - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/mpt/requirements.txt b/examples/models/contrib/mpt/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/mpt/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/opt/README.md b/examples/models/contrib/opt/README.md deleted file mode 100644 index 609c7ebfb910..000000000000 --- a/examples/models/contrib/opt/README.md +++ /dev/null @@ -1,252 +0,0 @@ -# OPT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [OPT](https://huggingface.co/docs/transformers/model_doc/opt) model using TensorRT LLM and run on a single GPU, a single node with -multiple GPUs or multiple nodes with multiple GPUs. - -- [OPT](#opt) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace Transformers](#1-download-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Summarization using the OPT model](#4-summarization-using-the-opt-model) - - [Fused MultiHead Attention (FMHA)](#fused-multihead-attention-fmha) - - [Tensor Parallelism for Embedding Lookup Table.](#tensor-parallelism-for-embedding-lookup-table) - - [1. Enable this feature](#1-enable-this-feature) - - [2. Choose the dimension for tensor parallelism](#2-choose-the-dimension-for-tensor-parallelism) - -## Overview - -The TensorRT LLM OPT implementation can be found in [`tensorrt_llm/models/opt/model.py`](../../tensorrt_llm/models/opt/model.py). The TensorRT LLM OPT example code is located in [`examples/models/contrib/opt`](./). There is one file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format - -In addition, there are two shared files in the parent folder [`examples`](../) for inference and evaluation: - -* [`../../../run.py`](../../../run.py) to run the inference on an input text; -* [`../../../summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * INT8 & INT4 Weight-Only - * Tensor Parallel - -## Usage - -The next two sections describe how to convert the weights from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) -format to the TensorRT LLM format. - -### 1. Download weights from HuggingFace Transformers - -You have to make sure `git-lfs` is properly installed to load the checkpoints. - -```bash -pip install -r requirements.txt && sudo apt-get install git-lfs -``` - -There are four different checkpoints available. Use one of the following commands to fetch the checkpoint you are interested in. - -```bash -# OPT-125M -git-lfs clone https://huggingface.co/facebook/opt-125m - -# OPT-350M -git-lfs clone https://huggingface.co/facebook/opt-350m - -# OPT-2.7B -git-lfs clone https://huggingface.co/facebook/opt-2.7b - -# OPT-66B -git-lfs clone https://huggingface.co/facebook/opt-66b -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -```bash -# OPT-125M -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --output_dir ./opt/125M/trt_ckpt/fp16/1-gpu/ - -# OPT-350M -python3 convert_checkpoint.py --model_dir ./opt-350m \ - --dtype float16 \ - --output_dir ./opt/350M/trt_ckpt/fp16/1-gpu/ - -# OPT-2.7B -python3 convert_checkpoint.py --model_dir ./opt-2.7b \ - --dtype float16 \ - --output_dir ./opt/2.7B/trt_ckpt/fp16/1-gpu/ - -# OPT-66B -python3 convert_checkpoint.py --model_dir ./opt-66b \ - --dtype float16 \ - --tp_size 4 \ - --output_dir ./opt/66B/trt_ckpt/fp16/4-gpu/ \ - --workers 2 -``` - -### 3. Build TensorRT engine(s) - -```bash -# OPT-125M -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/1-gpu/ - -# OPT-350M -trtllm-build --checkpoint_dir ./opt/350M/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/350M/trt_engines/fp16/1-gpu/ - -# OPT-2.7B -trtllm-build --checkpoint_dir ./opt/2.7B/trt_ckpt/fp16/1-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/2.7B/trt_engines/fp16/1-gpu/ - -# OPT-66B -trtllm-build --checkpoint_dir ./opt/66B/trt_ckpt/fp16/4-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/66B/trt_engines/fp16/4-gpu/ \ - --workers 2 -``` - -### 4. Summarization using the OPT model - -The following section describes how to run a TensorRT LLM OPT model to summarize the articles from the -[cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -The script can also perform the same summarization using the HF OPT model. - -```bash -# OPT-125M -python3 ../../../summarize.py --engine_dir ./opt/125M/trt_engines/fp16/1-gpu/ \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-125m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 - -# OPT-350M -python3 ../../../summarize.py --engine_dir ./opt/350M/trt_engines/fp16/1-gpu/ \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-350m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=20 - -# OPT-2.7B -python3 ../../../summarize.py --engine_dir ./opt/2.7B/trt_engines/fp16/1-gpu/ \ - --test_hf \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-2.7b \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=20 - -# OPT-66B -mpirun -n 4 --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./opt/66B/trt_engines/fp16/4-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-66b \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=20 -``` - -#### Fused MultiHead Attention (FMHA) - -You can enable the FMHA kernels for OPT by adding `--enable_context_fmha` to the invocation of `trtllm-build`. Note that it is disabled by default because of possible accuracy issues due to the use of Flash Attention. - -If you find that the default fp16 accumulation (`--context_fmha`) cannot meet the requirement, you can try to enable fp32 accumulation by adding `--enable_context_fmha_fp32_acc` to the inference command (`run.py` or `summarize.py`). However, it is expected to see performance drop. - -Note `--context_fmha` has to be used together with `--gpt_attention_plugin float16`. - -## Tensor Parallelism for Embedding Lookup Table. -Since the embedding lookup table can be several gigabytes in size. We can distribute this weight across multiple GPUs in order to reduce the memory consumption per GPU. - -### 1. Enable this feature -To enable this feature, add the flag `--use_parallel_embedding` to `trtllm-build`. - -### 2. Choose the dimension for tensor parallelism - -Assume the size of embedding lookup table is (vocab\_size \* hidden\_size), we can shard it along the vocab\_size (`--embedding_sharding_dim 0`) or hidden\_size (`--embedding_sharding_dim 1`) dimension. - -2.1 To shard the embedding lookup table along the hidden\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 1`. Here is an example: - -```Bash -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --output_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 1 -``` -2.2 To shard the embedding lookup table along the vocab\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 0`. - -Meanwhile, we provide a lookup plugin to support tensor parallelism on vocab\_size dimension. - -- An example of sharing along vocab\_size dimension with lookup plugin: - -```Bash -python3 convert_checkpoint.py --model_dir ./opt-125m \ - --dtype float16 \ - --output_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 - -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --workers 2 - -mpirun -n 2 --allow-run-as-root \ - python3 ../../../summarize.py --engine_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --batch_size 1 \ - --test_trt_llm \ - --hf_model_dir opt-125m \ - --data_type fp16 \ - --check_accuracy \ - --tensorrt_llm_rouge1_threshold=14 -``` - -- An example of sharing along vocab\_size dimension without lookup plugin: - -```Bash -trtllm-build --checkpoint_dir ./opt/125M/trt_ckpt/fp16/2-gpu/ \ - --gemm_plugin float16 \ - --max_batch_size 8 \ - --max_input_len 924 \ - --max_seq_len 1024 \ - --output_dir ./opt/125M/trt_engines/fp16/2-gpu/ \ - --workers 2 -``` diff --git a/examples/models/contrib/opt/convert_checkpoint.py b/examples/models/contrib/opt/convert_checkpoint.py deleted file mode 100644 index 37b3b610f24e..000000000000 --- a/examples/models/contrib/opt/convert_checkpoint.py +++ /dev/null @@ -1,361 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import safetensors -import torch -from transformers import AutoModelForCausalLM, Blip2ForConditionalGeneration - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.models.convert_utils import (get_weight, get_weight_and_bias, - split, split_matrix_tp, - split_qkv_bias_tp, split_qkv_tp) -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument( - '--model_type', - type=str, - default='opt', - choices=['opt', 'blip2'], - help= - 'Multimodal type when this script is used for multimodal conversion.') - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - args = parser.parse_args() - - return args - - -def split_embedding( - param: torch.Tensor, - tp_size: int, - tp_rank: int, - use_parallel_embedding: bool = False, - sharding_dim: int = 0, -) -> torch.Tensor: - if param is None: - return None - if not use_parallel_embedding: - return param - - vocab_size, hidden_size = param.size() - if sharding_dim == 0: - if vocab_size % tp_size != 0: - vocab_size_padded = pad_vocab_size(vocab_size, tp_size) - pad_width = vocab_size_padded - vocab_size - param = torch.nn.functional.pad(param, (0, 0, 0, pad_width), - value=0) - else: - assert hidden_size % tp_size == 0 - return split(param, tp_size, tp_rank, dim=sharding_dim) - - -def get_tllm_linear_weight(weight, - prefix, - bias=None, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[prefix + 'weight'] = processed_torch_weights - results[prefix + 'per_channel_scale'] = torch_weight_scales - else: - results[prefix + 'weight'] = weight.contiguous() - - if bias is not None: - results[prefix + 'bias'] = bias - - return results - - -def convert_hf_opt(hf_model, - rank=0, - tensor_parallel=1, - dtype='float32', - use_parallel_embedding=False, - sharding_dim=0, - use_weight_only=False, - plugin_weight_only_quant_type=torch.int8): - - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - dtype = getattr(torch, dtype) - do_layer_norm_before = hf_model.config.do_layer_norm_before - num_attention_heads = hf_model.config.num_attention_heads - hidden_size = hf_model.config.hidden_size - - for l in range(hf_model.config.num_hidden_layers): - prefix = f'model.decoder.layers.{l}.' - tllm_prex = f'transformer.layers.{l}.' - - q_weight, q_bias = get_weight_and_bias(model_params, - prefix + 'self_attn.q_proj', - dtype) - k_weight, k_bias = get_weight_and_bias(model_params, - prefix + 'self_attn.k_proj', - dtype) - v_weight, v_bias = get_weight_and_bias(model_params, - prefix + 'self_attn.v_proj', - dtype) - qkv_weight = torch.cat([q_weight, k_weight, v_weight], dim=0) - split_v = split_qkv_tp(qkv_weight, num_attention_heads, hidden_size, - tensor_parallel, rank) - qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=0) - bias = split_qkv_bias_tp(qkv_bias, num_attention_heads, hidden_size, - tensor_parallel, rank) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.qkv.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - - attn_dense_weight, attn_dense_bias = get_weight_and_bias( - model_params, prefix + 'self_attn.out_proj', dtype) - split_v = split_matrix_tp(attn_dense_weight, - tensor_parallel, - rank, - dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'attention.dense.', - attn_dense_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_fc_weight, mlp_fc_bias = get_weight_and_bias( - model_params, prefix + 'fc1', dtype) - split_v = split_matrix_tp(mlp_fc_weight, tensor_parallel, rank, dim=0) - bias = split_matrix_tp(mlp_fc_bias, tensor_parallel, rank, dim=0) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.fc.', bias, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight, mlp_proj_bias = get_weight_and_bias( - model_params, prefix + 'fc2', dtype) - split_v = split_matrix_tp(mlp_proj_weight, tensor_parallel, rank, dim=1) - weights.update( - get_tllm_linear_weight(split_v, tllm_prex + 'mlp.proj.', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight, input_ln_bias = get_weight_and_bias( - model_params, prefix + 'self_attn_layer_norm', dtype) - weights[tllm_prex + 'input_layernorm.weight'] = input_ln_weight - weights[tllm_prex + 'input_layernorm.bias'] = input_ln_bias - - post_ln_weight, post_ln_bias = get_weight_and_bias( - model_params, prefix + 'final_layer_norm', dtype) - weights[tllm_prex + 'post_layernorm.weight'] = post_ln_weight - weights[tllm_prex + 'post_layernorm.bias'] = post_ln_bias - - embed_w = get_weight(model_params, 'model.decoder.embed_tokens', dtype) - if 'model.decoder.project_in.weight' in model_params.keys(): - project_in = get_weight(model_params, 'model.decoder.project_in', dtype) - project_out = get_weight(model_params, 'model.decoder.project_out', - dtype) - lm_head_w = torch.matmul(embed_w.float(), project_out.float()).to(dtype) - embed_w = torch.matmul(embed_w.float(), - project_in.t().float()).to(dtype) - elif 'lm_head.weight' in model_params.keys(): - lm_head_w = get_weight(model_params, 'lm_head', dtype) - else: - lm_head_w = embed_w.clone() - - weights['lm_head.weight'] = split_matrix_tp(lm_head_w, - tensor_parallel, - rank, - dim=0) - - weights['transformer.vocab_embedding.weight'] = split_embedding( - embed_w, - tp_size=tensor_parallel, - tp_rank=rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - - embed_p = get_weight(model_params, 'model.decoder.embed_positions', dtype) - weights['transformer.position_embedding.weight'] = split_embedding( - embed_p[2:, :], - tp_size=tensor_parallel, - tp_rank=rank, - use_parallel_embedding=use_parallel_embedding, - sharding_dim=sharding_dim) - - if do_layer_norm_before: - ln_f_w, ln_f_b = get_weight_and_bias(model_params, - 'model.decoder.final_layer_norm', - dtype) - weights['transformer.ln_f.weight'] = ln_f_w - weights['transformer.ln_f.bias'] = ln_f_b - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - world_size = args.tp_size * args.pp_size - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_type == 'opt': - hf_model = AutoModelForCausalLM.from_pretrained(args.model_dir, - dtype="auto") - elif args.model_type == 'blip2': - hf_model = Blip2ForConditionalGeneration.from_pretrained( - args.model_dir, dtype="auto").language_model - - hf_config = hf_model.config - if hf_config.hidden_size != hf_config.word_embed_proj_dim: - args.use_parallel_embedding = False - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = QuantAlgo.W8A16 - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = QuantAlgo.W4A16 - - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'hidden_size': hf_config.hidden_size, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'learned_absolute', - 'max_position_embeddings': hf_config.max_position_embeddings, - 'hidden_act': hf_config.activation_function, - 'quantization': { - 'quant_algo': quant_algo - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'do_layer_norm_before': hf_config.do_layer_norm_before, - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def covert_and_save(rank): - weights = convert_hf_opt( - hf_model, - rank, - world_size, - dtype=args.dtype, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim) - safetensors.torch.save_file( - weights, os.path.join(args.output_dir, f'rank{rank}.safetensors')) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/contrib/opt/requirements.txt b/examples/models/contrib/opt/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/contrib/opt/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/contrib/stdit/README.md b/examples/models/contrib/stdit/README.md deleted file mode 100644 index 42116be7428e..000000000000 --- a/examples/models/contrib/stdit/README.md +++ /dev/null @@ -1,94 +0,0 @@ -# STDiT in OpenSoRA - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a STDiT in [OpenSoRA](https://github.com/hpcaitech/Open-Sora/tree/main) with TensorRT-LLM. - -## Overview - -The TensorRT LLM implementation of STDiT can be found in [tensorrt_llm/models/stdit/model.py](../../../../tensorrt_llm/models/stdit/model.py). The TensorRT LLM STDiT (OpenSoRA) example code is located in [`examples/models/contrib/stdit`](./). There are main files to build and run STDiT with TensorRT-LLM: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the STDiT model into TensorRT LLM checkpoint format. -* [`sample.py`](./sample.py) to run the pipeline with TensorRT engine(s) to generate videos. - -## Support Matrix - -- [x] TP -- [ ] CP -- [ ] FP8 - -## Usage - -The TensorRT LLM STDiT example code locates at [examples/models/contrib/stdit](./). It takes HuggingFace checkpoint as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Requirements - -Please install required packages first: - -```bash -pip install -r requirements.txt -# ColossalAI is also needed for text encoder. -pip install colossalai --no-deps -``` - -### Build STDiT TensorRT engine(s) - -This checkpoint will be converted to the TensorRT LLM checkpoint format by [`convert_checkpoint.py`](./convert_checkpoint.py). After that, we can build TensorRT engine(s) with the TensorRT LLM checkpoint. The pretrained checkpoint can be downloaded from [here](https://huggingface.co/hpcai-tech/OpenSora-STDiT-v3). - -```bash -# Convert to TRT-LLM -python convert_checkpoint.py --timm_ckpt= -# Build engine -trtllm-build --checkpoint_dir=tllm_checkpoint/ \ - --max_batch_size=2 \ - --gemm_plugin=float16 \ - --kv_cache_type=disabled \ - --remove_input_padding=enable \ - --gpt_attention_plugin=auto \ - --bert_attention_plugin=auto \ - --context_fmha=enable -``` - -After build, we can find a `./engine_output` directory, it is ready for running STDiT model with TensorRT LLM now. - -### Generate videos - -A [`sample.py`](./sample.py) is provided to generated videos with the optimized TensorRT engines. - -```bash -python sample.py "a beautiful waterfall" -``` - -And we can see a video named `sample_outputs/sample_0000.mp4` will be generated: - - - -### Tensor Parallel - -We can levaerage tensor parallel to further reduce latency and memory consumption on each GPU. - -```bash -# Convert to TRT-LLM -python convert_checkpoint.py --tp_size=2 --timm_ckpt= -# Build engines -trtllm-build --checkpoint_dir=tllm_checkpoint/ \ - --max_batch_size=2 \ - --gemm_plugin=float16 \ - --kv_cache_type=disabled \ - --remove_input_padding=enable \ - --gpt_attention_plugin=auto \ - --bert_attention_plugin=auto \ - --context_fmha=enable -# Run example -mpirun -n 2 --allow-run-as-root python sample.py "a beautiful waterfall" -``` - -### Context Parallel - -Not supported yet. diff --git a/examples/models/contrib/stdit/aspect.py b/examples/models/contrib/stdit/aspect.py deleted file mode 100644 index fe629a5da6ec..000000000000 --- a/examples/models/contrib/stdit/aspect.py +++ /dev/null @@ -1,526 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# 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. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/datasets/aspect.py - -import math - - -# computation -def get_h_w(a, ts, eps=1e-4): - h = (ts * a)**0.5 - h = h + eps - h = math.ceil(h) if math.ceil(h) % 2 == 0 else math.floor(h) - w = h / a - w = w + eps - w = math.ceil(w) if math.ceil(w) % 2 == 0 else math.floor(w) - return h, w - - -def get_aspect_ratios_dict(ars, ts=360 * 640): - est = {f"{a:.2f}": get_h_w(a, ts) for a in ars} - return est - - -def get_ar(ratio): - h, w = ratio.split(":") - return int(h) / int(w) - - -# H:W -ASPECT_RATIO_MAP = { - "3:8": "0.38", - "9:21": "0.43", - "12:25": "0.48", - "1:2": "0.50", - "9:17": "0.53", - "27:50": "0.54", - "9:16": "0.56", - "5:8": "0.62", - "2:3": "0.67", - "3:4": "0.75", - "1:1": "1.00", - "4:3": "1.33", - "3:2": "1.50", - "16:9": "1.78", - "17:9": "1.89", - "2:1": "2.00", - "50:27": "2.08", -} - -AR = [get_ar(ratio) for ratio in ASPECT_RATIO_MAP.keys()] - -# computed from above code -# S = 8294400 -ASPECT_RATIO_4K = { - "0.38": (1764, 4704), - "0.43": (1886, 4400), - "0.48": (1996, 4158), - "0.50": (2036, 4072), - "0.53": (2096, 3960), - "0.54": (2118, 3918), - "0.62": (2276, 3642), - "0.56": (2160, 3840), # base - "0.67": (2352, 3528), - "0.75": (2494, 3326), - "1.00": (2880, 2880), - "1.33": (3326, 2494), - "1.50": (3528, 2352), - "1.78": (3840, 2160), - "1.89": (3958, 2096), - "2.00": (4072, 2036), - "2.08": (4156, 1994), -} - -# S = 3686400 -ASPECT_RATIO_2K = { - "0.38": (1176, 3136), - "0.43": (1256, 2930), - "0.48": (1330, 2770), - "0.50": (1358, 2716), - "0.53": (1398, 2640), - "0.54": (1412, 2612), - "0.56": (1440, 2560), # base - "0.62": (1518, 2428), - "0.67": (1568, 2352), - "0.75": (1662, 2216), - "1.00": (1920, 1920), - "1.33": (2218, 1664), - "1.50": (2352, 1568), - "1.78": (2560, 1440), - "1.89": (2638, 1396), - "2.00": (2716, 1358), - "2.08": (2772, 1330), -} - -# S = 2073600 -ASPECT_RATIO_1080P = { - "0.38": (882, 2352), - "0.43": (942, 2198), - "0.48": (998, 2080), - "0.50": (1018, 2036), - "0.53": (1048, 1980), - "0.54": (1058, 1958), - "0.56": (1080, 1920), # base - "0.62": (1138, 1820), - "0.67": (1176, 1764), - "0.75": (1248, 1664), - "1.00": (1440, 1440), - "1.33": (1662, 1246), - "1.50": (1764, 1176), - "1.78": (1920, 1080), - "1.89": (1980, 1048), - "2.00": (2036, 1018), - "2.08": (2078, 998), -} - -# S = 921600 -ASPECT_RATIO_720P = { - "0.38": (588, 1568), - "0.43": (628, 1466), - "0.48": (666, 1388), - "0.50": (678, 1356), - "0.53": (698, 1318), - "0.54": (706, 1306), - "0.56": (720, 1280), # base - "0.62": (758, 1212), - "0.67": (784, 1176), - "0.75": (832, 1110), - "1.00": (960, 960), - "1.33": (1108, 832), - "1.50": (1176, 784), - "1.78": (1280, 720), - "1.89": (1320, 698), - "2.00": (1358, 680), - "2.08": (1386, 666), -} - -# S = 409920 -ASPECT_RATIO_480P = { - "0.38": (392, 1046), - "0.43": (420, 980), - "0.48": (444, 925), - "0.50": (452, 904), - "0.53": (466, 880), - "0.54": (470, 870), - "0.56": (480, 854), # base - "0.62": (506, 810), - "0.67": (522, 784), - "0.75": (554, 738), - "1.00": (640, 640), - "1.33": (740, 555), - "1.50": (784, 522), - "1.78": (854, 480), - "1.89": (880, 466), - "2.00": (906, 454), - "2.08": (924, 444), -} - -# S = 230400 -ASPECT_RATIO_360P = { - "0.38": (294, 784), - "0.43": (314, 732), - "0.48": (332, 692), - "0.50": (340, 680), - "0.53": (350, 662), - "0.54": (352, 652), - "0.56": (360, 640), # base - "0.62": (380, 608), - "0.67": (392, 588), - "0.75": (416, 554), - "1.00": (480, 480), - "1.33": (554, 416), - "1.50": (588, 392), - "1.78": (640, 360), - "1.89": (660, 350), - "2.00": (678, 340), - "2.08": (692, 332), -} - -# S = 102240 -ASPECT_RATIO_240P = { - "0.38": (196, 522), - "0.43": (210, 490), - "0.48": (222, 462), - "0.50": (226, 452), - "0.53": (232, 438), - "0.54": (236, 436), - "0.56": (240, 426), # base - "0.62": (252, 404), - "0.67": (262, 393), - "0.75": (276, 368), - "1.00": (320, 320), - "1.33": (370, 278), - "1.50": (392, 262), - "1.78": (426, 240), - "1.89": (440, 232), - "2.00": (452, 226), - "2.08": (462, 222), -} - -# S = 36864 -ASPECT_RATIO_144P = { - "0.38": (117, 312), - "0.43": (125, 291), - "0.48": (133, 277), - "0.50": (135, 270), - "0.53": (139, 262), - "0.54": (141, 260), - "0.56": (144, 256), # base - "0.62": (151, 241), - "0.67": (156, 234), - "0.75": (166, 221), - "1.00": (192, 192), - "1.33": (221, 165), - "1.50": (235, 156), - "1.78": (256, 144), - "1.89": (263, 139), - "2.00": (271, 135), - "2.08": (277, 132), -} - -# from PixArt -# S = 8294400 -ASPECT_RATIO_2880 = { - "0.25": (1408, 5760), - "0.26": (1408, 5568), - "0.27": (1408, 5376), - "0.28": (1408, 5184), - "0.32": (1600, 4992), - "0.33": (1600, 4800), - "0.34": (1600, 4672), - "0.40": (1792, 4480), - "0.42": (1792, 4288), - "0.47": (1920, 4096), - "0.49": (1920, 3904), - "0.51": (1920, 3776), - "0.55": (2112, 3840), - "0.59": (2112, 3584), - "0.68": (2304, 3392), - "0.72": (2304, 3200), - "0.78": (2496, 3200), - "0.83": (2496, 3008), - "0.89": (2688, 3008), - "0.93": (2688, 2880), - "1.00": (2880, 2880), - "1.07": (2880, 2688), - "1.12": (3008, 2688), - "1.21": (3008, 2496), - "1.28": (3200, 2496), - "1.39": (3200, 2304), - "1.47": (3392, 2304), - "1.70": (3584, 2112), - "1.82": (3840, 2112), - "2.03": (3904, 1920), - "2.13": (4096, 1920), - "2.39": (4288, 1792), - "2.50": (4480, 1792), - "2.92": (4672, 1600), - "3.00": (4800, 1600), - "3.12": (4992, 1600), - "3.68": (5184, 1408), - "3.82": (5376, 1408), - "3.95": (5568, 1408), - "4.00": (5760, 1408), -} - -# S = 4194304 -ASPECT_RATIO_2048 = { - "0.25": (1024, 4096), - "0.26": (1024, 3968), - "0.27": (1024, 3840), - "0.28": (1024, 3712), - "0.32": (1152, 3584), - "0.33": (1152, 3456), - "0.35": (1152, 3328), - "0.40": (1280, 3200), - "0.42": (1280, 3072), - "0.48": (1408, 2944), - "0.50": (1408, 2816), - "0.52": (1408, 2688), - "0.57": (1536, 2688), - "0.60": (1536, 2560), - "0.68": (1664, 2432), - "0.72": (1664, 2304), - "0.78": (1792, 2304), - "0.82": (1792, 2176), - "0.88": (1920, 2176), - "0.94": (1920, 2048), - "1.00": (2048, 2048), - "1.07": (2048, 1920), - "1.13": (2176, 1920), - "1.21": (2176, 1792), - "1.29": (2304, 1792), - "1.38": (2304, 1664), - "1.46": (2432, 1664), - "1.67": (2560, 1536), - "1.75": (2688, 1536), - "2.00": (2816, 1408), - "2.09": (2944, 1408), - "2.40": (3072, 1280), - "2.50": (3200, 1280), - "2.89": (3328, 1152), - "3.00": (3456, 1152), - "3.11": (3584, 1152), - "3.62": (3712, 1024), - "3.75": (3840, 1024), - "3.88": (3968, 1024), - "4.00": (4096, 1024), -} - -# S = 1048576 -ASPECT_RATIO_1024 = { - "0.25": (512, 2048), - "0.26": (512, 1984), - "0.27": (512, 1920), - "0.28": (512, 1856), - "0.32": (576, 1792), - "0.33": (576, 1728), - "0.35": (576, 1664), - "0.40": (640, 1600), - "0.42": (640, 1536), - "0.48": (704, 1472), - "0.50": (704, 1408), - "0.52": (704, 1344), - "0.57": (768, 1344), - "0.60": (768, 1280), - "0.68": (832, 1216), - "0.72": (832, 1152), - "0.78": (896, 1152), - "0.82": (896, 1088), - "0.88": (960, 1088), - "0.94": (960, 1024), - "1.00": (1024, 1024), - "1.07": (1024, 960), - "1.13": (1088, 960), - "1.21": (1088, 896), - "1.29": (1152, 896), - "1.38": (1152, 832), - "1.46": (1216, 832), - "1.67": (1280, 768), - "1.75": (1344, 768), - "2.00": (1408, 704), - "2.09": (1472, 704), - "2.40": (1536, 640), - "2.50": (1600, 640), - "2.89": (1664, 576), - "3.00": (1728, 576), - "3.11": (1792, 576), - "3.62": (1856, 512), - "3.75": (1920, 512), - "3.88": (1984, 512), - "4.00": (2048, 512), -} - -# S = 262144 -ASPECT_RATIO_512 = { - "0.25": (256, 1024), - "0.26": (256, 992), - "0.27": (256, 960), - "0.28": (256, 928), - "0.32": (288, 896), - "0.33": (288, 864), - "0.35": (288, 832), - "0.40": (320, 800), - "0.42": (320, 768), - "0.48": (352, 736), - "0.50": (352, 704), - "0.52": (352, 672), - "0.57": (384, 672), - "0.60": (384, 640), - "0.68": (416, 608), - "0.72": (416, 576), - "0.78": (448, 576), - "0.82": (448, 544), - "0.88": (480, 544), - "0.94": (480, 512), - "1.00": (512, 512), - "1.07": (512, 480), - "1.13": (544, 480), - "1.21": (544, 448), - "1.29": (576, 448), - "1.38": (576, 416), - "1.46": (608, 416), - "1.67": (640, 384), - "1.75": (672, 384), - "2.00": (704, 352), - "2.09": (736, 352), - "2.40": (768, 320), - "2.50": (800, 320), - "2.89": (832, 288), - "3.00": (864, 288), - "3.11": (896, 288), - "3.62": (928, 256), - "3.75": (960, 256), - "3.88": (992, 256), - "4.00": (1024, 256), -} - -# S = 65536 -ASPECT_RATIO_256 = { - "0.25": (128, 512), - "0.26": (128, 496), - "0.27": (128, 480), - "0.28": (128, 464), - "0.32": (144, 448), - "0.33": (144, 432), - "0.35": (144, 416), - "0.40": (160, 400), - "0.42": (160, 384), - "0.48": (176, 368), - "0.50": (176, 352), - "0.52": (176, 336), - "0.57": (192, 336), - "0.60": (192, 320), - "0.68": (208, 304), - "0.72": (208, 288), - "0.78": (224, 288), - "0.82": (224, 272), - "0.88": (240, 272), - "0.94": (240, 256), - "1.00": (256, 256), - "1.07": (256, 240), - "1.13": (272, 240), - "1.21": (272, 224), - "1.29": (288, 224), - "1.38": (288, 208), - "1.46": (304, 208), - "1.67": (320, 192), - "1.75": (336, 192), - "2.00": (352, 176), - "2.09": (368, 176), - "2.40": (384, 160), - "2.50": (400, 160), - "2.89": (416, 144), - "3.00": (432, 144), - "3.11": (448, 144), - "3.62": (464, 128), - "3.75": (480, 128), - "3.88": (496, 128), - "4.00": (512, 128), -} - - -def get_closest_ratio(height: float, width: float, ratios: dict): - aspect_ratio = height / width - closest_ratio = min(ratios.keys(), - key=lambda ratio: abs(float(ratio) - aspect_ratio)) - return closest_ratio - - -ASPECT_RATIOS = { - "144p": (36864, ASPECT_RATIO_144P), - "256": (65536, ASPECT_RATIO_256), - "240p": (102240, ASPECT_RATIO_240P), - "360p": (230400, ASPECT_RATIO_360P), - "512": (262144, ASPECT_RATIO_512), - "480p": (409920, ASPECT_RATIO_480P), - "720p": (921600, ASPECT_RATIO_720P), - "1024": (1048576, ASPECT_RATIO_1024), - "1080p": (2073600, ASPECT_RATIO_1080P), - "2k": (3686400, ASPECT_RATIO_2K), - "2048": (4194304, ASPECT_RATIO_2048), - "2880": (8294400, ASPECT_RATIO_2880), - "4k": (8294400, ASPECT_RATIO_4K), -} - - -def get_num_pixels(name): - return ASPECT_RATIOS[name][0] - - -def get_image_size(resolution, ar_ratio): - if ar_ratio in ASPECT_RATIO_MAP: - ar_key = ASPECT_RATIO_MAP[ar_ratio] - else: - ar_key = ar_ratio - rs_dict = ASPECT_RATIOS[resolution][1] - assert ar_key in rs_dict, f"Aspect ratio {ar_ratio} not found for resolution {resolution}" - return rs_dict[ar_key] - - -NUM_FRAMES_MAP = { - "1x": 51, - "2x": 102, - "4x": 204, - "8x": 408, - "16x": 816, - "2s": 51, - "4s": 102, - "8s": 204, - "16s": 408, - "32s": 816, -} - - -def get_num_frames(num_frames): - if num_frames in NUM_FRAMES_MAP: - return NUM_FRAMES_MAP[num_frames] - else: - return int(num_frames) diff --git a/examples/models/contrib/stdit/assets/a_beautiful_waterfall.mp4 b/examples/models/contrib/stdit/assets/a_beautiful_waterfall.mp4 deleted file mode 100644 index f423c02ad73e..000000000000 Binary files a/examples/models/contrib/stdit/assets/a_beautiful_waterfall.mp4 and /dev/null differ diff --git a/examples/models/contrib/stdit/convert_checkpoint.py b/examples/models/contrib/stdit/convert_checkpoint.py deleted file mode 100644 index 242dd4a9a141..000000000000 --- a/examples/models/contrib/stdit/convert_checkpoint.py +++ /dev/null @@ -1,252 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -from vae import get_vae - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import STDiT3Model - -PRETRAINED_STDIT_PATH = "hpcai-tech/OpenSora-STDiT-v3" - - -def pixel_size_to_latent_size(args): - vae = get_vae( - from_pretrained=args.vae_type, - micro_frame_size=args.vae_micro_frame_size, - micro_batch_size=args.vae_micro_batch_size, - ).eval() - spatial_patch_size = vae.spatial_vae.patch_size - temporal_patch_size = vae.temporal_vae.patch_size - vae_out_channels = vae.out_channels - pixel_size = (args.num_frames, args.height, args.width) - latent_size = vae.get_latent_size(pixel_size) - return { - 'in_channels': vae_out_channels, - 'latent_size': latent_size, - 'spatial_patch_size': spatial_patch_size, - 'temporal_patch_size': temporal_patch_size, - } - - -def size_str_to_list(repr): - return [int(it) for it in repr.split('x')] if 'x' in repr else [int(repr)] - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--timm_ckpt', type=str, default=None) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument('--caption_channels', - type=int, - default=4096, - help='The channel of input of caption embedder') - parser.add_argument('--depth', - type=int, - default=28, - help='The number of STDiT blocks') - parser.add_argument('--input_sq_size', - type=int, - default=512, - help='Base spatial position embedding size') - parser.add_argument('--stdit_type', - type=str, - default="STDiT3", - choices=["STDiT3"]) - parser.add_argument('--stdit_patch_size', - type=str, - default='1x2x2', - help='The patch size of stdit for patchify') - parser.add_argument('--width', - type=int, - default=640, - help='The width of image size') - parser.add_argument('--height', - type=int, - default=360, - help='The height of image size') - parser.add_argument('--num_frames', - type=int, - default=102, - help='The frames of generated video') - parser.add_argument('--vae_type', - type=str, - default="hpcai-tech/OpenSora-VAE-v1.2", - choices=["hpcai-tech/OpenSora-VAE-v1.2"]) - parser.add_argument('--vae_micro_frame_size', - type=int, - default=17, - help='The micro_frame_size for vae') - parser.add_argument('--vae_micro_batch_size', - type=int, - default=4, - help='The micro_batch_size for vae') - parser.add_argument('--hidden_size', - type=int, - default=1152, - help='The hidden size of STDiT') - parser.add_argument('--num_heads', - type=int, - default=16, - help='The number of heads of attention module') - parser.add_argument( - '--mlp_ratio', - type=float, - default=4.0, - help= - 'The ratio of hidden size compared to input hidden size in MLP layer') - parser.add_argument( - '--class_dropout_prob', - type=float, - default=0.1, - help='The probability to drop class token when training') - parser.add_argument('--model_max_length', - type=int, - default=300, - help='The max number of tokens (default: 300)') - parser.add_argument('--text_encoder_type', - type=str, - default="DeepFloyd/t5-v1_1-xxl", - choices=["DeepFloyd/t5-v1_1-xxl"]) - parser.add_argument('--learn_sigma', - type=bool, - default=True, - help='Whether the model learn sigma') - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--cp_size', - type=int, - default=1, - help='Context parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument('--disable_qk_norm', - action='store_true', - help='Disable norm for qk in attention') - parser.add_argument('--fp8', - action='store_true', - help='Whether use FP8 for layers') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - return args - - -def convert_and_save_model(args): - # [NOTE] PP is not supported yet. - world_size = args.tp_size * args.cp_size - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - cp_size=args.cp_size) - # Process args - runtime_config = { - 'architecture': "STDiT3", - 'checkpoint_path': os.path.abspath(args.timm_ckpt), - 'caption_channels': args.caption_channels, - 'num_hidden_layers': args.depth, - 'width': args.width, - 'height': args.height, - 'num_frames': args.num_frames, - 'hidden_size': args.hidden_size, - 'stdit_patch_size': size_str_to_list(args.stdit_patch_size), - 'input_sq_size': args.input_sq_size, - 'num_attention_heads': args.num_heads, - 'model_max_length': args.model_max_length, - 'mlp_ratio': args.mlp_ratio, - 'class_dropout_prob': args.class_dropout_prob, - 'learn_sigma': args.learn_sigma, - 'qk_norm': (not args.disable_qk_norm), - 'stdit_type': args.stdit_type, - 'vae_type': args.vae_type, - 'text_encoder_type': args.text_encoder_type, - } - runtime_config.update(pixel_size_to_latent_size(args)) - tik = time.time() - stdit = STDiT3Model.from_pretrained(os.path.dirname(args.timm_ckpt), - args.dtype, - mapping=mapping, - **runtime_config) - stdit.save_checkpoint(args.output_dir, save_config=True) - print(f'Total time of reading and converting: {time.time()-tik:.3f} s') - tik = time.time() - del stdit - print(f'Total time of saving checkpoint: {time.time()-tik:.3f} s') - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - assert args.pp_size == 1, "PP is not supported yet." - - tik = time.time() - - if args.timm_ckpt is None: - print( - f"No pretrained checkpoint provided, use default checkpoint from Huggingface instead." - ) - args.timm_ckpt = "./pretrained_ckpt/model.safetensors" - if not os.path.exists(args.timm_ckpt): - from huggingface_hub import snapshot_download - snapshot_download(repo_id=PRETRAINED_STDIT_PATH, - local_dir=os.path.dirname(args.timm_ckpt)) - - convert_and_save_model(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/contrib/stdit/pipeline_tllm.py b/examples/models/contrib/stdit/pipeline_tllm.py deleted file mode 100644 index 2e3ff8e346a3..000000000000 --- a/examples/models/contrib/stdit/pipeline_tllm.py +++ /dev/null @@ -1,607 +0,0 @@ -import os -import time -from functools import wraps -from typing import List - -import numpy as np -import torch -import torch.distributed -from cuda import cudart -from scheduler import timestep_transform -from utils import DataProcessor, print_progress_bar - -import tensorrt_llm -from tensorrt_llm._utils import (str_dtype_to_torch, str_dtype_to_trt, - torch_dtype_to_trt, trt_dtype_to_torch) -from tensorrt_llm.logger import logger -from tensorrt_llm.plugin.plugin import CustomAllReduceHelper -from tensorrt_llm.runtime.session import Session, TensorInfo - - -def CUASSERT(cuda_ret): - err = cuda_ret[0] - if err != cudart.cudaError_t.cudaSuccess: - raise RuntimeError( - f"CUDA ERROR: {err}, error code reference: https://nvidia.github.io/cuda-python/module/cudart.html#cuda.cudart.cudaError_t" - ) - if len(cuda_ret) > 1: - return cuda_ret[1:] - return None - - -class TllmSTDiT(): - - def __init__(self, - config, - tllm_model_dir, - debug_mode=True, - stream: torch.cuda.Stream = None): - self.dtype = config['pretrained_config']['dtype'] - - self.depth = config['pretrained_config']['num_hidden_layers'] - - rank = tensorrt_llm.mpi_rank() - world_size = config['pretrained_config']['mapping']['world_size'] - cp_size = config['pretrained_config']['mapping']['cp_size'] - tp_size = config['pretrained_config']['mapping']['tp_size'] - pp_size = config['pretrained_config']['mapping']['pp_size'] - gpus_per_node = config['pretrained_config']['mapping']['gpus_per_node'] - assert pp_size == 1 - self.mapping = tensorrt_llm.Mapping(world_size=world_size, - rank=rank, - cp_size=cp_size, - tp_size=tp_size, - pp_size=1, - gpus_per_node=gpus_per_node) - - local_rank = rank % self.mapping.gpus_per_node - self.device = torch.device(f'cuda:{local_rank}') - torch.cuda.set_device(self.device) - CUASSERT(cudart.cudaSetDevice(local_rank)) - - self.stream = stream - if self.stream is None: - self.stream = torch.cuda.Stream(self.device) - torch.cuda.set_stream(self.stream) - - engine_file = os.path.join(tllm_model_dir, f"rank{rank}.engine") - logger.info(f'Loading engine from {engine_file}') - with open(engine_file, "rb") as f: - engine_buffer = f.read() - - assert engine_buffer is not None - - self.session = Session.from_serialized_engine(engine_buffer) - - self.debug_mode = debug_mode - - self.inputs = {} - self.outputs = {} - self.buffer_allocated = False - - expected_tensor_names = [ - 'x', 'timestep', 'y', 'mask', 'x_mask', 'fps', 'height', 'width', - 'output' - ] - - if self.mapping.tp_size > 1: - self.buffer, self.all_reduce_workspace = CustomAllReduceHelper.allocate_workspace( - self.mapping, - CustomAllReduceHelper.max_workspace_size_auto( - self.mapping.tp_size)) - self.inputs['all_reduce_workspace'] = self.all_reduce_workspace - expected_tensor_names += ['all_reduce_workspace'] - - self.latent_size = config['pretrained_config']['latent_size'] - self.patch_size = config['pretrained_config']['stdit_patch_size'] - self.in_channels = config['pretrained_config']['in_channels'] - self.caption_channels = config['pretrained_config']['caption_channels'] - self.model_max_length = config['pretrained_config']['model_max_length'] - self.config = config['pretrained_config'] - - self.max_cattn_seq_len = int( - np.prod([ - np.ceil(d / p) - for d, p in zip(self.latent_size, self.patch_size) - ])) - - def _tensor_dtype(self, name): - # return torch dtype given tensor name for convenience - dtype = trt_dtype_to_torch(self.session.engine.get_tensor_dtype(name)) - return dtype - - def _get_extra_inputs_for_attention_plugin(self, batch_size, max_seq_len, - max_encoder_seq_len): - host_max_attention_window_sizes = torch.tensor([max_seq_len] * - self.depth, - dtype=torch.int32).cpu() - host_sink_token_length = torch.tensor([0], dtype=torch.int32).cpu() - context_lengths = torch.full([batch_size], - max_seq_len, - dtype=torch.int32).cuda() - host_context_lengths = torch.full([batch_size], - max_seq_len, - dtype=torch.int32).cpu() - host_request_types = torch.zeros([batch_size], dtype=torch.int32).cpu() - perf_knob_tensor_size = 16 - host_runtime_perf_knobs = torch.tensor([-1] * perf_knob_tensor_size, - dtype=torch.int64, - device='cpu') - host_context_progress = torch.tensor([0], - dtype=torch.int64, - device='cpu') - cross_encoder_input_lengths = torch.full([batch_size], - max_encoder_seq_len, - dtype=torch.int32, - device='cuda') - cross_encoder_max_input_length = torch.empty((max_encoder_seq_len, ), - dtype=torch.int32, - device='cuda') - - extra_inputs = { - 'host_max_attention_window_sizes': host_max_attention_window_sizes, - 'host_sink_token_length': host_sink_token_length, - 'context_lengths': context_lengths, - 'host_context_lengths': host_context_lengths, - 'encoder_input_lengths': cross_encoder_input_lengths, - 'encoder_max_input_length': cross_encoder_max_input_length, - 'host_request_types': host_request_types, - 'host_runtime_perf_knobs': host_runtime_perf_knobs, - 'host_context_progress': host_context_progress - } - return extra_inputs - - def _setup(self, batch_size, max_encoder_seq_len): - input_info = [ - TensorInfo(name='x', - dtype=str_dtype_to_trt(self.dtype), - shape=(batch_size, self.in_channels, *self.latent_size)), - TensorInfo(name='timestep', - dtype=str_dtype_to_trt(self.dtype), - shape=(batch_size, )), - TensorInfo(name='y', - dtype=str_dtype_to_trt(self.dtype), - shape=(batch_size, 1, self.model_max_length, - self.caption_channels)), - TensorInfo(name='mask', - dtype=str_dtype_to_trt('int32'), - shape=(1, self.model_max_length)), - TensorInfo(name='x_mask', - dtype=str_dtype_to_trt('bool'), - shape=(batch_size, self.latent_size[0])), - TensorInfo(name='fps', - dtype=str_dtype_to_trt(self.dtype), - shape=(1, )), - TensorInfo(name='height', - dtype=str_dtype_to_trt(self.dtype), - shape=(1, )), - TensorInfo(name='width', - dtype=str_dtype_to_trt(self.dtype), - shape=(1, )), - ] - - extra_inputs = self._get_extra_inputs_for_attention_plugin( - batch_size=batch_size, - max_seq_len=self.max_cattn_seq_len, - max_encoder_seq_len=max_encoder_seq_len, - ) - input_info += [ - tensorrt_llm.runtime.TensorInfo(name, - torch_dtype_to_trt(tensor.dtype), - tensor.shape) - for name, tensor in extra_inputs.items() - ] - - output_info = self.session.infer_shapes(input_info) - for t_info in output_info: - self.outputs[t_info.name] = torch.empty(tuple(t_info.shape), - dtype=trt_dtype_to_torch( - t_info.dtype), - device=self.device) - self.buffer_allocated = True - - def cuda_stream_guard(func): - """Sync external stream and set current stream to the one bound to the session. Reset on exit. - """ - - @wraps(func) - def wrapper(self, *args, **kwargs): - external_stream = torch.cuda.current_stream() - if external_stream != self.stream: - external_stream.synchronize() - torch.cuda.set_stream(self.stream) - ret = func(self, *args, **kwargs) - if external_stream != self.stream: - self.stream.synchronize() - torch.cuda.set_stream(external_stream) - return ret - - return wrapper - - @cuda_stream_guard - def forward(self, x: torch.Tensor, timestep: torch.Tensor, y: torch.Tensor, - mask: torch.Tensor, x_mask: torch.Tensor, fps: torch.Tensor, - height: torch.Tensor, width: torch.Tensor, y_lens: List[int]): - """ - Forward pass of STDiT. - x: (N, C, F, H, W) - timestep: (N,) - y: () - mask: () - x_mask: (N * 2, ) - fps: (1) - height: (1) - width: (1) - """ - self._setup(x.shape[0], max_encoder_seq_len=np.max(y_lens)) - if not self.buffer_allocated: - raise RuntimeError('Buffer not allocated, please call setup first!') - - inputs = { - 'x': x, - "timestep": timestep, - 'y': y, - "mask": mask, - "x_mask": x_mask, - "fps": fps, - "height": height, - "width": width, - } - self.inputs.update(**inputs) - extra_inputs = self._get_extra_inputs_for_attention_plugin( - batch_size=x.shape[0], - max_seq_len=self.max_cattn_seq_len, - max_encoder_seq_len=np.max(y_lens), - ) - self.inputs.update(**extra_inputs) - - self.session.set_shapes(self.inputs) - ok = self.session.run(self.inputs, self.outputs, - self.stream.cuda_stream) - - if not ok: - raise RuntimeError('Executing TRT engine failed!') - debug_tensors = {} - for name in list(self.outputs.keys()): - if name != 'output': - debug_tensors[name] = self.outputs.pop(name) - if len(debug_tensors) == 0: - return self.outputs['output'] - else: - return self.outputs['output'], debug_tensors - - -class TllmOpenSoraPipeline(): - - def __init__( - self, - stdit, - text_encoder, - vae, - scheduler, - num_sampling_steps=30, - num_timesteps=1000, - cfg_scale=4.0, - align=None, - aes=None, - flow=None, - camera_motion=None, - image_size=None, - resolution=None, - aspect_ratio=None, - num_frames=None, - fps=None, - save_fps=None, - diffusion_model_type=None, - condition_frame_length=None, - condition_frame_edit=None, - use_discrete_timesteps=False, - use_timestep_transform=False, - video_save_dir='./samples/', - dtype='float16', - device=torch.device('cuda'), - seed=None, - **kwargs, - ): - self.stdit = stdit - self.text_encoder = text_encoder - self.vae = vae - self.scheduler = scheduler - self.data_processor = DataProcessor(text_encoder=text_encoder, vae=vae) - - self.num_sampling_steps = num_sampling_steps - self.num_timesteps = num_timesteps - self.cfg_scale = cfg_scale - self.align = align - self.aes = aes - self.flow = flow - self.camera_motion = camera_motion - - if image_size is None: - assert ( - resolution is not None and aspect_ratio is not None - ), "resolution and aspect_ratio must be provided if image_size is not provided" - image_size = self.data_processor.get_image_size( - resolution, aspect_ratio) - self.image_size = image_size - self.num_frames = self.data_processor.get_num_frames(num_frames) - self.input_size = (num_frames, *image_size) - self.latent_size = self.data_processor.get_latent_size(self.input_size) - - self.fps = fps - self.save_fps = save_fps - self.diffusion_model_type = diffusion_model_type - self.condition_frame_length = condition_frame_length - self.condition_frame_edit = condition_frame_edit - - self.use_discrete_timesteps = use_discrete_timesteps - self.use_timestep_transform = use_timestep_transform - - self.video_save_dir = video_save_dir - self.seed = self.data_processor.set_random_seed(seed) - self.dtype = dtype - self.device = device - - self.prompt_as_path = kwargs.get("prompt_as_path", False) - self.watermark = kwargs.get("watermark", False) - - def __call__( - self, - prompts, - batch_size=1, - num_sample=1, - loop=1, - reference_path=None, - mask_strategy=None, - sample_name=None, - start_idx=0, - ): - reference_path = [""] * len( - prompts) if reference_path is None else reference_path - mask_strategy = [""] * len( - prompts) if mask_strategy is None else mask_strategy - assert len(reference_path) == len( - prompts), "Length of reference must be the same as prompts" - assert len(mask_strategy) == len( - prompts), "Length of mask_strategy must be the same as prompts" - - for i in range(0, len(prompts), batch_size): - # == prepare batch prompts == - batch_prompts = prompts[i:i + batch_size] - ms = mask_strategy[i:i + batch_size] - refs = reference_path[i:i + batch_size] - - # == get json from prompts == - batch_prompts, refs, ms = self.data_processor.extract_json_from_prompts( - batch_prompts, refs, ms) - original_batch_prompts = batch_prompts - - # == get reference for condition == - refs = self.data_processor.collect_references_batch( - refs, self.image_size) - - # == multi-resolution info == - model_args = self.data_processor.prepare_multi_resolution_info( - self.diffusion_model_type, len(batch_prompts), self.image_size, - self.num_frames, self.fps, self.device, - str_dtype_to_torch(self.dtype)) - - # == Iter over number of sampling for one prompt == - for k in range(num_sample): - # == prepare save paths == - save_paths = [ - self.data_processor.get_save_path_name( - self.video_save_dir, - sample_name=sample_name, - sample_idx=start_idx + idx, - prompt=original_batch_prompts[idx], - prompt_as_path=self.prompt_as_path, - num_sample=num_sample, - k=k, - ) for idx in range(len(batch_prompts)) - ] - - # NOTE: Skip if the sample already exists - # This is useful for resuming sampling VBench - if self.prompt_as_path and \ - all(os.path.exists(path) for path in save_paths): - continue - - # == process prompts step by step == - # 0. split prompt - # each element in the list is [prompt_segment_list, loop_idx_list] - batched_prompt_segment_list = [] - batched_loop_idx_list = [] - for prompt in batch_prompts: - prompt_segment_list, loop_idx_list = self.data_processor.split_prompt( - prompt) - batched_prompt_segment_list.append(prompt_segment_list) - batched_loop_idx_list.append(loop_idx_list) - - # [NOTE] Skip refining prompt by OpenAI - for idx, prompt_segment_list in enumerate( - batched_prompt_segment_list): - batched_prompt_segment_list[ - idx] = self.data_processor.append_score_to_prompts( - prompt_segment_list, - aes=self.aes, - flow=self.flow, - camera_motion=self.camera_motion, - ) - - # 3. clean prompt with T5 - for idx, prompt_segment_list in enumerate( - batched_prompt_segment_list): - batched_prompt_segment_list[idx] = [ - self.data_processor.text_preprocessing(prompt) - for prompt in prompt_segment_list - ] - - # 4. merge to obtain the final prompt - batch_prompts = [] - for prompt_segment_list, loop_idx_list in zip( - batched_prompt_segment_list, batched_loop_idx_list): - batch_prompts.append( - self.data_processor.merge_prompt( - prompt_segment_list, loop_idx_list)) - - # == Iter over loop generation == - video_clips = [] - for loop_i in range(loop): - # == get prompt for loop i == - batch_prompts_loop = self.data_processor.extract_prompts_loop( - batch_prompts, loop_i) - - # == add condition frames for loop == - if loop_i > 0: - refs, ms = self.data_processor.append_generated( - video_clips[-1], refs, ms, loop_i, - self.condition_frame_length, - self.condition_frame_edit) - - # == sampling == - torch.manual_seed(self.seed) - noise = torch.randn(len(batch_prompts), - self.vae.out_channels, - *self.latent_size, - device=self.device, - dtype=str_dtype_to_torch(self.dtype)) - masks = self.data_processor.apply_mask_strategy( - noise, refs, ms, loop_i, align=self.align) - samples = self.sample( - latent=noise, - prompts=batch_prompts_loop, - mask=masks, - additional_args=model_args, - ) - samples = self.vae.decode(samples.to( - str_dtype_to_torch(self.dtype)), - num_frames=self.num_frames) - video_clips.append(samples) - - self.save_video(video_clips, save_paths, len(batch_prompts), - loop) - start_idx += len(batch_prompts) - - def sample( - self, - latent, - prompts, - mask=None, - additional_args=None, - guidance_scale=None, - ): - # if no specific guidance scale is provided, use the default scale when initializing the scheduler - if guidance_scale is None: - guidance_scale = self.cfg_scale - - # text encoding - model_args = self.text_encoder.encode_with_null(prompts) - if additional_args is not None: - model_args.update(additional_args) - - # prepare timesteps - timesteps = [(1.0 - i / self.num_sampling_steps) * self.num_timesteps - for i in range(self.num_sampling_steps)] - if self.use_discrete_timesteps: - timesteps = [int(round(t)) for t in timesteps] - timesteps = [ - torch.tensor([t] * latent.shape[0], device=self.device) - for t in timesteps - ] - if self.use_timestep_transform: - timesteps = [ - timestep_transform(t, - additional_args, - num_timesteps=self.num_timesteps) - for t in timesteps - ] - - if mask is not None: - noise_added = torch.zeros_like(mask, dtype=torch.bool) - noise_added = noise_added | (mask == 1) - - for i, t in enumerate(timesteps): - print_progress_bar(i, self.num_sampling_steps) - # mask for adding noise - if mask is not None: - mask_t = mask * self.num_timesteps - x0 = latent.clone() - x_noise = self.scheduler.add_noise(x0, torch.randn_like(x0), t) - - mask_t_upper = mask_t >= t.unsqueeze(1) - model_args["x_mask"] = mask_t_upper.repeat(2, 1) - mask_add_noise = mask_t_upper & ~noise_added - - latent = torch.where(mask_add_noise[:, None, :, None, None], - x_noise, x0) - noise_added = mask_t_upper - - # classifier-free guidance - latent_in = torch.cat([latent, latent], 0) - t = torch.cat([t, t], 0) - for k in list(model_args.keys()): - if k not in ['y', 'mask', 'x_mask', 'fps', 'height', 'width']: - model_args.pop(k) - else: - if model_args[k].dtype == torch.float32: - model_args[k] = model_args[k].to(torch.float16) - latent_in = latent_in.to(torch.float16) - t = t.to(torch.float16) - model_args['mask'] = model_args['mask'].to(torch.int32) - model_args['y_lens'] = self._get_y_lens(model_args['y'], - model_args['mask']) - - pred = self.stdit.forward(x=latent_in, timestep=t, - **model_args).chunk(2, dim=1)[0] - pred_cond, pred_uncond = pred.chunk(2, dim=0) - v_pred = pred_uncond + guidance_scale * (pred_cond - pred_uncond) - - # update latent - dt = timesteps[i] - timesteps[i + 1] if i < len( - timesteps) - 1 else timesteps[i] - dt = dt / self.num_timesteps - latent = latent + v_pred * dt[:, None, None, None, None] - - if mask is not None: - latent = torch.where(mask_t_upper[:, None, :, None, None], - latent, x0) - print() - return latent - - def _get_y_lens(self, y, mask=None): - assert (len(y.shape) == 4) - if mask is not None: - if mask.shape[0] != y.shape[0]: - mask = mask.repeat(y.shape[0] // mask.shape[0], 1) - mask = mask.squeeze(1).squeeze(1) - y_lens = mask.sum(dim=1).tolist() - else: - y_lens = [y.shape[2]] * y.shape[0] - return y_lens - - def save_video( - self, - video_clips, - save_paths, - batch_size, - loop, - ): - os.makedirs(self.video_save_dir, exist_ok=True) - if tensorrt_llm.mpi_rank() == 0: - for idx in range(batch_size): - save_path = save_paths[idx] - video = [video_clips[i][idx] for i in range(loop)] - for i in range(1, loop): - video[i] = video[i][:, - self.data_processor.dframe_to_frame( - self.condition_frame_length):] - video = torch.cat(video, dim=1) - save_path = self.data_processor.save_sample( - video, - fps=self.save_fps, - save_path=save_path, - ) - if save_path.endswith(".mp4") and self.watermark: - time.sleep(1) # prevent loading previous generated video - self.data_processor.add_watermark(save_path) diff --git a/examples/models/contrib/stdit/requirements.txt b/examples/models/contrib/stdit/requirements.txt deleted file mode 100644 index bc9a3f3a24c8..000000000000 --- a/examples/models/contrib/stdit/requirements.txt +++ /dev/null @@ -1,13 +0,0 @@ -ftfy -av==12.3.0 -git+https://github.com/facebookresearch/xformers.git@v0.0.29#egg=xformers -rotary_embedding_torch==0.5.3 -git+https://github.com/hpcaitech/TensorNVMe.git -# For colossalai -fabric -contexttimer -ray -protobuf -bitsandbytes>=0.39.0 -rpyc==6.0.0 -galore_torch diff --git a/examples/models/contrib/stdit/sample.py b/examples/models/contrib/stdit/sample.py deleted file mode 100644 index 88890d8ee1d5..000000000000 --- a/examples/models/contrib/stdit/sample.py +++ /dev/null @@ -1,170 +0,0 @@ -import argparse -import glob -import json -import os - -import torch -from pipeline_tllm import TllmOpenSoraPipeline, TllmSTDiT -from safetensors.torch import load_file -from scheduler import RFlowScheduler -from text_encoder import CaptionEmbedder, T5Encoder -from utils import DataProcessor -from vae import get_vae - -import tensorrt_llm -from tensorrt_llm._utils import str_dtype_to_torch - - -class RuntimeConfig(dict): - - def __getattr__(self, key): - if key in self: - return self[key] - raise AttributeError( - f"'RuntimeConfig {self}' object has no attribute '{key}'") - - def __setattr__(self, key, value): - self[key] = value - - -def main(cfg): - tensorrt_llm.logger.set_level(cfg.log_level) - - torch.set_grad_enabled(False) - dtype = cfg.get("dtype", "float16") - device = torch.device('cuda') - - config_file = os.path.join(cfg.tllm_model_dir, 'config.json') - with open(config_file) as f: - config = json.load(f) - ## Build modules - model_config = config.get("pretrained_config") - stdit = TllmSTDiT(config, cfg.tllm_model_dir, debug_mode=cfg.debug_mode) - # HACK: for classifier-free guidance - ckpt_path = model_config['checkpoint_path'] - if os.path.isdir(ckpt_path): - ckpt_files = glob.glob(model_config['checkpoint_path'] + "/*") - else: - ckpt_files = [ckpt_path] - pretrained_weights = None - for file in ckpt_files: - if file.endswith('.safetensors'): - pretrained_weights = load_file(file) - elif file.endswith('.pt'): - pretrained_weights = dict(torch.load(file, weights_only=True)) - if pretrained_weights is not None: - break - if pretrained_weights is None: - raise FileNotFoundError - y_embedder = CaptionEmbedder( - in_channels=model_config.get('caption_channels'), - hidden_size=model_config.get('hidden_size'), - uncond_prob=model_config.get('class_dropout_prob'), - act_layer=torch.nn.GELU(approximate="tanh"), - token_num=model_config.get('model_max_length'), - ) - for name, param in y_embedder.named_parameters(): - param.data = pretrained_weights['y_embedder.' + name] - y_embedder.y_embedding = pretrained_weights['y_embedder.y_embedding'] - reuse_y_embedding = pretrained_weights['y_embedder.y_embedding'] - text_encoder = T5Encoder( - from_pretrained=cfg.text_encoder, - model_max_length=model_config.get('model_max_length'), - caption_channels=model_config.get('caption_channels'), - y_embedding=reuse_y_embedding, # HACK: for classifier-free guidance - hidden_size=model_config.get('hidden_size'), - device=device, - ) - vae = get_vae( - from_pretrained=cfg.vae, - micro_frame_size=17, - micro_batch_size=4, - ).to(dtype=str_dtype_to_torch(dtype), device=device).eval() - scheduler = RFlowScheduler( - use_timestep_transform=True, - num_timesteps=cfg.num_timesteps, - num_sampling_steps=cfg.num_sampling_steps, - ) - pipe = TllmOpenSoraPipeline( - stdit=stdit, - text_encoder=text_encoder, - vae=vae, - scheduler=scheduler, - num_sampling_steps=cfg.num_sampling_steps, - num_timesteps=cfg.num_timesteps, - cfg_scale=cfg.cfg_scale, - align=cfg.get("align", None), - aes=cfg.get("aes", None), - flow=cfg.get("flow", None), - camera_motion=cfg.get("camera_motion", None), - resolution=cfg.get("resolution", None), - aspect_ratio=cfg.get("aspect_ratio", None), - num_frames=cfg.get("num_frames", None), - fps=cfg.get("fps", 30), - save_fps=cfg.get("save_fps", - cfg.get("fps", 30) // cfg.get("frame_interval", 1)), - diffusion_model_type=cfg.get("diffusion_model_type", None), - condition_frame_length=cfg.get("condition_frame_length", 5), - condition_frame_edit=cfg.get("condition_frame_edit", 0.0), - use_timestep_transform=True, - video_save_dir=cfg.get("video_save_dir", "sample_outputs"), - dtype=dtype, - device=device, - seed=cfg.get('seed', 0), - ) - - # load prompts - prompts = cfg.get("prompt", None) - start_idx = cfg.get("start_index", 0) - if prompts is None: - if cfg.get("prompt_path", None) is not None: - prompts = DataProcessor.load_prompts(cfg.prompt_path, start_idx, - cfg.get("end_index", None)) - else: - prompts = [cfg.get("prompt_generator", "") - ] * 1_000_000 # endless loop - - pipe(prompts) - - -if __name__ == '__main__': - parser = argparse.ArgumentParser() - parser.add_argument('prompt', - nargs='*', - help="Text prompt(s) to guide video generation") - parser.add_argument("--batch-size", type=int, default=1) - parser.add_argument("--resolution", type=str, default="360p") - parser.add_argument("--aspect-ratio", type=str, default="9:16") - parser.add_argument("--num-timesteps", type=int, default=1000) - parser.add_argument("--num-sampling-steps", type=int, default=30) - parser.add_argument("--seed", type=int, default=2946901) - parser.add_argument("--tllm_model_dir", - type=str, - default='./engine_outputs/') - parser.add_argument("--video_save_dir", - type=str, - default='./sample_outputs/') - parser.add_argument("--gpus_per_node", type=int, default=8) - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument("--debug_mode", action='store_true') - args = parser.parse_args() - - runtime_config = dict( - num_frames=102, - fps=30, - frame_interval=1, - save_fps=30, - diffusion_model_type="STDiT3", - text_encoder="DeepFloyd/t5-v1_1-xxl", - vae="hpcai-tech/OpenSora-VAE-v1.2", - cfg_scale=7.0, - dtype="float16", - condition_frame_length=5, - align=5, - aes=6.5, - flow=None, - prompt="A scene from disaster movie.", - ) - runtime_config.update(vars(args)) - - main(RuntimeConfig(runtime_config)) diff --git a/examples/models/contrib/stdit/scheduler.py b/examples/models/contrib/stdit/scheduler.py deleted file mode 100644 index acd0cac58bdd..000000000000 --- a/examples/models/contrib/stdit/scheduler.py +++ /dev/null @@ -1,101 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# 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. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/schedulers/rf/rectified_flow.py - -import torch -from torch.distributions import LogisticNormal - - -def timestep_transform( - t, - model_kwargs, - base_resolution=512 * 512, - base_num_frames=1, - scale=1.0, - num_timesteps=1, -): - # Force fp16 input to fp32 to avoid nan output - for key in ["height", "width", "num_frames"]: - if model_kwargs[key].dtype == torch.float16: - model_kwargs[key] = model_kwargs[key].float() - t = t / num_timesteps - resolution = model_kwargs["height"] * model_kwargs["width"] - ratio_space = (resolution / base_resolution).sqrt() - # NOTE: currently, we do not take fps into account - # NOTE: temporal_reduction is hardcoded, this should be equal to the temporal reduction factor of the vae - if model_kwargs["num_frames"][0] == 1: - num_frames = torch.ones_like(model_kwargs["num_frames"]) - else: - num_frames = model_kwargs["num_frames"] // 17 * 5 - ratio_time = (num_frames / base_num_frames).sqrt() - ratio = ratio_space * ratio_time * scale - new_t = ratio * t / (1 + (ratio - 1) * t) - new_t = new_t * num_timesteps - return new_t - - -class RFlowScheduler: - - def __init__( - self, - num_timesteps=1000, - num_sampling_steps=10, - sample_method="uniform", - loc=0.0, - scale=1.0, - use_timestep_transform=False, - ): - self.num_timesteps = num_timesteps - self.num_sampling_steps = num_sampling_steps - - assert sample_method in ["uniform", "logit-normal"] - self.sample_method = sample_method - if sample_method == "logit-normal": - self.distribution = LogisticNormal(torch.tensor([loc]), - torch.tensor([scale])) - self.sample_t = lambda x: self.distribution.sample( - (x.shape[0], ))[:, 0].to(x.device) - self.use_timestep_transform = use_timestep_transform - - def add_noise( - self, - original_samples: torch.FloatTensor, - noise: torch.FloatTensor, - timesteps: torch.IntTensor, - ) -> torch.FloatTensor: - timepoints = timesteps.float() / self.num_timesteps - timepoints = 1 - timepoints # [1,1/1000] - # timepoint (bsz) noise: (bsz, 4, frame, w ,h) - # expand timepoint to noise shape - timepoints = timepoints.unsqueeze(1).unsqueeze(1).unsqueeze( - 1).unsqueeze(1) - timepoints = timepoints.repeat(1, noise.shape[1], noise.shape[2], - noise.shape[3], noise.shape[4]) - return timepoints * original_samples + (1 - timepoints) * noise diff --git a/examples/models/contrib/stdit/text_encoder.py b/examples/models/contrib/stdit/text_encoder.py deleted file mode 100644 index c48ea5b1cfca..000000000000 --- a/examples/models/contrib/stdit/text_encoder.py +++ /dev/null @@ -1,445 +0,0 @@ -# Copyright (C) 2023 PixArt-alpha/PixArt-alpha -# -# This program is free software: you can redistribute it and/or modify -# it under the terms of the GNU Affero General Public License as published -# by the Free Software Foundation, either version 3 of the License, or -# (at your option) any later version. -# -# This program is distributed in the hope that it will be useful, -# but WITHOUT ANY WARRANTY; without even the implied warranty of -# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the -# GNU Affero General Public License for more details. - -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# 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. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/models/text_encoder/t5.py - -import collections.abc -from functools import partial -from itertools import repeat - -import torch -from colossalai.shardformer import ShardConfig, ShardFormer -from colossalai.shardformer.modeling.jit import get_jit_fused_dropout_add_func -from colossalai.shardformer.modeling.t5 import ( - get_jit_fused_T5_layer_ff_forward, get_T5_layer_self_attention_forward) -from colossalai.shardformer.policies.base_policy import ( - Policy, SubModuleReplacementDescription) -from transformers import AutoTokenizer, T5EncoderModel -from transformers.models.t5.modeling_t5 import (T5LayerFF, T5LayerSelfAttention, - T5Stack) - - -def default(var, default_var): - return default_var if var is None else var - - -def _ntuple(n): - - def parse(x): - if isinstance(x, collections.abc.Iterable) and not isinstance(x, str): - return tuple(x) - return tuple(repeat(x, n)) - - return parse - - -to_1tuple = _ntuple(1) -to_2tuple = _ntuple(2) -to_3tuple = _ntuple(3) -to_4tuple = _ntuple(4) -to_ntuple = _ntuple - - -class T5LayerNorm(torch.nn.Module): - - def __init__(self, hidden_size, eps=1e-6): - """ - Construct a layernorm module in the T5 style. No bias and no subtraction of mean. - """ - super().__init__() - self.weight = torch.nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states): - # T5 uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean - # Square Layer Normalization https://arxiv.org/abs/1910.07467 thus variance is calculated - # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for - # half-precision inputs is done in fp32 - variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + - self.variance_epsilon) - # convert into half-precision if necessary - if self.weight.dtype in [torch.float16, torch.bfloat16]: - hidden_states = hidden_states.to(self.weight.dtype) - return self.weight * hidden_states - - @staticmethod - def from_native_module(module, *args, **kwargs): - assert module.__class__.__name__ == "FusedRMSNorm", ( - "Recovering T5LayerNorm requires the original layer to be apex's Fused RMS Norm." - "Apex's fused norm is automatically used by Hugging Face Transformers https://github.com/huggingface/transformers/blob/main/src/transformers/models/t5/modeling_t5.py#L265C5-L265C48" - ) - layer_norm = T5LayerNorm(module.normalized_shape, eps=module.eps) - layer_norm.weight.data.copy_(module.weight.data) - layer_norm = layer_norm.to(module.weight.device) - return layer_norm - - -class T5EncoderPolicy(Policy): - - def config_sanity_check(self): - assert not self.shard_config.enable_tensor_parallelism - assert not self.shard_config.enable_flash_attention - - def preprocess(self): - return self.model - - def module_policy(self): - policy = {} - # check whether apex is installed - try: - # recover hf from fused rms norm to T5 norm which is faster - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="layer_norm", - target_module=T5LayerNorm, - ), - policy=policy, - target_key=T5LayerFF, - ) - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="layer_norm", target_module=T5LayerNorm), - policy=policy, - target_key=T5LayerSelfAttention, - ) - self.append_or_create_submodule_replacement( - description=SubModuleReplacementDescription( - suffix="final_layer_norm", target_module=T5LayerNorm), - policy=policy, - target_key=T5Stack, - ) - except (ImportError, ModuleNotFoundError): - pass - # use jit operator - if self.shard_config.enable_jit_fused: - self.append_or_create_method_replacement( - description={ - "forward": get_jit_fused_T5_layer_ff_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=T5LayerFF, - ) - self.append_or_create_method_replacement( - description={ - "forward": get_T5_layer_self_attention_forward(), - "dropout_add": get_jit_fused_dropout_add_func(), - }, - policy=policy, - target_key=T5LayerSelfAttention, - ) - return policy - - def postprocess(self): - return self.model - - -class T5Embedder: - - def __init__( - self, - device, - from_pretrained=None, - *, - cache_dir=None, - hf_token=None, - use_text_preprocessing=True, - t5_model_kwargs=None, - torch_dtype=None, - use_offload_folder=None, - model_max_length=120, - local_files_only=False, - ): - self.device = torch.device(device) - self.torch_dtype = torch_dtype or torch.float16 - self.cache_dir = cache_dir - - if t5_model_kwargs is None: - t5_model_kwargs = { - "low_cpu_mem_usage": True, - "torch_dtype": self.torch_dtype, - } - - if use_offload_folder is not None: - t5_model_kwargs["offload_folder"] = use_offload_folder - t5_model_kwargs["device_map"] = { - "shared": self.device, - "encoder.embed_tokens": self.device, - "encoder.block.0": self.device, - "encoder.block.1": self.device, - "encoder.block.2": self.device, - "encoder.block.3": self.device, - "encoder.block.4": self.device, - "encoder.block.5": self.device, - "encoder.block.6": self.device, - "encoder.block.7": self.device, - "encoder.block.8": self.device, - "encoder.block.9": self.device, - "encoder.block.10": self.device, - "encoder.block.11": self.device, - "encoder.block.12": "disk", - "encoder.block.13": "disk", - "encoder.block.14": "disk", - "encoder.block.15": "disk", - "encoder.block.16": "disk", - "encoder.block.17": "disk", - "encoder.block.18": "disk", - "encoder.block.19": "disk", - "encoder.block.20": "disk", - "encoder.block.21": "disk", - "encoder.block.22": "disk", - "encoder.block.23": "disk", - "encoder.final_layer_norm": "disk", - "encoder.dropout": "disk", - } - else: - t5_model_kwargs["device_map"] = { - "shared": self.device, - "encoder": self.device, - } - - self.use_text_preprocessing = use_text_preprocessing - self.hf_token = hf_token - - self.tokenizer = AutoTokenizer.from_pretrained( - from_pretrained, - cache_dir=cache_dir, - local_files_only=local_files_only, - ) - self.model = T5EncoderModel.from_pretrained( - from_pretrained, - cache_dir=cache_dir, - local_files_only=local_files_only, - **t5_model_kwargs, - ).eval() - self.model_max_length = model_max_length - - def get_text_embeddings(self, texts): - text_tokens_and_mask = self.tokenizer( - texts, - max_length=self.model_max_length, - padding="max_length", - truncation=True, - return_attention_mask=True, - add_special_tokens=True, - return_tensors="pt", - ) - - input_ids = text_tokens_and_mask["input_ids"].to(self.device) - attention_mask = text_tokens_and_mask["attention_mask"].to(self.device) - with torch.no_grad(): - text_encoder_embs = self.model( - input_ids=input_ids, - attention_mask=attention_mask, - )["last_hidden_state"].detach() - return text_encoder_embs, attention_mask - - -class Mlp(torch.nn.Module): - - def __init__( - self, - in_features, - hidden_features=None, - out_features=None, - act_layer=torch.nn.GELU(), - norm_layer=None, - bias=True, - drop=0., - use_conv=False, - ): - super().__init__() - out_features = out_features or in_features - hidden_features = hidden_features or in_features - bias = to_2tuple(bias) - drop_probs = to_2tuple(drop) - linear_layer = partial(torch.nn.Conv2d, - kernel_size=1) if use_conv else torch.nn.Linear - - self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0]) - self.act = act_layer - self.drop1 = torch.nn.Dropout(drop_probs[0]) - self.norm = norm_layer( - hidden_features) if norm_layer is not None else torch.nn.Identity() - self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1]) - self.drop2 = torch.nn.Dropout(drop_probs[1]) - - def forward(self, x): - x = self.fc1(x) - x = self.act(x) - x = self.drop1(x) - x = self.norm(x) - x = self.fc2(x) - x = self.drop2(x) - return x - - -class CaptionEmbedder(torch.nn.Module): - - def __init__( - self, - in_channels, - hidden_size, - uncond_prob, - act_layer=torch.nn.GELU(approximate="tanh"), - token_num=120, - ): - super().__init__() - self.y_proj = Mlp( - in_features=in_channels, - hidden_features=hidden_size, - out_features=hidden_size, - act_layer=act_layer, - drop=0, - ) - self.register_buffer( - "y_embedding", - torch.randn(token_num, in_channels) / in_channels**0.5, - ) - self.uncond_prob = uncond_prob - - def token_drop(self, caption, force_drop_ids=None): - if force_drop_ids is None: - drop_ids = torch.rand(caption.shape[0]).cuda() < self.uncond_prob - else: - drop_ids = force_drop_ids == 1 - caption = torch.where(drop_ids[:, None, None, None], self.y_embedding, - caption) - return caption - - def forward(self, caption, train, force_drop_ids=None): - if train: - assert caption.shape[2:] == self.y_embedding.shape - use_dropout = self.uncond_prob > 0 - if (train and use_dropout) or (force_drop_ids is not None): - caption = self.token_drop(caption, force_drop_ids) - caption = self.y_proj(caption) - return caption - - -class T5Encoder: - - def __init__( - self, - from_pretrained=None, - model_max_length=120, - caption_channels=4096, - hidden_size=1152, - class_dropout_prob=0.1, - y_embedding=None, - device="cuda", - dtype=torch.float, - cache_dir=None, - shardformer=False, - local_files_only=False, - ): - assert from_pretrained is not None, "Please specify the path to the T5 model" - self.t5 = T5Embedder( - device=device, - torch_dtype=dtype, - from_pretrained=from_pretrained, - cache_dir=cache_dir, - model_max_length=model_max_length, - local_files_only=local_files_only, - ) - self.t5.model.to(dtype=dtype) - # [NOTE] disable y_embedder - if False: - self.y_embedder = CaptionEmbedder( - in_channels=caption_channels, - hidden_size=hidden_size, - uncond_prob=class_dropout_prob, - act_layer=torch.nn.GELU(approximate="tanh"), - token_num=model_max_length, - ) - else: - self.y_embedder = None - self.y_embedding = default( - y_embedding, - torch.randn(model_max_length, caption_channels) / - caption_channels**0.5).to(device) - - self.model_max_length = model_max_length - self.output_dim = self.t5.model.config.d_model - self.dtype = dtype - - if shardformer: - self.shardformer_t5() - - def shardformer_t5(self): - shard_config = ShardConfig( - tensor_parallel_process_group=None, - pipeline_stage_manager=None, - enable_tensor_parallelism=False, - enable_fused_normalization=False, - enable_flash_attention=False, - enable_jit_fused=True, - enable_sequence_parallelism=False, - enable_sequence_overlap=False, - ) - shard_former = ShardFormer(shard_config=shard_config) - optim_model, _ = shard_former.optimize(self.t5.model, - policy=T5EncoderPolicy()) - self.t5.model = optim_model.to(self.dtype) - - # ensure the weights are frozen - for p in self.t5.model.parameters(): - p.requires_grad = False - - def encode(self, text): - caption_embs, emb_masks = self.t5.get_text_embeddings(text) - caption_embs = caption_embs[:, None] - return dict(y=caption_embs, mask=emb_masks) - - def null(self, n): - if self.y_embedder is None: - null_y = self.y_embedding[None].repeat(n, 1, 1)[:, None] - else: - null_y = self.y_embedder.y_embedding[None].repeat(n, 1, 1)[:, None] - return null_y - - def encode_with_null(self, text): - batch_size = len(text) - encoded_outputs = self.encode(text) - y_null = self.null(batch_size) - encoded_outputs["y"] = torch.cat([encoded_outputs["y"], y_null], 0) - return encoded_outputs diff --git a/examples/models/contrib/stdit/utils.py b/examples/models/contrib/stdit/utils.py deleted file mode 100644 index 88a7269bc8e8..000000000000 --- a/examples/models/contrib/stdit/utils.py +++ /dev/null @@ -1,997 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# 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. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/utils/ckpt_utils.py - -import html -import json -import os -import random -import re -import sys -import urllib.parse as ul - -import ftfy -import numpy as np -import pandas as pd -import requests -import torch -import video_transforms -from aspect import get_image_size, get_num_frames -from bs4 import BeautifulSoup -from colossalai.checkpoint_io import GeneralCheckpointIO -from PIL import Image -from safetensors.torch import load_file -from torchvision import transforms -from torchvision.datasets.folder import IMG_EXTENSIONS, pil_loader -from torchvision.io import read_video, write_video -from torchvision.utils import save_image - -VID_EXTENSIONS = (".mp4", ".avi", ".mov", ".mkv") - -URL_REGEX = re.compile( - r"^(?:http|ftp)s?://" # http:// or https:// - r"(?:(?:[A-Z0-9](?:[A-Z0-9-]{0,61}[A-Z0-9])?\.)+(?:[A-Z]{2,6}\.?|[A-Z0-9-]{2,}\.?)|" # domain... - r"localhost|" # localhost... - r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})" # ...or ip - r"(?::\d+)?" # optional port - r"(?:/?|[/?]\S+)$", - re.IGNORECASE, -) - -HF_ENDPOINT = os.environ.get("HF_ENDPOINT") -if HF_ENDPOINT is None: - HF_ENDPOINT = "https://huggingface.co" -PRETRAINED_MODELS = { - "DiT-XL-2-512x512.pt": - "https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-512x512.pt", - "DiT-XL-2-256x256.pt": - "https://dl.fbaipublicfiles.com/DiT/models/DiT-XL-2-256x256.pt", - "Latte-XL-2-256x256-ucf101.pt": - HF_ENDPOINT + "/maxin-cn/Latte/resolve/main/ucf101.pt", - "PixArt-XL-2-256x256.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-256x256.pth", - "PixArt-XL-2-SAM-256x256.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-SAM-256x256.pth", - "PixArt-XL-2-512x512.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-512x512.pth", - "PixArt-XL-2-1024-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-alpha/resolve/main/PixArt-XL-2-1024-MS.pth", - "OpenSora-v1-16x256x256.pth": - HF_ENDPOINT + - "/hpcai-tech/Open-Sora/resolve/main/OpenSora-v1-16x256x256.pth", - "OpenSora-v1-HQ-16x256x256.pth": - HF_ENDPOINT + - "/hpcai-tech/Open-Sora/resolve/main/OpenSora-v1-HQ-16x256x256.pth", - "OpenSora-v1-HQ-16x512x512.pth": - HF_ENDPOINT + - "/hpcai-tech/Open-Sora/resolve/main/OpenSora-v1-HQ-16x512x512.pth", - "PixArt-Sigma-XL-2-256x256.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-256x256.pth", - "PixArt-Sigma-XL-2-512-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-512-MS.pth", - "PixArt-Sigma-XL-2-1024-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-1024-MS.pth", - "PixArt-Sigma-XL-2-2K-MS.pth": - HF_ENDPOINT + - "/PixArt-alpha/PixArt-Sigma/resolve/main/PixArt-Sigma-XL-2-2K-MS.pth", -} - - -def is_img(path): - ext = os.path.splitext(path)[-1].lower() - return ext in IMG_EXTENSIONS - - -def is_vid(path): - ext = os.path.splitext(path)[-1].lower() - return ext in VID_EXTENSIONS - - -def is_url(url): - return re.match(URL_REGEX, url) is not None - - -def read_file(input_path): - if input_path.endswith(".csv"): - return pd.read_csv(input_path) - elif input_path.endswith(".parquet"): - return pd.read_parquet(input_path) - else: - raise NotImplementedError(f"Unsupported file format: {input_path}") - - -def download_url(input_path): - output_dir = "cache" - os.makedirs(output_dir, exist_ok=True) - base_name = os.path.basename(input_path) - output_path = os.path.join(output_dir, base_name) - try: - img_data = requests.get(input_path).content - except requests.exceptions.RequestException as e: - print(f"Error downloading URL {input_path}: {e}") - return None - with open(output_path, "wb") as handler: - handler.write(img_data) - print(f"URL {input_path} downloaded to {output_path}") - return output_path - - -def reparameter(ckpt, name=None, model=None): - model_name = name - name = os.path.basename(name) - if not torch.distributed.is_initialized() or torch.distributed.get_rank( - ) == 0: - print("loading pretrained model: %s", model_name) - if name in ["DiT-XL-2-512x512.pt", "DiT-XL-2-256x256.pt"]: - ckpt["x_embedder.proj.weight"] = ckpt[ - "x_embedder.proj.weight"].unsqueeze(2) - del ckpt["pos_embed"] - if name in ["Latte-XL-2-256x256-ucf101.pt"]: - ckpt = ckpt["ema"] - ckpt["x_embedder.proj.weight"] = ckpt[ - "x_embedder.proj.weight"].unsqueeze(2) - del ckpt["pos_embed"] - del ckpt["temp_embed"] - if name in [ - "PixArt-XL-2-256x256.pth", - "PixArt-XL-2-SAM-256x256.pth", - "PixArt-XL-2-512x512.pth", - "PixArt-XL-2-1024-MS.pth", - "PixArt-Sigma-XL-2-256x256.pth", - "PixArt-Sigma-XL-2-512-MS.pth", - "PixArt-Sigma-XL-2-1024-MS.pth", - "PixArt-Sigma-XL-2-2K-MS.pth", - ]: - ckpt = ckpt["state_dict"] - ckpt["x_embedder.proj.weight"] = ckpt[ - "x_embedder.proj.weight"].unsqueeze(2) - if "pos_embed" in ckpt: - del ckpt["pos_embed"] - if name in [ - "PixArt-1B-2.pth", - ]: - ckpt = ckpt["state_dict"] - if "pos_embed" in ckpt: - del ckpt["pos_embed"] - # no need pos_embed - if "pos_embed_temporal" in ckpt: - del ckpt["pos_embed_temporal"] - if "pos_embed" in ckpt: - del ckpt["pos_embed"] - # different text length - if "y_embedder.y_embedding" in ckpt: - if ckpt["y_embedder.y_embedding"].shape[ - 0] < model.y_embedder.y_embedding.shape[0]: - print( - "Extend y_embedding from %s to %s", - ckpt["y_embedder.y_embedding"].shape[0], - model.y_embedder.y_embedding.shape[0], - ) - additional_length = model.y_embedder.y_embedding.shape[0] - ckpt[ - "y_embedder.y_embedding"].shape[0] - new_y_embedding = torch.zeros(additional_length, - model.y_embedder.y_embedding.shape[1]) - new_y_embedding[:] = ckpt["y_embedder.y_embedding"][-1] - ckpt["y_embedder.y_embedding"] = torch.cat( - [ckpt["y_embedder.y_embedding"], new_y_embedding], dim=0) - elif ckpt["y_embedder.y_embedding"].shape[ - 0] > model.y_embedder.y_embedding.shape[0]: - print( - "Shrink y_embedding from %s to %s", - ckpt["y_embedder.y_embedding"].shape[0], - model.y_embedder.y_embedding.shape[0], - ) - ckpt["y_embedder.y_embedding"] = ckpt[ - "y_embedder.y_embedding"][:model.y_embedder.y_embedding. - shape[0]] - # stdit3 special case - if type(model).__name__ == "STDiT3" and "PixArt-Sigma" in name: - ckpt_keys = list(ckpt.keys()) - for key in ckpt_keys: - if "blocks." in key: - ckpt[key.replace("blocks.", "spatial_blocks.")] = ckpt[key] - del ckpt[key] - return ckpt - - -def download_model(model_name=None, local_path=None, url=None): - """ - Downloads a pre-trained DiT model from the web. - """ - if model_name is not None: - assert model_name in PRETRAINED_MODELS - local_path = f"pretrained_models/{model_name}" - web_path = PRETRAINED_MODELS[model_name] - else: - assert local_path is not None - assert url is not None - web_path = url - if not os.path.isfile(local_path): - os.makedirs("pretrained_models", exist_ok=True) - dir_name = os.path.dirname(local_path) - file_name = os.path.basename(local_path) - download_url(web_path, dir_name, file_name) - model = torch.load(local_path, map_location=lambda storage, loc: storage) - return model - - -def find_model(model_name, model=None): - """ - Finds a pre-trained DiT model, downloading it if necessary. Alternatively, loads a model from a local path. - """ - if model_name in PRETRAINED_MODELS: # Find/download our pre-trained DiT checkpoints - model_ckpt = download_model(model_name) - model_ckpt = reparameter(model_ckpt, model_name, model=model) - else: # Load a custom DiT checkpoint: - assert os.path.isfile( - model_name), f"Could not find DiT checkpoint at {model_name}" - model_ckpt = torch.load(model_name, - map_location=lambda storage, loc: storage) - model_ckpt = reparameter(model_ckpt, model_name, model=model) - return model_ckpt - - -def load_from_sharded_state_dict(model, - ckpt_path, - model_name="model", - strict=False): - ckpt_io = GeneralCheckpointIO() - ckpt_io.load_model(model, - os.path.join(ckpt_path, model_name), - strict=strict) - - -def load_checkpoint(model, - ckpt_path, - save_as_pt=False, - model_name="model", - strict=False): - if ckpt_path.endswith(".pt") or ckpt_path.endswith(".pth"): - state_dict = find_model(ckpt_path, model=model) - missing_keys, unexpected_keys = model.load_state_dict(state_dict, - strict=strict) - print("Missing keys: %s", missing_keys) - print("Unexpected keys: %s", unexpected_keys) - elif ckpt_path.endswith(".safetensors"): - state_dict = load_file(ckpt_path) - missing_keys, unexpected_keys = model.load_state_dict(state_dict, - strict=False) - print(f"Missing keys: {missing_keys}") - print(f"Unexpected keys: {unexpected_keys}") - elif os.path.isdir(ckpt_path): - load_from_sharded_state_dict(model, - ckpt_path, - model_name, - strict=strict) - print("Model checkpoint loaded from %s", ckpt_path) - if save_as_pt: - save_path = os.path.join(ckpt_path, model_name + "_ckpt.pt") - torch.save(model.state_dict(), save_path) - print("Model checkpoint saved to %s", save_path) - else: - raise ValueError(f"Invalid checkpoint path: {ckpt_path}") - - -def print_progress_bar(iteration, total, length=40): - iteration += 1 - filled_length = int(length * iteration // total) - bar = '█' * filled_length + '-' * (length - filled_length) - sys.stdout.write(f'\rDenoising steps: |{bar}| {iteration}/{total}') - sys.stdout.flush() - - -class PromptProcessor(): - - @staticmethod - def load_prompts(prompt_path, start_idx=None, end_idx=None): - with open(prompt_path, "r") as f: - prompts = [line.strip() for line in f.readlines()] - prompts = prompts[start_idx:end_idx] - return prompts - - @staticmethod - def extract_json_from_prompts(prompts, reference, mask_strategy): - ret_prompts = [] - for i, prompt in enumerate(prompts): - parts = re.split(r"(?=[{])", prompt) - assert len(parts) <= 2, f"Invalid prompt: {prompt}" - ret_prompts.append(parts[0]) - if len(parts) > 1: - additional_info = json.loads(parts[1]) - for key in additional_info: - assert key in ["reference_path", - "mask_strategy"], f"Invalid key: {key}" - if key == "reference_path": - reference[i] = additional_info[key] - elif key == "mask_strategy": - mask_strategy[i] = additional_info[key] - return ret_prompts, reference, mask_strategy - - @staticmethod - def split_prompt(prompt_text): - if prompt_text.startswith("|0|"): - # this is for prompts which look like - # |0| a beautiful day |1| a sunny day |2| a rainy day - # we want to parse it into a list of prompts with the loop index - prompt_list = prompt_text.split("|")[1:] - text_list = [] - loop_idx = [] - for i in range(0, len(prompt_list), 2): - start_loop = int(prompt_list[i]) - text = prompt_list[i + 1].strip() - text_list.append(text) - loop_idx.append(start_loop) - return text_list, loop_idx - else: - return [prompt_text], None - - @staticmethod - def merge_prompt(text_list, loop_idx_list=None): - if loop_idx_list is None: - return text_list[0] - else: - prompt = "" - for i, text in enumerate(text_list): - prompt += f"|{loop_idx_list[i]}|{text}" - return prompt - - @staticmethod - def extract_prompts_loop(prompts, num_loop): - ret_prompts = [] - for prompt in prompts: - if prompt.startswith("|0|"): - prompt_list = prompt.split("|")[1:] - text_list = [] - for i in range(0, len(prompt_list), 2): - start_loop = int(prompt_list[i]) - text = prompt_list[i + 1] - end_loop = int(prompt_list[i + 2]) if i + 2 < len( - prompt_list) else num_loop + 1 - text_list.extend([text] * (end_loop - start_loop)) - prompt = text_list[num_loop] - ret_prompts.append(prompt) - return ret_prompts - - @staticmethod - def append_score_to_prompts(prompts, - aes=None, - flow=None, - camera_motion=None): - new_prompts = [] - for prompt in prompts: - new_prompt = prompt - if aes is not None and "aesthetic score:" not in prompt: - new_prompt = f"{new_prompt} aesthetic score: {aes:.1f}." - if flow is not None and "motion score:" not in prompt: - new_prompt = f"{new_prompt} motion score: {flow:.1f}." - if camera_motion is not None and "camera motion:" not in prompt: - new_prompt = f"{new_prompt} camera motion: {camera_motion}." - new_prompts.append(new_prompt) - return new_prompts - - @classmethod - def basic_clean(cls, text): - text = ftfy.fix_text(text) - text = html.unescape(html.unescape(text)) - return text.strip() - - @classmethod - def clean_caption(cls, caption): - BAD_PUNCT_REGEX = re.compile(r"[" + "#®•©™&@·º½¾¿¡§~" + r"\)" + r"\(" + - r"\]" + r"\[" + r"\}" + r"\{" + r"\|" + - "\\" + "\/" + "\*" + r"]{1,}") # noqa - caption = str(caption) - caption = ul.unquote_plus(caption) - caption = caption.strip().lower() - caption = re.sub("", "person", caption) - # urls: - caption = re.sub( - r"\b((?:https?:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - caption = re.sub( - r"\b((?:www:(?:\/{1,3}|[a-zA-Z0-9%])|[a-zA-Z0-9.\-]+[.](?:com|co|ru|net|org|edu|gov|it)[\w/-]*\b\/?(?!@)))", # noqa - "", - caption, - ) # regex for urls - # html: - caption = BeautifulSoup(caption, features="html.parser").text - # @ - caption = re.sub(r"@[\w\d]+\b", "", caption) - # 31C0—31EF CJK Strokes - # 31F0—31FF Katakana Phonetic Extensions - # 3200—32FF Enclosed CJK Letters and Months - # 3300—33FF CJK Compatibility - # 3400—4DBF CJK Unified Ideographs Extension A - # 4DC0—4DFF Yijing Hexagram Symbols - # 4E00—9FFF CJK Unified Ideographs - caption = re.sub(r"[\u31c0-\u31ef]+", "", caption) - caption = re.sub(r"[\u31f0-\u31ff]+", "", caption) - caption = re.sub(r"[\u3200-\u32ff]+", "", caption) - caption = re.sub(r"[\u3300-\u33ff]+", "", caption) - caption = re.sub(r"[\u3400-\u4dbf]+", "", caption) - caption = re.sub(r"[\u4dc0-\u4dff]+", "", caption) - caption = re.sub(r"[\u4e00-\u9fff]+", "", caption) - ####################################################### - - # все виды тире / all types of dash --> "-" - caption = re.sub( - r"[\u002D\u058A\u05BE\u1400\u1806\u2010-\u2015\u2E17\u2E1A\u2E3A\u2E3B\u2E40\u301C\u3030\u30A0\uFE31\uFE32\uFE58\uFE63\uFF0D]+", # noqa - "-", - caption, - ) - # кавычки к одному стандарту - caption = re.sub(r"[`´«»“”¨]", '"', caption) - caption = re.sub(r"[‘’]", "'", caption) - # " - caption = re.sub(r""?", "", caption) - # & - caption = re.sub(r"&", "", caption) - # ip addresses: - caption = re.sub(r"\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}", " ", caption) - # article ids: - caption = re.sub(r"\d:\d\d\s+$", "", caption) - # \n - caption = re.sub(r"\\n", " ", caption) - # "#123" - caption = re.sub(r"#\d{1,3}\b", "", caption) - # "#12345.." - caption = re.sub(r"#\d{5,}\b", "", caption) - # "123456.." - caption = re.sub(r"\b\d{6,}\b", "", caption) - # filenames: - caption = re.sub(r"[\S]+\.(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)", - "", caption) - # - caption = re.sub(r"[\"\']{2,}", r'"', caption) # """AUSVERKAUFT""" - caption = re.sub(r"[\.]{2,}", r" ", caption) # """AUSVERKAUFT""" - caption = re.sub(BAD_PUNCT_REGEX, r" ", - caption) # ***AUSVERKAUFT***, #AUSVERKAUFT - caption = re.sub(r"\s+\.\s+", r" ", caption) # " . " - # this-is-my-cute-cat / this_is_my_cute_cat - regex2 = re.compile(r"(?:\-|\_)") - if len(re.findall(regex2, caption)) > 3: - caption = re.sub(regex2, " ", caption) - caption = cls.basic_clean(caption) - caption = re.sub(r"\b[a-zA-Z]{1,3}\d{3,15}\b", "", caption) # jc6640 - caption = re.sub(r"\b[a-zA-Z]+\d+[a-zA-Z]+\b", "", caption) # jc6640vc - caption = re.sub(r"\b\d+[a-zA-Z]+\d+\b", "", caption) # 6640vc231 - caption = re.sub(r"(worldwide\s+)?(free\s+)?shipping", "", caption) - caption = re.sub(r"(free\s)?download(\sfree)?", "", caption) - caption = re.sub(r"\bclick\b\s(?:for|on)\s\w+", "", caption) - caption = re.sub( - r"\b(?:png|jpg|jpeg|bmp|webp|eps|pdf|apk|mp4)(\simage[s]?)?", "", - caption) - caption = re.sub(r"\bpage\s+\d+\b", "", caption) - caption = re.sub(r"\b\d*[a-zA-Z]+\d+[a-zA-Z]+\d+[a-zA-Z\d]*\b", r" ", - caption) # j2d1a2a... - caption = re.sub(r"\b\d+\.?\d*[xх×]\d+\.?\d*\b", "", caption) - caption = re.sub(r"\b\s+\:\s+", r": ", caption) - caption = re.sub(r"(\D[,\./])\b", r"\1 ", caption) - caption = re.sub(r"\s+", " ", caption) - caption.strip() - caption = re.sub(r"^[\"\']([\w\W]+)[\"\']$", r"\1", caption) - caption = re.sub(r"^[\'\_,\-\:;]", r"", caption) - caption = re.sub(r"[\'\_,\-\:\-\+]$", r"", caption) - caption = re.sub(r"^\.\S+$", "", caption) - return caption.strip() - - @classmethod - def text_preprocessing(cls, text, use_text_preprocessing: bool = True): - if use_text_preprocessing: - # The exact text cleaning as was in the training stage: - text = cls.clean_caption(text) - text = cls.clean_caption(text) - return text - else: - return text.lower().strip() - - -class VideoProcessor(): - MASK_DEFAULT = ["0", "0", "0", "0", "1", "0"] - - @staticmethod - def get_image_size(resolution, ar_ratio): - return get_image_size(resolution=resolution, ar_ratio=ar_ratio) - - @staticmethod - def get_num_frames(num_frames): - return get_num_frames(num_frames=num_frames) - - @staticmethod - def get_latent_size(vae, input_size): - return vae.get_latent_size(input_size=input_size) - - @staticmethod - def center_crop_arr(pil_image, image_size): - """ - Center cropping implementation from ADM. - https://github.com/openai/guided-diffusion/blob/8fb3ad9197f16bbc40620447b2742e13458d2831/guided_diffusion/image_datasets.py#L126 - """ - while min(*pil_image.size) >= 2 * image_size: - pil_image = pil_image.resize(tuple(x // 2 for x in pil_image.size), - resample=Image.BOX) - - scale = image_size / min(*pil_image.size) - pil_image = pil_image.resize(tuple( - round(x * scale) for x in pil_image.size), - resample=Image.BICUBIC) - - arr = np.array(pil_image) - crop_y = (arr.shape[0] - image_size) // 2 - crop_x = (arr.shape[1] - image_size) // 2 - return Image.fromarray(arr[crop_y:crop_y + image_size, - crop_x:crop_x + image_size]) - - @staticmethod - def resize_crop_to_fill(pil_image, image_size): - w, h = pil_image.size # PIL is (W, H) - th, tw = image_size - rh, rw = th / h, tw / w - if rh > rw: - sh, sw = th, round(w * rh) - image = pil_image.resize((sw, sh), Image.BICUBIC) - i = 0 - j = int(round((sw - tw) / 2.0)) - else: - sh, sw = round(h * rw), tw - image = pil_image.resize((sw, sh), Image.BICUBIC) - i = int(round((sh - th) / 2.0)) - j = 0 - arr = np.array(image) - assert i + th <= arr.shape[0] and j + tw <= arr.shape[1] - return Image.fromarray(arr[i:i + th, j:j + tw]) - - @classmethod - def get_transforms_image(cls, name="center", image_size=(256, 256)): - if name is None: - return None - elif name == "center": - assert image_size[0] == image_size[ - 1], "Image size must be square for center crop" - transform = transforms.Compose([ - transforms.Lambda(lambda pil_image: cls.center_crop_arr( - pil_image, image_size[0])), - # transforms.RandomHorizontalFlip(), - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - elif name == "resize_crop": - transform = transforms.Compose([ - transforms.Lambda(lambda pil_image: cls.resize_crop_to_fill( - pil_image, image_size)), - transforms.ToTensor(), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - else: - raise NotImplementedError(f"Transform {name} not implemented") - return transform - - @classmethod - def get_transforms_video(cls, name="center", image_size=(256, 256)): - if name is None: - return None - elif name == "center": - assert image_size[0] == image_size[ - 1], "image_size must be square for center crop" - transform_video = transforms.Compose([ - video_transforms.ToTensorVideo(), # TCHW - # video_transforms.RandomHorizontalFlipVideo(), - video_transforms.UCFCenterCropVideo(image_size[0]), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - elif name == "resize_crop": - transform_video = transforms.Compose([ - video_transforms.ToTensorVideo(), # TCHW - video_transforms.ResizeCrop(image_size), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - else: - raise NotImplementedError(f"Transform {name} not implemented") - return transform_video - - @classmethod - def read_image_from_path(cls, - path, - transform=None, - transform_name="center", - num_frames=1, - image_size=(256, 256)): - image = pil_loader(path) - if transform is None: - transform = cls.get_transforms_image(image_size=image_size, - name=transform_name) - image = transform(image) - video = image.unsqueeze(0).repeat(num_frames, 1, 1, 1) - video = video.permute(1, 0, 2, 3) - return video - - @classmethod - def read_video_from_path(cls, - path, - transform=None, - transform_name="center", - image_size=(256, 256)): - vframes, aframes, info = read_video(filename=path, - pts_unit="sec", - output_format="TCHW") - if transform is None: - transform = cls.get_transforms_video(image_size=image_size, - name=transform_name) - video = transform(vframes) # T C H W - video = video.permute(1, 0, 2, 3) - return video - - @classmethod - def read_from_path(cls, path, image_size, transform_name="center"): - if is_url(path): - path = download_url(path) - ext = os.path.splitext(path)[-1].lower() - if ext.lower() in VID_EXTENSIONS: - return cls.read_video_from_path(path, - image_size=image_size, - transform_name=transform_name) - else: - assert ext.lower( - ) in IMG_EXTENSIONS, f"Unsupported file format: {ext}" - return cls.read_image_from_path(path, - image_size=image_size, - transform_name=transform_name) - - @classmethod - def collect_references_batch(cls, vae, reference_paths, image_size): - refs_x = [] # refs_x: [batch, ref_num, C, T, H, W] - for reference_path in reference_paths: - if reference_path == "": - refs_x.append([]) - continue - ref_path = reference_path.split(";") - ref = [] - for r_path in ref_path: - r = cls.read_from_path(r_path, - image_size, - transform_name="resize_crop") - r_x = vae.encode(r.unsqueeze(0).to(vae.device, vae.dtype)) - r_x = r_x.squeeze(0) - ref.append(r_x) - refs_x.append(ref) - return refs_x - - @staticmethod - def append_generated(vae, generated_video, refs_x, mask_strategy, loop_i, - condition_frame_length, condition_frame_edit): - ref_x = vae.encode(generated_video) - for j, refs in enumerate(refs_x): - if refs is None: - refs_x[j] = [ref_x[j]] - else: - refs.append(ref_x[j]) - if mask_strategy[j] is None or mask_strategy[j] == "": - mask_strategy[j] = "" - else: - mask_strategy[j] += ";" - mask_strategy[ - j] += f"{loop_i},{len(refs)-1},-{condition_frame_length},0,{condition_frame_length},{condition_frame_edit}" - return refs_x, mask_strategy - - @classmethod - def parse_mask_strategy(cls, mask_strategy): - mask_batch = [] - if mask_strategy == "" or mask_strategy is None: - return mask_batch - mask_strategy = mask_strategy.split(";") - for mask in mask_strategy: - mask_group = mask.split(",") - num_group = len(mask_group) - assert num_group >= 1 and num_group <= 6, f"Invalid mask strategy: {mask}" - mask_group.extend(cls.MASK_DEFAULT[num_group:]) - for i in range(5): - mask_group[i] = int(mask_group[i]) - mask_group[5] = float(mask_group[5]) - mask_batch.append(mask_group) - return mask_batch - - @classmethod - def find_nearest_point(cls, value, point, max_value): - t = value // point - if value % point > point / 2 and t < max_value // point - 1: - t += 1 - return t * point - - @classmethod - def apply_mask_strategy(cls, z, refs_x, mask_strategys, loop_i, align=None): - masks = [] - no_mask = True - for i, mask_strategy in enumerate(mask_strategys): - no_mask = False - mask = torch.ones(z.shape[2], dtype=torch.float, device=z.device) - mask_strategy = cls.parse_mask_strategy(mask_strategy) - for mst in mask_strategy: - loop_id, m_id, m_ref_start, m_target_start, m_length, edit_ratio = mst - if loop_id != loop_i: - continue - ref = refs_x[i][m_id] - - if m_ref_start < 0: - # ref: [C, T, H, W] - m_ref_start = ref.shape[1] + m_ref_start - if m_target_start < 0: - # z: [B, C, T, H, W] - m_target_start = z.shape[2] + m_target_start - if align is not None: - m_ref_start = cls.find_nearest_point( - m_ref_start, align, ref.shape[1]) - m_target_start = cls.find_nearest_point( - m_target_start, align, z.shape[2]) - m_length = min(m_length, z.shape[2] - m_target_start, - ref.shape[1] - m_ref_start) - z[i, :, m_target_start:m_target_start + - m_length] = ref[:, m_ref_start:m_ref_start + m_length] - mask[m_target_start:m_target_start + m_length] = edit_ratio - masks.append(mask) - if no_mask: - return None - masks = torch.stack(masks) - return masks - - @staticmethod - def dframe_to_frame(num): - assert num % 5 == 0, f"Invalid num: {num}" - return num // 5 * 17 - - @staticmethod - def save_sample(x, - save_path=None, - fps=8, - normalize=True, - value_range=(-1, 1), - force_video=False, - verbose=True): - """ - Args: - x (Tensor): shape [C, T, H, W] - """ - assert x.ndim == 4 - if not force_video and x.shape[1] == 1: # T = 1: save as image - save_path += ".png" - x = x.squeeze(1) - save_image([x], - save_path, - normalize=normalize, - value_range=value_range) - else: - save_path += ".mp4" - if normalize: - low, high = value_range - x.clamp_(min=low, max=high) - x.sub_(low).div_(max(high - low, 1e-5)) - x = x.mul(255).add_(0.5).clamp_(0, 255).permute(1, 2, 3, 0).to( - "cpu", torch.uint8) - write_video(save_path, x, fps=fps, video_codec="h264") - if verbose: - print(f"Saved to {save_path}") - return save_path - - @staticmethod - def add_watermark( - input_video_path, - watermark_image_path="./assets/images/watermark/watermark.png", - output_video_path=None): - # execute this command in terminal with subprocess - # return if the process is successful - if output_video_path is None: - output_video_path = input_video_path.replace( - ".mp4", "_watermark.mp4") - cmd = f'ffmpeg -y -i {input_video_path} -i {watermark_image_path}' - cmd += f'-filter_complex "[1][0]scale2ref=oh*mdar:ih*0.1[logo][video];[video][logo]overlay" {output_video_path}' - import subprocess - exit_code = subprocess.call(cmd, shell=True) - is_success = exit_code == 0 - return is_success - - -class DataProcessor(): - - def __init__(self, text_encoder, vae): - self._text_encoder = text_encoder - self._vae = vae - - @staticmethod - def set_random_seed(seed: int): - random.seed(seed) - np.random.seed(seed) - torch.manual_seed(seed) - if torch.cuda.is_available(): - torch.cuda.manual_seed_all(seed) - return seed - - @staticmethod - def prepare_multi_resolution_info(model_type, batch_size, image_size, - num_frames, fps, device, dtype): - IMG_FPS = 120 - ret = {} - if model_type in ["STDiT3", "OpenSora"]: - fps = fps if num_frames > 1 else IMG_FPS - fps = torch.tensor([fps], device=device, - dtype=dtype).repeat(batch_size) - height = torch.tensor([image_size[0]], device=device, - dtype=dtype).repeat(batch_size) - width = torch.tensor([image_size[1]], device=device, - dtype=dtype).repeat(batch_size) - num_frames = torch.tensor([num_frames], device=device, - dtype=dtype).repeat(batch_size) - ar = torch.tensor([image_size[0] / image_size[1]], - device=device, - dtype=dtype).repeat(batch_size) - ret = dict(height=height, - width=width, - num_frames=num_frames, - ar=ar, - fps=fps) - else: - raise NotImplementedError(f"Model type is {model_type}") - return ret - - @staticmethod - def get_save_path_name( - save_dir, - sample_name=None, # prefix - sample_idx=None, # sample index - prompt=None, # used prompt - prompt_as_path=False, # use prompt as path - num_sample=1, # number of samples to generate for one prompt - k=None, # kth sample - ): - if sample_name is None: - sample_name = "" if prompt_as_path else "sample" - sample_name_suffix = prompt if prompt_as_path else f"_{sample_idx:04d}" - save_path = os.path.join(save_dir, f"{sample_name}{sample_name_suffix}") - if num_sample != 1: - save_path = f"{save_path}-{k}" - return save_path - - @staticmethod - def load_prompts(prompt_path, start_idx=None, end_idx=None): - return PromptProcessor.load_prompts(prompt_path=prompt_path, - start_idx=start_idx, - end_idx=end_idx) - - @staticmethod - def extract_json_from_prompts(prompts, reference, mask_strategy): - return PromptProcessor.extract_json_from_prompts( - prompts=prompts, reference=reference, mask_strategy=mask_strategy) - - @staticmethod - def split_prompt(prompt_text): - return PromptProcessor.split_prompt(prompt_text=prompt_text) - - @staticmethod - def merge_prompt(text_list, loop_idx_list=None): - return PromptProcessor.merge_prompt(text_list=text_list, - loop_idx_list=loop_idx_list) - - @staticmethod - def extract_prompts_loop(prompts, num_loop): - return PromptProcessor.extract_prompts_loop(prompts=prompts, - num_loop=num_loop) - - @staticmethod - def append_score_to_prompts(prompts, - aes=None, - flow=None, - camera_motion=None): - return PromptProcessor.append_score_to_prompts( - prompts=prompts, aes=aes, flow=flow, camera_motion=camera_motion) - - @staticmethod - def text_preprocessing(text, use_text_preprocessing: bool = True): - return PromptProcessor.text_preprocessing( - text=text, use_text_preprocessing=use_text_preprocessing) - - @staticmethod - def get_image_size(resolution, ar_ratio): - return VideoProcessor.get_image_size(resolution=resolution, - ar_ratio=ar_ratio) - - @staticmethod - def get_num_frames(num_frames): - return VideoProcessor.get_num_frames(num_frames=num_frames) - - def get_latent_size(self, input_size): - return VideoProcessor.get_latent_size(self._vae, input_size=input_size) - - def collect_references_batch(self, reference_paths, image_size): - return VideoProcessor.collect_references_batch( - self._vae, reference_paths=reference_paths, image_size=image_size) - - def append_generated(self, generated_video, refs_x, mask_strategy, loop_i, - condition_frame_length, condition_frame_edit): - return VideoProcessor.append_generated( - self._vae, - generated_video=generated_video, - refs_x=refs_x, - mask_strategy=mask_strategy, - loop_i=loop_i, - condition_frame_length=condition_frame_length, - condition_frame_edit=condition_frame_edit) - - @staticmethod - def apply_mask_strategy(z, refs_x, mask_strategys, loop_i, align=None): - return VideoProcessor.apply_mask_strategy(z=z, - refs_x=refs_x, - mask_strategys=mask_strategys, - loop_i=loop_i, - align=align) - - @staticmethod - def dframe_to_frame(num): - return VideoProcessor.dframe_to_frame(num=num) - - @staticmethod - def save_sample(x, - save_path=None, - fps=8, - normalize=True, - value_range=(-1, 1), - force_video=False, - verbose=True): - return VideoProcessor.save_sample(x=x, - save_path=save_path, - fps=fps, - normalize=normalize, - value_range=value_range, - force_video=force_video, - verbose=verbose) - - @staticmethod - def add_watermark( - input_video_path, - watermark_image_path="./assets/images/watermark/watermark.png", - output_video_path=None): - return VideoProcessor.add_watermark( - input_video_path=input_video_path, - watermark_image_path=watermark_image_path, - output_video_path=output_video_path) diff --git a/examples/models/contrib/stdit/vae.py b/examples/models/contrib/stdit/vae.py deleted file mode 100644 index cd193b8eb75a..000000000000 --- a/examples/models/contrib/stdit/vae.py +++ /dev/null @@ -1,802 +0,0 @@ -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# 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. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/models/vae/vae.py - -import os -from typing import Tuple, Union - -import torch -import torch.nn.functional as F -from diffusers.models import AutoencoderKL -from einops import rearrange -from transformers import PretrainedConfig, PreTrainedModel -from utils import load_checkpoint - - -class VideoAutoencoderKL(torch.nn.Module): - - def __init__( - self, - from_pretrained=None, - micro_batch_size=None, - cache_dir=None, - local_files_only=False, - subfolder=None, - scaling_factor=0.18215, - ): - super().__init__() - self.module = AutoencoderKL.from_pretrained( - from_pretrained, - cache_dir=cache_dir, - local_files_only=local_files_only, - subfolder=subfolder, - ) - self.out_channels = self.module.config.latent_channels - self.patch_size = (1, 8, 8) - self.micro_batch_size = micro_batch_size - self.scaling_factor = scaling_factor - - def encode(self, x): - # x: (B, C, T, H, W) - B = x.shape[0] - x = rearrange(x, "B C T H W -> (B T) C H W") - - if self.micro_batch_size is None: - x = self.module.encode(x).latent_dist.sample().mul_( - self.scaling_factor) - else: - # NOTE: cannot be used for training - bs = self.micro_batch_size - x_out = [] - for i in range(0, x.shape[0], bs): - x_bs = x[i:i + bs] - x_bs = self.module.encode(x_bs).latent_dist.sample().mul_( - self.scaling_factor) - x_out.append(x_bs) - x = torch.cat(x_out, dim=0) - x = rearrange(x, "(B T) C H W -> B C T H W", B=B) - return x - - def decode(self, x, **kwargs): - # x: (B, C, T, H, W) - B = x.shape[0] - x = rearrange(x, "B C T H W -> (B T) C H W") - if self.micro_batch_size is None: - x = self.module.decode(x / self.scaling_factor).sample - else: - # NOTE: cannot be used for training - bs = self.micro_batch_size - x_out = [] - for i in range(0, x.shape[0], bs): - x_bs = x[i:i + bs] - x_bs = self.module.decode(x_bs / self.scaling_factor).sample - x_out.append(x_bs) - x = torch.cat(x_out, dim=0) - x = rearrange(x, "(B T) C H W -> B C T H W", B=B) - return x - - def get_latent_size(self, input_size): - latent_size = [] - for i in range(3): - # assert ( - # input_size[i] is None or input_size[i] % self.patch_size[i] == 0 - # ), "Input size must be divisible by patch size" - latent_size.append( - input_size[i] // - self.patch_size[i] if input_size[i] is not None else None) - return latent_size - - @property - def device(self): - return next(self.parameters()).device - - @property - def dtype(self): - return next(self.parameters()).dtype - - -def cast_tuple(t, length=1): - return t if isinstance(t, tuple) else ((t, ) * length) - - -def divisible_by(num, den): - return (num % den) == 0 - - -def is_odd(n): - return not divisible_by(n, 2) - - -def pad_at_dim(t, pad, dim=-1): - dims_from_right = (-dim - 1) if dim < 0 else (t.ndim - dim - 1) - zeros = (0, 0) * dims_from_right - return F.pad(t, (*zeros, *pad), mode="constant") - - -def exists(v): - return v is not None - - -class DiagonalGaussianDistribution: - """Stripped version of https://github.com/richzhang/PerceptualSimilarity/tree/master/models""" - - def __init__( - self, - parameters, - deterministic=False, - ): - self.parameters = parameters - self.mean, self.logvar = torch.chunk(parameters, 2, dim=1) - self.logvar = torch.clamp(self.logvar, -30.0, 20.0) - self.deterministic = deterministic - self.std = torch.exp(0.5 * self.logvar) - self.var = torch.exp(self.logvar) - if self.deterministic: - self.var = self.std = torch.zeros_like(self.mean).to( - device=self.parameters.device, dtype=self.mean.dtype) - - def sample(self): - # torch.randn: standard normal distribution - x = self.mean + self.std * torch.randn(self.mean.shape).to( - device=self.parameters.device, dtype=self.mean.dtype) - return x - - def kl(self, other=None): - if self.deterministic: - return torch.Tensor([0.0]) - else: - if other is None: # SCH: assumes other is a standard normal distribution - return 0.5 * torch.sum( - torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, - dim=[1, 2, 3, 4]) - else: - return 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var + - self.var / other.var - 1.0 - self.logvar + other.logvar, - dim=[1, 2, 3, 4], - ) - - def nll(self, sample, dims=[1, 2, 3, 4]): - if self.deterministic: - return torch.Tensor([0.0]) - logtwopi = torch.log(torch.Tensor([2.0 * torch.pi])) - return 0.5 * torch.sum(logtwopi + self.logvar + - torch.pow(sample - self.mean, 2) / self.var, - dim=dims) - - def mode(self): - return self.mean - - -class CausalConv3d(torch.nn.Module): - - def __init__( - self, - chan_in, - chan_out, - kernel_size: Union[int, Tuple[int, int, int]], - pad_mode="constant", - strides=None, # allow custom stride - **kwargs, - ): - super().__init__() - kernel_size = cast_tuple(kernel_size, 3) - time_kernel_size, height_kernel_size, width_kernel_size = kernel_size - assert is_odd(height_kernel_size) and is_odd(width_kernel_size) - dilation = kwargs.pop("dilation", 1) - stride = strides[0] if strides is not None else kwargs.pop("stride", 1) - self.pad_mode = pad_mode - time_pad = dilation * (time_kernel_size - 1) + (1 - stride) - height_pad = height_kernel_size // 2 - width_pad = width_kernel_size // 2 - self.time_pad = time_pad - self.time_causal_padding = (width_pad, width_pad, height_pad, - height_pad, time_pad, 0) - stride = strides if strides is not None else (stride, 1, 1) - dilation = (dilation, 1, 1) - self.conv = torch.nn.Conv3d(chan_in, - chan_out, - kernel_size, - stride=stride, - dilation=dilation, - **kwargs) - - def forward(self, x): - x = F.pad(x, self.time_causal_padding, mode=self.pad_mode) - x = self.conv(x) - return x - - -class ResBlock(torch.nn.Module): - - def __init__( - self, - in_channels, # SCH: added - filters, - conv_fn, - activation_fn=torch.nn.SiLU, - use_conv_shortcut=False, - num_groups=32, - ): - super().__init__() - self.in_channels = in_channels - self.filters = filters - self.activate = activation_fn() - self.use_conv_shortcut = use_conv_shortcut - - # SCH: MAGVIT uses GroupNorm by default - self.norm1 = torch.nn.GroupNorm(num_groups, in_channels) - self.conv1 = conv_fn(in_channels, - self.filters, - kernel_size=(3, 3, 3), - bias=False) - self.norm2 = torch.nn.GroupNorm(num_groups, self.filters) - self.conv2 = conv_fn(self.filters, - self.filters, - kernel_size=(3, 3, 3), - bias=False) - if in_channels != filters: - if self.use_conv_shortcut: - self.conv3 = conv_fn(in_channels, - self.filters, - kernel_size=(3, 3, 3), - bias=False) - else: - self.conv3 = conv_fn(in_channels, - self.filters, - kernel_size=(1, 1, 1), - bias=False) - - def forward(self, x): - residual = x - x = self.norm1(x) - x = self.activate(x) - x = self.conv1(x) - x = self.norm2(x) - x = self.activate(x) - x = self.conv2(x) - if self.in_channels != self.filters: # SCH: ResBlock X->Y - residual = self.conv3(residual) - return x + residual - - -def get_activation_fn(activation): - if activation == "relu": - activation_fn = torch.nn.ReLU - elif activation == "swish": - activation_fn = torch.nn.SiLU - else: - raise NotImplementedError - return activation_fn - - -class Encoder(torch.nn.Module): - """Encoder Blocks.""" - - def __init__( - self, - in_out_channels=4, - latent_embed_dim=512, # num channels for latent vector - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(False, True, True), - num_groups=32, # for nn.GroupNorm - activation_fn="swish", - ): - super().__init__() - self.filters = filters - self.num_res_blocks = num_res_blocks - self.num_blocks = len(channel_multipliers) - self.channel_multipliers = channel_multipliers - self.temporal_downsample = temporal_downsample - self.num_groups = num_groups - self.embedding_dim = latent_embed_dim - - self.activation_fn = get_activation_fn(activation_fn) - self.activate = self.activation_fn() - self.conv_fn = CausalConv3d - self.block_args = dict( - conv_fn=self.conv_fn, - activation_fn=self.activation_fn, - use_conv_shortcut=False, - num_groups=self.num_groups, - ) - - # first layer conv - self.conv_in = self.conv_fn( - in_out_channels, - filters, - kernel_size=(3, 3, 3), - bias=False, - ) - - # ResBlocks and conv downsample - self.block_res_blocks = torch.nn.ModuleList([]) - self.conv_blocks = torch.nn.ModuleList([]) - - filters = self.filters - prev_filters = filters # record for in_channels - for i in range(self.num_blocks): - filters = self.filters * self.channel_multipliers[i] - block_items = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - block_items.append( - ResBlock(prev_filters, filters, **self.block_args)) - prev_filters = filters # update in_channels - self.block_res_blocks.append(block_items) - - if i < self.num_blocks - 1: - if self.temporal_downsample[i]: - t_stride = 2 if self.temporal_downsample[i] else 1 - s_stride = 1 - self.conv_blocks.append( - self.conv_fn(prev_filters, - filters, - kernel_size=(3, 3, 3), - strides=(t_stride, s_stride, s_stride))) - prev_filters = filters # update in_channels - else: - # if no t downsample, don't add since this does nothing for pipeline models - self.conv_blocks.append( - torch.nn.Identity(prev_filters)) # Identity - prev_filters = filters # update in_channels - - # last layer res block - self.res_blocks = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - self.res_blocks.append( - ResBlock(prev_filters, filters, **self.block_args)) - prev_filters = filters # update in_channels - - # MAGVIT uses Group Normalization - self.norm1 = torch.nn.GroupNorm(self.num_groups, prev_filters) - - self.conv2 = self.conv_fn(prev_filters, - self.embedding_dim, - kernel_size=(1, 1, 1), - padding="same") - - def forward(self, x): - x = self.conv_in(x) - - for i in range(self.num_blocks): - for j in range(self.num_res_blocks): - x = self.block_res_blocks[i][j](x) - if i < self.num_blocks - 1: - x = self.conv_blocks[i](x) - for i in range(self.num_res_blocks): - x = self.res_blocks[i](x) - - x = self.norm1(x) - x = self.activate(x) - x = self.conv2(x) - return x - - -class Decoder(torch.nn.Module): - """Decoder Blocks.""" - - def __init__( - self, - in_out_channels=4, - latent_embed_dim=512, - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(False, True, True), - num_groups=32, # for nn.GroupNorm - activation_fn="swish", - ): - super().__init__() - self.filters = filters - self.num_res_blocks = num_res_blocks - self.num_blocks = len(channel_multipliers) - self.channel_multipliers = channel_multipliers - self.temporal_downsample = temporal_downsample - self.num_groups = num_groups - self.embedding_dim = latent_embed_dim - self.s_stride = 1 - - self.activation_fn = get_activation_fn(activation_fn) - self.activate = self.activation_fn() - self.conv_fn = CausalConv3d - self.block_args = dict( - conv_fn=self.conv_fn, - activation_fn=self.activation_fn, - use_conv_shortcut=False, - num_groups=self.num_groups, - ) - - filters = self.filters * self.channel_multipliers[-1] - prev_filters = filters - - # last conv - self.conv1 = self.conv_fn(self.embedding_dim, - filters, - kernel_size=(3, 3, 3), - bias=True) - - # last layer res block - self.res_blocks = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - self.res_blocks.append(ResBlock(filters, filters, - **self.block_args)) - - # ResBlocks and conv upsample - self.block_res_blocks = torch.nn.ModuleList([]) - self.num_blocks = len(self.channel_multipliers) - self.conv_blocks = torch.nn.ModuleList([]) - # reverse to keep track of the in_channels, but append also in a reverse direction - for i in reversed(range(self.num_blocks)): - filters = self.filters * self.channel_multipliers[i] - # resblock handling - block_items = torch.nn.ModuleList([]) - for _ in range(self.num_res_blocks): - block_items.append( - ResBlock(prev_filters, filters, **self.block_args)) - prev_filters = filters # SCH: update in_channels - self.block_res_blocks.insert(0, block_items) # SCH: append in front - - # conv blocks with upsampling - if i > 0: - if self.temporal_downsample[i - 1]: - t_stride = 2 if self.temporal_downsample[i - 1] else 1 - # SCH: T-Causal Conv 3x3x3, f -> (t_stride * 2 * 2) * f, depth to space t_stride x 2 x 2 - self.conv_blocks.insert( - 0, - self.conv_fn(prev_filters, - prev_filters * t_stride * self.s_stride * - self.s_stride, - kernel_size=(3, 3, 3)), - ) - else: - self.conv_blocks.insert( - 0, - torch.nn.Identity(prev_filters), - ) - - self.norm1 = torch.nn.GroupNorm(self.num_groups, prev_filters) - - self.conv_out = self.conv_fn(filters, in_out_channels, 3) - - def forward(self, x): - x = self.conv1(x) - for i in range(self.num_res_blocks): - x = self.res_blocks[i](x) - for i in reversed(range(self.num_blocks)): - for j in range(self.num_res_blocks): - x = self.block_res_blocks[i][j](x) - if i > 0: - t_stride = 2 if self.temporal_downsample[i - 1] else 1 - x = self.conv_blocks[i - 1](x) - x = rearrange( - x, - "B (C ts hs ws) T H W -> B C (T ts) (H hs) (W ws)", - ts=t_stride, - hs=self.s_stride, - ws=self.s_stride, - ) - - x = self.norm1(x) - x = self.activate(x) - x = self.conv_out(x) - return x - - -class VAE_Temporal(torch.nn.Module): - - def __init__( - self, - in_out_channels=4, - latent_embed_dim=4, - embed_dim=4, - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(True, True, False), - num_groups=32, # for nn.GroupNorm - activation_fn="swish", - ): - super().__init__() - - self.time_downsample_factor = 2**sum(temporal_downsample) - # self.time_padding = self.time_downsample_factor - 1 - self.patch_size = (self.time_downsample_factor, 1, 1) - self.out_channels = in_out_channels - - # NOTE: following MAGVIT, conv in bias=False in encoder first conv - self.encoder = Encoder( - in_out_channels=in_out_channels, - latent_embed_dim=latent_embed_dim * 2, - filters=filters, - num_res_blocks=num_res_blocks, - channel_multipliers=channel_multipliers, - temporal_downsample=temporal_downsample, - num_groups=num_groups, # for nn.GroupNorm - activation_fn=activation_fn, - ) - self.quant_conv = CausalConv3d(2 * latent_embed_dim, 2 * embed_dim, 1) - - self.post_quant_conv = CausalConv3d(embed_dim, latent_embed_dim, 1) - self.decoder = Decoder( - in_out_channels=in_out_channels, - latent_embed_dim=latent_embed_dim, - filters=filters, - num_res_blocks=num_res_blocks, - channel_multipliers=channel_multipliers, - temporal_downsample=temporal_downsample, - num_groups=num_groups, # for nn.GroupNorm - activation_fn=activation_fn, - ) - - def get_latent_size(self, input_size): - latent_size = [] - for i in range(3): - if input_size[i] is None: - lsize = None - elif i == 0: - time_padding = (0 if - (input_size[i] % self.time_downsample_factor - == 0) else self.time_downsample_factor - - input_size[i] % self.time_downsample_factor) - lsize = (input_size[i] + time_padding) // self.patch_size[i] - else: - lsize = input_size[i] // self.patch_size[i] - latent_size.append(lsize) - return latent_size - - def encode(self, x): - time_padding = x.shape[2] % self.time_downsample_factor - if time_padding != 0: - time_padding = self.time_downsample_factor - time_padding - x = pad_at_dim(x, (time_padding, 0), dim=2) - encoded_feature = self.encoder(x) - moments = self.quant_conv(encoded_feature).to(x.dtype) - posterior = DiagonalGaussianDistribution(moments) - return posterior - - def decode(self, z, num_frames=None): - time_padding = num_frames % self.time_downsample_factor - if time_padding != 0: - time_padding = self.time_downsample_factor - time_padding - z = self.post_quant_conv(z) - x = self.decoder(z) - x = x[:, :, time_padding:] - return x - - def forward(self, x, sample_posterior=True): - posterior = self.encode(x) - if sample_posterior: - z = posterior.sample() - else: - z = posterior.mode() - recon_video = self.decode(z, num_frames=x.shape[2]) - return recon_video, posterior, z - - -VAE_MODELS = { - "VideoAutoencoderKL": VideoAutoencoderKL, - "VAE_Temporal_SD": VAE_Temporal, -} - - -class VideoAutoencoderPipelineConfig(PretrainedConfig): - model_type = "VideoAutoencoderPipeline" - - def __init__( - self, - spatial_vae_config=None, - temporal_vae_config=None, - from_pretrained=None, - freeze_vae_2d=False, - cal_loss=False, - micro_frame_size=None, - shift=0.0, - scale=1.0, - **kwargs, - ): - self.spatial_vae_config = spatial_vae_config - self.temporal_vae_config = temporal_vae_config - self.from_pretrained = from_pretrained - self.freeze_vae_2d = freeze_vae_2d - self.cal_loss = cal_loss - self.micro_frame_size = micro_frame_size - self.shift = shift - self.scale = scale - super().__init__(**kwargs) - - -class VideoAutoencoderPipeline(PreTrainedModel): - config_class = VideoAutoencoderPipelineConfig - - def __init__(self, config: VideoAutoencoderPipelineConfig): - super().__init__(config=config) - vae_type = config.spatial_vae_config.pop('type') - self.spatial_vae = VAE_MODELS[vae_type](**config.spatial_vae_config) - - vae_type = config.temporal_vae_config.pop('type') - pretrained_path = config.temporal_vae_config.pop( - 'from_pretrained' - ) if 'from_pretrained' in config.temporal_vae_config else None - self.temporal_vae = VAE_MODELS[vae_type](**config.temporal_vae_config) - if pretrained_path is not None: - load_checkpoint(self.temporal_vae, pretrained_path) - - self.cal_loss = config.cal_loss - self.micro_frame_size = config.micro_frame_size - self.micro_z_frame_size = self.temporal_vae.get_latent_size( - [config.micro_frame_size, None, None])[0] - if config.freeze_vae_2d: - for param in self.spatial_vae.parameters(): - param.requires_grad = False - self.out_channels = self.temporal_vae.out_channels - # normalization parameters - scale = torch.tensor(config.scale) - shift = torch.tensor(config.shift) - if len(scale.shape) > 0: - scale = scale[None, :, None, None, None] - if len(shift.shape) > 0: - shift = shift[None, :, None, None, None] - self.register_buffer("scale", scale) - self.register_buffer("shift", shift) - - def encode(self, x): - x_z = self.spatial_vae.encode(x) - - if self.micro_frame_size is None: - posterior = self.temporal_vae.encode(x_z) - z = posterior.sample() - else: - z_list = [] - for i in range(0, x_z.shape[2], self.micro_frame_size): - x_z_bs = x_z[:, :, i:i + self.micro_frame_size] - posterior = self.temporal_vae.encode(x_z_bs) - z_list.append(posterior.sample()) - z = torch.cat(z_list, dim=2) - - if self.cal_loss: - return z, posterior, x_z - else: - return (z - self.shift) / self.scale - - def decode(self, z, num_frames=None): - if not self.cal_loss: - z = z * self.scale.to(z.dtype) + self.shift.to(z.dtype) - - if self.micro_frame_size is None: - x_z = self.temporal_vae.decode(z, num_frames=num_frames) - x = self.spatial_vae.decode(x_z) - else: - x_z_list = [] - for i in range(0, z.size(2), self.micro_z_frame_size): - z_bs = z[:, :, i:i + self.micro_z_frame_size] - x_z_bs = self.temporal_vae.decode(z_bs, - num_frames=min( - self.micro_frame_size, - num_frames)) - x_z_list.append(x_z_bs) - num_frames -= self.micro_frame_size - x_z = torch.cat(x_z_list, dim=2) - x = self.spatial_vae.decode(x_z) - - if self.cal_loss: - return x, x_z - else: - return x - - def forward(self, x): - assert self.cal_loss, "This method is only available when cal_loss is True" - z, posterior, x_z = self.encode(x) - x_rec, x_z_rec = self.decode(z, num_frames=x_z.shape[2]) - return x_rec, x_z_rec, z, posterior, x_z - - def get_latent_size(self, input_size): - if self.micro_frame_size is None or input_size[0] is None: - return self.temporal_vae.get_latent_size( - self.spatial_vae.get_latent_size(input_size)) - else: - sub_input_size = [ - self.micro_frame_size, input_size[1], input_size[2] - ] - sub_latent_size = self.temporal_vae.get_latent_size( - self.spatial_vae.get_latent_size(sub_input_size)) - sub_latent_size[0] = sub_latent_size[0] * (input_size[0] // - self.micro_frame_size) - remain_temporal_size = [ - input_size[0] % self.micro_frame_size, None, None - ] - if remain_temporal_size[0] > 0: - remain_size = self.temporal_vae.get_latent_size( - remain_temporal_size) - sub_latent_size[0] += remain_size[0] - return sub_latent_size - - def get_temporal_last_layer(self): - return self.temporal_vae.decoder.conv_out.conv.weight - - @property - def device(self): - return next(self.parameters()).device - - @property - def dtype(self): - return next(self.parameters()).dtype - - -def get_vae( - micro_batch_size=4, - micro_frame_size=17, - from_pretrained=None, - local_files_only=False, - freeze_vae_2d=False, - cal_loss=False, - force_huggingface=False, -): - spatial_vae_config = dict( - type="VideoAutoencoderKL", - from_pretrained="PixArt-alpha/pixart_sigma_sdxlvae_T5_diffusers", - subfolder="vae", - micro_batch_size=micro_batch_size, - local_files_only=local_files_only, - ) - temporal_vae_config = dict( - type="VAE_Temporal_SD", - from_pretrained=None, - in_out_channels=4, - latent_embed_dim=4, - embed_dim=4, - filters=128, - num_res_blocks=4, - channel_multipliers=(1, 2, 2, 4), - temporal_downsample=(False, True, True), - ) - shift = (-0.10, 0.34, 0.27, 0.98) - scale = (3.85, 2.32, 2.33, 3.06) - kwargs = dict( - spatial_vae_config=spatial_vae_config, - temporal_vae_config=temporal_vae_config, - freeze_vae_2d=freeze_vae_2d, - cal_loss=cal_loss, - micro_frame_size=micro_frame_size, - shift=shift, - scale=scale, - ) - - if force_huggingface or (from_pretrained is not None - and not os.path.exists(from_pretrained)): - model = VideoAutoencoderPipeline.from_pretrained( - from_pretrained, **kwargs) - else: - config = VideoAutoencoderPipelineConfig(**kwargs) - model = VideoAutoencoderPipeline(config) - if from_pretrained: - load_checkpoint(model, from_pretrained) - return model diff --git a/examples/models/contrib/stdit/video_transforms.py b/examples/models/contrib/stdit/video_transforms.py deleted file mode 100644 index 3cabd640c5c6..000000000000 --- a/examples/models/contrib/stdit/video_transforms.py +++ /dev/null @@ -1,581 +0,0 @@ -# Copyright 2024 Vchitect/Latte - -# 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.# Modified from Latte - -# Copyright 2024 HPC-AI Technology Inc. - -# 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. - -# 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. - -# reference: https://github.com/hpcaitech/Open-Sora/blob/main/opensora/datasets/video_transforms.py - -import numbers -import random - -import numpy as np -import torch - - -def _is_tensor_video_clip(clip): - if not torch.is_tensor(clip): - raise TypeError("clip should be Tensor. Got %s" % type(clip)) - - if not clip.ndimension() == 4: - raise ValueError("clip should be 4D. Got %dD" % clip.dim()) - - return True - - -def crop(clip, i, j, h, w): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - """ - if len(clip.size()) != 4: - raise ValueError("clip should be a 4D tensor") - return clip[..., i:i + h, j:j + w] - - -def resize(clip, target_size, interpolation_mode): - if len(target_size) != 2: - raise ValueError( - f"target size should be tuple (height, width), instead got {target_size}" - ) - return torch.nn.functional.interpolate(clip, - size=target_size, - mode=interpolation_mode, - align_corners=False) - - -def resize_scale(clip, target_size, interpolation_mode): - if len(target_size) != 2: - raise ValueError( - f"target size should be tuple (height, width), instead got {target_size}" - ) - H, W = clip.size(-2), clip.size(-1) - scale_ = target_size[0] / min(H, W) - return torch.nn.functional.interpolate(clip, - scale_factor=scale_, - mode=interpolation_mode, - align_corners=False) - - -def resized_crop(clip, i, j, h, w, size, interpolation_mode="bilinear"): - """ - Do spatial cropping and resizing to the video clip - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - i (int): i in (i,j) i.e coordinates of the upper left corner. - j (int): j in (i,j) i.e coordinates of the upper left corner. - h (int): Height of the cropped region. - w (int): Width of the cropped region. - size (tuple(int, int)): height and width of resized clip - Returns: - clip (torch.tensor): Resized and cropped clip. Size is (T, C, H, W) - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - clip = crop(clip, i, j, h, w) - clip = resize(clip, size, interpolation_mode) - return clip - - -def center_crop(clip, crop_size): - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - th, tw = crop_size - if h < th or w < tw: - raise ValueError("height and width must be no smaller than crop_size") - - i = int(round((h - th) / 2.0)) - j = int(round((w - tw) / 2.0)) - return crop(clip, i, j, th, tw) - - -def center_crop_using_short_edge(clip): - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - if h < w: - th, tw = h, h - i = 0 - j = int(round((w - tw) / 2.0)) - else: - th, tw = w, w - i = int(round((h - th) / 2.0)) - j = 0 - return crop(clip, i, j, th, tw) - - -def resize_crop_to_fill(clip, target_size): - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - th, tw = target_size[0], target_size[1] - rh, rw = th / h, tw / w - if rh > rw: - sh, sw = th, round(w * rh) - clip = resize(clip, (sh, sw), "bilinear") - i = 0 - j = int(round(sw - tw) / 2.0) - else: - sh, sw = round(h * rw), tw - clip = resize(clip, (sh, sw), "bilinear") - i = int(round(sh - th) / 2.0) - j = 0 - assert i + th <= clip.size(-2) and j + tw <= clip.size(-1) - return crop(clip, i, j, th, tw) - - -def random_shift_crop(clip): - """ - Slide along the long edge, with the short edge as crop size - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - h, w = clip.size(-2), clip.size(-1) - - if h <= w: - short_edge = h - else: - short_edge = w - - th, tw = short_edge, short_edge - - i = torch.randint(0, h - th + 1, size=(1, )).item() - j = torch.randint(0, w - tw + 1, size=(1, )).item() - return crop(clip, i, j, th, tw) - - -def to_tensor(clip): - """ - Convert tensor data type from uint8 to float, divide value by 255.0 and - permute the dimensions of clip tensor - Args: - clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) - Return: - clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) - """ - _is_tensor_video_clip(clip) - if not clip.dtype == torch.uint8: - raise TypeError("clip tensor should have data type uint8. Got %s" % - str(clip.dtype)) - # return clip.float().permute(3, 0, 1, 2) / 255.0 - return clip.float() / 255.0 - - -def normalize(clip, mean, std, inplace=False): - """ - Args: - clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W) - mean (tuple): pixel RGB mean. Size is (3) - std (tuple): pixel standard deviation. Size is (3) - Returns: - normalized clip (torch.tensor): Size is (T, C, H, W) - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - if not inplace: - clip = clip.clone() - mean = torch.as_tensor(mean, dtype=clip.dtype, device=clip.device) - # print(mean) - std = torch.as_tensor(std, dtype=clip.dtype, device=clip.device) - clip.sub_(mean[:, None, None, None]).div_(std[:, None, None, None]) - return clip - - -def hflip(clip): - """ - Args: - clip (torch.tensor): Video clip to be normalized. Size is (T, C, H, W) - Returns: - flipped clip (torch.tensor): Size is (T, C, H, W) - """ - if not _is_tensor_video_clip(clip): - raise ValueError("clip should be a 4D torch.tensor") - return clip.flip(-1) - - -class ResizeCrop: - - def __init__(self, size): - if isinstance(size, numbers.Number): - self.size = (int(size), int(size)) - else: - self.size = size - - def __call__(self, clip): - clip = resize_crop_to_fill(clip, self.size) - return clip - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size})" - - -class RandomCropVideo: - - def __init__(self, size): - if isinstance(size, numbers.Number): - self.size = (int(size), int(size)) - else: - self.size = size - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: randomly cropped video clip. - size is (T, C, OH, OW) - """ - i, j, h, w = self.get_params(clip) - return crop(clip, i, j, h, w) - - def get_params(self, clip): - h, w = clip.shape[-2:] - th, tw = self.size - - if h < th or w < tw: - raise ValueError( - f"Required crop size {(th, tw)} is larger than input image size {(h, w)}" - ) - - if w == tw and h == th: - return 0, 0, h, w - - i = torch.randint(0, h - th + 1, size=(1, )).item() - j = torch.randint(0, w - tw + 1, size=(1, )).item() - - return i, j, th, tw - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size})" - - -class CenterCropResizeVideo: - """ - First use the short side for cropping length, - center crop video, then resize to the specified size - """ - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: scale resized / center cropped video clip. - size is (T, C, crop_size, crop_size) - """ - clip_center_crop = center_crop_using_short_edge(clip) - clip_center_crop_resize = resize( - clip_center_crop, - target_size=self.size, - interpolation_mode=self.interpolation_mode) - return clip_center_crop_resize - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" - - -class UCFCenterCropVideo: - """ - First scale to the specified size in equal proportion to the short edge, - then center cropping - """ - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: scale resized / center cropped video clip. - size is (T, C, crop_size, crop_size) - """ - clip_resize = resize_scale(clip=clip, - target_size=self.size, - interpolation_mode=self.interpolation_mode) - clip_center_crop = center_crop(clip_resize, self.size) - return clip_center_crop - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" - - -class KineticsRandomCropResizeVideo: - """ - Slide along the long edge, with the short edge as crop size. And resie to the desired size. - """ - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - clip_random_crop = random_shift_crop(clip) - clip_resize = resize(clip_random_crop, self.size, - self.interpolation_mode) - return clip_resize - - -class CenterCropVideo: - - def __init__( - self, - size, - interpolation_mode="bilinear", - ): - if isinstance(size, tuple): - if len(size) != 2: - raise ValueError( - f"size should be tuple (height, width), instead got {size}") - self.size = size - else: - self.size = (size, size) - - self.interpolation_mode = interpolation_mode - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Video clip to be cropped. Size is (T, C, H, W) - Returns: - torch.tensor: center cropped video clip. - size is (T, C, crop_size, crop_size) - """ - clip_center_crop = center_crop(clip, self.size) - return clip_center_crop - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(size={self.size}, interpolation_mode={self.interpolation_mode}" - - -class NormalizeVideo: - """ - Normalize the video clip by mean subtraction and division by standard deviation - Args: - mean (3-tuple): pixel RGB mean - std (3-tuple): pixel RGB standard deviation - inplace (boolean): whether do in-place normalization - """ - - def __init__(self, mean, std, inplace=False): - self.mean = mean - self.std = std - self.inplace = inplace - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): video clip must be normalized. Size is (C, T, H, W) - """ - return normalize(clip, self.mean, self.std, self.inplace) - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(mean={self.mean}, std={self.std}, inplace={self.inplace})" - - -class ToTensorVideo: - """ - Convert tensor data type from uint8 to float, divide value by 255.0 and - permute the dimensions of clip tensor - """ - - def __init__(self): - pass - - def __call__(self, clip): - """ - Args: - clip (torch.tensor, dtype=torch.uint8): Size is (T, C, H, W) - Return: - clip (torch.tensor, dtype=torch.float): Size is (T, C, H, W) - """ - return to_tensor(clip) - - def __repr__(self) -> str: - return self.__class__.__name__ - - -class RandomHorizontalFlipVideo: - """ - Flip the video clip along the horizontal direction with a given probability - Args: - p (float): probability of the clip being flipped. Default value is 0.5 - """ - - def __init__(self, p=0.5): - self.p = p - - def __call__(self, clip): - """ - Args: - clip (torch.tensor): Size is (T, C, H, W) - Return: - clip (torch.tensor): Size is (T, C, H, W) - """ - if random.random() < self.p: - clip = hflip(clip) - return clip - - def __repr__(self) -> str: - return f"{self.__class__.__name__}(p={self.p})" - - -# ------------------------------------------------------------ -# --------------------- Sampling --------------------------- -# ------------------------------------------------------------ -class TemporalRandomCrop(object): - """Temporally crop the given frame indices at a random location. - - Args: - size (int): Desired length of frames will be seen in the model. - """ - - def __init__(self, size): - self.size = size - - def __call__(self, total_frames): - rand_end = max(0, total_frames - self.size - 1) - begin_index = random.randint(0, rand_end) - end_index = min(begin_index + self.size, total_frames) - return begin_index, end_index - - -if __name__ == "__main__": - import os - - import numpy as np - import torchvision.io as io - from torchvision import transforms - from torchvision.utils import save_image - - vframes, aframes, info = io.read_video(filename="./v_Archery_g01_c03.avi", - pts_unit="sec", - output_format="TCHW") - - trans = transforms.Compose([ - ToTensorVideo(), - RandomHorizontalFlipVideo(), - UCFCenterCropVideo(512), - # NormalizeVideo(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5], inplace=True), - transforms.Normalize(mean=[0.5, 0.5, 0.5], - std=[0.5, 0.5, 0.5], - inplace=True), - ]) - - target_video_len = 32 - frame_interval = 1 - total_frames = len(vframes) - print(total_frames) - - temporal_sample = TemporalRandomCrop(target_video_len * frame_interval) - - # Sampling video frames - start_frame_ind, end_frame_ind = temporal_sample(total_frames) - # print(start_frame_ind) - # print(end_frame_ind) - assert end_frame_ind - start_frame_ind >= target_video_len - frame_indice = np.linspace(start_frame_ind, - end_frame_ind - 1, - target_video_len, - dtype=int) - print(frame_indice) - - select_vframes = vframes[frame_indice] - print(select_vframes.shape) - print(select_vframes.dtype) - - select_vframes_trans = trans(select_vframes) - print(select_vframes_trans.shape) - print(select_vframes_trans.dtype) - - select_vframes_trans_int = ((select_vframes_trans * 0.5 + 0.5) * - 255).to(dtype=torch.uint8) - print(select_vframes_trans_int.dtype) - print(select_vframes_trans_int.permute(0, 2, 3, 1).shape) - - io.write_video("./test.avi", - select_vframes_trans_int.permute(0, 2, 3, 1), - fps=8) - - for i in range(target_video_len): - save_image(select_vframes_trans[i], - os.path.join("./test000", "%04d.png" % i), - normalize=True, - value_range=(-1, 1)) diff --git a/examples/models/core/bert/.gitignore b/examples/models/core/bert/.gitignore deleted file mode 100644 index 4b9ff316e843..000000000000 --- a/examples/models/core/bert/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -bert* -*.log diff --git a/examples/models/core/bert/README.md b/examples/models/core/bert/README.md deleted file mode 100644 index 05c9e16ca4a3..000000000000 --- a/examples/models/core/bert/README.md +++ /dev/null @@ -1,117 +0,0 @@ -# BERT and BERT Variants - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the BERT family, specifically [BERT](https://huggingface.co/docs/transformers/model_doc/bert) and [RoBERTa](https://huggingface.co/docs/transformers/model_doc/roberta) model using TensorRT-LLM. It also describes how to run on a single GPU and two GPUs. - -## Overview - -The TensorRT LLM BERT family implementation can be found in [`tensorrt_llm/models/bert/model.py`](../../../../tensorrt_llm/models/bert/model.py). -The TensorRT LLM BERT family example code is located in [`examples/models/core/bert`](./). There are two main files in that folder: - - * [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the BERT model into TensorRT LLM checkpoint format. - * [`run.py`](./run.py) to run the inference on an input text, - -## Convert Weights - -The `convert_checkpoint.py` script converts weights from HuggingFace format to TRT-LLM format. You need to prepare HuggingFace checkpoint files before you run the convert script. - -Use `--model_dir` to specify the HuggingFace checkpoint directory. - -Supported `model_name` options include: BertModel, BertForQuestionAnswering, BertForSequenceClassification. Use `--model` to specify the target BERT model. Please note that if you choose BertModel, [`convert_checkpoint.py`](./convert_checkpoint.py) will ignore the BERT model class specified in HuggingFace config file, and throw a warning to remind you. - -Use `--output_dir` to specify the converted checkpoint and configuration directory. The default value is `./tllm_checkpoint`. This directory will be used for next engine building phase. - -Take BertForQuestionAnswering for example, - -```bash -export hf_model_dir= -export model_name='bertqa' -export model='BertForQuestionAnswering' -export dtype='float16' - -# convert -python convert_checkpoint.py \ ---model $model \ ---model_dir $hf_model_dir \ ---output_dir ${model_name}_${dtype}_tllm_checkpoint \ ---dtype $dtype - -# convert tp=2 -python convert_checkpoint.py \ ---model $model \ ---model_dir $hf_model_dir \ ---output_dir ${model_name}_${dtype}_tllm_checkpoint \ ---dtype $dtype \ ---tp_size 2 - -``` - -## Build TensorRT engine(s) - -TensorRT LLM converts HuggingFace BERT family models into TensorRT engine(s). -To build the TensorRT engine, the basic command is: - -```bash -trtllm-build --checkpoint_dir ./${model_name}_${dtype}_tllm_checkpoint \ ---output_dir ${model_name}_engine_outputs \ -``` -Beside the basic engine build, TensorRT LLM provides these features by adding these flags with basic build command: - -- To use `bert_attention_plugin`, add `--bert_attention_plugin` to command. - -- To use remove input padding, add `--remove_input_padding=enable` and `--bert_attention_plugin` to command. Please note that remove input padding feature has to come with bert attention plugin. - -- To use FMHA kernels, add `--context_fmha=enable` or `--bert_context_fmha_fp32_acc=enable`(to enable FP32 accumulation). Note that these two flags should be used with `--bert_attention_plugin` - -Continue the BertForQuestionAnswering example: -```bash -# Build TensorRT engine for BertForQuestionAnswering model, with remove_input_padding enabled. -# TP=1 and TP=2 share the same build command -trtllm-build --checkpoint_dir ./${model_name}_${dtype}_tllm_checkpoint \ ---output_dir=${model_name}_engine_outputs \ ---remove_input_padding=enable \ ---bert_attention_plugin=${dtype} \ ---max_batch_size 8 \ ---max_input_len 512 -``` - -## Run TensorRT engine(s) -Run a TensorRT LLM BERT model using the engines generated by build command mentioned above. -Note that during model deployment, only the TensorRT engine files are needed. Previously downloaded model checkpoints and converted weights can be removed. - -[`run.py`](./run.py) provides an example for performing the inference and decoding the output. By default, it will use the task specific datasets as input text, for example, ['squad_v2'](https://huggingface.co/datasets/rajpurkar/squad_v2) for BertForQuestionAnswering. - -To run the TensorRT engine, the basic command is: - -```bash -run.py --engine_dir ./${model_name}_engine_outputs \ ---hf_model_dir $hf_model_dir \ # used for loading tokenizer -``` -Please note that: - -- To use remove input padding, add `--remove_input_padding` to command. This flag is used to tell the runtime how to process the input and decode the output. - -- To compare the result with HuggingFace model, add `--run_hf_test` to command. The runtime will load the HF model from `hf_model_dir` and compare the result. Refer to [`run.py`](./run.py) for more details. - -Continue the BertForQuestionAnswering example: -```bash -# Run TensorRT engine -python run.py \ ---engine_dir ./${model_name}_engine_outputs \ ---hf_model_dir=$hf_model_dir \ ---remove_input_padding \ ---run_hf_test - -# Run TP=2 inference -mpirun -n 2 \ -python run.py \ ---engine_dir ./${model_name}_engine_outputs \ ---hf_model_dir=$hf_model_dir \ ---remove_input_padding \ ---run_hf_test -``` diff --git a/examples/models/core/bert/__init__.py b/examples/models/core/bert/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/examples/models/core/bert/convert_checkpoint.py b/examples/models/core/bert/convert_checkpoint.py deleted file mode 100644 index f17624e8334e..000000000000 --- a/examples/models/core/bert/convert_checkpoint.py +++ /dev/null @@ -1,192 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Union - -from transformers import AutoConfig - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import (BertForQuestionAnswering, - BertForSequenceClassification, BertModel, - RobertaForQuestionAnswering, - RobertaForSequenceClassification, RobertaModel) -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model', - required=True, - choices=[ - 'BertModel', - 'BertForQuestionAnswering', - 'BertForSequenceClassification', - 'RobertaModel', - 'RobertaForQuestionAnswering', - 'RobertaForSequenceClassification', - ]) - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'float16']) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - # Quantization args - parser.add_argument("--use_fp8", - action="store_true", - default=False, - help="Enable FP8 per-tensor quantization") - parser.add_argument( - '--quant_ckpt_path', - type=str, - default=None, - help='Path of a quantized model checkpoint in .safetensors format') - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - - parser.add_argument('--log_level', type=str, default='info') - - args = parser.parse_args() - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - - if args.use_fp8: - quant_config.quant_algo = QuantAlgo.FP8 - return quant_config - - -def convert_and_save_hf(args): - model_dir = args.model_dir - - world_size = args.tp_size * args.pp_size - #TODO: add override_fields if needed - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - - #TODO: add fp8 support later - quant_config = args_to_quant_config(args) - - hf_config = AutoConfig.from_pretrained(model_dir, trust_remote_code=True) - assert hf_config is not None, "Failed to load huggingface config, please check!" - - def convert_and_save_rank(args, rank, tllm_class: Union[ - BertModel, - RobertaModel, - BertForQuestionAnswering, - RobertaForQuestionAnswering, - BertForSequenceClassification, - RobertaForSequenceClassification, - ]): - mapping = Mapping( - world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size, - ) - tik = time.time() - tllm_bert = tllm_class.from_hugging_face( - model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - ) - print(f'Total time of reading and converting {time.time()-tik} s') - tik = time.time() - tllm_bert.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del tllm_bert - print(f'Total time of saving checkpoint {time.time()-tik} s') - - tllm_class = globals()[f'{args.model}'] - if not args.model == hf_config.architectures[0]: - logger.warning( - "The model doesn't match the architecture in huggingface config.") - - execute(args.workers, [convert_and_save_rank] * world_size, args, - tllm_class) - release_gc() - - -def execute(workers, func, args, - tllm_class: Union[BertModel, RobertaModel, BertForQuestionAnswering, - RobertaForQuestionAnswering, - BertForSequenceClassification, - RobertaForSequenceClassification]): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank, tllm_class) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [ - p.submit(f, args, rank, tllm_class) - for rank, f in enumerate(func) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - assert ((args.tp_size <= 2) - and (args.pp_size == 1)), "For now we only support TP = 2!" - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/bert/run.py b/examples/models/core/bert/run.py deleted file mode 100644 index 3796d2a28337..000000000000 --- a/examples/models/core/bert/run.py +++ /dev/null @@ -1,338 +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. -import argparse -import json -import os - -# isort: off -import torch -import tensorrt as trt -# isort: on - -from transformers import AutoConfig, AutoTokenizer -from utils import (compare_bertcls_result, compare_bertqa_result, - decode_bertcls_output, decode_bertqa_output, get_engine_name, - intermediate_check, prepare_text_inputs, process_input, - temporary_datasets_config) - -import tensorrt_llm -from tensorrt_llm import logger -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import trt_dtype_to_torch -from tensorrt_llm.runtime import Session, TensorInfo - -from transformers import BertConfig, BertPreTrainedModel, BertForQuestionAnswering, BertForSequenceClassification, BertModel # isort:skip -from transformers import RobertaConfig, RobertaPreTrainedModel, RobertaForQuestionAnswering, RobertaForSequenceClassification, RobertaModel # isort:skip - -OUTPUT_NAME_MAPPING = { - 'BertModel': 'hidden_states', - 'BertForQuestionAnswering': 'logits', - 'BertForSequenceClassification': 'logits', - 'RobertaModel': 'hidden_states', - 'RobertaForQuestionAnswering': 'logits', - 'RobertaForSequenceClassification': 'logits' -} - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument('--engine_dir', type=str, default='bert_outputs') - parser.add_argument('--hf_model_dir', type=str, required=True) - parser.add_argument('--run_hf_test', action='store_true') - parser.add_argument('--remove_input_padding', action='store_true') - parser.add_argument('--debug', action='store_true') - - return parser.parse_args() - - -if __name__ == '__main__': - emit_engine_arch_deprecation("run.py") - args = parse_arguments() - - tensorrt_llm.logger.set_level(args.log_level) - - config_path = os.path.join(args.engine_dir, 'config.json') - with open(config_path, 'r') as f: - config = json.load(f) - - remove_padding = config['build_config']['plugin_config'][ - 'remove_input_padding'] - assert args.remove_input_padding == remove_padding, \ - f"The engine is build with remove_input_padding={remove_padding}, \ - but the inference runtime is performed with remove_input_padding={args.remove_input_padding}!" - - world_size = config['pretrained_config']['mapping']['world_size'] - assert world_size == tensorrt_llm.mpi_world_size(), \ - f'Engine world size ({world_size}) != Runtime world size ({tensorrt_llm.mpi_world_size()})' - - model_name = config['pretrained_config']['architecture'] - - # Roberta doesn't have token_type_ids, use all zeros to replace - is_roberta = "Roberta" in model_name - - runtime_rank = tensorrt_llm.mpi_rank() if world_size > 1 else 0 - - runtime_mapping = tensorrt_llm.Mapping(world_size, - runtime_rank, - tp_size=world_size) - torch.cuda.set_device(runtime_rank % runtime_mapping.gpus_per_node) - - serialize_path = get_engine_name(runtime_rank) - serialize_path = os.path.join(args.engine_dir, serialize_path) - - stream = torch.cuda.current_stream().cuda_stream - logger.info(f'Loading engine from {serialize_path}') - with open(serialize_path, 'rb') as f: - engine_buffer = f.read() - logger.info(f'Creating session from engine') - session = Session.from_serialized_engine(engine_buffer) - if args.debug: session._print_engine_info() - - #NOTE: prepare input - with temporary_datasets_config(HF_DATASETS_OFFLINE=False): - test_inputs = prepare_text_inputs(model_name) - hf_tokenizer = AutoTokenizer.from_pretrained(args.hf_model_dir) - if args.remove_input_padding: - #NOTE:Remove padding - inputs_without_padding = hf_tokenizer(**test_inputs) - input_ids_list = [ - torch.tensor(ids).int().cuda() \ - for ids in inputs_without_padding['input_ids'] - ] - # attention_mask_list = inputs_without_padding['attention_mask'], - if is_roberta: - token_type_ids_list = [ - torch.zeros_like(torch.tensor(ids)).int().cuda() \ - for ids in inputs_without_padding['input_ids'] - ] - else: - token_type_ids_list = [ - torch.tensor(ids).int().cuda() \ - for ids in inputs_without_padding['token_type_ids'] - ] - - input_ids, input_lengths, token_type_ids, position_ids, max_input_length = \ - process_input(input_ids_list=input_ids_list, - token_type_ids_list=token_type_ids_list, - is_roberta=is_roberta, - padding_idx=config['pretrained_config']['pad_token_id']) - - else: - - #NOTE:Padding: pad to longest seq len - inputs_with_padding = hf_tokenizer( - **test_inputs, - padding=True, - ) - - inputs_without_padding = hf_tokenizer(**test_inputs) - input_ids = torch.tensor(inputs_with_padding['input_ids']).int().cuda() - input_lengths = [len(x) for x in inputs_without_padding['input_ids']] - input_lengths = torch.tensor(input_lengths, - device=input_ids.device, - dtype=torch.int32) - attention_mask = torch.tensor(inputs_with_padding['attention_mask'], - device=input_ids.device, - dtype=torch.int32) - if is_roberta: - token_type_ids = torch.zeros_like(torch.tensor( - inputs_with_padding['input_ids']), - device=input_ids.device, - dtype=torch.int32) - else: - token_type_ids = torch.tensor(inputs_with_padding['token_type_ids'], - device=input_ids.device, - dtype=torch.int32) - - # NOTE: TRT-LLM perform inference - output_name = OUTPUT_NAME_MAPPING[model_name] - if args.remove_input_padding: - # NOTE: Remove padding: - inputs = { - "input_ids": input_ids, - "input_lengths": input_lengths, - "token_type_ids": token_type_ids, - "position_ids": position_ids, - "max_input_length": max_input_length - } - output_info = session.infer_shapes([ - TensorInfo("input_ids", trt.DataType.INT32, input_ids.shape), - TensorInfo("input_lengths", trt.DataType.INT32, - input_lengths.shape), - TensorInfo("token_type_ids", trt.DataType.INT32, - token_type_ids.shape), - TensorInfo("position_ids", trt.DataType.INT32, position_ids.shape), - TensorInfo("max_input_length", trt.DataType.INT32, - max_input_length.shape) - ]) - else: - #NOTE: Padding: - inputs = { - 'input_ids': input_ids, - 'input_lengths': input_lengths, - 'token_type_ids': token_type_ids, - } - output_info = session.infer_shapes([ - TensorInfo('input_ids', trt.DataType.INT32, input_ids.shape), - TensorInfo('input_lengths', trt.DataType.INT32, - input_lengths.shape), - TensorInfo('token_type_ids', trt.DataType.INT32, - token_type_ids.shape) - ]) - - outputs = { - t.name: - torch.empty(tuple(t.shape), - dtype=trt_dtype_to_torch(t.dtype), - device='cuda') - for t in output_info - } - assert output_name in outputs, f'{output_name} not found in outputs, check if build.py set output name correctly' - - logger.info(f"Rank{runtime_rank} is running inference...") - ok = session.run(inputs=inputs, outputs=outputs, stream=stream) - assert ok, "Runtime execution failed" - torch.cuda.synchronize() - res = outputs[output_name] - if args.debug: logger.info(f"Outputs:{outputs.keys()}") - - # NOTE: load hf model and perform inference as reference (only on rank0) - if tensorrt_llm.mpi_rank() == 0: - logger.info(f"Rank{runtime_rank} is generating HF reference...") - if args.run_hf_test: - hf_bert = globals()[f'{model_name}'].from_pretrained( - args.hf_model_dir).cuda().to(torch.float16).eval() - hf_inputs = hf_tokenizer(**test_inputs, - padding=True, - return_tensors="pt") - hf_inputs = hf_inputs.to(hf_bert.device) - with torch.no_grad(): - hf_outputs = hf_bert.forward(output_hidden_states=args.debug, - **hf_inputs) - - torch.cuda.synchronize() - # NOTE: Decode output (only on rank0) - - if tensorrt_llm.mpi_rank() == 0: - logger.info(f"Rank{runtime_rank} is comparing with HF reference...") - if model_name == "BertModel" or model_name == "RobertaModel": - if args.remove_input_padding: - # reshape result back to [batch_size, ...], - # and then "padding" so we could compare the tensor - from torch.nn.utils.rnn import pad_sequence - res = torch.split(res, input_lengths.tolist(), dim=0) - res = pad_sequence(list(res), batch_first=True, padding_value=0) - else: - # applied attention mask on trtllm res - attention_mask_tmp = attention_mask.unsqueeze(-1) - res = res * attention_mask_tmp - if args.run_hf_test: - ref = hf_outputs.last_hidden_state - ref = ref * hf_inputs['attention_mask'].unsqueeze(-1) - - if args.debug: - intermediate_check(outputs, hf_outputs['hidden_states'], - attention_mask_tmp, logger) - - if world_size == 1: - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=1.5e-2, - atol=1.5e-2) - else: - # the arithmetic order of TP>1 is different from HF ref, which is always TP=1 for convenience. - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=4e-2, - atol=2e-2) - print(f"{model_name} result is all close to HF reference!") - - if model_name == 'BertForQuestionAnswering' or model_name == 'RobertaForQuestionAnswering': - if args.remove_input_padding: - - # [num_tokens, 2] -> [num_tokens, 1] - res_start_logits, res_end_logits = torch.split(res, 1, -1) - - # reshape result back to [batch_size, ...] - res_start_logits = torch.split(res_start_logits, - input_lengths.tolist(), - dim=0) - res_start_logits = tuple(t.squeeze() for t in res_start_logits) - res_end_logits = torch.split(res_end_logits, - input_lengths.tolist(), - dim=0) - res_end_logits = tuple(t.squeeze() for t in res_end_logits) - - else: - #NOTE: Padding - # [B, Padding_len, 2] -> [B, Padding_len, 1] - res_start_logits, res_end_logits = torch.split(res, 1, -1) - # [B, Padding_len, 1] -> [B, Padding_len] - res_start_logits = res_start_logits.squeeze() - res_end_logits = res_end_logits.squeeze() - res_start_logits = res_start_logits * attention_mask - res_end_logits = res_end_logits * attention_mask - - res_start_logits = torch.split(res_start_logits, 1, dim=0) - res_start_logits = tuple(t.squeeze(0) for t in res_start_logits) - res_end_logits = torch.split(res_end_logits, 1, dim=0) - res_end_logits = tuple(t.squeeze(0) for t in res_end_logits) - - - decode_res = decode_bertqa_output( inputs_text=test_inputs, \ - hf_tokenizer=hf_tokenizer, start_logits=res_start_logits, \ - end_logits=res_end_logits) - - if args.run_hf_test: - ref_start_logits = hf_outputs.start_logits - ref_end_logits = hf_outputs.end_logits - # when we use_plugin and have real-data model_dir and input - # We do not need to care about the output of padding positions: - ref_start_logits = ref_start_logits * hf_inputs['attention_mask'] - ref_end_logits = ref_end_logits * hf_inputs['attention_mask'] - - decode_ref = decode_bertqa_output( inputs_text=test_inputs, \ - hf_tokenizer=hf_tokenizer, start_logits=ref_start_logits, \ - end_logits=ref_end_logits) - compare_bertqa_result(inputs_text=test_inputs, - res_answers=decode_res, - ref_answers=decode_ref) - - elif model_name == 'BertForSequenceClassification' or model_name == 'RobertaForSequenceClassification': - hf_config = AutoConfig.from_pretrained(args.hf_model_dir) - decode_res = decode_bertcls_output(logits=res, - hf_model_config=hf_config, - inputs_text=test_inputs) - - if args.run_hf_test: - ref = hf_outputs.logits - if world_size == 1: - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=1.5e-2, - atol=1.5e-2) - else: - # the arithmetic order of TP>1 is different from HF ref, which is always TP=1 for convenience. - torch.testing.assert_close(actual=res.half(), - expected=ref, - rtol=4e-2, - atol=2e-2) - decode_ref = decode_bertcls_output(logits=hf_outputs.logits, - hf_model_config=hf_config, - inputs_text=test_inputs) - compare_bertcls_result(inputs_text=test_inputs, - res_answers=decode_res, - ref_answers=decode_res) diff --git a/examples/models/core/bert/utils.py b/examples/models/core/bert/utils.py deleted file mode 100644 index b3714e523e87..000000000000 --- a/examples/models/core/bert/utils.py +++ /dev/null @@ -1,217 +0,0 @@ -from contextlib import contextmanager -from typing import Dict, List, Tuple - -# isort: off -import torch -# isort: on - -import os -from pathlib import Path -from typing import Optional - -import datasets -from datasets import load_dataset - -from transformers import BertConfig, BertPreTrainedModel, BertForQuestionAnswering, BertForSequenceClassification, BertModel # isort:skip -from transformers import RobertaConfig, RobertaPreTrainedModel, RobertaForQuestionAnswering, RobertaForSequenceClassification, RobertaModel # isort:skip - - -# NOTE: This routine is copied from from tests/unittests/utils/llm_data.py -def llm_models_root(check=False) -> Optional[Path]: - root = Path("/home/scratch.trt_llm_data_ci/llm-models/") - - if "LLM_MODELS_ROOT" in os.environ: - root = Path(os.environ.get("LLM_MODELS_ROOT")) - - if not root.exists(): - root = Path("/scratch.trt_llm_data/llm-models/") - - if check: - assert root.exists(), \ - "You shall set LLM_MODELS_ROOT env or be able to access /home/scratch.trt_llm_data_ci to run this test" - - return root if root.exists() else None - - -def llm_datasets_root() -> Path: - return llm_models_root(check=True) / "datasets" - - -def prepare_text_inputs(model_name, batch_size=8): - print( - f"HF_DATASETS_OFFLINE inside function: {datasets.config.HF_DATASETS_OFFLINE}" - ) - if model_name == "BertForQuestionAnswering" or model_name == "RobertaForQuestionAnswering": - squad_dataset_root = str(llm_datasets_root( - )) + "/" if datasets.config.HF_DATASETS_OFFLINE else "" - squad_dataset_path = squad_dataset_root + "squad_v2" - squad_dataset = load_dataset(squad_dataset_path, trust_remote_code=True) - val_dataset = squad_dataset["validation"] - samples = val_dataset.select(range(batch_size)) - - qa_real_test_inputs = { - 'text': samples["question"], - 'text_pair': samples["context"] - } - return qa_real_test_inputs - elif model_name == "BertForSequenceClassification" or model_name == "RobertaForSequenceClassification": - yelp_dataset_root = str(llm_datasets_root( - )) + "/" if datasets.config.HF_DATASETS_OFFLINE else "fancyzhx/" - yelp_dataset_path = yelp_dataset_root + "yelp_polarity" - yelp_dataset = load_dataset(yelp_dataset_path, trust_remote_code=True) - val_dataset = yelp_dataset["test"] - samples = val_dataset.select(range(batch_size)) - - seqcls_real_test_inputs = {'text': samples['text']} - return seqcls_real_test_inputs - elif model_name == "BertModel" or model_name == "RobertaModel": - #NOTE: For BertModel, it is used as an encoder, so we use dummy input here, - # you can choose whatevert you like, but the numerical accuracy might vary. - test_input = 'To be or not to be: that is the question' - input_strings = [test_input for _ in range(batch_size)] - base_real_test_inputs = {'text': input_strings} - return base_real_test_inputs - - else: - raise NotImplementedError(f"Unknown model {model_name}") - - -def get_engine_name(rank): - return 'rank{}.engine'.format(rank) - - -def decode_bertqa_output(inputs_text, hf_tokenizer, - start_logits: Tuple[torch.Tensor], - end_logits: Tuple[torch.Tensor]): - question, context = inputs_text['text'], inputs_text['text_pair'] - assert len(context) == len(question) - batch_size = len(context) - - # regenerate inputs_ids because it is flatten for remove_input_padding=True - inputs = hf_tokenizer(**inputs_text, padding=True, return_tensors='pt') - inputs_ids = inputs['input_ids'] - answer_start_index = [logit.argmax(dim=0) for logit in start_logits] - answer_end_index = [logit.argmax(dim=0) for logit in end_logits] - decode_answer = [] - for i in range(batch_size): - predict_answer_tokens = inputs_ids[ - i, answer_start_index[i]:answer_end_index[i] + 1] - predict_text = hf_tokenizer.decode(predict_answer_tokens, - skip_special_tokens=True) - decode_answer.append(predict_text) - return decode_answer - - -def compare_bertqa_result(inputs_text, res_answers, ref_answers): - from difflib import SequenceMatcher - question, context = inputs_text['text'], inputs_text['text_pair'] - assert len(res_answers) == len(ref_answers) - batch_size = len(res_answers) - for i in range(batch_size): - print(f"Context: {context[i]}\nQuestion: {question[i]}") - print(f"Ref Answer: {ref_answers[i]}") - print(f"Res Answer: {res_answers[i]}") - match_rate = SequenceMatcher(None, "\n".join(res_answers[i]), - "\n".join(ref_answers[i])).ratio() - assert match_rate > 0.95 - print( - f"TRT-LLM results match HF results with literal match rate {match_rate}" - ) - - -def decode_bertcls_output(logits: torch.Tensor, hf_model_config, inputs_text): - text = inputs_text['text'] - id2label = hf_model_config.id2label - class_ids = logits.argmax(dim=1) - decode_answer = [] - batch_size = len(text) - for i in range(batch_size): - predicted_class_id = class_ids[i].item() - predicted_label = id2label[predicted_class_id] - decode_answer.append(predicted_label) - return decode_answer - - -def compare_bertcls_result(inputs_text, res_answers, ref_answers): - from difflib import SequenceMatcher - text = inputs_text['text'] - batch_size = len(text) - for i in range(batch_size): - print(f"Context: {text[i]}") - print(f"Ref Label: {ref_answers[i]}") - print(f"Res Label: {res_answers[i]}") - match_rate = SequenceMatcher(None, "\n".join(res_answers[i]), - "\n".join(ref_answers[i])).ratio() - assert match_rate > 0.95 - print( - f"TRT-LLM results match HF results with literal match rate {match_rate}" - ) - - -def process_input(input_ids_list: List[torch.Tensor], - token_type_ids_list: List[torch.Tensor], - is_roberta=False, - padding_idx=1): - input_lengths = [] - position_ids_list = [] - max_input_length = 0 - for i, input_ids in enumerate(input_ids_list): - input_len = len(input_ids) - assert input_len == len(token_type_ids_list[i]), f"sample {i}: len(input_ids)={len(input_ids)}, " \ - f"len(token_type_ids)={len(token_type_ids_list[i])}, not equal" - input_lengths.append(input_len) - position_ids = torch.arange(0, input_len, dtype=torch.int32) - if is_roberta: - position_ids = position_ids + 1 + padding_idx - - position_ids_list.append(position_ids) - max_input_length = max(max_input_length, input_len) - - # [num_tokens] - input_ids = torch.concat(input_ids_list).int().cuda() - token_type_ids = torch.concat(token_type_ids_list).int().cuda() - position_ids = torch.concat(position_ids_list).int().cuda() - - input_lengths = torch.tensor(input_lengths).int().cuda() # [batch_size] - max_input_length = torch.empty((max_input_length, )).int().cuda() - return input_ids, input_lengths, token_type_ids, position_ids, max_input_length - - -def intermediate_check(tllm_inter: Dict, hf_ref: Tuple[torch.Tensor], attn_mask, - logger): - - def apply_mask(x): - return x * attn_mask - - # minus one because there is an embedding output - num_layers = len(hf_ref) - 1 - - res = tllm_inter['embedding_output'] - res = apply_mask(res) - ref = hf_ref[0] - ref = apply_mask(ref) - torch.testing.assert_close(actual=res, expected=ref, rtol=1e-2, atol=1e-2) - logger.debug("Embedding are all close") - - for i in range(num_layers - 1): - res = tllm_inter[f'layer_{i}_output'] - res = apply_mask(res) - ref = hf_ref[i + 1] - ref = apply_mask(ref) - is_close = torch.allclose(res, ref, rtol=1e-2, atol=1e-2) - logger.debug(f'BertEncoderLayer_{i}_output is close: {is_close}') - - -@contextmanager -def temporary_datasets_config(**kwargs): - # Save original settings - original_settings = {} - for key, value in kwargs.items(): - original_settings[key] = getattr(datasets.config, key) - setattr(datasets.config, key, value) - try: - yield - finally: - # Restore original settings - for key, value in original_settings.items(): - setattr(datasets.config, key, value) diff --git a/examples/models/core/commandr/README.md b/examples/models/core/commandr/README.md deleted file mode 100644 index 12c978d15a0a..000000000000 --- a/examples/models/core/commandr/README.md +++ /dev/null @@ -1,213 +0,0 @@ -# Command R - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [C4AI Command-R](https://huggingface.co/CohereForAI/c4ai-command-r-v01), [C4AI Command R+](https://huggingface.co/CohereForAI/c4ai-command-r-plus), [Aya-23-8B](https://huggingface.co/CohereForAI/aya-23-8B), [Aya-23-35B](https://huggingface.co/CohereForAI/aya-23-35B) models using TensorRT LLM and run on a single GPU or a single node with multiple GPUs. - -- [Command R](#Command-R) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download repo and weights from HuggingFace Transformers](#1-download-repo-and-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [4. Run inference](#4-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multi GPU](#single-node-multi-gpu) - - [5. Run summarization task](#5-run-summarization-task) - - [Weight Only quantization](#weight-only-quantization) - - -## Overview - -The TensorRT LLM Command-R implementation can be found in [`tensorrt_llm/models/commandr/model.py`](../../../../tensorrt_llm/models/commandr/model.py). -The TensorRT LLM Command-R example code is located in [`examples/models/core/commandr`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - - * FP16 - * INT8 & INT4 Weight-Only - * Tensor Parallel - -## Usage - -The next section describe how to build the engine and run the inference demo. - -### 1. Download repo and weights from HuggingFace Transformers - -```bash -pip install -r requirements.txt -apt-get update -apt-get install git-lfs - -# clone one or more models we want to build -git clone https://huggingface.co/CohereForAI/c4ai-command-r-v01 command_r_v01 -git clone https://huggingface.co/CohereForAI/c4ai-command-r-plus command_r_plus -git clone https://huggingface.co/CohereForAI/aya-23-8B aya_23_8B -git clone https://huggingface.co/CohereForAI/aya-23-35B aya_23_35B -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# Command-R: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir command_r_v01 --output_dir trt_ckpt/command_r_v01/fp16/1-gpu - -# Command-R+: 4-way tensor parallelism -python3 convert_checkpoint.py --model_dir command_r_plus --tp_size 4 --output_dir trt_ckpt/command_r_plus/fp16/4-gpu - -# Aya-23-8B: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir aya_23_8B --output_dir trt_ckpt/aya_23_8B/fp16/1-gpu - -# Aya-23-35B: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir aya_23_35B --output_dir trt_ckpt/aya_23_35B/fp16/1-gpu -``` - -### 3. Build TensorRT engine(s) - -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is also same to the number of GPUs used to run inference. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# Command-R: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/command_r_v01/fp16/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/command_r_v01/fp16/1-gpu - -# Command-R+: 4-way tensor parallelism -trtllm-build --checkpoint_dir trt_ckpt/command_r_plus/fp16/4-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/command_r_plus/fp16/4-gpu - -# Command-R: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/aya_23_8B/fp16/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/aya_23_8B/fp16/1-gpu - -# Command-R: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/aya_23_35B/fp16/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/aya_23_35B/fp16/1-gpu -``` - -If the engines are built successfully, you will see output like (Command-R as the example): - -```txt -...... -[09/19/2024-03:34:30] [TRT] [I] Engine generation completed in 26.9495 seconds. -[09/19/2024-03:34:30] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 4 MiB, GPU 70725 MiB -[09/19/2024-03:34:55] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 176260 MiB -[09/19/2024-03:34:55] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:00:52 -[09/19/2024-03:34:55] [TRT] [I] Serialized 26 bytes of code generator cache. -[09/19/2024-03:34:55] [TRT] [I] Serialized 315007 bytes of compilation cache. -[09/19/2024-03:34:55] [TRT] [I] Serialized 12 timing cache entries -[09/19/2024-03:34:55] [TRT-LLM] [I] Timing cache serialized to model.cache -[09/19/2024-03:34:55] [TRT-LLM] [I] Build phase peak memory: 176257.29 MB, children: 17.65 MB -[09/19/2024-03:34:55] [TRT-LLM] [I] Serializing engine to trt_engines/command_r_v01/fp16/1-gpu/rank0.engine... -[09/19/2024-03:35:20] [TRT-LLM] [I] Engine serialized. Total time: 00:00:25 -[09/19/2024-03:35:23] [TRT-LLM] [I] Total time of building all engines: 00:01:47 -``` - -### 4. Run inference - -#### Single node, single GPU - -```bash -# Run the default engine of Command-R on single GPU. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/fp16/1-gpu - -# Run the default engine of Command-R on single GPU, using streaming output. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/fp16/1-gpu \ - --streaming - -# Run the default engine of Aya-23-8B on single GPU. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir aya_23_8B \ - --engine_dir trt_engines/aya_23_8B/fp16/1-gpu - -# Run the default engine of Aya-23-35B on single GPU. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir aya_23_35B \ - --engine_dir trt_engines/aya_23_35B/fp16/1-gpu -``` - -#### Single node, multi GPU - -```bash -# Run the Tensor Parallel 4 engine of Command-R+ on 4 GPUs. -mpirun -n 4 \ - python ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_plus \ - --engine_dir trt_engines/command_r_plus/fp16/4-gpu -``` - -If the engines are run successfully, you will see output like (Command-R as the example): - -```txt -...... -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef in Paris and worked in the kitchens of the French royal family. He came to England in 1814 and worked in a number of London hotels and restaurants, including the Reform Club and the London Tavern. He also opened his own restaurant" -``` - -### 5. Run summarization task - -```bash -# Run the summarization of Command-R task. -python3 ../../../summarize.py --test_trt_llm \ - --hf_model_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/fp16/1-gpu -``` - -If the engines are run successfully, you will see output like (Command-R as the example): - -```txt -...... -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM (total latency: 81.05689692497253 sec) -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2000) -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM (tokens per second: 24.67402621952367) -[01/26/2024-02:51:56] [TRT-LLM] [I] TensorRT LLM beam 0 result -[01/26/2024-02:51:56] [TRT-LLM] [I] rouge1 : 24.06804397902119 -[01/26/2024-02:51:56] [TRT-LLM] [I] rouge2 : 6.456513335555016 -[01/26/2024-02:51:56] [TRT-LLM] [I] rougeL : 16.77644999660741 -[01/26/2024-02:51:56] [TRT-LLM] [I] rougeLsum : 20.57359472317834 -``` - -### Weight Only quantization - -Use `--use_weight_only` to enable INT8-Weight-Only quantization, this will significantly lower the latency and memory footprint. Furthermore, use `--weight_only_precision int8` or `--weight_only_precision int4` to configure the data type of the weights. - -```bash -# Command-R: single gpu, int8 weight only quantization -python3 convert_checkpoint.py --model_dir command_r_v01 \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir trt_ckpt/command_r_v01/int8_wo/1-gpu - -# Command-R: single-gpu engine with int8 weight only quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/command_r_v01/int8_wo/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/command_r_v01/int8_wo/1-gpu - -# Run inference. -python3 ../../../run.py --max_output_len 50 \ - --tokenizer_dir command_r_v01 \ - --engine_dir trt_engines/command_r_v01/int8_wo/1-gpu -``` diff --git a/examples/models/core/commandr/convert_checkpoint.py b/examples/models/core/commandr/convert_checkpoint.py deleted file mode 100644 index 666803143e02..000000000000 --- a/examples/models/core/commandr/convert_checkpoint.py +++ /dev/null @@ -1,187 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import CohereForCausalLM -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--disable_weight_only_quant_plugin', - default=False, - action="store_true", - help= - 'By default, using plugin implementation for weight quantization. Enabling disable_weight_only_quant_plugin flag will use ootb implementation instead of plugin.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument("--load_model_on_cpu", action="store_true") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - - args = parser.parse_args() - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - quant_config.quant_algo = QuantAlgo.W4A16 - - return quant_config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'disable_weight_only_quant_plugin': - args.disable_weight_only_quant_plugin, - 'load_model_on_cpu': args.load_model_on_cpu, - } - - -def convert_and_save_hf(args): - model_dir = args.model_dir - world_size = args.tp_size * args.pp_size - # Need to convert the cli args to the kay-value pairs and override them in the generate config dict. - # Ideally these fields will be moved out of the config and pass them into build API, keep them here for compatibility purpose for now, - # before the refactor is done. - override_fields = {} - override_fields.update(args_to_build_options(args)) - - quant_config = args_to_quant_config(args) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - cohere = CohereForCausalLM.from_hugging_face( - model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - **override_fields, - ) - cohere.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del cohere - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - assert args.model_dir is not None - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/commandr/requirements.txt b/examples/models/core/commandr/requirements.txt deleted file mode 100644 index 2f8713d8658a..000000000000 --- a/examples/models/core/commandr/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score diff --git a/examples/models/core/enc_dec/README.md b/examples/models/core/enc_dec/README.md deleted file mode 100644 index 92ba03cd04c2..000000000000 --- a/examples/models/core/enc_dec/README.md +++ /dev/null @@ -1,444 +0,0 @@ -# Encoder-Decoder - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run an Encoder-Decoder (Enc-Dec) model in TensorRT LLM on NVIDIA GPUs. - -## Table of Contents - -- [Encoder-Decoder](#encoder-decoder) - - [Table of Contents](#table-of-contents) - - [Overview](#overview) - - [Usage](#usage) - - [Encoder-Decoder Model Support](#encoder-decoder-model-support) - - [Download weights from HuggingFace Transformers](#download-weights-from-huggingface-transformers) - - [Convert and Split Weights](#convert-and-split-weights) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [Run](#run) - - [Run C++ runtime](#run-c-runtime) - - [Run with Triton Backend](#run-with-triton-backend) - - [Run Python runtime](#run-python-runtime) - - [Benchmark](#benchmark) - - [Benchmark C++ runtime](#benchmark-c-runtime) - - [Run BART with LoRA](#run-bart-with-lora) - - [Reminders](#reminders) - - [Attention Scaling Factors](#attention-scaling-factors) - - [Run FairSeq NMT (Neural Machine Translation) models](#run-fairseq-nmt-neural-machine-translation-models) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Get quantized checkpoint with ModelOpt](#get-quantized-checkpoint-with-modelopt) - -## Overview - -The TensorRT LLM Enc-Dec implementation can be found in [tensorrt_llm/models/enc_dec/model.py](../../../../tensorrt_llm/models/enc_dec/model.py). The TensorRT LLM Enc-Dec example code is located in [`examples/models/core/enc_dec`](./): - - * `trtllm-build` to build the [TensorRT](https://developer.nvidia.com/tensorrt) engine(s) needed to run the Enc-Dec model, - * [`run.py`](./run.py) to run the inference on an example input text. - * Enc-Dec models can have specific implementations, such as the popular T5 family (T5, mT5, Flan-T5, ByT5), BART family (BART, mBART), and FairSeq family (WMTs). They are now merged into a single convert script: - * [`convert_checkpoint.py`](./convert_checkpoint.py) to convert weights from HuggingFace or FairSeq format to TRT-LLM format, and split weights for multi-GPU inference, -## Usage - -The TensorRT LLM Enc-Dec example code locates at [examples/models/core/enc_dec](./). It takes HuggingFace or FairSeq model name as input, and builds the corresponding TensorRT engines. On each GPU, there will be two TensorRT engines, one for Encoder and one for Decoder. - -## Encoder-Decoder Model Support - -The implementation is designed to support generic encoder-decoder models by abstracting the common and derivative components of different model architectures, such as: -- [T5](https://huggingface.co/docs/transformers/main/en/model_doc/t5) -- [T5v1.1](https://huggingface.co/docs/transformers/model_doc/t5v1.1) and [Flan-T5](https://huggingface.co/docs/transformers/model_doc/flan-t5) -- [mT5](https://huggingface.co/docs/transformers/model_doc/mt5) -- [BART](https://huggingface.co/docs/transformers/model_doc/bart) -- [mBART](https://huggingface.co/docs/transformers/model_doc/mbart) -- [FairSeq NMT](https://pytorch.org/hub/pytorch_fairseq_translation/) -- [ByT5](https://huggingface.co/docs/transformers/main/en/model_doc/byt5) -- [UL2 (coming)](https://huggingface.co/docs/transformers/model_doc/ul2) and [Flan-UL2 (coming)](https://huggingface.co/docs/transformers/model_doc/flan-ul2) - -It also supports full Tensor Parallelism (TP), Pipeline Parallelism (PP), and a hybrid of the two. Currently, Fused Multi-Head Attention (FMHA) is not yet enabled for T5 family due to its relative attention design. - -In this example, we use T5 (`t5-small`) and Flan-T5 (`google/flan-t5-small`) to showcase TRT-LLM support on Enc-Dec models. BART models and FairSeq models can follow very similar steps by just replacing the model name. - -### Download weights from HuggingFace Transformers - -```bash -git clone https://huggingface.co/t5-small tmp/hf_models/t5-small -git clone https://huggingface.co/google/flan-t5-small tmp/hf_models/flan-t5-small -git clone https://huggingface.co/facebook/bart-large-cnn tmp/hf_models/bart-large-cnn -git clone https://huggingface.co/facebook/mbart-large-50-many-to-one-mmt tmp/hf_models/mbart-large-50-many-to-one-mmt -git clone https://huggingface.co/google/byt5-small tmp/hf_models/byt5-small -``` - -### Convert and Split Weights -The `convert_checkpoint.py` script converts weights from HuggingFace or FairSeq format to TRT-LLM format, and splits weights for multi-GPU inference. `--tp_size` specifies the number of GPUs for tensor parallelism during inference. Pipeline Parallelism size can be set with `--pp_size` for distributed inference. - -The HuggingFace or Fairseq checkpoints of the enc-dec models mentioned in this Readme are all float32 precision. Use `--dtype` to set the target inference precision during the weight conversion. - -After weight conversion, TensorRT LLM converted weights and model configuration will be saved under `/` directory, which is the `--checkpoint_dir` input path you should give to the **next** engine building phase. - -Take T5 for example: - -```bash -# Example: build t5-small using 4-way tensor parallelism on a node with 8 GPUs (but only use 4 of them, for demonstration purpose), BF16, enabling beam search up to width=1. -export MODEL_NAME="t5-small" # or "flan-t5-small" -export MODEL_TYPE="t5" -export INFERENCE_PRECISION="bfloat16" -export TP_SIZE=4 -export PP_SIZE=1 -export WORLD_SIZE=4 -export MAX_BEAM_WIDTH=1 -python convert_checkpoint.py --model_type ${MODEL_TYPE} \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION} \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} \ - --dtype ${INFERENCE_PRECISION} -``` - -### Build TensorRT engine(s) - -TensorRT LLM builds TensorRT engine(s) with flexible controls on different types of optimizations. Note that these are just examples to demonstrate multi-GPU inference. For small models like T5-small, single GPU is usually sufficient. - -After engine building, TensorRT engines will be saved under `/` directory, which is the `--engine_dir` path you should give to the next engine running phase. It is recommended to have `/` in the output path where `Y` is number of total GPU ranks in a multi-node, multi-GPU setup, because the same `Y` number GPUs could be executed with different TP (Tensor Parallelism) and PP (Pipeline Parallelism) combinations. - -We should distinguish between `X` - TP size and `Y` - total number of GPU ranks: -* When `X = Y`, only TP is enabled -* When `X < Y`, both TP and PP are enabled. In such case, please make sure you have completed weight conversion step for `TP=X`. - -The default value of `--max_input_len` is 1024. When building DecoderModel, specify decoder input length with `--max_input_len=1` for encoder-decoder model to start generation from decoder_start_token_id of length 1. If the start token is a single token (the default behavior of T5/BART/etc.), you should set `--max_input_len` as 1; if you want the decoder-only type of generation, set `--max_input_len` above 1 to get similar behavior as HF's `decoder_forced_input_ids`. - -EncoderModel does not generate prompt. `--max_seq_len` should be the same as `--max_input_len`. `--max_seq_len` would be set as `--max_input_len` if not specified. - -DecoderModel takes `--max_encoder_input_len` and `--max_input_len` as model inputs, `--max_encoder_input_len` is set to 1024 as default since `--max_input_len` is 1024 for EncoderModel. - -To be noted: -1. For T5, add `--context_fmha disable`. FMHA with T5's relative attention bias is not implemented. Add `--use_implicit_relative_attention` when `--max_seq_len` is extremely large, causing decoder engine size to be too large to fit in memory. Compute relative attention on-the-fly (implicitly, without pre-computation) instead. -2. `--bert_attention_plugin`, `--gpt_attention_plugin`, `--remove_input_padding`, `--gemm_plugin` require explicit disabling and setting, or else they'll be set to default value in `trtllm-build`. - -```bash -# --gpt_attention_plugin is necessary in Enc-Dec. -# Try --gemm_plugin to prevent accuracy issue. -# It is recommended to use --remove_input_padding along with --gpt_attention_plugin for better performance -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable \ - --context_fmha disable - -# For decoder, refer to the above content and set --max_input_len correctly -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable \ - --context_fmha disable - -``` - -For BART, `--context_fmha` can be enabled. `trtllm-build` has the default setting to enable it. - -```bash -# Example: build bart-large-cnn using a single GPU, FP32, running greedy search -export MODEL_NAME="bart-large-cnn" # or "mbart-large-50-many-to-one-mmt" -export MODEL_TYPE="bart" -export INFERENCE_PRECISION="float32" -export TP_SIZE=1 -export PP_SIZE=1 -export WORLD_SIZE=1 -export MAX_BEAM_WIDTH=1 -python convert_checkpoint.py --model_type ${MODEL_TYPE} \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --output_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION} \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} \ - --dtype ${INFERENCE_PRECISION} - -# Note: non-T5 models can enable FMHA for the encoder part, for FP16/BF16, the default is enabled -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable - # --context_fmha disable should be removed - -# Use the same command for decoder engine -trtllm-build --checkpoint_dir tmp/trt_models/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width ${MAX_BEAM_WIDTH} \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding enable - # --context_fmha disable should be removed - -``` - -### Run - -Run a TensorRT LLM Enc-Dec model using the engines generated by build.py. -Note that during model deployment, only the TensorRT engine files are needed. Previously downloaded model checkpoints and converted weights can be removed. - -Different types of runtime are provided for encoder-decoder models. Following an order of serving performance and good usability, we recommend: -- (NEW) Python binding of C++ runtime w/ Paged KV Cache and Inflight Batching (IFB) -- Python runtime w/ Static Batching -- (NEW) C++ runtime w/ Paged KV Cache and Inflight Batching - -Please refer to the documentation for the details of [paged kv cache](../../../../docs/source/legacy/advanced/gpt-attention.md#paged-kv-cache) and [inflight batching](../../../../docs/source/legacy/advanced/gpt-attention.md#inflight-batching). - -#### Run C++ runtime -**Note: to use inflight batching and paged kv cache features in C++ runtime, please make sure you have set `--paged_kv_cache enable` (which is by default enabled) in the `trtllm-build` command of the decoder. Meanwhile, if using Python runtime, it is recommended to disable this flag by `--paged_kv_cache disable` to avoid any unnecessary overhead.** - -Note that for C++ runtime and Triton backend, Pipeline Parallelism (PP) is not supported yet, because PP usage is relatively rare for encoder-decoder models. If PP is really needed, it is recommended to use the Python runtime instead. - -For good usability, Python binding of the C++ runtime is provided. You can use the high-level C++ `ModelRunner` under the `examples/` root folder. - -```python -# Inferencing via python binding of C++ runtime with inflight batching (IFB) -python3 ../../../run.py --engine_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION} --tokenizer_dir tmp/hf_models/${MODEL_NAME} --max_output_len 64 --num_beams=1 --input_text "translate English to German: The house is wonderful." -``` - -You can specify `--kv_cache_free_gpu_memory_fraction` to control the percentage of free GPU memory to be used by KV cache (by default 0.9), and `--cross_kv_cache_fraction` to control the percentage of KV cache to be used by cross attention (by default 0.5, and rest of the KV cache will be used by self attention). - -For pure C++ runtime, there is no example given yet. Please check the [`Executor`](../../../../cpp/include/tensorrt_llm/executor/executor.h) API to implement your own end-to-end workflow. It is highly recommended to leverage more encapsulated solutions such as the above C++ Python binding or [Triton backend](https://github.com/triton-inference-server/tensorrtllm_backend). - -#### Run with Triton Backend -[Triton backend](https://github.com/triton-inference-server/tensorrtllm_backend/blob/main/docs/encoder_decoder.md) contains the tutorial on how to run encoder-decoder engines with Tritonserver. - -#### Run Python runtime - -For pure Python runtime, you can still use the encoder-decoder specific script under `examples/models/core/enc_dec/`. - -```bash -# Inferencing w/ single GPU greedy search, compare results with HuggingFace FP32 -python3 run.py --engine_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION} --engine_name ${MODEL_NAME} --model_name tmp/hf_models/${MODEL_NAME} --max_new_token=64 --num_beams=1 --compare_hf_fp32 - -# Inferencing w/ 4 GPUs (4-way TP, as configured during the engine building step), greedy search, compare results with HuggingFace FP32 -mpirun --allow-run-as-root -np ${WORLD_SIZE} python3 run.py --engine_dir tmp/trt_engines/${MODEL_NAME}/${INFERENCE_PRECISION} --engine_name ${MODEL_NAME} --model_name tmp/hf_models/${MODEL_NAME} --max_new_token=64 --num_beams=1 --compare_hf_fp32 -``` - -### Benchmark - -#### Benchmark C++ runtime - -The tutorial for encoder-decoder C++ runtime benchmark can be found in [`benchmarks/cpp`](../../benchmarks/cpp/README.md#2-launch-c-benchmarking-inflightv1-batching) - - -### Run BART with LoRA - -* Download the base model and lora model from HF: - -```bash -git clone https://huggingface.co/facebook/bart-large-cnn tmp/hf_models/bart-large-cnn -git clone https://huggingface.co/sooolee/bart-large-cnn-samsum-lora tmp/hf_models/bart-large-cnn-samsum-lora -``` - -If using customize models, just put both the base model and lora model dirs into `tmp/hf_models`. - -* Convert and Split Weights, setting `--hf_lora_dir`. - -```bash -export INFERENCE_PRECISION="float16" -python convert_checkpoint.py --model_type bart \ - --model_dir tmp/hf_models/bart-large-cnn \ - --output_dir tmp/trt_models/bart-large-cnn/${INFERENCE_PRECISION} \ - --tp_size 1 \ - --pp_size 1 \ - --dtype ${INFERENCE_PRECISION} -``` - -* Build engine, setting `--use_lora_plugin`. - -```bash - -trtllm-build --checkpoint_dir tmp/trt_models/bart-large-cnn/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable \ - --lora_plugin ${INFERENCE_PRECISION} \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ \ - --lora_target_modules attn_q attn_v - -trtllm-build --checkpoint_dir tmp/trt_models/bart-large-cnn/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --gemm_plugin ${INFERENCE_PRECISION} \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable \ - --lora_plugin ${INFERENCE_PRECISION} \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ \ - --lora_target_modules attn_q cross_attn_q attn_v cross_attn_v -``` - -* Run the engine, setting `--lora_dir` and `--lora_task_uids`. `--lora_task_uids` should be set as a list of uids which length equals to batch size. The following example is for batch size = 3: - -```bash -python run.py \ - --engine_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/ \ - --engine_name bart-large-cnn \ - --model_name tmp/hf_models/bart-large-cnn \ - --max_new_token=64 \ - --num_beams=1 \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ \ - --lora_task_uids 0 0 0 -``` - -* Run with multi-loRA, append `--lora_dir` with other lora directories and set `--lora_task_uids` according to the index of the lora directories. Set to "-1" to run with the base model: - -```bash -python run.py \ - --engine_dir tmp/trt_engines/bart-large-cnn/${INFERENCE_PRECISION}/ \ - --engine_name bart-large-cnn \ - --model_name tmp/hf_models/bart-large-cnn \ - --max_new_token=64 \ - --num_beams=1 \ - --lora_dir tmp/hf_models/bart-large-cnn-samsum-lora/ ... \ - --lora_task_uids 0 -1 -1 0 0 -1 -``` - -### Reminders - -- Flan-T5 models have known issues regarding FP16 precision and using BF16 precision is recommended, regardless of TRT-LLM. Please stay with FP32 or BF16 precision for Flan-T5 family. -- For T5 and Flan-T5 family that have relative attention bias design, the relative attention table is split along `num_heads` dimension in Tensor Parallelism mode. Therefore, `num_heads` must be divisible by `tp_size`. Please be aware of this when setting the TP parameter. -- For mBART, models that can control output languages (e.g. [`mbart-large-50-many-to-many-mmt`](https://huggingface.co/facebook/mbart-large-50-many-to-many-mmt)) are not currently supported, as the script does not support `ForcedBOSTokenLogitsProcessor` to control output languages. - -### Attention Scaling Factors - -The `q_scaling` convention in the TRT-LLM plugin is defined as follows: -``` -norm_factor = 1.f / (q_scaling * sqrt(head_size)) -``` -In the Multi-Head Attention (MHA) mechanism, the output of the `Q*K^T` product is scaled by this constant value `norm_factor` as `norm_factor * (Q*K^T)` for `softmax`. This scaling factor can be adjusted or neutralized based on the model's requirements. - -Handling in Different Models: -- BART/FairSeq NMT: For the BART models, `q_scaling` is set to `1.f`. Therefore, the `norm_factor` for BART becomes `1.f / sqrt(head_size)`. TRT-LLM uses the default value `q_scaling = 1.f`. Similar to FairSeq NMT models. -- T5: For the T5 models, `q_scaling` is `1.f/sqrt(head_size)`, leading to a `norm_factor` of `1.f`. This is handled in T5 by the TRT-LLM's `get_offset_q_scaling()` function, which reads `head_size` from the T5 model configuration and sets `q_scaling = 1.f/sqrt(head_size)` to effectively offset the `norm_factor` to `1.f`. - -### Run FairSeq NMT (Neural Machine Translation) models - -FairSeq model download and library dependency are different from HuggingFace ones. Especially if you are following the recommended docker container setup in [README](../../README.md), it has a custom PyTorch build but FairSeq installation will force upgrade the PyTorch version. As a workaround, we skip the `torch` and `torchaudio` dependencies in FairSeq to make everything work nicely inside the TRT-LLM container. - -```bash -# Download weights from HuggingFace Transformers -# Instructions from: https://github.com/facebookresearch/fairseq/blob/main/examples/translation/README.md#example-usage-cli-tools. Public model checkpoints are also listed there. Here we use WMT'14 Transformer model as an example. -mkdir -p tmp/fairseq_models && curl https://dl.fbaipublicfiles.com/fairseq/models/wmt14.en-fr.joined-dict.transformer.tar.bz2 | tar xvjf - -C tmp/fairseq_models --one-top-level=wmt14 --strip-components 1 --no-same-owner - -# Install FairSeq dependency -# avoid base torch to be upgraded by fairseq -pushd tmp && (git clone https://github.com/facebookresearch/fairseq.git || true) && pushd fairseq && sed -i '/torch>=/d;/torchaudio>=/d' setup.py && pip install -e . && pip install sacremoses subword_nmt && popd && popd - -# Convert and Split Weights, single GPU example -export TP_SIZE=1 -export PP_SIZE=1 -export WORLD_SIZE=1 -export INFERENCE_PRECISION="float32" -python convert_checkpoint.py --model_type nmt \ - --model_dir tmp/fairseq_models/wmt14 \ - --output_dir tmp/trt_models/wmt14/${INFERENCE_PRECISION} \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} \ - --dtype ${INFERENCE_PRECISION} - -# Build TensorRT engine(s) -# Note: non-T5 models can enable FMHA for the encoder part, although only FP16/BF16 precisions are valid -trtllm-build --checkpoint_dir tmp/trt_models/wmt14/${INFERENCE_PRECISION}/encoder \ - --output_dir tmp/trt_engines/wmt14/${INFERENCE_PRECISION}/encoder \ - --paged_kv_cache disable \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1024 \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable - -trtllm-build --checkpoint_dir tmp/trt_models/wmt14/${INFERENCE_PRECISION}/decoder \ - --output_dir tmp/trt_engines/wmt14/${INFERENCE_PRECISION}/decoder \ - --moe_plugin disable \ - --max_beam_width 1 \ - --max_batch_size 8 \ - --max_input_len 1 \ - --max_seq_len 201 \ - --max_encoder_input_len 1024 \ - --bert_attention_plugin ${INFERENCE_PRECISION} \ - --gpt_attention_plugin ${INFERENCE_PRECISION} \ - --remove_input_padding disable -# Run -mpirun --allow-run-as-root -np ${WORLD_SIZE} python3 run.py --engine_dir tmp/trt_engines/wmt14/${INFERENCE_PRECISION} --engine_name wmt14 --model_name tmp/fairseq_models/wmt14/${INFERENCE_PRECISION} --max_new_token=24 --num_beams=1 -``` - -### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit `nvidia-modelopt>=0.22.1` is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)). - -> [!NOTE] -> Modelopt 0.22.1 is not yet released. - -#### Get quantized checkpoint with ModelOpt -Currently supported conversion are `bart-large-cnn` and `T5` family. For `bart`, please set `--dtype float16`; for `T5` family, please set `--dtype float32` due to known bug with apex+HF mentioned in [transformer:issue/34264](https://github.com/huggingface/transformers/issues/34264). - -```bash -# Example: quantize bart-large-cnn using 4-way tensor parallelism on a node with 8 GPUs (but only use 4 of them, for demonstration purpose) into FP8 weight and convert to TRTLLM checkpoint. -export MODEL_NAME="bart-large-cnn" -export MODEL_TYPE="bart" -export INFERENCE_PRECISION="float16" -export TP_SIZE=4 -export PP_SIZE=1 -export WORLD_SIZE=4 -export MAX_BEAM_WIDTH=1 -python ../quantization/quantize.py \ - --model_dir tmp/hf_models/${MODEL_NAME} \ - --dtype ${INFERENCE_PRECISION} \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir tmp/trt_models/${MODEL_NAME}/fp8 \ - --calib_size 512 \ - --batch_size 16 \ - --tp_size ${TP_SIZE} \ - --pp_size ${PP_SIZE} -``` - -The rest may follow the same command in [Build TensorRT engine(s)](#build-tensorrt-engines), with some notes: -* For `bart`, please add `--use_fp8_context_fmha enable` for fp8 context fmha support. For `t5`, context fmha is not supported due to relative attention bias. -* Please ensure `--paged_kv_cache enable` for decoder for fp8 paged kv cache. -* Please use `--gemm_plugin auto`, `--bert_attention_plugin auto`, `--gpt_attention_plugin auto` instead of setting precision to these plugins. -* Please use CPP runtime for better performance. diff --git a/examples/models/core/enc_dec/__init__.py b/examples/models/core/enc_dec/__init__.py deleted file mode 100644 index e69de29bb2d1..000000000000 diff --git a/examples/models/core/enc_dec/convert_checkpoint.py b/examples/models/core/enc_dec/convert_checkpoint.py deleted file mode 100755 index 883a2fdebfff..000000000000 --- a/examples/models/core/enc_dec/convert_checkpoint.py +++ /dev/null @@ -1,1958 +0,0 @@ -import argparse -import configparser -import copy -import json -import logging -import os -import types -from ast import literal_eval -from collections import defaultdict -from datetime import datetime -from pathlib import Path - -import safetensors -from helper import (convert_weight_to_dtype, fairseq_sin_pos_embedding, - fuse_qkv_one_layer, reshape, split) -from transformers import (AutoModelForSeq2SeqLM, Blip2ForConditionalGeneration, - MBartForConditionalGeneration, NougatProcessor, - Pix2StructForConditionalGeneration, - T5ForConditionalGeneration, VisionEncoderDecoderModel) - -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import pad_vocab_size -from tensorrt_llm.functional import (LayerNormPositionType, LayerNormType, - MLPType) -from tensorrt_llm.layers import LanguageAdapterConfig -from tensorrt_llm.models import PretrainedConfig - -dir_path = os.path.dirname(os.path.realpath(__file__)) -LOGGER = logging.getLogger(__name__) - -layernorm_type_map = {i.name: i.value for i in LayerNormType} -layernorm_position_map = {i.name: i.value for i in LayerNormPositionType} -mlp_type_map = {i.name: i.value for i in MLPType} - -# Constants for specific model configurations -ECLAIR_RADIO_MAX_POSITION_EMBEDDINGS = 20000 - - -def copy_args_to_component_config(component_config, args): - for arg in vars(args): - setattr(component_config, arg, getattr(args, arg)) - return component_config - - -def parse_t5_config(args, hf_model): - config = configparser.ConfigParser() - - config["encoder"] = {} - for key, val in hf_model.encoder.config.to_dict().items(): - config["encoder"][key] = f"{val}" - - # manually set q_scaling to offset attention scaling's effect. - # TODO: modify kernels to control whether to disable attention scaling - def get_offset_q_scaling(config): - scaling = 1 / config.head_size**.5 - return scaling - - config["decoder"] = {} - for key, val in hf_model.decoder.config.to_dict().items(): - config["decoder"][key] = f"{val}" - - config["structure"] = dict() - config["structure"]["t5_with_bias"] = "false" - config["structure"]["use_gated_activation"] = str( - hf_model.encoder.config.is_gated_act) - config["structure"]["position_embedding_type"] = "relative" - config["structure"]["model_type"] = args.model_type - - def parse_t5_config_by_component(config, component, args): - component_config = types.SimpleNamespace() - component_config = copy_args_to_component_config(component_config, args) - component_config.n_head = config.getint(component, 'num_heads') - component_config.head_size = config.getint(component, 'd_kv') - component_config.hidden_size = config.getint(component, 'd_model') - component_config.ffn_hidden_size = config.getint(component, 'd_ff') - component_config.vocab_size = config.getint(component, 'vocab_size') - component_config.n_positions = config.getint(component, - 'n_positions', - fallback=512) - component_config.has_position_embedding = config.getboolean( - component, 'has_position_embedding', - fallback=False) # TODO: hardcoded here - - component_config.has_token_type_embedding = config.getboolean( - component, 'has_token_type_embedding', fallback=False) - component_config.has_embedding_layernorm = config.getboolean( - component, 'has_embedding_layernorm', fallback=False) - component_config.has_embedding_scale = config.getboolean( - component, 'has_embedding_scale', fallback=False) - component_config.q_scaling = get_offset_q_scaling(component_config) - component_config.has_attention_qkvo_bias = config.getboolean( - component, 'has_attention_qkvo_bias', - fallback=False) # TODO: hardcoded here - component_config.has_mlp_bias = config.getboolean(component, - 'has_mlp_bias', - fallback=False) - component_config.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm', fallback=True) - component_config.layernorm_eps = config.getfloat( - component, 'layer_norm_epsilon') - component_config.layernorm_position = layernorm_position_map[config.get( - component, 'layernorm_position', - fallback='pre_layernorm')] # TODO: hardcoded here - component_config.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type', fallback='RmsNorm')] - component_config.hidden_act = config.get(component, 'dense_act_fn') - component_config.gated_act = config.getboolean(component, - 'is_gated_act') - component_config.mlp_type = mlp_type_map['GatedMLP' if component_config. - gated_act else 'MLP'] - component_config.num_buckets = config.getint( - component, 'relative_attention_num_buckets') - component_config.max_distance = config.getint( - component, 'relative_attention_max_distance') - component_config.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - component_config.logits_dtype = config.get(component, - 'logits_dtype', - fallback='float32') - - if component == 'encoder': - component_config.n_layer = config.getint(component, 'num_layers') - - component_config.relative_attention = config.get( - 'structure', 'position_embedding_type') == 'relative' - - elif component == 'decoder': - component_config.n_layer = config.getint(component, - 'num_decoder_layers') - component_config.has_lm_head_bias = config.getboolean( - component, # TODO: T5 with bias - 'has_lm_head_bias', - fallback=False) - component_config.relative_attention = config.getboolean( - component, 'relative_attention', fallback=True) - component_config.rescale_before_lm_head = config.getboolean( - component, 'tie_word_embeddings' - ) # default is True (for T5), but False for Flan-T5 - component_config.encoder_hidden_size = config.getint( - 'encoder', 'd_model') - component_config.encoder_num_heads = config.getint( - 'encoder', 'num_heads') - component_config.encoder_head_size = config.getint( - 'encoder', 'd_kv') - component_config.decoder_start_token_id = config.getint( - 'decoder', 'decoder_start_token_id', fallback=0) - component_config.eos_token_id = config.getint('decoder', - 'eos_token_id', - fallback=1) - bos_token_id = config.get('decoder', 'bos_token_id', fallback=None) - # T5 does not have bos_token_id - component_config.bos_token_id = int( - bos_token_id) if bos_token_id not in (None, "None") else None - component_config.pad_token_id = config.getint('decoder', - 'pad_token_id', - fallback=0) - - else: - assert False, 'Unsupported component!' - - return component_config - - encoder_config = parse_t5_config_by_component(config, "encoder", args) - decoder_config = parse_t5_config_by_component(config, "decoder", args) - - return encoder_config, decoder_config - - -def convert_t5_weights_to_tllm_safetensors(config, component, params): - weights = {} - - mapping = config.mapping - - convert_weight_to_dtype(params, config.dtype) - hidden_size = config.hidden_size - ffn_hidden_size = config.intermediate_size - num_layers = config.num_hidden_layers - n_head = config.num_attention_heads - head_size = config.head_size - attention_hidden_size = n_head * head_size # head size * num_heads not necessarily equals hidden_dim, such as Flan-T5 - - hf_param_prefix = f'{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - hf_component_idx = 1 if component == 'encoder' else 2 - - def get_attn_module_name(component, block, layer, attn_type): - return f'{component}.block.{int(block)}.layer.{int(layer)}.{attn_type}' - - weights['transformer.vocab_embedding.weight'] = reshape( - params['shared.weight'].clone(), - None) if not config.use_parallel_embedding else reshape( - split(params['shared.weight'].clone(), mapping.tp_size, - mapping.tp_rank, 0), None) - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - hf_layer_name_prefix = f'{hf_param_prefix}.block.{layer_idx}' - - hidden_layer_name_split = { - f'{hf_layer_name_prefix}.layer.0.SelfAttention.o.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wo.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi_0.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - } - - hidden_layer_name_no_split = { - f'{hf_layer_name_prefix}.layer.0.layer_norm.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.layer_norm.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp_layernorm.weight', - "shape": None - }, - } - - if config.gated_act: - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi2.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - f'{hf_layer_name_prefix}.layer.{hf_component_idx}.DenseReluDense.wi_1.weight': - { - "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - }) - - if component == 'decoder': - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.layer.1.EncDecAttention.o.weight': { - "name": - f'{trtllm_layer_name_prefix}.cross_attention.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - f'{hf_layer_name_prefix}.layer.1.layer_norm.weight': { - "name": - f'{trtllm_layer_name_prefix}.cross_attention_layernorm.weight', - "shape": None - }, - }) - self_attn_module_name = get_attn_module_name( - component, layer_idx, "1", 'EncDecAttention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - self_attn_module_name = get_attn_module_name(component, layer_idx, "0", - 'SelfAttention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - weights[ - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.rel_attn_table'] = reshape( - split( - params[ - f'{component}.block.0.layer.0.SelfAttention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - split(params[hf_weight_name], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - params[hf_weight_name].clone(), shape=weight_info["shape"]) - - weights['transformer.ln_f.weight'] = reshape( - params[f'{component}.final_layer_norm.weight'].clone(), None) - - if component == 'decoder': - weights['lm_head.weight'] = reshape( - split(params['lm_head.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - if not config.use_implicit_relative_attention: - weights['rel_attn_table'] = reshape( - split( - params[ - f'{component}.block.0.layer.0.SelfAttention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - return weights - - -convert_blip2_weights_to_tllm_safetensors = convert_t5_weights_to_tllm_safetensors # func alias - - -def parse_nmt_config_by_component(config, component, args): - assert component in ('encoder', 'decoder'), 'Unsupported component!' - component_config = types.SimpleNamespace() - component_config = copy_args_to_component_config(component_config, args) - component_config.n_layer = config.getint(component, f'{component}_layers') - component_config.n_head = config.getint(component, - f'{component}_attention_heads') - component_config.hidden_size = config.getint( - component, f'{component}_embed_dim') # fairseq naming - component_config.head_size = config.getint( - component, - 'd_kv', - fallback=component_config.hidden_size // component_config.n_head) - component_config.ffn_hidden_size = config.getint( - component, f'{component}_ffn_embed_dim') # fairseq naming - component_config.vocab_size = config.getint(component, 'vocab_size') - component_config.n_positions = config.getint( - component, 'max_source_positions') # fairseq naming - component_config.has_position_embedding = not config.getboolean( - component, 'no_token_positional_embeddings', - fallback=False) # fairseq naming - component_config.has_token_type_embedding = config.getboolean( - component, 'has_token_type_embedding', fallback=False) - component_config.has_embedding_layernorm = config.getboolean( - component, 'layernorm_embedding', fallback=True) # fairseq naming - component_config.has_embedding_scale = not config.getboolean( - component, 'no_scale_embedding') # fairseq naming - component_config.q_scaling = config.getfloat(component, - 'q_scaling', - fallback=1.0) - component_config.has_attention_qkvo_bias = config.getboolean('structure', - 't5_with_bias', - fallback=True) - component_config.has_mlp_bias = config.getboolean('structure', - 't5_with_bias', - fallback=True) - component_config.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm') - component_config.layernorm_eps = config.getfloat( - component, 'layer_norm_epsilon', fallback=1e-5) # fairseq naming - - normalize_before = config.getboolean( - component, f'{component}_normalize_before') # fairseq naming - component_config.layernorm_position = layernorm_position_map[ - 'pre_layernorm' if normalize_before else 'post_layernorm'] - - component_config.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type', fallback='LayerNorm')] - component_config.hidden_act = config.get(component, - 'activation_fn') # fairseq naming - component_config.gated_act = config.getboolean(component, - 'is_gated_act', - fallback=False) - component_config.mlp_type = mlp_type_map['GatedMLP' if component_config. - gated_act else 'MLP'] - component_config.relative_attention = config.get( - 'structure', 'position_embedding_type') == 'relative' - - component_config.num_buckets = config.getint( - component, 'relative_attention_num_buckets', fallback=0) - component_config.max_distance = config.getint( - component, 'relative_attention_max_distance', fallback=0) - component_config.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - component_config.logits_dtype = config.get(component, - 'logits_dtype', - fallback='float32') - if component == 'decoder': - component_config.rescale_before_lm_head = config.getboolean( - component, 'rescale_before_lm_head') - - component_config.encoder_hidden_size = config.getint( - 'encoder', 'encoder_embed_dim') # fairseq naming - component_config.encoder_num_heads = config.getint( - 'encoder', 'encoder_attention_heads') - component_config.encoder_head_size = config.getint( - 'encoder', - 'd_kv', - fallback=component_config.encoder_hidden_size // - component_config.encoder_num_heads) - component_config.decoder_start_token_id = None - component_config.eos_token_id = None - component_config.bos_token_id = None - component_config.pad_token_id = None - - return component_config - - -def parse_nmt_config(args, model): - config = configparser.ConfigParser() - fairseq_config = vars(model.cfg.model) # Namespace --> dict - - config['encoder'] = dict() - for key, val in fairseq_config.items(): - config["encoder"][key] = f"{val}" - config["encoder"]["q_scaling"] = '1' - # NMT has final layernorm for pre-norm model architecture. - config['encoder']['has_model_final_layernorm'] = config['encoder'][ - 'encoder_normalize_before'] - config['encoder']['vocab_size'] = str(len(model.src_dict)) # fairseq naming - - config['decoder'] = dict() - for key, val in fairseq_config.items(): - config["decoder"][key] = f"{val}" - config["decoder"]["q_scaling"] = '1' - config["decoder"]["rescale_before_lm_head"] = 'false' - config['decoder']['has_model_final_layernorm'] = str( - config['decoder'].getboolean('decoder_normalize_before', False) - and not config['decoder'].getboolean('no_decoder_final_norm', False)) - config['decoder']['vocab_size'] = str(len(model.tgt_dict)) # fairseq naming - - config["structure"] = dict() - config["structure"]["t5_with_bias"] = "true" - config["structure"]["use_gated_activation"] = "false" - config["structure"][ - "position_embedding_type"] = "learned_absolute" # "sinusoid" - config["structure"]["model_type"] = args.model_type - - encoder_config = parse_nmt_config_by_component(config, "encoder", args) - decoder_config = parse_nmt_config_by_component(config, "decoder", args) - - return encoder_config, decoder_config - - -def convert_nmt_weights_to_tllm_safetensors(config, component, params, - sin_pos_embedding): - weights = {} - - mapping = config.mapping - - hidden_size = config.hidden_size - - convert_weight_to_dtype(params, config.dtype) - ffn_hidden_size = config.intermediate_size - vocab_size = config.vocab_size - - hf_param_prefix = f'models.0.{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - - hidden_layer_name_split = { - 'self_attn.out_proj.weight': { - "name": f'{trtllm_attn_layer_name}.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - 'fc1.weight': { - "name": 'mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - 'fc1.bias': { - "name": 'mlp.fc.bias', - "shape": (ffn_hidden_size // mapping.tp_size), - "split_dim": 0 - }, - 'fc2.weight': { - "name": 'mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - } - - hidden_layer_name_no_split = { - 'self_attn.out_proj.bias': { - "name": f'{trtllm_attn_layer_name}.dense.bias', - "shape": (hidden_size) - }, - 'self_attn_layer_norm.weight': { - "name": f'{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - 'self_attn_layer_norm.bias': { - "name": f'{trtllm_attn_layernorm_name}.bias', - "shape": None - }, - 'fc2.bias': { - "name": 'mlp.proj.bias', - "shape": (hidden_size) - }, - 'final_layer_norm.weight': { - "name": 'mlp_layernorm.weight', - "shape": None - }, - 'final_layer_norm.bias': { - "name": 'mlp_layernorm.bias', - "shape": None - }, - } - - if component == "decoder": - hidden_layer_name_split.update({ - 'encoder_attn.out_proj.weight': { - "name": 'cross_attention.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - 'encoder_attn.out_proj.bias': { - "name": 'cross_attention.dense.bias', - "shape": (hidden_size) - }, - 'encoder_attn_layer_norm.weight': { - "name": 'cross_attention_layernorm.weight', - "shape": None, - }, - 'encoder_attn_layer_norm.bias': { - "name": 'cross_attention_layernorm.bias', - "shape": None - }, - }) - - def get_attn_module_name(component, layer, attn_type): - return f'models.0.{component}.layers.{int(layer)}.{attn_type}' - - weights["transformer.vocab_embedding.weight"] = reshape( - params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - (vocab_size, -1)) if not config.use_parallel_embedding else reshape( - split(params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - mapping.tp_size, mapping.tp_rank, 0), - (vocab_size // mapping.tp_size, -1)) - weights["transformer.position_embedding.weight"] = reshape( - sin_pos_embedding, (config.max_position_embeddings, hidden_size)) - - num_layers = config.num_hidden_layers - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - hf_layer_name_prefix = f'{hf_param_prefix}.layers.{layer_idx}' - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - weights[ - f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape( - split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}' - hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}' - weights[trtllm_layer_fullname] = reshape( - params[hf_layer_fullname].clone(), shape=weight_info["shape"]) - - self_attn_module_name = get_attn_module_name(component, layer_idx, - 'self_attn') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - if component == 'decoder': - cross_attn_module_name = get_attn_module_name( - component, layer_idx, 'encoder_attn') - weights.update( - fuse_qkv_one_layer( - params, cross_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - - if component == 'decoder': - weights['lm_head.weight'] = reshape( - split(params[f'{hf_param_prefix}.output_projection.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - - if config.has_model_final_layernorm: - weights['transformer.ln_f.weight'] = params[ - f'{hf_param_prefix}.layer_norm.weight'].clone() - weights['transformer.ln_f.bias'] = params[ - f'{hf_param_prefix}.layer_norm.bias'].clone() - - return weights - - -def parse_bart_config(args, hf_model): - - config = configparser.ConfigParser() - - config['decoder'] = dict() - if args.eclair_radio: - for key, val in hf_model.config.to_dict().items(): - config["decoder"][key] = f"{val}" - else: - for key, val in hf_model.model.decoder.config.to_dict().items(): - config["decoder"][key] = f"{val}" - config["decoder"]["q_scaling"] = '1' - config["decoder"]["rescale_before_lm_head"] = str(False) - config['decoder']['has_model_final_layernorm'] = str( - args.nougat or args.eclair_radio - or isinstance(hf_model, MBartForConditionalGeneration)) - - if args.nougat or args.eclair_radio: - # These flags are true for mbart decoders, but missing in HF config - config['decoder']['normalize_before'] = str(True) - config['decoder']['normalize_embeddings'] = str(True) - - config['encoder'] = dict() - # Init few encoder configs, needed by build, from decoder config - encoder_config_keys = [ - "encoder_ffn_dim", "encoder_layers", "encoder_attention_heads", - "encoder_layerdrop", "d_model" - ] - for key in encoder_config_keys: - config['encoder'][key] = config['decoder'][key] - else: - config['encoder'] = dict() - for key, val in hf_model.model.encoder.config.to_dict().items(): - config["encoder"][key] = f"{val}" - config["encoder"]["q_scaling"] = '1' - - # mBART has final layernorm, BART does not - config['encoder']['has_model_final_layernorm'] = str( - isinstance(hf_model, MBartForConditionalGeneration)) - - config["structure"] = dict() - config["structure"]["t5_with_bias"] = "true" - config["structure"]["use_gated_activation"] = "false" - config["structure"]["position_embedding_type"] = "learned_absolute" - config["structure"]["model_type"] = args.model_type - - def parse_bart_config_by_component(config, component, args): - assert component in ('encoder', 'decoder'), 'Unsupported component!' - component_config = types.SimpleNamespace() - component_config = copy_args_to_component_config(component_config, args) - component_config.n_layer = config.getint(component, - f'{component}_layers') - component_config.n_head = config.getint(component, - f'{component}_attention_heads') - component_config.hidden_size = config.getint(component, 'd_model') - component_config.head_size = config.getint( - component, - 'd_kv', - fallback=component_config.hidden_size // component_config.n_head) - component_config.ffn_hidden_size = config.getint( - component, f'{component}_ffn_dim') - component_config.vocab_size = config.getint(component, 'vocab_size') - component_config.n_positions = config.getint(component, - 'max_position_embeddings') - component_config.has_position_embedding = config.getboolean( - component, 'has_position_embedding', - fallback=True) # TODO: hardcoded here - component_config.has_token_type_embedding = config.getboolean( - component, 'has_token_type_embedding', fallback=False) - component_config.has_embedding_layernorm = config.getboolean( - component, 'has_embedding_layernorm', fallback=True) - component_config.has_embedding_scale = config.getboolean( - component, 'scale_embedding') - component_config.q_scaling = config.getfloat(component, - 'q_scaling', - fallback=1.0) - component_config.has_attention_qkvo_bias = config.getboolean( - 'structure', 't5_with_bias', fallback=True) - component_config.has_mlp_bias = config.getboolean('structure', - 't5_with_bias', - fallback=True) - component_config.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm') - component_config.layernorm_eps = config.getfloat(component, - 'layer_norm_epsilon', - fallback=False) - - normalize_before = config.getboolean(component, 'normalize_before') - component_config.layernorm_position = layernorm_position_map[ - 'pre_layernorm' if normalize_before else 'post_layernorm'] - - component_config.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type', fallback='LayerNorm')] - component_config.hidden_act = config.get(component, - 'activation_function') - component_config.gated_act = config.getboolean(component, - 'is_gated_act', - fallback=False) - component_config.mlp_type = mlp_type_map['GatedMLP' if component_config. - gated_act else 'MLP'] - component_config.relative_attention = config.get( - 'structure', 'position_embedding_type') == 'relative' - - component_config.num_buckets = config.getint( - component, 'relative_attention_num_buckets', fallback=0) - component_config.max_distance = config.getint( - component, 'relative_attention_max_distance', fallback=0) - component_config.max_lora_rank = config.getint(component, - 'max_lora_rank', - fallback=0) - component_config.lora_target_modules = literal_eval( - config.get(component, 'lora_target_modules', fallback="[]")) - component_config.hf_modules_to_trtllm_modules = literal_eval( - config.get(component, 'hf_modules_to_trtllm_modules', - fallback="{}")) - component_config.trtllm_modules_to_hf_modules = literal_eval( - config.get(component, 'trtllm_modules_to_hf_modules', - fallback="{}")) - component_config.logits_dtype = config.get(component, - 'logits_dtype', - fallback='float32') - component_config.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - - if component == 'decoder': - component_config.rescale_before_lm_head = config.getboolean( - component, 'rescale_before_lm_head') - - component_config.encoder_hidden_size = config.getint( - 'encoder', 'd_model') - component_config.encoder_num_heads = config.getint( - 'encoder', 'encoder_attention_heads') - component_config.encoder_head_size = config.getint( - 'encoder', - 'd_kv', - fallback=component_config.encoder_hidden_size // - component_config.encoder_num_heads) - - # nougat has decoder_start_token_id = None, special handling - decoder_start_token_id = config.get('decoder', - 'decoder_start_token_id') - component_config.decoder_start_token_id = int( - decoder_start_token_id - ) if decoder_start_token_id != "None" else None - component_config.eos_token_id = config.getint( - 'decoder', 'eos_token_id') - component_config.bos_token_id = config.getint( - 'decoder', 'bos_token_id') - component_config.pad_token_id = config.getint( - 'decoder', 'pad_token_id') - - return component_config - - encoder_config = None - if not (args.nougat or args.eclair_radio): - encoder_config = parse_bart_config_by_component(config, "encoder", args) - decoder_config = parse_bart_config_by_component(config, "decoder", args) - - # Override n_positions for eclair_radio model - if args.eclair_radio: - decoder_config.n_positions = ECLAIR_RADIO_MAX_POSITION_EMBEDDINGS - - return encoder_config, decoder_config - - -def convert_bart_weights_to_tllm_safetensors(config, component, params): - weights = {} - - mapping = config.mapping - - hidden_size = config.hidden_size - - convert_weight_to_dtype(params, config.dtype) - ffn_hidden_size = config.intermediate_size - vocab_size = config.vocab_size - - hf_param_prefix = f'model.{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - embedding_layer_names = { - 'embed_tokens.weight': { - "name": 'transformer.vocab_embedding.weight', - "shape": (vocab_size, -1) - }, - 'embed_positions.weight': { - "name": 'transformer.position_embedding.weight', - "shape": (config.max_position_embeddings, hidden_size) - }, - 'layernorm_embedding.weight': { - "name": 'transformer.ln_embed.weight', - "shape": None - }, - 'layernorm_embedding.bias': { - "name": 'transformer.ln_embed.bias', - "shape": None - }, - } - - hidden_layer_name_split = { - 'self_attn.out_proj.weight': { - "name": f'{trtllm_attn_layer_name}.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - 'fc1.weight': { - "name": 'mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - 'fc1.bias': { - "name": 'mlp.fc.bias', - "shape": (ffn_hidden_size // mapping.tp_size), - "split_dim": 0 - }, - 'fc2.weight': { - "name": 'mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - } - - hidden_layer_name_no_split = { - 'self_attn.out_proj.bias': { - "name": f'{trtllm_attn_layer_name}.dense.bias', - "shape": (hidden_size) - }, - 'self_attn_layer_norm.weight': { - "name": f'{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - 'self_attn_layer_norm.bias': { - "name": f'{trtllm_attn_layernorm_name}.bias', - "shape": None - }, - 'fc2.bias': { - "name": 'mlp.proj.bias', - "shape": (hidden_size) - }, - 'final_layer_norm.weight': { - "name": 'mlp_layernorm.weight', - "shape": None - }, - 'final_layer_norm.bias': { - "name": 'mlp_layernorm.bias', - "shape": None - }, - } - - if config.model_type == 'mbart': - hidden_layer_name_split['layer_norm.weight'] = { - "name": 'transformer.ln_f.weight', - "shape": None, - "split_dim": 0 - } - hidden_layer_name_no_split['layer_norm.bias'] = { - "name": 'transformer.ln_f.bias', - "shape": None, - "split_dim": 0 - } - - if component == "decoder": - hidden_layer_name_split.update({ - 'encoder_attn.out_proj.weight': { - "name": 'cross_attention.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - } - }) - hidden_layer_name_no_split.update({ - 'encoder_attn.out_proj.bias': { - "name": 'cross_attention.dense.bias', - "shape": (hidden_size) - }, - 'encoder_attn_layer_norm.weight': { - "name": 'cross_attention_layernorm.weight', - "shape": None - }, - 'encoder_attn_layer_norm.bias': { - "name": 'cross_attention_layernorm.bias', - "shape": None - }, - }) - - def get_attn_module_name(component, layer, attn_type): - return f'model.{component}.layers.{int(layer)}.{attn_type}' - - for hf_weight_name, weight_info in embedding_layer_names.items(): - if 'position' in hf_weight_name: - weights[weight_info["name"]] = params[ - f'{hf_param_prefix}.{hf_weight_name}'][2:].clone() - else: - weights[weight_info["name"]] = params[ - f'{hf_param_prefix}.{hf_weight_name}'].clone() - weights[weight_info["name"]] = reshape(weights[weight_info["name"]], - weight_info["shape"]) - - weights["embedding.vocab_embedding.weight"] = reshape( - params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - (vocab_size, -1)) if not config.use_parallel_embedding else reshape( - split(params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - mapping.tp_size, mapping.tp_rank, 0), - (vocab_size // mapping.tp_size, -1)) - - num_layers = config.num_hidden_layers - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - hf_layer_name_prefix = f'{hf_param_prefix}.layers.{layer_idx}' - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - weights[ - f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape( - split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}' - hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}' - weights[trtllm_layer_fullname] = reshape( - params[hf_layer_fullname].clone(), shape=weight_info["shape"]) - - self_attn_module_name = get_attn_module_name(component, layer_idx, - 'self_attn') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - if component == 'decoder': - cross_attn_module_name = get_attn_module_name( - component, layer_idx, 'encoder_attn') - weights.update( - fuse_qkv_one_layer( - params, cross_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - - if component == 'decoder': - import torch - lm_head_weights = params['lm_head.weight'].clone().detach() - vocab_size = config.vocab_size - if params['lm_head.weight'].shape[0] % mapping.tp_size != 0: - vocab_size_padded = pad_vocab_size(config.vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - config.vocab_size - - lm_head_weights = torch.nn.functional.pad(lm_head_weights, - (0, 0, 0, pad_width), - 'constant', - value=0) - vocab_size = vocab_size_padded - weights['lm_head.weight'] = reshape( - split(lm_head_weights, mapping.tp_size, mapping.tp_rank, dim=0), - (vocab_size // mapping.tp_size, hidden_size)) - - if config.has_model_final_layernorm: - weights['transformer.ln_f.weight'] = params[ - f'{hf_param_prefix}.layer_norm.weight'].clone() - weights['transformer.ln_f.bias'] = params[ - f'{hf_param_prefix}.layer_norm.bias'].clone() - - return weights - - -def parse_pix2struct_config(args, hf_model): - # manually set q_scaling to offset attention scaling's effect. - # TODO: modify kernels to control whether to disable attention scaling - config = configparser.ConfigParser() - - def get_offset_q_scaling(config) -> str: - d_model = config.hidden_size - num_heads = config.num_heads - head_size = d_model / num_heads - scaling = 1 / head_size**.5 - return str(scaling) - - config["decoder"] = {} - for key, val in hf_model.decoder.config.to_dict().items(): - config["decoder"][key] = f"{val}" - - config["decoder"]["q_scaling"] = get_offset_q_scaling( - hf_model.decoder.config) - - config["structure"] = dict() - config["structure"]["pix2struct_with_bias"] = "false" - config["structure"]["use_gated_activation"] = "false" - config["structure"]["position_embedding_type"] = "relative" - config["structure"]["model_type"] = args.model_type - - def parse_pix2struct_config_by_component(config, component, args): - if component == 'decoder': - args.n_layer = config.getint(component, 'num_layers') - args.n_head = config.getint(component, 'num_heads') - args.head_size = config.getint(component, 'd_kv') - args.hidden_size = config.getint(component, 'hidden_size') - args.ffn_hidden_size = config.getint(component, 'd_ff') - args.vocab_size = config.getint(component, 'vocab_size') - args.n_positions = config.getint(component, - 'n_positions', - fallback=512) - args.has_position_embedding = config.getboolean( - component, 'has_position_embedding', - fallback=False) # TODO: hardcoded here - args.has_token_type_embedding = config.getboolean( - component, 'has_token_type_embedding', fallback=False) - args.has_embedding_layernorm = config.getboolean( - component, 'has_embedding_layernorm', fallback=False) - args.has_embedding_scale = config.getboolean(component, - 'has_embedding_scale', - fallback=False) - args.q_scaling = config.getfloat(component, - 'q_scaling', - fallback=1.0) - args.has_attention_qkvo_bias = config.getboolean( - component, 'has_attention_qkvo_bias', fallback=False) - args.has_mlp_bias = config.getboolean(component, - 'has_mlp_bias', - fallback=False) - args.has_model_final_layernorm = config.getboolean( - component, 'has_model_final_layernorm', fallback=True) - args.layernorm_eps = config.getfloat(component, - 'layer_norm_epsilon') - args.layernorm_position = layernorm_position_map[config.get( - component, 'layernorm_position', - fallback='pre_layernorm')] # TODO: hardcoded here - args.layernorm_type = layernorm_type_map[config.get( - component, 'layernorm_type', fallback='RmsNorm')] - args.hidden_act = config.get(component, 'dense_act_fn') - args.gated_act = True - args.mlp_type = mlp_type_map['GatedMLP' if args. - gated_act else 'MLP'] - args.has_lm_head_bias = config.getboolean( - component, # TODO: T5 with bias - 'has_lm_head_bias', - fallback=False) - args.relative_attention = config.getboolean(component, - 'relative_attention', - fallback=True) - args.num_buckets = config.getint(component, - 'relative_attention_num_buckets') - args.max_distance = config.getint( - component, 'relative_attention_max_distance') - args.logits_dtype = config.get(component, - 'logits_dtype', - fallback='float32') - args.rescale_before_lm_head = config.getboolean( - component, 'tie_word_embeddings' - ) # default is True (for T5), but False for Flan-T5 - args.encoder_hidden_size = config.getint('decoder', 'hidden_size') - args.encoder_num_heads = config.getint('decoder', 'num_heads') - args.encoder_head_size = config.getint('decoder', 'd_kv') - args.position_embedding_type = config.get( - 'structure', 'position_embedding_type') - args.decoder_start_token_id = config.getint( - 'decoder', 'decoder_start_token_id') - args.eos_token_id = config.getint('decoder', 'eos_token_id') - bos_token_id = config.get('decoder', 'bos_token_id') - # pix2struct does not have bos_token_id - args.bos_token_id = int( - bos_token_id) if bos_token_id != "None" else None - args.pad_token_id = config.getint('decoder', 'pad_token_id') - - else: - assert False, 'Unsupported component!' - return args - - decoder_args = parse_pix2struct_config_by_component(config, "decoder", args) - return None, decoder_args - - -def convert_pix2struct_weights_to_tllm_safetensors(config, component, params): - weights = {} - - mapping = config.mapping - - convert_weight_to_dtype(params, config.dtype) - hidden_size = config.hidden_size - ffn_hidden_size = config.intermediate_size - num_layers = config.num_hidden_layers - n_head = config.num_attention_heads - head_size = config.head_size - attention_hidden_size = n_head * head_size # head size * num_heads not necessarily equals hidden_dim, such as Flan-T5 - - hf_param_prefix = f'{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' - - def get_attn_module_name(component, layer, attn_type): - return f'{component}.layer.{int(layer)}.{attn_type}.attention' - - weights['transformer.vocab_embedding.weight'] = reshape( - params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - None) if not config.use_parallel_embedding else reshape( - split(params[f'{hf_param_prefix}.embed_tokens.weight'].clone(), - mapping.tp_size, mapping.tp_rank, 0), None) - - layers_range = mapping.pp_layers(num_layers) - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - hf_layer_name_prefix = f'{hf_param_prefix}.layer.{layer_idx}' - - hidden_layer_name_split = { - f'{hf_layer_name_prefix}.self_attention.attention.output.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.mlp.DenseReluDense.wo.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{hf_layer_name_prefix}.mlp.DenseReluDense.wi_0.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - } - - hidden_layer_name_no_split = { - f'{hf_layer_name_prefix}.self_attention.layer_norm.weight': { - "name": - f'{trtllm_layer_name_prefix}.{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - f'{hf_layer_name_prefix}.mlp.layer_norm.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp_layernorm.weight', - "shape": None - }, - } - - if config.gated_act: - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.mlp.DenseReluDense.wi_1.weight': { - "name": f'{trtllm_layer_name_prefix}.mlp.gate.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - }) - - hidden_layer_name_split.update({ - f'{hf_layer_name_prefix}.encoder_decoder_attention.attention.output.weight': - { - "name": - f'{trtllm_layer_name_prefix}.cross_attention.dense.weight', - "shape": - (hidden_size, attention_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - f'{hf_layer_name_prefix}.encoder_decoder_attention.layer_norm.weight': - { - "name": - f'{trtllm_layer_name_prefix}.cross_attention_layernorm.weight', - "shape": None - }, - }) - self_attn_module_name = get_attn_module_name( - component, layer_idx, 'encoder_decoder_attention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', mapping.tp_size, - mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - self_attn_module_name = get_attn_module_name(component, layer_idx, - 'self_attention') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (attention_hidden_size * 3 // mapping.tp_size, hidden_size), - None)) - - weights[ - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}.rel_attn_table'] = reshape( - split( - params[ - f'{component}.layer.0.self_attention.attention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - split(params[hf_weight_name], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - if hf_weight_name in params.keys(): - weights[weight_info["name"]] = reshape( - params[hf_weight_name].clone(), shape=weight_info["shape"]) - - weights[f'transformer.ln_f.weight'] = reshape( - params[f'{component}.final_layer_norm.weight'].clone(), None) - - weights['lm_head.weight'] = reshape( - split(params[f'{component}.lm_head.weight'], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - if not config.use_implicit_relative_attention: - weights[f'rel_attn_table'] = reshape( - split( - params[ - f'{component}.layer.0.self_attention.attention.relative_attention_bias.weight'] - .T, mapping.tp_size, mapping.tp_rank, 0), - (n_head // mapping.tp_size, config.num_buckets)) - - return weights - - -def parse_language_adapter_config(args, model): - config = configparser.ConfigParser() - config.read(args.model_dir + "/config.ini") - # rename from "sinusoid" to "learned_absolute" - config["structure"]["position_embedding_type"] = "learned_absolute" - - locale_list = ['de_DE', 'en_US', 'es_ES', 'fr_FR'] # to be changed by user - - encoder_config = parse_nmt_config_by_component(config, "encoder", args) - encoder_config.residual_scaling = config.getfloat("encoder", - 'residual_scaling', - fallback=1.0) - encoder_config.language_adapter_config = LanguageAdapterConfig( - num_languages=config.getint("encoder", 'adapter_langs', fallback=None), - ffn_hidden_size=config.getint("encoder", - 'encoder_adapter_embed_dim', - fallback=None), - language_list=locale_list).to_dict() - - decoder_config = parse_nmt_config_by_component(config, "decoder", args) - decoder_config.residual_scaling = config.getfloat("decoder", - 'residual_scaling', - fallback=1.0) - decoder_config.language_adapter_config = LanguageAdapterConfig( - num_languages=config.getint("decoder", 'adapter_langs', fallback=None), - ffn_hidden_size=config.getint("decoder", - 'decoder_adapter_embed_dim', - fallback=None), - language_list=locale_list).to_dict() - - decoder_config.decoder_start_token_id = 2 - decoder_config.eos_token_id = 2 - decoder_config.bos_token_id = 2 - decoder_config.pad_token_id = 0 - - return encoder_config, decoder_config - - -def convert_language_adapter_weights_to_tllm_safetensors( - config, component, params): - weights = {} - - mapping = config.mapping - - convert_weight_to_dtype(params, config.dtype) - - param_prefix = f'{component}' - trtllm_layer_name = f'transformer.layers' - trtllm_attn_layer_name = 'attention' if component == 'encoder' else 'self_attention' - trtllm_attn_layernorm_name = 'self_attention_layernorm' if component == 'decoder' else 'attention_layernorm' - mlp_param_prefix = '' if f'{param_prefix}.0.fc1.weight' in params else 'mlp.' - - hidden_size = params[ - f'{param_prefix}.layers.0.self_attn.out_proj.weight'].shape[0] - ffn_hidden_size = params[ - f'{param_prefix}.layers.0.{mlp_param_prefix}fc1.weight'].shape[0] - - hidden_layer_name_split = { - 'self_attn.out_proj.weight': { - "name": f'{trtllm_attn_layer_name}.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - f'{mlp_param_prefix}fc1.weight': { - "name": 'mlp.fc.weight', - "shape": (ffn_hidden_size // mapping.tp_size, hidden_size), - "split_dim": 0 - }, - f'{mlp_param_prefix}fc1.bias': { - "name": 'mlp.fc.bias', - "shape": (ffn_hidden_size // mapping.tp_size), - "split_dim": 0 - }, - f'{mlp_param_prefix}fc2.weight': { - "name": 'mlp.proj.weight', - "shape": (hidden_size, ffn_hidden_size // mapping.tp_size), - "split_dim": -1 - }, - } - - hidden_layer_name_no_split = { - 'self_attn.out_proj.bias': { - "name": f'{trtllm_attn_layer_name}.dense.bias', - "shape": (hidden_size) - }, - 'self_attn_layer_norm.weight': { - "name": f'{trtllm_attn_layernorm_name}.weight', - "shape": None - }, - 'self_attn_layer_norm.bias': { - "name": f'{trtllm_attn_layernorm_name}.bias', - "shape": None - }, - f'{mlp_param_prefix}fc2.bias': { - "name": 'mlp.proj.bias', - "shape": (hidden_size) - }, - 'final_layer_norm.weight': { - "name": 'mlp_layernorm.weight', - "shape": None - }, - 'final_layer_norm.bias': { - "name": 'mlp_layernorm.bias', - "shape": None - }, - } - - if component == "decoder": - hidden_layer_name_split.update({ - 'encoder_attn.out_proj.weight': { - "name": 'cross_attention.dense.weight', - "shape": (hidden_size, hidden_size // mapping.tp_size), - "split_dim": -1 - }, - }) - hidden_layer_name_no_split.update({ - 'encoder_attn.out_proj.bias': { - "name": 'cross_attention.dense.bias', - "shape": (hidden_size) - }, - 'encoder_attn_layer_norm.weight': { - "name": 'cross_attention_layernorm.weight', - "shape": None, - }, - 'encoder_attn_layer_norm.bias': { - "name": 'cross_attention_layernorm.bias', - "shape": None - }, - }) - - def get_attn_module_name(layer, attn_type): - return f'{param_prefix}.layers.{int(layer)}.{attn_type}' - - # support MostlyFreezedEmbedding in 5.5B model - embed_tokens_weight_name = f'{param_prefix}.embed_tokens.weight' - if embed_tokens_weight_name not in params: - embed_tokens_weight_name = f'{param_prefix}.embed_tokens.weight_' - - weights['transformer.vocab_embedding.weight'] = reshape( - params[embed_tokens_weight_name].clone(), - None) if not config.use_parallel_embedding else reshape( - split(params[embed_tokens_weight_name].clone(), mapping.tp_size, - mapping.tp_rank, 0), None) - - weights[ - 'transformer.position_embedding.weight'] = fairseq_sin_pos_embedding( - config.max_position_embeddings, - params[embed_tokens_weight_name].shape[1]) - - num_layers = config.num_hidden_layers - layers_range = mapping.pp_layers(num_layers) - - for layer_idx in layers_range: - local_layer_idx = layer_idx - layers_range[0] - hf_layer_name_prefix = f'{param_prefix}.layers.{layer_idx}' - trtllm_layer_name_prefix = f'{trtllm_layer_name}.{local_layer_idx}' - - for hf_weight_name, weight_info in hidden_layer_name_split.items(): - weights[ - f'{trtllm_layer_name_prefix}.{weight_info["name"]}'] = reshape( - split(params[f'{hf_layer_name_prefix}.{hf_weight_name}'], - mapping.tp_size, - mapping.tp_rank, - dim=weight_info["split_dim"]), weight_info["shape"]) - - for hf_weight_name, weight_info in hidden_layer_name_no_split.items(): - trtllm_layer_fullname = f'{trtllm_layer_name_prefix}.{weight_info["name"]}' - hf_layer_fullname = f'{hf_layer_name_prefix}.{hf_weight_name}' - weights[trtllm_layer_fullname] = reshape( - params[hf_layer_fullname].clone(), shape=weight_info["shape"]) - - self_attn_module_name = get_attn_module_name(layer_idx, 'self_attn') - weights.update( - fuse_qkv_one_layer( - params, self_attn_module_name, - f'{trtllm_layer_name_prefix}.{trtllm_attn_layer_name}', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - - if component == 'decoder': - cross_attn_module_name = get_attn_module_name( - layer_idx, 'encoder_attn') - weights.update( - fuse_qkv_one_layer( - params, cross_attn_module_name, - f'{trtllm_layer_name_prefix}.cross_attention', - mapping.tp_size, mapping.tp_rank, config.model_type, - (hidden_size * 3 // mapping.tp_size, hidden_size), - (hidden_size * 3 // mapping.tp_size))) - assert len(config.language_adapter_config['language_list']) > 0 - - language_adapter_weights = defaultdict(list) - language_adapter_weight_info = { - 'fc1.weight': { - "name": f'{trtllm_layer_name_prefix}.adapter.layers.fc.weight', - "shape": None - }, - 'fc1.bias': { - "name": f'{trtllm_layer_name_prefix}.adapter.layers.fc.bias', - "shape": None - }, - 'fc2.weight': { - "name": - f'{trtllm_layer_name_prefix}.adapter.layers.proj.weight', - "shape": None - }, - 'fc2.bias': { - "name": f'{trtllm_layer_name_prefix}.adapter.layers.proj.bias', - "shape": None - }, - } - - for language in config.language_adapter_config['language_list']: - for key in language_adapter_weight_info.keys(): - language_adapter_weights[key].append(params[ - f'{param_prefix}.layers.{layer_idx}.adapter.{language}.{key}'] - .unsqueeze(0)) - - import torch - for key, weight_info in language_adapter_weight_info.items(): - weights[weight_info["name"]] = torch.cat( - language_adapter_weights[key], dim=0) - - weights[ - f'{trtllm_layer_name_prefix}.adapter_layer_norm.weight'] = params[ - f'{param_prefix}.layers.{layer_idx}.adapter_layer_norm.weight'] - weights[f'{trtllm_layer_name_prefix}.adapter_layer_norm.bias'] = params[ - f'{param_prefix}.layers.{layer_idx}.adapter_layer_norm.bias'] - - if component == 'decoder': - - # share_decoder_input_output_embed=True, output_proj = embed_tokens.transpose() - lm_head_weight_name = f'{param_prefix}.output_projection.weight' - if lm_head_weight_name not in params: - lm_head_weight_name = embed_tokens_weight_name - weights['lm_head.weight'] = reshape( - split(params[lm_head_weight_name], - mapping.tp_size, - mapping.tp_rank, - dim=0), (config.vocab_size // mapping.tp_size, hidden_size)) - - return weights - - -def get_model(args): - if args.model_type == "t5": - model = T5ForConditionalGeneration.from_pretrained(args.model_dir) - elif args.model_type == "nmt": - from fairseq.models.transformer import TransformerModel - model = TransformerModel.from_pretrained(args.model_dir) - elif args.model_type == "bart": - if args.nougat: - model = VisionEncoderDecoderModel.from_pretrained(args.model_dir) - model = model.get_decoder() - elif args.eclair_radio: - import torch - - class RadioWithNeck(torch.nn.Module): - - def __init__(self): - super().__init__() - - self.model_encoder = torch.hub.load("NVlabs/RADIO", - "radio_model", - version="radio_v2.5-h") - self.model_encoder.summary_idxs = torch.tensor(4) - - self.conv1 = torch.nn.Conv1d(1280, 1024, 1) - self.layer_norm1 = torch.nn.LayerNorm( - 1024, eps=1e-6, elementwise_affine=True) - self.conv2 = torch.nn.Conv2d(1024, - 1024, - kernel_size=(1, 4), - stride=(1, 4), - padding=0, - bias=False) - self.layer_norm2 = torch.nn.LayerNorm( - 1024, eps=1e-6, elementwise_affine=True) - - def forward(self, pixel_values): - _, feature = self.model_encoder(pixel_values) - output = self.conv1(feature.permute(0, 2, - 1)).permute(0, 2, 1) - output = self.layer_norm1(output).permute(0, 2, 1) - - b, d, _ = output.shape - h = pixel_values.shape[-2] // 16 - w = pixel_values.shape[-1] // 16 - output = self.conv2(output.reshape(b, d, h, w)) - output = output.flatten(-2, -1).permute(0, 2, 1) - output = self.layer_norm2(output) - return output - - def get_processor(): - processor = NougatProcessor.from_pretrained( - "facebook/nougat-base") - - special_tokens = { - "output_plain_index": "", - "output_markdown_index": "", - "output_no_text_index": "", - "output_ocr_index": "", - "predict_bbox_index": "", - "no_bbox_index": "", - "bbox_start_index": "", # not used but can keep - # "bbox_end_index": "", # not used but can keep - "no_class_index": "", - "predict_classes_index": "", - } - for key, special_t in special_tokens.items(): - processor.tokenizer.add_special_tokens( - {"additional_special_tokens": [special_t]}) - setattr(processor.tokenizer, key, - processor.tokenizer.encode(special_t)[1]) - - # Add regular tokens for boxes - processor.tokenizer.add_tokens( - [f"" for x_i in range(1024)]) - processor.tokenizer.add_tokens( - [f"" for y_i in range(1280)]) - # Add regular tokens for classes - #"" - possible_classes = [ - "Text", "Title", "Section-header", "List-item", "TOC", - "Bibliography", "Footnote", "Page-header", "Page-footer", - "Picture", "Formula", "Page-number", "Table", "Caption" - ] - processor.tokenizer.add_tokens( - [f"" for cls in possible_classes]) - return processor - - processor = get_processor() - model = VisionEncoderDecoderModel.from_pretrained( - "facebook/nougat-base") - model.encoder = RadioWithNeck() - model.decoder.resize_token_embeddings(len(processor.tokenizer), - pad_to_multiple_of=64) - model.config.decoder_start_token_id = processor.tokenizer.eos_token_id # 2 - model.config.pad_token_id = processor.tokenizer.pad_token_id # 1 - from transformers.models.mbart.modeling_mbart import \ - MBartLearnedPositionalEmbedding - _, d_model = model.device, model.config.decoder.d_model - - with torch.inference_mode(): - # Inspect checkpoint shapes - safetensors.torch.load_model(model, - os.path.join( - args.model_dir, - "model.safetensors"), - strict=False) - model.decoder.model.decoder.embed_positions = MBartLearnedPositionalEmbedding( - ECLAIR_RADIO_MAX_POSITION_EMBEDDINGS, d_model) - model.decoder.model.decoder.embed_positions.weight.data.zero_() - model.decoder.model.decoder.embed_positions.weight.requires_grad_( - True) - model.decoder.lm_head.weight = model.decoder.get_input_embeddings( - ).weight - - model.eval() - model = model.get_decoder() - - else: - model = AutoModelForSeq2SeqLM.from_pretrained(args.model_dir) - elif args.model_type == "pix2struct": - model = Pix2StructForConditionalGeneration.from_pretrained( - args.model_dir) - elif args.model_type == "blip2": - model = Blip2ForConditionalGeneration.from_pretrained( - args.model_dir).language_model - elif args.model_type == "language_adapter": - import torch - - class DummyTorchModel: - - def __init__(self, model) -> None: - self.model = model - - def state_dict(self): - return self.model['model'] - - model = torch.load(args.model_dir + "/model.pt", weights_only=False) - return DummyTorchModel(model) - - return model - - -def convert_checkpoint(args): - - model = get_model(args) - - saved_dir = Path(args.output_dir) - saved_dir.mkdir(parents=True, exist_ok=True) - - encoder_saved_dir = saved_dir / "encoder" - encoder_saved_dir.mkdir(parents=True, exist_ok=True) - decoder_saved_dir = saved_dir / "decoder" - decoder_saved_dir.mkdir(parents=True, exist_ok=True) - - world_size = args.tp_size * args.pp_size - - kv_cache_quant_algo = None - quant_algo = None - - model_type = args.model_type if args.model_type != "blip2" else "t5" - parse_config_mapper = { - 't5': parse_t5_config, - 'pix2struct': parse_pix2struct_config, - 'blip2': parse_t5_config, # blip2 uses t5 config parser - 'language_adapter': parse_language_adapter_config, - 'nmt': parse_nmt_config, - 'bart': parse_bart_config, - } - encoder_config, decoder_config = parse_config_mapper[model_type](args, - model) - - additional_settings = ["gated_act"] - if model_type == 'language_adapter': - additional_settings += ["residual_scaling", "language_adapter_config"] - - if not (args.nougat - or args.eclair_radio) and args.model_type != "pix2struct": - tllm_encoder_config = { - 'architecture': "EncoderModel", - 'dtype': args.dtype, - 'logits_dtype': encoder_config.logits_dtype, - 'num_hidden_layers': encoder_config.n_layer, - 'num_attention_heads': encoder_config.n_head, - 'hidden_size': encoder_config.hidden_size, - 'norm_epsilon': encoder_config.layernorm_eps, - 'vocab_size': encoder_config.vocab_size, - 'position_embedding_type': encoder_config.position_embedding_type, - 'hidden_act': encoder_config.hidden_act, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_position_embeddings': encoder_config.n_positions, - 'num_key_value_heads': encoder_config.n_head, - 'head_size': encoder_config.head_size, - 'has_position_embedding': encoder_config.has_position_embedding, - 'layernorm_type': encoder_config.layernorm_type, - 'has_attention_qkvo_bias': encoder_config.has_attention_qkvo_bias, - 'has_mlp_bias': encoder_config.has_mlp_bias, - 'has_model_final_layernorm': - encoder_config.has_model_final_layernorm, - 'has_embedding_layernorm': encoder_config.has_embedding_layernorm, - 'has_embedding_scale': encoder_config.has_embedding_scale, - 'intermediate_size': encoder_config.ffn_hidden_size, - 'q_scaling': encoder_config.q_scaling, - 'layernorm_position': encoder_config.layernorm_position, - 'mlp_type': encoder_config.mlp_type, - 'relative_attention': encoder_config.relative_attention, - 'max_distance': encoder_config.max_distance, - 'num_buckets': encoder_config.num_buckets, - 'model_type': encoder_config.model_type, - } - - for additional_setting in additional_settings: - if hasattr(encoder_config, additional_setting): - tllm_encoder_config.update({ - additional_setting: - getattr(encoder_config, additional_setting) - }) - - with (encoder_saved_dir / "config.json").open('w') as f: - json.dump(tllm_encoder_config, f, indent=4) - - encoder_convert_args = dict(params=model.state_dict(), - component="encoder") - tllm_decoder_config = { - 'architecture': "DecoderModel", - 'dtype': args.dtype, - 'logits_dtype': decoder_config.logits_dtype, - 'num_hidden_layers': decoder_config.n_layer, - 'num_attention_heads': decoder_config.n_head, - 'hidden_size': decoder_config.hidden_size, - 'norm_epsilon': decoder_config.layernorm_eps, - 'vocab_size': decoder_config.vocab_size, - 'position_embedding_type': decoder_config.position_embedding_type, - 'hidden_act': decoder_config.hidden_act, - 'quantization': { - 'quant_algo': quant_algo, - 'kv_cache_quant_algo': kv_cache_quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'max_position_embeddings': decoder_config.n_positions, - 'head_size': decoder_config.head_size, - 'has_position_embedding': decoder_config.has_position_embedding, - 'layernorm_type': decoder_config.layernorm_type, - 'has_attention_qkvo_bias': decoder_config.has_attention_qkvo_bias, - 'has_mlp_bias': decoder_config.has_mlp_bias, - 'has_model_final_layernorm': decoder_config.has_model_final_layernorm, - 'has_embedding_layernorm': decoder_config.has_embedding_layernorm, - 'has_embedding_scale': decoder_config.has_embedding_scale, - 'intermediate_size': decoder_config.ffn_hidden_size, - 'q_scaling': decoder_config.q_scaling, - 'layernorm_position': decoder_config.layernorm_position, - 'mlp_type': decoder_config.mlp_type, - 'relative_attention': decoder_config.relative_attention, - 'max_distance': decoder_config.max_distance, - 'num_buckets': decoder_config.num_buckets, - 'model_type': decoder_config.model_type, - 'rescale_before_lm_head': decoder_config.rescale_before_lm_head, - 'encoder_hidden_size': decoder_config.encoder_hidden_size, - 'encoder_num_heads': decoder_config.encoder_num_heads, - 'encoder_head_size': decoder_config.encoder_head_size, - 'skip_cross_kv': args.skip_cross_kv, - 'use_implicit_relative_attention': args.use_implicit_relative_attention, - 'decoder_start_token_id': decoder_config.decoder_start_token_id, - 'eos_token_id': decoder_config.eos_token_id, - 'bos_token_id': decoder_config.bos_token_id, - 'pad_token_id': decoder_config.pad_token_id, - } - for additional_setting in additional_settings: - if hasattr(decoder_config, additional_setting): - tllm_decoder_config.update({ - additional_setting: - getattr(decoder_config, additional_setting) - }) - - with (decoder_saved_dir / "config.json").open('w') as f: - json.dump(tllm_decoder_config, f, indent=4) - - decoder_convert_args = dict(params=model.state_dict(), component="decoder") - - if args.model_type == "nmt": - fairseq_config = vars(model.cfg.model) # Namespace --> dict - num_embeddings = fairseq_config['max_source_positions'] - embedding_dim = fairseq_config['encoder_embed_dim'] - padding_idx = model.models[0].encoder.embed_tokens.padding_idx # 1 - - sin_pos_embedding = model.models[ - 0].encoder.embed_positions.get_embedding( - padding_idx + 1 + num_embeddings, - embedding_dim, - padding_idx=padding_idx) # [2 + num_embeddings, embed_dim] - sin_pos_embedding = sin_pos_embedding[2:, :] # remove offset embeddings - - encoder_convert_args["sin_pos_embedding"] = sin_pos_embedding - decoder_convert_args["sin_pos_embedding"] = sin_pos_embedding - - if args.workers == 1: - if not (args.nougat - or args.eclair_radio) and args.model_type != "pix2struct": - convert(0, world_size, args, tllm_encoder_config, - encoder_convert_args, encoder_saved_dir) - convert(0, world_size, args, tllm_decoder_config, decoder_convert_args, - decoder_saved_dir) - else: - if args.workers > world_size: - args.workers = world_size - LOGGER.info(f'Convert checkpoint using {args.workers} workers.') - import torch.multiprocessing as mp - if not (args.nougat - or args.eclair_radio) and args.model_type != "pix2struct": - mp.spawn(convert, - nprocs=args.workers, - args=(world_size, args, tllm_encoder_config, - encoder_convert_args, encoder_saved_dir)) - mp.spawn(convert, - nprocs=args.workers, - args=(world_size, args, tllm_decoder_config, - decoder_convert_args, decoder_saved_dir)) - - -def convert(worker_rank, world_size, args, model_config, convert_args, - saved_dir): - for rank in range(worker_rank, world_size, args.workers): - rank_config = copy.deepcopy(PretrainedConfig.from_dict(model_config)) - rank_config.set_rank(rank) - weights = globals( - )[f'convert_{rank_config.model_type}_weights_to_tllm_safetensors']( - config=rank_config, **convert_args) - safetensors.torch.save_file(weights, - f'{saved_dir}/rank{rank}.safetensors') - - -if __name__ == "__main__": - emit_engine_arch_deprecation("convert_checkpoint.py") - parser = argparse.ArgumentParser( - formatter_class=argparse.RawTextHelpFormatter) - parser.add_argument( - '--model_type', - type=str, - default='t5', - choices=[ - 't5', 'nmt', 'bart', 'pix2struct', 'blip2', 'language_adapter' - ], - help= - 'Multimodal type when this script is used for multimodal conversion.') - - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument("--model_dir", - "-i", - type=str, - help="Path to the framework checkpoint file", - required=True) - parser.add_argument("--output_dir", - "-o", - type=str, - help="Path to the converted TRT-LLM model weight file", - required=True) - parser.add_argument( - "--workers", - type=int, - help="How many workers to spawn for conversion (default: 4)", - default=4) - parser.add_argument("--nougat", - action="store_true", - help="Model which uses vision encoder + mbart decoder") - parser.add_argument("--eclair_radio", - action="store_true", - help="Model which uses vision encoder + mbart decoder") - parser.add_argument("--verbose", - action="store_true", - help="Provide verbose messages") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharding is only enabled when embedding_sharding_dim = 0' - ) - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--dtype', - type=str, - default='float16', - choices=['float16', 'float32', 'bfloat16'], - help= - 'Target inference dtype. Weights and Computation will be in this dtype, no matter what original dtype the weight checkpoint has.' - ) - parser.add_argument( - '--skip_cross_kv', - action='store_true', - help= - 'Skip redundant cross qkv computation by using TensorRT IfConditional switch (experimental).' - ) - parser.add_argument( - '--use_implicit_relative_attention', - action='store_true', - help= - 'Compute relative attention bias on the fly instead of pre-compute a relative attention bias table.' - ) - args = parser.parse_args() - log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s" - logging.basicConfig(level=logging.DEBUG if args.verbose else logging.INFO, - format=log_format) - LOGGER.info("\n=============== Argument ===============") - for key in vars(args): - LOGGER.info(f"{key}: {vars(args)[key]}") - LOGGER.info("========================================") - - start_time = datetime.now() - convert_checkpoint(args) - stop_time = datetime.now() - run_time = (stop_time - start_time) - LOGGER.info("Spend {} (h:m:s) to convert the model".format(run_time)) diff --git a/examples/models/core/enc_dec/helper.py b/examples/models/core/enc_dec/helper.py deleted file mode 100755 index ed3628c1bbbd..000000000000 --- a/examples/models/core/enc_dec/helper.py +++ /dev/null @@ -1,121 +0,0 @@ -import math -import typing -from typing import Union - -import numpy as np -import torch # pytype: disable=import-error - -from tensorrt_llm._utils import str_dtype_to_torch - - -def split(v: Union[np.ndarray, torch.Tensor], - tp_size: int, - tp_rank: int, - dim=0): - if tp_size == 1: - if isinstance(v, np.ndarray): - return np.ascontiguousarray(v.copy()) - else: - return v.clone().detach() - assert len(v.shape) > 1 or dim == 0 - if isinstance(v, np.ndarray): - return np.ascontiguousarray( - np.split(v, tp_size, axis=dim)[tp_rank].copy()) - else: - assert v.shape[dim] % tp_size == 0, \ - 'Unable to split: shape={v.shape} (dim={dim}) tp_size={tp_size}.' - split_size = v.shape[dim] // tp_size - return v.split(split_size, dim=dim)[tp_rank].clone().detach() - - -def reshape(v: torch.Tensor, shape=None): - if shape is None: - return v.contiguous() - else: - return v.reshape(shape).contiguous() - - -def fuse_qkv_one_layer(params, attn_module_name, trtllm_layer_name, tp_size, - tp_rank, model_type, weight_shape, bias_shape): - - qkv_module_names = get_qkv_module_name(model_type) - - weight = {} - - # fuse weights of q, k, v - q_w = params[f'{attn_module_name}.{qkv_module_names["q"]}.weight'] - k_w = params[f'{attn_module_name}.{qkv_module_names["k"]}.weight'] - v_w = params[f'{attn_module_name}.{qkv_module_names["v"]}.weight'] - - # fuse qkv weight - shape = q_w.shape # (do, din) - qkv_w = torch.cat([q_w, k_w, v_w], - dim=0).reshape([3, shape[0], shape[1]]) # (3, do, din) - qkv_w = split(qkv_w, tp_size, tp_rank, dim=1) - weight[f'{trtllm_layer_name}.qkv.weight'] = reshape(qkv_w, - shape=weight_shape) - - # fuse qkv biases if present - if f'{attn_module_name}.{qkv_module_names["q"]}.bias' in params.keys( - ) and params[f'{attn_module_name}.{qkv_module_names["q"]}.bias'] is not None: - q_b = params[f'{attn_module_name}.{qkv_module_names["q"]}.bias'] - k_b = params[f'{attn_module_name}.{qkv_module_names["k"]}.bias'] - v_b = params[f'{attn_module_name}.{qkv_module_names["v"]}.bias'] - shape = q_b.shape[0] # (do,) - qkv_b = torch.cat([q_b, k_b, v_b], dim=0).reshape([3, shape]) # (3, do) - qkv_b = split(qkv_b, tp_size, tp_rank, dim=1) - weight[f'{trtllm_layer_name}.qkv.bias'] = reshape(qkv_b, - shape=bias_shape) - return weight - - -def get_qkv_module_name(model_type): - if model_type in ["t5", "blip2"]: - q = "q" - k = "k" - v = "v" - elif model_type == "bart" or model_type == "nmt" or model_type == "language_adapter": - q = "q_proj" - k = "k_proj" - v = "v_proj" - elif model_type == "pix2struct": - q = "query" - k = "key" - v = "value" - return {"q": q, "k": k, "v": v} - - -def convert_weight_to_dtype(params: typing.Dict[str, torch.Tensor], - dtype: typing.Optional[np.dtype] = None): - if dtype is not None: - assert isinstance(dtype, - str), f"dtype must be str, but get type {type(dtype)}" - for name in params.keys(): - params[name] = params[name].to(str_dtype_to_torch(dtype)) - - -def fairseq_sin_pos_embedding(num_embeddings: int, embedding_dim: int): - ''' - generate fairseq specific sinusoidal position embedding [sin, sin, ... cos, cos...] - https://github.com/facebookresearch/fairseq/blob/main/fairseq/modules/sinusoidal_positional_embedding.py - ''' - padding_offset = 2 - half_dim = embedding_dim // 2.0 - emb = math.log(10000) / (half_dim - 1) - emb = torch.exp(torch.arange(half_dim, dtype=torch.float16) * -emb) - emb = torch.arange(num_embeddings + padding_offset, - dtype=torch.float16).unsqueeze(1) * emb.unsqueeze(0) - emb = torch.cat([torch.sin(emb), torch.cos(emb)], - dim=1).view(num_embeddings + padding_offset, -1) - if embedding_dim % 2 == 1: - # zero pad - emb = torch.cat( - [emb, torch.zeros(num_embeddings + padding_offset, 1)], - dim=1, - dtype=torch.float16) - ''' - remove first 2 column to match position_id setup difference between fairseq & trt - fairseq position_id starts with 2, ex: [2, 3, 4 ..] - trt position_id starts with 0 [0, 1, 2] - ''' - return emb[padding_offset:, :] diff --git a/examples/models/core/enc_dec/run.py b/examples/models/core/enc_dec/run.py deleted file mode 100644 index 8ae81960e6c5..000000000000 --- a/examples/models/core/enc_dec/run.py +++ /dev/null @@ -1,512 +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. -import argparse -import json -import time - -import numpy as np -import torch -from transformers import (AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, - BartForConditionalGeneration, - MBartForConditionalGeneration, - T5ForConditionalGeneration) - -import tensorrt_llm -from tensorrt_llm import logger -from tensorrt_llm.runtime import EncDecModelRunner - - -def print_tensor(tensor_name, tensor, num_elements=10): - if tensor.dtype in (torch.int32, torch.int64): - tensor = tensor.to(dtype=float) - print( - f'{tensor_name}: mean={tensor.abs().mean().item():.3f}, sum={tensor.abs().sum().item():.3f}, max={tensor.abs().max().item():.3f}' - ) - # Pass num_elements=-1 will print the whole tensor - if num_elements < 0: - num_elements = torch.numel(tensor) - print(f'{tensor.flatten()[:num_elements]}') - print("Tensor Shape: ", tensor.size()) - print("") - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument("--max_new_tokens", type=int, default=64) - parser.add_argument("--log_level", type=str, default="error") - parser.add_argument("--engine_dir", "-i", type=str, default="trt_engines") - parser.add_argument("--engine_name", type=str, default="enc_dec") - parser.add_argument("--model_name", - type=str, - help="HuggingFace model name or FairSeq model path", - default="t5-small") - parser.add_argument("--num_beams", - type=int, - help="Use beam search if num_beams >1", - default=1) - parser.add_argument("--debug_mode", - help="Whether or not to turn on the debug mode", - action='store_true') - parser.add_argument("--compare_hf_fp32", - help="Compare results with HuggingFace FP32", - action='store_true') - parser.add_argument('--lora_dir', type=str, default=None, nargs="+") - parser.add_argument('--lora_task_uids', type=str, default=None, nargs="+") - parser.add_argument("--output_npy", - type=str, - default=None, - help="Store input/output tensors C++ runtime testing") - return parser.parse_args() - - -def test_fairseq_models(args): - ## Note: NMT is the only FairSeq model. Adding FairSeq dependency is too heavy for the CI workflow, hence we used fixed input/output ids for correctness check and leave FairSeq code in comments. Users can follow Encoder-Decoder's README to install FairSeq and test locally. - ''' - from fairseq.models.transformer import TransformerModel - - fairseq_model = TransformerModel.from_pretrained(model_name_or_path=args.model_name, data_name_or_path=args.model_name, bpe='subword_nmt', tokenizer='moses').cuda() - - input_text = "Good Morning! How are you doing today?" - input_ids = fairseq_model.encode(input_text) - - tik = time.time() - # Note: FairSeq sampling=True results are not deterministic, disable during accuracy check - fairseq_output_ids = fairseq_model.generate(input_ids, beam=1, sampling=False) # - tik = time.time() - - fairseq_output_ids = fairseq_output_ids[0]['tokens'] - fairseq_output_text = fairseq_model.decode(fairseq_output_ids) - - print("--------------------------------------") - print("input text: ", input_text) - print("input ids: ", input_ids) # [9938, 5384, 9328, 812, 3619, 53, 181, 3829, 1735, 171, 2] - print("fairseq_output ids: ", fairseq_output_ids) # [9804, 391, 4, 4625, 167, 25, 1003, 5123, 17, 167, 1466, 1234, 171, 2] - print("fairseq_output text: ", fairseq_output_text) # "Bonjour, Comment vous en tirez-vous aujourd'hui ?" - print(f"FairSeq E2E time {(tok-tik)*1000}ms") - print("--------------------------------------") - ''' - - max_new_tokens = args.max_new_tokens - bos_token_id = 2 - pad_token_id = 0 - eos_token_id = 2 - decoder_start_token_id = bos_token_id - - input_ids = torch.tensor( - [9938, 5384, 9328, 812, 3619, 53, 181, 3829, 1735, 171, 2]) - fairseq_output_ids = torch.tensor( - [9804, 391, 4, 4625, 167, 25, 1003, 5123, 17, 167, 1466, 1234, 171, 2]) - input_ids = torch.tensor([input_ids.tolist()]).type(torch.IntTensor).cuda() - decoder_input_ids = torch.IntTensor([[decoder_start_token_id]]).cuda() - decoder_input_ids = decoder_input_ids.repeat((input_ids.shape[0], 1)) - - tllm_model = EncDecModelRunner.from_engine(args.engine_name, - args.engine_dir, - debug_mode=args.debug_mode) - - inference_dtype = tllm_model.encoder_model_config.dtype - - return_dict = False # when set return_dict=True, get outputs by key - tik = time.time() - tllm_output = tllm_model.generate( - encoder_input_ids=input_ids, - decoder_input_ids=decoder_input_ids, - max_new_tokens=max_new_tokens, - num_beams=args.num_beams, - bos_token_id=bos_token_id, - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - debug_mode=args.debug_mode, - ) - torch.cuda.synchronize() - tok = time.time() - - if return_dict: - tllm_output_ids = tllm_output['output_ids'] - else: - tllm_output_ids = tllm_output - - if tensorrt_llm.mpi_rank() == 0: - output_ids = tllm_output_ids[:, 0, :] - output_ids = output_ids[output_ids != eos_token_id] - fairseq_output_ids = fairseq_output_ids[fairseq_output_ids != - eos_token_id] - - print("--------------------------------------") - print("TRT-LLM output_ids: ", output_ids) - print(f"TRT-LLM E2E time {(tok-tik)*1000}ms") - print("Precision:", inference_dtype) - print("--------------------------------------") - - assert output_ids.tolist() == fairseq_output_ids.tolist( - ), f"TRT-LLM output ids {output_ids} does not match Fairseq ids {fairseq_output_ids}" - - -def test_language_adapter_models(args): - # TRT-LLM runtime - tllm_model = EncDecModelRunner.from_engine(args.engine_name, - args.engine_dir, - debug_mode=args.debug_mode) - inference_dtype = tllm_model.encoder_model_config.dtype - - tokenized_inputs = [[ - 34901, 3048, 3011, 123250, 9517, 3018, 45732, 3048, 3003, 6553, 3781, - 383416, 33356, 3032, 97339, 3382, 3003, 19143, 3022, 169460, 3001, - 87966, 35848, 2996, 3002, 6358, 7387, 25864, 3032, 3011, 4570, 3022, - 7235, 182168, 2992, 3003, 2991, 39861, 2997, 26629, 98419, 5339, 2993, - 423511, 2544, 2 - ], - [ - 34901, 3048, 3011, 123250, 9517, 3018, 45732, 3048, - 3003, 6553, 3781, 383416, 33356, 3032, 97339, 3382, - 3003, 19143, 3022, 169460, 3001, 87966, 35848, 2996, - 3002, 6358, 7387, 25864, 3032, 3011, 4570, 3022, - 7235, 182168, 2992, 3003, 2991, 39861, 2997, 26629, - 98419, 5339, 2993, 423512, 2712, 2 - ]] - language_task_uids = [2, 3] - - target_outputs = [[ - 4094, 82383, 3501, 3073, 12672, 3535, 45217, 3018, 45732, 3158, 3116, - 400231, 3010, 7212, 12398, 52837, 3046, 391725, 3164, 3116, 40625, 2994, - 204507, 3001, 402406, 35848, 2996, 3002, 3003, 8317, 2994, 3007, 80104, - 55333, 3046, 3073, 4755, 2994, 7235, 182168, 2992, 3030, 4005, 2994, - 63261, 60932, 3010, 2991, 39861, 2993 - ], - [ - 62366, 3099, 14803, 3056, 9517, 3056, 3495, 36942, - 3975, 292422, 3262, 3315, 3010, 53857, 41472, 9823, - 3010, 6493, 26179, 151498, 3062, 286084, 3453, 3315, - 45059, 2994, 286488, 3001, 53771, 16240, 35848, 2996, - 3002, 22161, 3072, 3315, 25864, 51019, 3062, 3072, - 3063, 2999, 10657, 2994, 7235, 182168, 2992, 3030, - 7109, 3077, 2999, 109181, 51563, 3366, 2991, 39861, - 2993 - ]] - - max_new_tokens = args.max_new_tokens - input_ids = torch.IntTensor(tokenized_inputs) - - with open(f"{args.engine_dir}/decoder/config.json", "r") as f: - model_config = json.load(f) - decoder_start_token_id = model_config['pretrained_config'][ - 'decoder_start_token_id'] - decoder_input_ids = torch.IntTensor([[decoder_start_token_id]]).to('cuda') - decoder_input_ids = decoder_input_ids.repeat((input_ids.shape[0], 1)) - - def get_language_adapter_routings(language_uids, input_ids): - language_adapter_routing_masks = torch.tensor(language_uids, - dtype=torch.int32) - language_adapter_routings = [] - - for i, input_id in enumerate(input_ids): - mask = language_adapter_routing_masks[i].repeat(len(input_id), 1) - language_adapter_routings.append(mask) - - return torch.cat(language_adapter_routings) - - encoder_language_adapter_routings = get_language_adapter_routings( - language_task_uids, input_ids) - decoder_language_adapter_routings = get_language_adapter_routings( - language_task_uids, decoder_input_ids) - - bos_token_id = 2 - pad_token_id = 0 - eos_token_id = 2 - - return_dict = True # when set return_dict=True, get outputs by key - tik = time.time() - tllm_output = tllm_model.generate( - encoder_input_ids=input_ids, - decoder_input_ids=decoder_input_ids, - max_new_tokens=max_new_tokens, - num_beams=args.num_beams, - bos_token_id=bos_token_id, - pad_token_id=pad_token_id, - eos_token_id=eos_token_id, - debug_mode=args.debug_mode, - return_dict=return_dict, - encoder_language_adapter_routings=encoder_language_adapter_routings, - decoder_language_adapter_routings=decoder_language_adapter_routings, - return_encoder_output=args.output_npy and tensorrt_llm.mpi_rank() == 0) - torch.cuda.synchronize() - tok = time.time() - - if tensorrt_llm.mpi_rank() == 0: - tllm_output_ids = tllm_output['output_ids'] - - output_ids = tllm_output_ids[:, 0, :] - output_ids_list = [ - output_id[output_id != eos_token_id].tolist() - for output_id in output_ids - ] - - decoder_input_lengths = (decoder_input_ids != pad_token_id).sum(dim=1) - output_gen_lengths = (output_ids != eos_token_id).sum( - dim=1) - decoder_input_lengths - - print( - f"------ TRT-LLM beam = {args.num_beams} --------------------------------" - ) - if 'encoder_output' in tllm_output: - encoder_output = tllm_output['encoder_output'] - print_tensor('TRT-LLM encoder_output:', encoder_output) - print("TRT-LLM output_ids: ", output_ids) - print("TRT-LLM output generated lengths: ", output_gen_lengths) - print(f"TRT-LLM E2E time {(tok-tik)*1000}ms") - print("Precision:", inference_dtype) - print("--------------------------------------") - - assert output_ids_list == target_outputs, f"TRT-LLM output ids {output_ids_list} does not match Fairseq ids {target_outputs}" - - if args.output_npy: - output_npy(args, tokenized_inputs, tllm_output, output_ids) - - -def output_npy(args, tokenized_inputs, tllm_output, output_ids): - os.makedirs(args.output_npy, exist_ok=True) - - if hasattr(tokenized_inputs, "attention_mask"): - input_lengths = tokenized_inputs.attention_mask.sum(dim=1).type( - torch.IntTensor) - input_ids = tokenized_inputs.input_ids.type(torch.IntTensor) - else: - input_lengths = torch.IntTensor( - [len(input_ids) for input_ids in tokenized_inputs]) - input_ids = torch.IntTensor(tokenized_inputs) - - input_ids_flatten = torch.cat( - [input_ids[i][:input_lengths[i]] for i in range(len(input_lengths))]) - encoder_output = tllm_output['encoder_output'].type(torch.float16) - - def save_npy(tensor, name): - np.save(os.path.join(args.output_npy, f'{name}.npy'), - tensor.cpu().numpy()) - - print( - f"Saving input/output tensors to {args.output_npy} for C++ runtime testing" - ) - save_npy(input_ids_flatten, 'input_ids') # [num_tokens] - save_npy(input_lengths, 'input_lengths') # [batch_size] - save_npy(encoder_output, 'encoder_output') # [num_tokens, hidden_size] - save_npy( - output_ids, f'output_ids_beam{args.num_beams}' - ) # [batch_size, max_output_tokens], max_output_tokens = decoder_input_tokens + max_new_tokens - - -if __name__ == "__main__": - import os - - os.environ["TOKENIZERS_PARALLELISM"] = "false" - args = parse_arguments() - logger.set_level(args.log_level) - - # FairSeq NMT test logic is different from HuggingFace models - if 'wmt' in args.model_name: - test_fairseq_models(args) - exit() - - # language adapter test logic is different from HuggingFace models - if 'language_adapter' in args.engine_name: - test_language_adapter_models(args) - exit() - - test_remove_padding = True - if not test_remove_padding: - if 't5' in args.model_name: - input_text = "translate English to German: The house is wonderful, radiating timeless charm and offering a warm, inviting interior with beautiful details and a serene backyard." - elif 'bart' in args.model_name: - input_text = "The tower is 324 metres (1,063 ft) tall, about the same height as an 81-storey building, and the tallest structure in Paris. Its base is square, measuring 125 metres (410 ft) on each side. During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world, a title it held for 41 years until the Chrysler Building in New York City was finished in 1930. It was the first structure to reach a height of 300 metres. Due to the addition of a broadcasting aerial at the top of the tower in 1957, it is now taller than the Chrysler Building by 5.2 metres (17 ft). Excluding transmitters, the Eiffel Tower is the second tallest free-standing structure in France after the Millau Viaduct." - else: - raise RuntimeError('Unsupported model type!') - - else: - input_text = [ - "translate English to German: The house is wonderful.", - "summarize: I am a high-performance inference optimizer and runtime.", - "During its construction, the Eiffel Tower surpassed the Washington Monument to become the tallest man-made structure in the world", - ] - - # TRT-LLM runtime - tllm_model = EncDecModelRunner.from_engine(args.engine_name, - args.engine_dir, - args.lora_dir, - args.lora_task_uids, - debug_mode=args.debug_mode) - - inference_dtype = tllm_model.encoder_model_config.dtype - if inference_dtype == 'float32': - if "byt5" in args.model_name: - print( - "ByT5 models tokenize input by bytes instead of words, causing the input text in this example to be longer than the default value during build stage. Please adjust --max_input_len during trtllm-build to select the right length limit for ByT5 models." - ) - else: - input_text.append( - "Summarize this article in one sentence.\n\nKristine Watts (Molie Weeks) is broken apart, missing her lover; she is not able to overcome her love for him that is lost in the past. She hires a stranger (Douglas Davis) and gives a list of her mistakes to him with things to fix. But time is irreversible and sometimes the cure for the pain is a tragic end.\n\nThe first point that impresses in \"The Cure\" is the stylish cinematography that alternates black and white with color. The concise and sharp screenplay is capable to develop a tragic and bleak tale of love with an unexpected plot point in the very end in less than eight minutes. The soundtrack is beautiful but the volume is a little loud and associated to the fact that English is not my native language, in some moments I needed to repeat some words whispered by the narrator. The unknown lead actress has magnificent performance and is extremely gorgeous. I hope to have a chance to see her again on the screen. Last but not the least, the debut of the director and writer Ryan Jafri could not be better. My vote is nine.\n\nTitle (Brazil): Not Available", - ) - - tokenizer = AutoTokenizer.from_pretrained( - args.model_name) # TODO: use model path instead - tokenized_inputs = tokenizer(input_text, return_tensors="pt", padding=True) - - max_new_tokens = args.max_new_tokens - input_ids = tokenized_inputs.input_ids.type(torch.IntTensor).to( - 'cuda') # [batch_size, padded_length] - # by default int64, must cast to int32! otherwise C++ kernel will interpret as [a, 0, b, 0, c, 0, ...] - - if tensorrt_llm.mpi_rank() == 0: - print("--------------------------------------") - print( - f"BOS={tokenizer.bos_token_id}, PAD={tokenizer.pad_token_id}, EOS={tokenizer.eos_token_id}" - ) - print("input text: ", input_text) - print("input ids: ", input_ids) - print("input lengths: ", tokenized_inputs.attention_mask.sum(dim=1)) - print("--------------------------------------") - - model_config = AutoConfig.from_pretrained(args.model_name) - - # start_id for decoder (could add more input_ids as forced_decoder_ids) - decoder_input_ids = torch.IntTensor([[model_config.decoder_start_token_id] - ]).to('cuda') - decoder_input_ids = decoder_input_ids.repeat((input_ids.shape[0], 1)) - - # simple comparison with HF on FP32 - if args.compare_hf_fp32: - if tensorrt_llm.mpi_rank() == 0: - hf_model = AutoModelForSeq2SeqLM.from_pretrained( - args.model_name, # TODO: use model path instead - # dtype=torch.float16 if '16' in dtype else torch.float32, # TODO: use matched torch dtype - ).to('cuda').eval() # TODO: create config model path instead - assert type(hf_model) in ( - T5ForConditionalGeneration, BartForConditionalGeneration, - MBartForConditionalGeneration), 'Unsupported model!' - - if args.lora_dir is not None: - assert len(args.lora_dir - ) >= 1, "At least one lora model dir is required" - # we can only test single lora with HF - from peft import PeftModel - hf_model = PeftModel.from_pretrained( - hf_model, args.lora_dir[0]).to('cuda').eval() - - tik = time.time() - hf_gen_output = hf_model.generate( - input_ids=input_ids, - decoder_input_ids=decoder_input_ids, - max_new_tokens=max_new_tokens, - num_beams=args.num_beams, - bos_token_id=tokenizer.bos_token_id, - pad_token_id=tokenizer.pad_token_id, - eos_token_id=tokenizer.eos_token_id, - use_cache=True, - # control logits processors - no_repeat_ngram_size=0, # disable no repeat post-processor - forced_bos_token_id=None, # disable forced first/last token - forced_eos_token_id=None, - min_length=0, - # for debug - output_scores=True, - output_hidden_states=True, - return_dict_in_generate=True) - # get hf output scores - hf_output_ids = hf_gen_output.sequences - # convert to logits - torch.cuda.synchronize() - tok = time.time() - - output_ids = hf_output_ids.squeeze(dim=1) - hf_output_text = tokenizer.batch_decode(output_ids, - skip_special_tokens=True) - decoder_input_lengths = (decoder_input_ids - != tokenizer.pad_token_id).sum(dim=1) - output_gen_lengths = (output_ids != tokenizer.eos_token_id).sum( - dim=1) - decoder_input_lengths - print( - f"------ HF beam = {args.num_beams} --------------------------------" - ) - print("HF output_ids: ", output_ids) - print("HF output text: ", hf_output_text) - print("HF output generated lengths: ", output_gen_lengths) - print(f"HF E2E time {(tok-tik)*1000}ms") - print("--------------------------------------") - - return_dict = True # when set return_dict=True, get outputs by key - tik = time.time() - tllm_output = tllm_model.generate( - encoder_input_ids=input_ids, - decoder_input_ids=decoder_input_ids, - max_new_tokens=max_new_tokens, - num_beams=args.num_beams, - bos_token_id=tokenizer.bos_token_id, - pad_token_id=tokenizer.pad_token_id, - eos_token_id=tokenizer.eos_token_id, - debug_mode=args.debug_mode, - return_dict=return_dict, - attention_mask=tokenized_inputs.attention_mask, - time_encoder=True, - return_encoder_output=args.output_npy and tensorrt_llm.mpi_rank() == 0) - torch.cuda.synchronize() - tok = time.time() - - if tensorrt_llm.mpi_rank() == 0: - if return_dict: - tllm_output_ids = tllm_output['output_ids'] - else: - tllm_output_ids = tllm_output - - output_ids = tllm_output_ids[:, 0, :] - output_text = tokenizer.batch_decode(output_ids, - skip_special_tokens=True) - decoder_input_lengths = (decoder_input_ids - != tokenizer.pad_token_id).sum(dim=1) - output_gen_lengths = (output_ids != tokenizer.eos_token_id).sum( - dim=1) - decoder_input_lengths - - print( - f"------ TRT-LLM beam = {args.num_beams} --------------------------------" - ) - if 'encoder_output' in tllm_output: - encoder_output = tllm_output['encoder_output'] - print_tensor('TRT-LLM encoder_output:', encoder_output) - print("TRT-LLM output_ids: ", output_ids) - print("TRT-LLM output text: ", output_text) - print("TRT-LLM output generated lengths: ", output_gen_lengths) - print(f"TRT-LLM E2E time {(tok-tik)*1000}ms") - print("Precision:", inference_dtype) - print("--------------------------------------") - - # save input/output tensors for C++ runtime testing - if args.output_npy: - output_npy(args, tokenized_inputs, tllm_output, output_ids) - - # simple accuracy check - if args.compare_hf_fp32: - from difflib import SequenceMatcher - match_rate = SequenceMatcher(None, "\n".join(output_text), - "\n".join(hf_output_text)).ratio() - print(output_text) - print(hf_output_text) - if inference_dtype != "float32": - print("") - print( - f"[CAVEAT] Comparing TRT-LLM {inference_dtype} results with HF float32 results. Close match are not expected!" - ) - assert match_rate > 0.8, f"Incorrect results! Match rate {match_rate}" - else: - assert match_rate > 0.95, f"Incorrect results! Match rate {match_rate}" - print( - f"TRT-LLM results match HF FP32 results with literal match rate {match_rate}" - ) diff --git a/examples/models/core/gemma/README.md b/examples/models/core/gemma/README.md deleted file mode 100644 index 7c94f7ca52ff..000000000000 --- a/examples/models/core/gemma/README.md +++ /dev/null @@ -1,895 +0,0 @@ -# Run Gemma on TensorRT-LLM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -## Table Of Contents - -- [Run Gemma on TensorRT-LLM](#run-gemma-on-tensorrt-llm) - - [Table Of Contents](#table-of-contents) - - [Support Matrix](#support-matrix) - - [Common scripts](#common-scripts) - - [Convert checkpoint](#convert-checkpoint) - - [Build engine](#build-engine) - - [Run inference](#run-inference) - - [Specific commands](#specific-commands) - - [Run Gemma 2B](#run-gemma-2b) - - [Run inference under bfloat16 for HF checkpoint](#run-inference-under-bfloat16-for-hf-checkpoint) - - [Run inference under FP8 for keras checkpoint](#run-inference-under-fp8-for-keras-checkpoint) - - [Run 2B inference under SmoothQuant for jax checkpoint](#run-2b-inference-under-smoothquant-for-jax-checkpoint) - - [Run inference under weight only for jax checkpoint](#run-inference-under-weight-only-for-jax-checkpoint) - - [Run inference under INT8 KV caches for jax checkpoint](#run-inference-under-int8-kv-caches-for-jax-checkpoint) - - [Run Gemma 7B](#run-gemma-7b) - - [Run inference under bfloat16 for torch checkpoint](#run-inference-under-bfloat16-for-torch-checkpoint) - - [Run inference under FP8 for jax checkpoint](#run-inference-under-fp8-for-jax-checkpoint) - - [Run 7B inference under SmoothQuant for jax checkpoint](#run-7b-inference-under-smoothquant-for-jax-checkpoint) - - [Run inference under weight only for keras checkpoint](#run-inference-under-weight-only-for-keras-checkpoint) - - [Run inference under INT8 KV caches for keras checkpoint](#run-inference-under-int8-kv-caches-for-keras-checkpoint) - - [Run Gemma 2](#run-gemma-2) - - [Run inference under bfloat16 for torch checkpoint](#run-inference-under-bfloat16-for-torch-checkpoint-1) - - [Run Gemma 3](#run-gemma-3) - - [Run inference under bfloat16 for HF checkpoint](#run-inference-under-bfloat16-for-hf-checkpoint-1) - - [Disaggregated Serving](#disaggregated-serving) - - [Dynamo](#dynamo) - - [Run Gemma 4](#run-gemma-4) - - [Serve with `trtllm-serve` (OpenAI-compatible API)](#serve-with-trtllm-serve-openai-compatible-api) - - [Accuracy evaluation with `trtllm-eval`](#accuracy-evaluation-with-trtllm-eval) - - [Run Modelopt Quantization](#run-modelopt-quantization) - - [Requirements](#requirements) - - [Quantize Checkpoints](#quantize-checkpoints) - - [Build Engines](#build-engines) - - [Accuracy Results on MMLU](#accuracy-results-on-mmlu) - -## Support Matrix - * FP32/FP16/BF16/INT8 Weight-Only/INT4 AWQ/SmoothQuant/FP8 - * For SmoothQuant, TRT-LLM only supports FP16 higher precision now. - * checkpoint type: Jax, Torch, Keras, Huggingface (HF) - * STRONGLY TYPED - * python runtime and triton backend - -## Common scripts - -### Convert checkpoint - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -Users can use `convert_checkpoint.py` to convert the different source checkpoint to unified TensorRT LLM checkpoint format. Users could set `--dtype` to determine the inference data type, and set the quantization options like `--enable_fp8`, `--fp8_kv_cache` `--use_smooth_quant`, `--calibrate_kv_cache` (for INT8 kv cache) and `--use-weight-only-with-precision` (weight only). Users could also control the source checkpoint type by `--ckpt-type`. Currently, supported checkpoint types are `jax`, `torch` and `keras`. - -```bash -CKPT_PATH=/tmp/models/gemma_nv/checkpoints/tmp_2b_it -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/bf16/tp1/ - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} -``` - -### Build engine - -After getting checkpoint, we can use `trtllm-build` command to build TensorRT LLM engines from TensorRT LLM checkpoints. - -```bash -ENGINE_PATH=/tmp/gemma/2B/bf16/1-gpu/ -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} -``` - -### Run inference - -We provide three examples to run inference `run.py`, `summarize.py` and `mmlu.py`. `run.py` only run inference with `input_text` and show the output. - -`summarize.py` runs summarization on [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset and evaluate the model by [ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. - -`mmlu.py` runs MMLU to evaluate the model by accuracy. - -Note that we need to download the dataset of MMLU first and the evaluation of MMLU requires more time. - -* run.py - -```bash -VOCAB_FILE_PATH=/tmp/models/gemma_nv/checkpoints/tmp_vocab.model -python3 ../../../run.py --engine_dir ${ENGINE_PATH} \ - --max_output_len 30 \ - --vocab_file ${VOCAB_FILE_PATH} - -[TensorRT-LLM] TensorRT LLM version: 0.9.0.dev2024020600Input [Text 0]: " Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: "chef in the renowned kitchens of Lyon. After honing his skills in various Michelin-starred establishments, he embarked on a solo venture, establishing his own restaurant" -``` - -* summarize.py - -```bash -python3 ../../../summarize.py --test_trt_llm \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 \ - --vocab_file ${VOCAB_FILE_PATH} - -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM (total latency: 3.2821836471557617 sec) -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1989) -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM (tokens per second: 605.9989975648089) -[02/06/2024-10:08:54] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/06/2024-10:08:55] [TRT-LLM] [I] rouge1 : 26.376388677070615 -[02/06/2024-10:08:55] [TRT-LLM] [I] rouge2 : 7.468157586877296 -[02/06/2024-10:08:55] [TRT-LLM] [I] rougeL : 17.953060795106556 -[02/06/2024-10:08:55] [TRT-LLM] [I] rougeLsum : 22.410938121151652 -``` - -* mmlu.py - -Download the dataset first - -```bash -mkdir data -wget https://people.eecs.berkeley.edu/~hendrycks/data.tar -O data/mmlu.tar -tar -xf data/mmlu.tar -C data -mv data/data data/mmlu -``` - -Evaluate on MMLU dataset. - -```bash -python3 ../../../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.358 - social sciences -Average accuracy 0.359 - other (business, health, misc.) -Average accuracy: 0.329 -``` - -## Specific commands - -In this section, we demonstrate the scripts to convert checkpoint, building engine and run inference on different settings. We will not demonstrate all combinations here because there are too many cases. We choose some important cases to demonstrate. - -### Run Gemma 2B - -#### Run inference under bfloat16 for HF checkpoint - -```bash -git clone git@hf.co:google/gemma-2b -CKPT_PATH=gemma-2b/ -UNIFIED_CKPT_PATH=/tmp/ckpt/hf/gemma/2b/1-gpu/ -ENGINE_PATH=/tmp/engines/gemma/2B/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-2b/ - -python3 ./convert_checkpoint.py \ - --ckpt-type hf \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM (total latency: 3.0897433757781982 sec) -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2141) -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM (tokens per second: 692.9378073221881) -[03/05/2024-02:24:39] [TRT-LLM] [I] TensorRT LLM beam 0 result -[03/05/2024-02:24:39] [TRT-LLM] [I] rouge1 : 21.042873132085678 -[03/05/2024-02:24:39] [TRT-LLM] [I] rouge2 : 6.322669223228836 -[03/05/2024-02:24:39] [TRT-LLM] [I] rougeL : 16.450116567540338 -[03/05/2024-02:24:39] [TRT-LLM] [I] rougeLsum : 18.836567173262736 -``` - -#### Run inference under FP8 for keras checkpoint - -WARNING: This way of running FP8 will introduce noticeable accuracy drop. To avoid that, use Modelopt quantization mentioned in this readme. - -In this example, we demonstrate how to run FP8 inference on Gemma. Note that `convert_checkpoint.py` only uses identity activation scales, so the accuracy might be little worse than higher precision in some cases, but it is still very good because we don't do any calibration. This also shows the stability of FP8 compared to INT8. - -```bash -git clone git@hf.co:google/gemma-2b-it-keras -GIT_LFS_SKIP_SMUDGE=1 git clone git@hf.co:google/gemma-2b-it-flax # clone tokenizer model -cd gemma-2b-it-flax -git lfs pull -I tokenizer.model - -CKPT_PATH=gemma-2b-it-keras -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_en_tensorrt_llm/fp8/tp1/ -ENGINE_PATH=/tmp/gemma/2B/fp8/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type keras \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --enable_fp8 \ - --fp8_kv_cache \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM (total latency: 3.116227149963379 sec) -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2419) -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM (tokens per second: 776.259201781368) -[02/08/2024-10:37:15] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-10:37:15] [TRT-LLM] [I] rouge1 : 20.206082692133098 -[02/08/2024-10:37:15] [TRT-LLM] [I] rouge2 : 5.902141189518428 -[02/08/2024-10:37:15] [TRT-LLM] [I] rougeL : 15.403458457907643 -[02/08/2024-10:37:15] [TRT-LLM] [I] rougeLsum : 17.44535527417846 - -python3 ../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.390 - social sciences -Average accuracy 0.405 - other (business, health, misc.) -Average accuracy: 0.356 -``` - -#### Run 2B inference under SmoothQuant for jax checkpoint - -```bash -git clone git@hf.co:google/gemma-2b-it-flax -CKPT_PATH=gemma-2b-it-flax/2b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/sq/tp1 -ENGINE_PATH=/tmp/gemma/2B/int8_sq/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype float16 \ - --use_smooth_quant_plugin 0.5 \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM (total latency: 3.460859775543213 sec) -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1786) -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM (tokens per second: 516.0567361385428) -[02/08/2024-04:42:06] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-04:42:06] [TRT-LLM] [I] rouge1 : 22.534044843245525 -[02/08/2024-04:42:06] [TRT-LLM] [I] rouge2 : 5.940093176022924 -[02/08/2024-04:42:06] [TRT-LLM] [I] rougeL : 16.258991712579736 -[02/08/2024-04:42:06] [TRT-LLM] [I] rougeLsum : 19.60977626046262 -``` - -#### Run inference under weight only for jax checkpoint - -Available precisions: `int8` and `int4` - -Note that `int4-weight-only` might not be able to keep the accuracies on all models. If users want to use int4 to run inference, we recommend using `int4_awq`. - -* `int8` - -```bash -git clone git@hf.co:google/gemma-2b-it-flax -CKPT_PATH=gemma-2b-it-flax/2b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/w8_a16/tp1/ -ENGINE_PATH=/tmp/gemma/2B/w8_a16/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --use-weight-only-with-precision int8 \ - --dtype bfloat16 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM (total latency: 3.5987987518310547 sec) -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1797) -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM (tokens per second: 499.3332842203787) -[02/08/2024-04:44:54] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-04:44:54] [TRT-LLM] [I] rouge1 : 24.48521318679745 -[02/08/2024-04:44:54] [TRT-LLM] [I] rouge2 : 7.240543314565931 -[02/08/2024-04:44:54] [TRT-LLM] [I] rougeL : 17.857921729984078 -[02/08/2024-04:44:54] [TRT-LLM] [I] rougeLsum : 21.214162155642896 -``` - -#### Run inference under INT8 KV caches for jax checkpoint - -```bash -git clone git@hf.co:google/gemma-2b-it-flax -CKPT_PATH=gemma-2b-it-flax/2b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_2b_it_tensorrt_llm/int8kv/tp1 -ENGINE_PATH=/tmp/gemma/2B/int8kv/1-gpu/ -VOCAB_FILE_PATH=gemma-2b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --world-size 1 \ - --dtype bfloat16 \ - --calibrate_kv_cache \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM (total latency: 3.5348474979400635 sec) -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM (total output tokens: 1819) -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM (tokens per second: 514.5907994786265) -[02/08/2024-04:52:22] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-04:52:22] [TRT-LLM] [I] rouge1 : 24.0397941580232 -[02/08/2024-04:52:22] [TRT-LLM] [I] rouge2 : 7.325311340360227 -[02/08/2024-04:52:22] [TRT-LLM] [I] rougeL : 17.54210044633271 -[02/08/2024-04:52:22] [TRT-LLM] [I] rougeLsum : 20.627861723682177 -``` - -### Run Gemma 7B - -#### Run inference under bfloat16 for torch checkpoint - -Since torch model does not have model config, we need to add it manually in `CKPT_PATH` with file name `config.json`. - -```bash -git clone git@hf.co:google/gemma-7b-pytorch - -CKPT_PATH=gemma-7b-pytorch/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/bf16/tp1/ -ENGINE_PATH=/tmp/gemma/7B/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-7b-pytorch/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type torch \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -python3 ../../../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.739 - social sciences -Average accuracy 0.697 - other (business, health, misc.) -Average accuracy: 0.630 -``` - -#### Run inference under FP8 for jax checkpoint - -WARNING: This way of running FP8 will introduce noticeable accuracy drop. To avoid that, use Modelopt quantization mentioned in this readme. - -In this example, we demonstrate how to run FP8 inference on Gemma. Note that `convert_checkpoint.py` only uses identity activation scales, so the accuracy might be little worse than higher precision in some cases, but it is still very good because we don't do any calibration. This also shows the stability of FP8 compared to INT8. - -```bash -CKPT_PATH=/tmp/models/gemma_nv/checkpoints/tmp_7b_it -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/fp8/tp1/ -ENGINE_PATH=/tmp/gemma/7B/fp8/1-gpu/ -VOCAB_FILE_PATH=/tmp/models/gemma_nv/checkpoints/tmp_vocab.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --enable_fp8 \ - --fp8_kv_cache \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM (total latency: 5.884302377700806 sec) -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2694) -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM (tokens per second: 457.8282737830064) -[02/08/2024-06:42:13] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-06:42:13] [TRT-LLM] [I] rouge1 : 27.18633861010837 -[02/08/2024-06:42:13] [TRT-LLM] [I] rouge2 : 7.734928823230158 -[02/08/2024-06:42:13] [TRT-LLM] [I] rougeL : 19.32537431798716 -[02/08/2024-06:42:13] [TRT-LLM] [I] rougeLsum : 22.82522575944535 -``` - -#### Run 7B inference under SmoothQuant for jax checkpoint - -```bash -git clone git@hf.co:google/gemma-7b-it-flax -CKPT_PATH=gemma-7b-it-flax/7b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/sq/tp1 -ENGINE_PATH=/tmp/gemma/7B/int8_sq/1-gpu/ -VOCAB_FILE_PATH=gemma-7b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type jax \ - --model-dir ${CKPT_PATH} \ - --dtype float16 \ - --use_smooth_quant_plugin 0.5 \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/19/2024-10:02:53] [TRT-LLM] [I] --------------------------------------------------------- -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM (total latency: 13.65670919418335 sec) -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM (total output tokens: 8351) -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM (tokens per second: 611.494312521266) -[02/19/2024-10:03:09] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/19/2024-10:03:09] [TRT-LLM] [I] rouge1 : 28.8107815115074 -[02/19/2024-10:03:09] [TRT-LLM] [I] rouge2 : 8.623835512061866 -[02/19/2024-10:03:09] [TRT-LLM] [I] rougeL : 19.7277195532959 -[02/19/2024-10:03:09] [TRT-LLM] [I] rougeLsum : 23.434950511855114 -``` - -#### Run inference under weight only for keras checkpoint - -Available precisions: `int8` and `int4` - -Note that `int4-weight-only` might not be able to keep the accuracies on all models. If users want to use int4 to run inference, we recommend using `int4_awq`. - -* `int8` - -```bash -git clone git@hf.co:google/gemma-7b-it-keras -GIT_LFS_SKIP_SMUDGE=1 git clone git@hf.co:google/gemma-7b-it-flax # clone tokenizer model -cd gemma-7b-it-flax -git lfs pull -I tokenizer.model - -CKPT_PATH=gemma-7b-it-keras -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/w8_a16/tp1/ -ENGINE_PATH=/tmp/gemma/7B/w8_a16/1-gpu/ -VOCAB_FILE_PATH=gemma-7b-it-flax/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type keras \ - --model-dir ${CKPT_PATH} \ - --use-weight-only-with-precision int8 \ - --dtype bfloat16 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM (total latency: 8.49835753440857 sec) -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2654) -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM (tokens per second: 312.2956393931832) -[02/08/2024-07:38:15] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-07:38:16] [TRT-LLM] [I] rouge1 : 20.396209981234687 -[02/08/2024-07:38:16] [TRT-LLM] [I] rouge2 : 5.73302850102211 -[02/08/2024-07:38:16] [TRT-LLM] [I] rougeL : 16.001683776127507 -[02/08/2024-07:38:16] [TRT-LLM] [I] rougeLsum : 18.36957526315223 -``` - -#### Run inference under INT8 KV caches for keras checkpoint - -```bash -CKPT_PATH=/tmp/models/gemma_keras/keras/gemma_7b_en/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_7b_it_tensorrt_llm/int8kv/tp1 -ENGINE_PATH=/tmp/gemma/7B/int8kv/1-gpu/ -VOCAB_FILE_PATH=/tmp/models/gemma_nv/checkpoints/tmp_vocab.model - -python3 ./convert_checkpoint.py \ - --ckpt-type keras \ - --model-dir ${CKPT_PATH} \ - --world-size 1 \ - --dtype bfloat16 \ - --calibrate_kv_cache \ - --tokenizer_dir ${VOCAB_FILE_PATH} \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 32 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM (total latency: 8.73880124092102 sec) -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM (total output tokens: 2771) -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM (tokens per second: 317.09154649544956) -[02/08/2024-07:51:11] [TRT-LLM] [I] TensorRT LLM beam 0 result -[02/08/2024-07:51:11] [TRT-LLM] [I] rouge1 : 20.934864626327627 -[02/08/2024-07:51:11] [TRT-LLM] [I] rouge2 : 4.954721611692932 -[02/08/2024-07:51:11] [TRT-LLM] [I] rougeL : 15.307592049634444 -[02/08/2024-07:51:11] [TRT-LLM] [I] rougeLsum : 17.94213019528988 -``` - - -### Run Gemma 2 - -Gemma 2 currently has following limitations: - - Only HF style checkpoints are supported. - - The maximum sequence length allowed is 4096. -#### Run inference under bfloat16 for torch checkpoint -```bash -variant=9b # 27b -git clone git@hf.co:google/gemma-2-$variant-it - -CKPT_PATH=gemma-2-$variant-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_$variant_it_tensorrt_llm/bf16/tp1/ -ENGINE_PATH=/tmp/gemma2/$variant/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-2-$variant-it/tokenizer.model - -python3 ./examples/models/core/gemma/convert_checkpoint.py \ - --ckpt-type hf \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 8 \ - --max_ite 5 - -python3 ../../../mmlu.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} - -Average accuracy 0.739 - social sciences -Average accuracy 0.697 - other (business, health, misc.) -Average accuracy: 0.630 -``` - -### Run Gemma 3 - -Gemma 3's text generation model from HuggingFace is supported. Gemma3 1B model interleaves 5 local layers between each global layer. While local layers use sliding window attention with a short span of 512 tokens, global layers attend to the long context. TRTLLM support layerwise sliding-window attention and the sliding window size for each layer could be passed in using the `--max_attention_window_size` parameter at runtime. If a subpattern is provided, TRTLLM can extrapolate the complete pattern and the extrapolation logic is printed to terminal. - -#### Run inference under bfloat16 for HF checkpoint -```bash -git clone https://huggingface.co/google/gemma-3-1b-it - -CKPT_PATH=gemma-3-1b-it/ -UNIFIED_CKPT_PATH=/tmp/checkpoints/tmp_1b_it_tensorrt_llm/bf16/tp1/ -ENGINE_PATH=/tmp/gemma3/1b/bf16/1-gpu/ -VOCAB_FILE_PATH=gemma-3-1b-it/tokenizer.model - -python3 ./convert_checkpoint.py \ - --ckpt-type hf \ - --model-dir ${CKPT_PATH} \ - --dtype bfloat16 \ - --world-size 1 \ - --output-model-dir ${UNIFIED_CKPT_PATH} - -trtllm-build --checkpoint_dir ${UNIFIED_CKPT_PATH} \ - --gemm_plugin auto \ - --max_batch_size 8 \ - --max_input_len 3000 \ - --max_seq_len 3100 \ - --output_dir ${ENGINE_PATH} - -python3 ../../../summarize.py --test_trt_llm \ - --vocab_file ${VOCAB_FILE_PATH} \ - --engine_dir ${ENGINE_PATH} \ - --batch_size 1 \ - --max_ite 5 \ - --max_attention_window_size 512 512 512 512 512 3100 - -... -[TensorRT-LLM][INFO] TRTGptModel mMaxAttentionWindowSize: (512, 512, 512, 512, 512, 3100) * 4 + (512, 512) -... -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM (total latency: 1.6197962760925293 sec) -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM (total output tokens: 475) -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM (tokens per second: 293.2467539349165) -[04/09/2025-18:28:26] [TRT-LLM] [I] TensorRT LLM beam 0 result -[04/09/2025-18:28:26] [TRT-LLM] [I] rouge1: 22.780314381954003 -[04/09/2025-18:28:26] [TRT-LLM] [I] rouge2: 4.331099231480823 -[04/09/2025-18:28:26] [TRT-LLM] [I] rougeL: 15.26751867562475 -[04/09/2025-18:28:26] [TRT-LLM] [I] rougeLsum: 20.14696930976001 -``` - -#### Disaggregated Serving - -To serve the model in disaggregated mode, you should launch context and generation servers using `trtllm-serve`. - -For example, you can launch a single context server on port 8001 with: - -```bash -export TRTLLM_USE_UCX_KVCACHE=1 - -cat >./ctx_config.yml < output_ctx_8001 & -``` - -Then launch a single generation server on port 8002 with: - -```bash -cat >./gen_config.yml < output_gen_8002 & -``` - -Finally, you can launch the disaggregated server which will accept requests from the client and do -the orchestration between the context and generation servers with: - -```bash -cat >./disagg-config.yaml < argparse.Namespace: - parser = argparse.ArgumentParser() - parser.add_argument("--ckpt-type", - type=CheckpointType, - choices=list(CheckpointType)) - parser.add_argument("--model-dir", "--model_dir", type=Path, required=True) - parser.add_argument("--output-model-dir", - "--output_dir", - type=Path, - required=True) - parser.add_argument("--world-size", - type=int, - default=1, - help="world size, only support tensor parallelism now") - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - "--use-weight-only-with-precision", - choices=["int8", "int4", "w4a8_awq", "w4a16_awq"], - help= - "help='Quantize weights for the various GEMMs to INT4/INT8. Define the precision for the weights.", - ) - parser.add_argument( - "--use-int8-weight-only-embedding", - action="store_true", - help= - "Use weight only on embedding table and lm_head. (Only supported on Hopper GPU)", - ) - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument( - "--enable_fp8", - action="store_true", - help="Use FP8 Linear layer for Attention QKV/Dense and MLP.") - parser.add_argument( - "--fp8_kv_cache", - action="store_true", - help= - "By default, we use dtype for KV cache. fp8_kv_cache chooses fp8 quantization for KV", - ) - parser.add_argument( - "--quant_ckpt_path", - default=None, - help= - "Path of a directory to quantized model checkpoints in .safetensors format or \ - path of a quantized model checkpoint in .npz format") - parser.add_argument( - '--calib_dataset', - type=str, - default='ccdv/cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument('--use_smooth_quant', - action="store_true", - help="Use smooth quant.") - parser.add_argument( - "--int8_kv_cache", - "--calibrate_kv_cache", - "-kv", - action="store_true", - help= - "Generate scaling factors for KV cache. Used for storing KV cache in int8." - ) - parser.add_argument( - '--per_channel', - action="store_true", - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - action="store_true", - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - "--smoothquant", - "--use_smooth_quant_plugin", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument( - '--tokenizer_dir', - default=None, - help='tokenizer path; defaults to jax_model_dir if left unspecified') - parser.add_argument("--load_model_on_cpu", action="store_true") - args = parser.parse_args() - - if args.use_weight_only: - assert args.weight_only_precision is not None - return args - - -CKPT_PARSER: Dict[CheckpointType, Type[Parsers]] = { - CheckpointType.jax: JAXParser, - CheckpointType.keras: KerasParser, - CheckpointType.torch: TorchParser, - CheckpointType.hf: HfParser -} - - -def compute_quant_algo(args: argparse.Namespace) -> Optional[QuantAlgo]: - if args.weight_only_precision: - return { - "int8": QuantAlgo.W8A16, - "int4": QuantAlgo.W4A16, - "w4a8_awq": QuantAlgo.W4A8_AWQ, - "w4a16_awq": QuantAlgo.W4A16_AWQ, - }[args.weight_only_precision] - elif args.enable_fp8: - return QuantAlgo.FP8 - if args.use_smooth_quant: - return QuantAlgo.W8A8_SQ_PER_CHANNEL - elif args.smoothquant is not None: - if args.per_token and args.per_channel: - return QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - elif not args.per_token and not args.per_channel: - return QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - elif not args.per_token and args.per_channel: - return QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - elif args.per_token and not args.per_channel: - return QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - return None - - -def create_quant_config(args: argparse.Namespace) -> QuantConfig: - quant_algo = compute_quant_algo(args) - GemmaForCausalLM.assert_valid_quant_algo(quant_algo) - quant_config = QuantConfig(quant_algo=quant_algo) - if args.smoothquant is not None: - quant_config.smoothquant_val = args.smoothquant - - if args.fp8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - if args.int8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.INT8 - - if args.weight_only_precision: - use_awq = args.weight_only_precision.endswith("awq") - use_int4 = args.weight_only_precision.endswith("int4") - if use_awq: - quant_config.group_size = 128 - - if use_awq or use_int4 or not args.use_int8_weight_only_embedding: - quant_config.has_zero_point = False - quant_config.pre_quant_scale = True - else: - quant_config.exclude_modules = ['router'] - - return quant_config - - -def main() -> None: - emit_engine_arch_deprecation("convert_checkpoint.py") - args = parse_arguments() - tik = time.time() - quant_config = create_quant_config(args) - - ckpt_parser = CKPT_PARSER[args.ckpt_type]() - - mapping = Mapping( - world_size=args.world_size, - tp_size=args.world_size, - pp_size=1, - ) - """We don't support pipeline parallelism yet for Gemma.""" - - if isinstance(ckpt_parser, HfParser): - trt_llm_config = GemmaConfig.from_hugging_face( - args.model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - ) - else: - print(f"Loading source parameters from {args.model_dir.absolute()}") - ckpt_params = ckpt_parser.load_parameters(args.model_dir) - input_embedding_weights = ckpt_parser.embedding_weights(ckpt_params) - num_embed, _ = input_embedding_weights.shape - ckpt_params_dtype = str(input_embedding_weights.dtype).split(".")[ - -1] # np.bfloat16 -> bfloat16 - ckpt_config = ckpt_parser.get_config(args.model_dir, ckpt_params, - num_embed) - # 2B TransformerConfig(num_layers=18, num_embed=256128, embed_dim=2048, hidden_dim=16384, num_heads=8, head_dim=256, num_kv_heads=1) - # 7B TransformerConfig(...) - - del ckpt_params - - print(f"Source configuration determined from parameters: {ckpt_config}") - - trt_llm_config = tensorrt_llm.models.GemmaConfig( - architecture=GEMMA_ARCHITECTURE, - dtype=infer_dtype(args.dtype, ckpt_params_dtype), - logits_dtype="float32", - vocab_size=ckpt_config.num_embed, - max_position_embeddings=8192, - hidden_size=ckpt_config.embed_dim, - num_hidden_layers=ckpt_config.num_layers, - num_attention_heads=ckpt_config.num_heads, - num_key_value_heads=ckpt_config.num_kv_heads, - head_size=ckpt_config.head_dim, - hidden_act="gelu", - intermediate_size=ckpt_config.hidden_dim, - norm_epsilon=1e-6, # hard-coded in RMSNorm from gemma/layers.py - position_embedding_type="rope_gpt_neox", - mapping=mapping, - gpus_per_node=8, - quantization=quant_config, - use_parallel_embedding=mapping.tp_size > 1, - ) - if hasattr(ckpt_config, - "model_type") and ckpt_config.model_type == "gemma2": - trt_llm_config.inter_layernorms = True - trt_llm_config.final_logit_softcapping = ckpt_config.final_logit_softcapping - trt_llm_config.attn_logit_softcapping = ckpt_config.attn_logit_softcapping - trt_llm_config.query_pre_attn_scalar = ckpt_config.query_pre_attn_scalar - - trt_llm_config_dict = trt_llm_config.to_dict() - print(f"Determined TensorRT LLM configuration {trt_llm_config_dict}") - - save_config(trt_llm_config, output_dir=args.output_model_dir, log=True) - - for config in trt_llm_config.for_each_rank(): - hf_weights = load_gemma_weights( - parameters_or_model_dir=args.model_dir, - trt_llm_config=config, - ckpt_parser=ckpt_parser, - load_model_on_cpu=args.load_model_on_cpu) - ranked_weights = non_modelopt_quantize_if_needed( - hf_weights, - model_dir=args.model_dir, - quantize_modifiers=QuantizeModifiers.from_args(args), - trt_llm_config=config) - save_checkpoint(output_dir=args.output_model_dir, - weights=ranked_weights, - rank=config.mapping.rank) - - elapsed = time.strftime("%H:%M:%S", time.gmtime(time.time() - tik)) - print(f"Total time of converting checkpoints: {elapsed}") - - -if __name__ == "__main__": - main() diff --git a/examples/models/core/gemma/requirements.txt b/examples/models/core/gemma/requirements.txt deleted file mode 100644 index a1bbed25b68a..000000000000 --- a/examples/models/core/gemma/requirements.txt +++ /dev/null @@ -1,14 +0,0 @@ --f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --c ../../../constraints.txt -# WAR the new posting of "nvidia-cudnn-cu12~=9.0". -# "jax[cuda12_pip]~=0.4.19" specifies "nvidia-cudnn-cu12>=8.9" but actually requires "nvidia-cudnn-cu12~=8.9". -nvidia-cudnn-cu12~=8.9; platform_machine == "x86_64" -tensorrt_llm>=0.0.0.dev0 -flax~=0.8.0 -# jax[cuda12_pip]~=0.4.19 -safetensors~=0.4.1 -sentencepiece>=0.1.99 -h5py~=3.12.1 -rouge_score -nltk -datasets==3.1.0 diff --git a/examples/models/core/glm-4-9b/README.md b/examples/models/core/glm-4-9b/README.md deleted file mode 100644 index b2789d59f92f..000000000000 --- a/examples/models/core/glm-4-9b/README.md +++ /dev/null @@ -1,292 +0,0 @@ -# ChatGLM - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [glm-4-9b](https://huggingface.co/THUDM/glm-4-9b) models using TensorRT LLM and run on a single GPU, a single node with multiple GPUs or multiple nodes with multiple GPUs. - -- [glm-4-9b](#glm-4-9b) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Model comparison](#model-comparison) - - [Tokenizer and special tokens comparison](#tokenizer-and-special-tokens-comparison) - - [Usage](#usage) - - [1. Download repo and weights from HuggingFace Transformers](#1-download-repo-and-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [Enable plugins](#enable-plugins) - - [In-flight batching](#in-flight-batching) - - [4. Run inference](#4-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multi GPU](#single-node-multi-gpu) - - [5. Run summarization task](#5-run-summarization-task) - - [Weight Only quantization](#weight-only-quantization) - - [Smooth Quantization (SQ)](#smooth-quantization-sq) - - [Activation-aware Weight Quantization (AWQ)](#activation-aware-weight-quantization-awq) - - [FP8 Quantization](#fp8-quantization) - - [Benchmark](#benchmark) - - -## Overview - -The TensorRT LLM ChatGLM implementation can be found in [`tensorrt_llm/models/chatglm/model.py`](../../../../tensorrt_llm/models/chatglm/model.py). -The TensorRT LLM ChatGLM example code is located in [`examples/models/core/glm-4-9b`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - -| Model Name | FP16 | FMHA | WO | SQ | AWQ | FP8 | TP | PP | ST | C++ | benchmark | IFB | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :-------: | :---: | -| glm_4_9b | Y | Y | Y | | | | Y | | | Y | | | -| glm_4_9b_chat | Y | Y | Y | | | | Y | | | Y | | | - -* Model Name: the name of the model, the same as the name on HuggingFace -* FMHA: Fused MultiHead Attention (see introduction below) -* WO: Weight Only Quantization (int8 / int4) -* SQ: Smooth Quantization (int8) -* AWQ: Activation Aware Weight Quantization (int4) -* FP8: FP8 Quantization -* TP: Tensor Parallel -* PP: Pipeline Parallel -* ST: Strongly Typed -* C++: C++ Runtime -* benchmark: benchmark by python / C++ Runtime -* IFB: In-flight Batching (see introduction below) - -## Model comparison - -| Name | nL | nAH | nKH | nHW | nH | nF | nMSL | nV | bP2D | bBQKV | bBDense | Comments | -| :--------------: | :---: | :---: | :---: | :---: | :---: | :---: | :---: | :----: | :---: | :---: | :-----: | :----------------------------------------------------------------- | -| glm_4_9b | 40 | 32 | 2 | 128 | 4096 | 13696 | 8192 | 151552 | N | Y | N | | -| glm_4_9b_chat | 40 | 32 | 2 | 128 | 4096 | 13696 | 8192 | 151552 | N | Y | N | | - -* nL: number of layers -* nAH: number of attention heads -* nKH: number of kv heads (less than nAH if multi_query_attention is used) -* nHW: head width -* nH: hidden size -* nF: FFN hidden size -* nMSL: max sequence length (input + output) -* nV: vocabulary size -* bP2D: use position_encoding_2d (Y: Yes, N: No) -* bBQKV: use bias for QKV multiplication in self-attention -* bBDense: use bias for Dense multiplication in self-attention - -## Tokenizer and special tokens comparison - -| Name | Tokenizer | bos | eos | pad | cls | startofpiece | endofpiece | mask | smask | gmask | -| :--------------: | :--------------: | :----: | :----: | :---: | :---: | :----------: | :--------: | :----: | :---: | :----: | -| glm_4_9b | ChatGLM4Tokenizer | | 151329 | 151329 | | 151333 | | | | | -| glm_4_9b_chat | ChatGLM4Tokenizer | | 151329 | 151329 | | 151333 | | | | | - -## Usage - -The next section describe how to build the engine and run the inference demo. - -### 1. Download repo and weights from HuggingFace Transformers - -```bash -pip install -r requirements.txt -apt-get update -apt-get install git-lfs - -# clone one or more models we want to build -git clone https://huggingface.co/THUDM/glm-10b glm_10b -git clone https://huggingface.co/THUDM/glm-4-9b glm_4_9b -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format - -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# GLM-4-9B: single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir glm_4_9b --output_dir trt_ckpt/glm_4_9b/fp16/1-gpu -``` - -### 3. Build TensorRT engine(s) - -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The number of engine files is also same to the number of GPUs used to run inference. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# GLM-4-9B: single-gpu engine with dtype float16, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/fp16/1-gpu \ - --gemm_plugin bfloat16 \ - --output_dir trt_engines/glm_4_9b/fp16/1-gpu -``` - -If the engines are run successfully, you will see output like (glm-4-9b as the example): - -```txt -...... -[01/26/2024-02:40:36] [TRT] [I] Engine generation completed in 136.52 seconds. -[01/26/2024-02:40:36] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 1016 MiB, GPU 11909 MiB -[01/26/2024-02:40:36] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +11909, now: CPU 0, GPU 11909 (MiB) -[01/26/2024-02:40:40] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 29706 MiB -[01/26/2024-02:40:40] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:02:20 -[01/26/2024-02:40:42] [TRT-LLM] [I] Serializing engine to trt_engines/glm_4_9b/fp16/1-gpu/rank0.engine... -[01/26/2024-02:42:29] [TRT-LLM] [I] Engine serialized. Total time: 00:01:47 -[01/26/2024-02:42:30] [TRT-LLM] [I] Total time of building all engines: 00:05:19 -``` - -#### Enable plugins - -* Use `--gpt_attention_plugin ` to configure GPT Attention plugin (default as float16) -* Use `--gemm_plugin ` to configure GEMM plugin (default as float16) -* Use `--context_fmha enable` to enable FMHA kernels, which can provide better performance and low GPU memory occupation. - * `--gpt_attention_plugin float16` must be used when using FMHA. - -#### In-flight batching - -* The engine(s) must be built accordingly if [in-flight batching in C++ runtime](../../docs/in_flight_batching.md) will be used. -* Use `--gpt_attention_plugin float16`, `--paged_kv_cache enable`, `--remove_input_padding enable` to build engine(s) supporting In-flight Batching. - * It is possible to use `--gpt_attention_plugin float32` In-flight Batching. - * The size of the block in paged KV cache can be controlled additionally by using `--tokens_per_block=N`. - -### 4. Run inference - -#### Single node, single GPU - -```bash -# Run the default engine of GLM-4-9B on single GPU, other model name is available if built. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp16/1-gpu -``` - -#### Single node, multi GPU - -```bash -# Run the Tensor Parallel 2 engine of glm_4_9b on two GPU, other model name is available if built. -mpirun -n 2 \ - python ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp16/2-gpu -``` - -* `--allow-run-as-root` might be needed if using `mpirun` as root. -* `trtllm-build` flag `--context_fmha enable` uses FP16 accumulator, which might cause low accuracy. In this case, add `--enable_context_fmha_fp32_acc` to the inference command should be used to protect accuracy at a cost of small performance drop. - -If the engines are run successfully, you will see output like (ChatGLM3-6B as the example): - -```txt -...... -Input [Text 0]: "[gMASK]sop What's new between ChatGLM3-6B and ChatGLM2-6B?" -Output [Text 0 Beam 0]: "There is no new information provided in the official documentation, but I found some differences in the code. The ChatGLM3-6B has an additional parameter called 'config', which is not present in ChatGLM2-6B. Additionally" -``` - -### 5. Run summarization task - -```bash -# Run the summarization of glm_4_9b task, other model name is available if built. -python3 ../../../summarize.py --test_trt_llm \ - --hf_model_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp16/1-gpu -``` - -### Weight Only quantization - -Use `--use_weight_only` to enable INT8-Weight-Only quantization, this will significantly lower the latency and memory footprint. Furthermore, use `--weight_only_precision int8` or `--weight_only_precision int4` to configure the data type of the weights. - -```bash -# glm_4_9b: single gpu, int8 weight only quantization -python3 convert_checkpoint.py --model_dir glm_4_9b \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir trt_ckpt/glm_4_9b/int8_wo/1-gpu - -# glm_4_9b: single-gpu engine with int8 weight only quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/int8_wo/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/int8_wo/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/int8_wo/1-gpu -``` - -### Smooth Quantization (SQ) - -Use `--smoothquant` to enable smooth quantization. - -```bash -# glm_4_9b: single gpu, int8 smooth quantization -python3 convert_checkpoint.py --model_dir glm_4_9b \ - --smoothquant 0.5 \ - --per_channel \ - --per_token \ - --output_dir trt_ckpt/glm_4_9b/sq/1-gpu - -# glm_4_9b: single-gpu engine with int8 smooth quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/sq/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/sq/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/sq/1-gpu -``` - -### Activation-aware Weight Quantization (AWQ) - -The [`quantize.py`](../../../quantization/quantize.py) script can be used to quantize the models and export TensorRT LLM checkpoints. - -```bash -# glm_4_9b: single gpu, int4 awq quantization -python ../../../quantization/quantize.py --model_dir glm_4_9b \ - --dtype float16 \ - --qformat int4_awq \ - --output_dir trt_ckpt/glm_4_9b/int4_awq/1-gpu - -# glm_4_9b: single-gpu engine with int4 awq quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/int4_awq/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/int4_awq/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/int4_awq/1-gpu -``` - -### FP8 Quantization - -The [`quantize.py`](../../../quantization/quantize.py) script can be used to quantize the models and export TensorRT LLM checkpoints. - -```bash -# glm_4_9b: single gpu, fp8 quantization -python ../../../quantization/quantize.py --model_dir glm_4_9b \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir trt_ckpt/glm_4_9b/fp8/1-gpu - -# glm_4_9b: single-gpu engine with fp8 quantization, GPT Attention plugin, Gemm plugin -trtllm-build --checkpoint_dir trt_ckpt/glm_4_9b/fp8/1-gpu \ - --gemm_plugin float16 \ - --output_dir trt_engines/glm_4_9b/fp8/1-gpu - -# Run inference. -python3 ../../../run.py --input_text "What's new between ChatGLM3-6B and ChatGLM2-6B?" \ - --max_output_len 50 \ - --tokenizer_dir glm_4_9b \ - --engine_dir trt_engines/glm_4_9b/fp8/1-gpu -``` diff --git a/examples/models/core/glm-4-9b/convert_checkpoint.py b/examples/models/core/glm-4-9b/convert_checkpoint.py deleted file mode 100644 index 4f7f67e03790..000000000000 --- a/examples/models/core/glm-4-9b/convert_checkpoint.py +++ /dev/null @@ -1,271 +0,0 @@ -import argparse -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.logger import logger -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import ChatGLMForCausalLM -from tensorrt_llm.models.chatglm.config import GLM_VERSIONS -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument( - '--chatglm_version', - default=None, - choices=[None] + GLM_VERSIONS, - help= - "By default the script will try to infer the chatglm_version from model_dir. " - "Or users may overwrite chatglm_version by explicitly passing the version." - ) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--cp_size', - type=int, - default=1, - help='N-way context parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument('--logits_dtype', - type=str, - default='float32', - choices=['float16', 'float32']) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument( - '--calib_dataset', - type=str, - default='cnn_dailymail', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument( - '--per_channel', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - action="store_true", - default=False, - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--device', - help= - "The device to run calibration; effective for HuggingFace model only.", - default='cuda', - choices=['cuda', 'cpu']) - parser.add_argument("--load_model_on_cpu", action="store_true") - args = parser.parse_args() - - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - quant_config.quant_algo = QuantAlgo.W4A16 - elif args.smoothquant: - quant_config.smoothquant_val = args.smoothquant - if args.per_channel: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.INT8 - - return quant_config - - -def args_to_build_options(args): - return { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'logits_dtype': args.logits_dtype, - } - - -def execute(workers, func): - if workers == 1: - for rank, f in enumerate(func): - f(rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_hf(args): - world_size = args.tp_size * args.pp_size - quant_config = args_to_quant_config(args) - override_fields = args_to_build_options(args) - - if args.smoothquant is not None or args.int8_kv_cache: - mapping = Mapping(world_size=world_size, - tp_size=args.tp_size, - pp_size=args.pp_size, - cp_size=args.cp_size) - ChatGLMForCausalLM.quantize(args.model_dir, - args.output_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - device=args.device, - calib_dataset=args.calib_dataset, - **override_fields) - else: - - def convert_and_save_rank(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - glm = ChatGLMForCausalLM.from_hugging_face( - args.model_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - chatglm_version=args.chatglm_version, - load_model_on_cpu=args.load_model_on_cpu, - **override_fields) - glm.config.mapping.cp_size = args.cp_size - glm.config.mapping.attn_tp_size = -1 - glm.config.mapping.attn_cp_size = -1 - glm.config.mapping.world_size *= args.cp_size - glm.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del glm - - execute(args.workers, [convert_and_save_rank] * world_size) - release_gc() - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - print(tensorrt_llm.__version__) - args = parse_arguments() - logger.set_level(args.log_level) - - assert args.pp_size == 1, "Pipeline parallelism is not supported." - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - convert_and_save_hf(args) - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/glm-4-9b/requirements.txt b/examples/models/core/glm-4-9b/requirements.txt deleted file mode 100644 index cdc65bf2bb38..000000000000 --- a/examples/models/core/glm-4-9b/requirements.txt +++ /dev/null @@ -1,8 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -protobuf -rouge_score -sentencepiece -tiktoken diff --git a/examples/models/core/glm-4-9b/tokenization_chatglm.py b/examples/models/core/glm-4-9b/tokenization_chatglm.py deleted file mode 100644 index 75495af5f7f7..000000000000 --- a/examples/models/core/glm-4-9b/tokenization_chatglm.py +++ /dev/null @@ -1,343 +0,0 @@ -import base64 -import json -import os -from typing import Any, Dict, List, Optional, Union - -import regex as re -import tiktoken -from torch import TensorType -from transformers import PreTrainedTokenizer -from transformers.tokenization_utils_base import BatchEncoding, EncodedInput -from transformers.utils import PaddingStrategy - - -class ChatGLM4Tokenizer(PreTrainedTokenizer): - vocab_files_names = {"vocab_file": "tokenizer.model"} - model_input_names = ["input_ids", "attention_mask", "position_ids"] - - def __init__(self, - vocab_file, - padding_side="left", - clean_up_tokenization_spaces=False, - encode_special_tokens=False, - **kwargs): - self.name = "GLM4Tokenizer" - self.vocab_file = vocab_file - pat_str = "(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\\r\\n\\p{L}\\p{N}]?\\p{L}+|\\p{N}{1,3}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+" - self.pat_str = re.compile(pat_str) - self.encode_special_tokens = encode_special_tokens - - mergeable_ranks = {} - with open(vocab_file) as f: - for line in f: - token, rank = line.strip().split() - rank = int(rank) - token = base64.b64decode(token) - mergeable_ranks[token] = rank - - self.mergeable_ranks = mergeable_ranks - - self.tokenizer = tiktoken.Encoding(name="my_tokenizer", - pat_str=pat_str, - mergeable_ranks=mergeable_ranks, - special_tokens={}) - self.decoder = {rank: token for token, rank in mergeable_ranks.items()} - self.n_words = len(self.decoder) - - super().__init__( - padding_side=padding_side, - clean_up_tokenization_spaces=clean_up_tokenization_spaces, - **kwargs) - - @property - def vocab_size(self): - return self.n_words - - def get_vocab(self): - """ Returns vocab as a dict """ - vocab = { - self._convert_id_to_token(i): i - for i in range(self.vocab_size) - } - vocab.update(self.added_tokens_encoder) - return vocab - - def convert_tokens_to_string(self, tokens: List[Union[bytes, str, - int]]) -> str: - """ - Converts a sequence of tokens in a single string. - """ - text = "" - temp = b"" - for t in tokens: - if isinstance(t, int): - t = chr(t) - if isinstance(t, str): - if temp: - text += temp.decode("utf-8", errors="replace") - elif isinstance(t, bytes): - temp += t - else: - raise TypeError( - "token should only be of type int, bytes or str") - if temp: - text += temp.decode("utf-8", errors="replace") - return text - - def _tokenize(self, text, **kwargs): - tokens = [] - ids = self.tokenizer.encode(text) - for t in ids: - tokens.append(self.decoder[t]) - return tokens - - def _convert_token_to_id(self, token): - """ Converts a token (str) in an id using the vocab. """ - return self.mergeable_ranks[token] - - def _convert_id_to_token(self, index): - """Converts an index (integer) in a token (str) using the vocab.""" - return self.decoder.get(index, "") - - def save_vocabulary(self, save_directory, filename_prefix=None): - """ - Save the vocabulary and special tokens file to a directory. - - Args: - save_directory (`str`): - The directory in which to save the vocabulary. - filename_prefix (`str`, *optional*): - An optional prefix to add to the named of the saved files. - - Returns: - `Tuple(str)`: Paths to the files saved. - """ - if os.path.isdir(save_directory): - vocab_file = os.path.join(save_directory, - self.vocab_files_names["vocab_file"]) - else: - vocab_file = save_directory - - with open(self.vocab_file, 'rb') as fin: - proto_str = fin.read() - - with open(vocab_file, "wb") as writer: - writer.write(proto_str) - - return (vocab_file, ) - - def get_prefix_tokens(self): - prefix_tokens = [ - self.convert_tokens_to_ids("[gMASK]"), - self.convert_tokens_to_ids("") - ] - return prefix_tokens - - def build_single_message(self, role, metadata, message, tokenize=True): - assert role in ["system", "user", "assistant", "observation"], role - if tokenize: - role_tokens = [self.convert_tokens_to_ids(f"<|{role}|>") - ] + self.tokenizer.encode(f"{metadata}\n", - disallowed_special=()) - message_tokens = self.tokenizer.encode(message, - disallowed_special=()) - tokens = role_tokens + message_tokens - return tokens - else: - return str(f"<|{role}|>{metadata}\n{message}") - - def apply_chat_template( - self, - conversation: Union[List[Dict[str, str]], List[List[Dict[str, str]]], - "Conversation"], - add_generation_prompt: bool = False, - tokenize: bool = True, - padding: bool = False, - truncation: bool = False, - max_length: Optional[int] = None, - return_tensors: Optional[Union[str, TensorType]] = None, - return_dict: bool = False, - tokenizer_kwargs: Optional[Dict[str, Any]] = None, - add_special_tokens: bool = True, - **kwargs, - ) -> Union[str, List[int], List[str], List[List[int]], BatchEncoding]: - - if return_dict and not tokenize: - raise ValueError( - "`return_dict=True` is incompatible with `tokenize=False`, because there is no dict " - "of tokenizer outputs to return.") - - def handle_single_conversation(conversation): - input_ids = self.get_prefix_tokens() if add_special_tokens else [] - input_message = "[gMASK]" if add_special_tokens else "" - for item in conversation: - if item.get("tools"): - tools = item["tools"] - content = "你是一个名为 GhatGLM 的人工智能助手。你是基于智谱AI训练的语言模型 GLM-4 模型开发的，你的任务是针对用户的问题和要求提供适当的答复和支持。" - content += "\n\n# 可用工具" - for tool in tools: - if tool["type"] == "function": - function = tool["function"] - content += f"\n\n## {function['name']}\n\n{json.dumps(function, ensure_ascii=False, indent=4)}" - content += "\n在调用上述函数时，请使用 Json 格式表示调用的参数。" - elif tool["type"] == "python": - content += "\n\n## python\n\n当你向 `python` 发送包含 Python 代码的消息时，该代码将会在一个有状态的 Jupyter notebook 环境中执行。\n`python` 返回代码执行的输出，或在执行 60 秒后返回超时。\n`/mnt/data` 将会持久化存储你的文件。在此会话中，`python` 无法访问互联网。不要使用 `python` 进行任何网络请求或者在线 API 调用，这些在线内容的访问将不会成功。" - elif tool["type"] == "simple_browser": - content += "\n\n## simple_browser\n\n你可以使用 `simple_browser` 工具。该工具支持以下函数：\n`search(query: str, recency_days: int)`：使用搜索引擎进行查询并显示结果，可以使用 `recency_days` 参数控制搜索内容的时效性。\n`mclick(ids: list[int])`：获取一系列指定 id 的页面内容。每次调用时，须选择3-10个页面。选择多个角度的页面，同时尽可能选择可信任的信息来源。考虑到部分页面是无法加载的，你也可以多打开一些可能有用的页面而不用担心内容过多。\n`open_url(url: str)`：打开指定的 URL。\n\n使用 `【{引用 id}†{引用文本}】` 来引用内容。\n\n操作步骤：1. 使用 `search` 来获得信息列表; 2. 使用 `mclick` 来获取指定 ID 页面的内容; 3. 根据获得的内容进行回复。在回复中应当引用信息来源。\n 如果用户提供了 URL，也可以用 `open_url` 直接打开页面。\n如果初次搜索结果没有找到合适的信息，也可以再次使用 `search` 进行搜索。" - elif tool["type"] == "cogview": - content += "\n\n## cogview\n\n如果用户的请求中包含了对图像的描述，你可以使用 `cogview` 来生成图像并展示给用户。你需要向 `cogview` 发送图像描述，规则：\n- 发送给 `cogview` 的消息必须使用英语。用户的中文描述必须完全翻译为英语。\n- 应当尽可能详细地描述图像生成的需求，需求描述约 100 英文单词。\n- 保持用户原始描述的意图。不要虚构内容或者没见过的人物。\n- 如无特殊说明，所在地为中国，持有中国立场并遵循中国社会主义价值观。" - else: - raise NotImplementedError( - f"Unknown tool type {tool['type']}") - input = self.build_single_message("system", - "", - content, - tokenize=tokenize) - if tokenize: - input_ids.extend(input) - else: - input_message += input - if item["content"]: - input = self.build_single_message(item["role"], - item.get("metadata", ""), - item["content"], - tokenize=tokenize) - if tokenize: - input_ids.extend(input) - else: - input_message += input - if add_generation_prompt: - if tokenize: - input_ids.extend( - [self.convert_tokens_to_ids("<|assistant|>")]) - else: - input_message += "<|assistant|>" - return input_ids if tokenize else input_message - - # Main logic to handle different conversation formats - if isinstance(conversation, list) and all( - isinstance(i, dict) for i in conversation): - result = handle_single_conversation(conversation) - elif isinstance(conversation, list) and all( - isinstance(i, list) for i in conversation): - result = [handle_single_conversation(c) for c in conversation] - elif hasattr(conversation, "messages"): - result = handle_single_conversation(conversation.messages) - else: - raise ValueError("Invalid conversation format") - - if tokenize: - output = self.batch_encode_plus( - [result] if isinstance(result[0], int) else result, - padding=padding, - truncation=truncation, - max_length=max_length, - return_tensors=return_tensors, - is_split_into_words=True, - add_special_tokens=False) - if return_dict: - return output - else: - return output["input_ids"] - else: - return result - - def build_inputs_with_special_tokens( - self, - token_ids_0: List[int], - token_ids_1: Optional[List[int]] = None) -> List[int]: - """ - Build model inputs from a sequence or a pair of sequence for sequence classification tasks by concatenating and - adding special tokens. A BERT sequence has the following format: - - - single sequence: `[CLS] X [SEP]` - - pair of sequences: `[CLS] A [SEP] B [SEP]` - - Args: - token_ids_0 (`List[int]`): - List of IDs to which the special tokens will be added. - token_ids_1 (`List[int]`, *optional*): - Optional second list of IDs for sequence pairs. - - Returns: - `List[int]`: List of [input IDs](../glossary#input-ids) with the appropriate special tokens. - """ - prefix_tokens = self.get_prefix_tokens() - token_ids_0 = prefix_tokens + token_ids_0 - if token_ids_1 is not None: - token_ids_0 = token_ids_0 + token_ids_1 + [ - self.convert_tokens_to_ids("") - ] - return token_ids_0 - - def _pad( - self, - encoded_inputs: Union[Dict[str, EncodedInput], BatchEncoding], - max_length: Optional[int] = None, - padding_strategy: PaddingStrategy = PaddingStrategy.DO_NOT_PAD, - pad_to_multiple_of: Optional[int] = None, - return_attention_mask: Optional[bool] = None, - padding_side: str = "left", # Fix for new transformers - ) -> dict: - """ - Pad encoded inputs (on left/right and up to predefined length or max length in the batch) - - Args: - encoded_inputs: - Dictionary of tokenized inputs (`List[int]`) or batch of tokenized inputs (`List[List[int]]`). - max_length: maximum length of the returned list and optionally padding length (see below). - Will truncate by taking into account the special tokens. - padding_strategy: PaddingStrategy to use for padding. - - - PaddingStrategy.LONGEST Pad to the longest sequence in the batch - - PaddingStrategy.MAX_LENGTH: Pad to the max length (default) - - PaddingStrategy.DO_NOT_PAD: Do not pad - The tokenizer padding sides are defined in self.padding_side: - - - 'left': pads on the left of the sequences - - 'right': pads on the right of the sequences - pad_to_multiple_of: (optional) Integer if set will pad the sequence to a multiple of the provided value. - This is especially useful to enable the use of Tensor Core on NVIDIA hardware with compute capability - `>= 7.5` (Volta). - return_attention_mask: - (optional) Set to False to avoid returning attention mask (default: set to model specifics) - """ - # Load from model defaults - assert self.padding_side == "left" - - required_input = encoded_inputs[self.model_input_names[0]] - seq_length = len(required_input) - - if padding_strategy == PaddingStrategy.LONGEST: - max_length = len(required_input) - - if max_length is not None and pad_to_multiple_of is not None and ( - max_length % pad_to_multiple_of != 0): - max_length = ( - (max_length // pad_to_multiple_of) + 1) * pad_to_multiple_of - - needs_to_be_padded = padding_strategy != PaddingStrategy.DO_NOT_PAD and len( - required_input) != max_length - - # Initialize attention mask if not present. - if "attention_mask" not in encoded_inputs: - encoded_inputs["attention_mask"] = [1] * seq_length - - if "position_ids" not in encoded_inputs: - encoded_inputs["position_ids"] = list(range(seq_length)) - - if needs_to_be_padded: - difference = max_length - len(required_input) - - if "attention_mask" in encoded_inputs: - encoded_inputs["attention_mask"] = [ - 0 - ] * difference + encoded_inputs["attention_mask"] - if "position_ids" in encoded_inputs: - encoded_inputs["position_ids"] = [ - 0 - ] * difference + encoded_inputs["position_ids"] - encoded_inputs[self.model_input_names[ - 0]] = [self.pad_token_id] * difference + required_input - - return encoded_inputs diff --git a/examples/models/core/gpt/.gitignore b/examples/models/core/gpt/.gitignore deleted file mode 100644 index 4c73534033eb..000000000000 --- a/examples/models/core/gpt/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -gpt* -*.log -c-model -*.safetensors -*.json diff --git a/examples/models/core/gpt/README.md b/examples/models/core/gpt/README.md deleted file mode 100644 index f50f45e75d67..000000000000 --- a/examples/models/core/gpt/README.md +++ /dev/null @@ -1,786 +0,0 @@ -# GPT - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document explains how to build the [GPT](https://huggingface.co/gpt2) model using TensorRT LLM and run on a single GPU, a single node with multiple GPUs or multiple nodes with multiple GPUs. - -- [GPT](#gpt) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [1. Download weights from HuggingFace Transformers](#1-download-weights-from-huggingface-transformers) - - [2. Convert weights from HF Transformers to TensorRT LLM format](#2-convert-weights-from-hf-transformers-to-tensorrt-llm-format) - - [3. Build TensorRT engine(s)](#3-build-tensorrt-engines) - - [Fused MultiHead Attention (FMHA)](#fused-multihead-attention-fmha) - - [In-flight batching and paged KV cache](#in-flight-batching-and-paged-kv-cache) - - [4. Build TensorRT engine(s) with Random Weights](#4-build-tensorrt-engines-with-random-weights) - - [5. Run inference](#5-run-inference) - - [Single node, single GPU](#single-node-single-gpu) - - [Single node, multiple GPUs](#single-node-multiple-gpus) - - [Multiple nodes, multiple GPUs using Slurm](#multiple-nodes-multiple-gpus-using-slurm) - - [Quantization](#quantization) - - [SmoothQuant](#smoothquant) - - [Model Transformation](#model-transformation) - - [INT8 inference](#int8-inference) - - [Usage](#usage-1) - - [INT8 KV Cache](#int8-kv-cache) - - [Weight Only Quantization](#weight-only-quantization) - - [FP8 Quantization](#fp8-quantization) - - [Embedding Parallelism](#embedding-parallelism) - - [1. Embedding parallelism](#1-embedding-parallelism) - - [2. The sharding dimension for embedding parallelism](#2-the-sharding-dimension-for-embedding-parallelism) - - [GPT Variant - Granite(20B and 34B)](#gpt-variant---granite) - - [GPT Variant - SantaCoder](#gpt-variant---santacoder) - - [GPT Variant - StarCoder (v1 and v2)](#gpt-variant---starcoder-v1-and-v2) - - [Run StarCoder2 with LoRA](#run-starcoder2-with-lora) - - [GPT-Next](#gpt-next) - - [Prompt-tuning](#prompt-tuning) - - [MultiLoRA with the Nemo checkpoint](#multilora-with-the-nemo-checkpoint) - -## Overview - -The TensorRT LLM GPT implementation can be found in [`tensorrt_llm/models/gpt/model.py`](../../../../tensorrt_llm/models/gpt/model.py). The TensorRT LLM GPT example code is located in [`examples/models/core/gpt`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert a checkpoint from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) format to the TensorRT LLM format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 - * FP8 - * Inflight Batching - * PAGED_KV_CACHE - * FP8 KV CACHE - * Tensor Parallel - * Pipeline Parallel - * STRONGLY TYPED - * INT8 SmoothQuant - * INT8 weight only - * INT4 weight only - -## Usage - -The next two sections describe how to convert the weights from the [HuggingFace (HF) Transformers](https://github.com/huggingface/transformers) -format to the TensorRT LLM format. - -### 1. Download weights from HuggingFace Transformers - -Please install required packages first: - -```bash -pip install -r requirements.txt -``` - -```bash -# Download hf gpt2 model -rm -rf gpt2 && git clone https://huggingface.co/gpt2-medium gpt2 -pushd gpt2 && rm pytorch_model.bin model.safetensors && wget -q https://huggingface.co/gpt2-medium/resolve/main/pytorch_model.bin && popd -``` - -### 2. Convert weights from HF Transformers to TensorRT LLM format -The [`convert_checkpoint.py`](./convert_checkpoint.py) script converts HF weights to TensorRT LLM checkpoints. The number of checkpoint files (in .safetensors format) is same to the number of GPUs used to run inference. - -```bash -# single gpu, dtype float16 -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --output_dir gpt2/trt_ckpt/fp16/1-gpu - -# 2-way tensor parallelism -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --output_dir gpt2/trt_ckpt/fp16/2-gpu - -# 2-way tensor parallelism and 2-way pipeline parallelism -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --pp_size 2 \ - --output_dir gpt2/trt_ckpt/fp16/4-gpu -``` - -### 3. Build TensorRT engine(s) -The `trtllm-build` command builds TensorRT LLM engines from TensorRT LLM checkpoints. The checkpoint directory provides the model's weights and architecture configuration. The number of engine files is also same to the number of GPUs used to run inference. - -`trtllm-build` command has a variety of options. In particular, the plugin-related options have two categories: -* Plugin options that requires a data type (e.g., `gpt_attention_plugin`), you can - * explicitly specify `float16`/`bfloat16`/`float32`, so that the plugins are enabled with the specified precision; - * implicitly specify `auto`, so that the plugins are enabled with the precision automatically inferred from model dtype (i.e., the dtype specified in weight conversion); or - * disable the plugin by `disable`. -* Other features that requires a boolean (e.g., `context_fmha`, `paged_kv_cache`, `remove_input_padding`), you can - * enable/disable the feature by specifying `enable`/`disable`. - -The defaults have been carefully tuned for better performance. For example, `gpt_attention_plugin`, `context_fmha`, `paged_kv_cache` and `remove_input_padding` are enabled by default. See more details by `trtllm-build --help`. - -Normally, the `trtllm-build` command only requires a single GPU, but you can enable parallel building by passing the number of GPUs to the `--workers` argument. - -```bash -# Build a single-GPU float16 engine from TensorRT LLM checkpoint. -# gpt_attention_plugin (the special TensorRT LLM GPT Attention plugin) and remove_input_padding are enabled by default for runtime performance. -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/1-gpu \ - --output_dir gpt2/trt_engines/fp16/1-gpu - -# Build 2-way tensor parallelism engines from TensorRT LLM checkpoint. -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/2-gpu \ - --output_dir gpt2/trt_engines/fp16/2-gpu - -# Build 2-way tensor parallelism and 2-way pipeline parallelism engines from TensorRT LLM checkpoint. -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/4-gpu \ - --output_dir gpt2/trt_engines/fp16/4-gpu -``` - -If the engines are built successfully, you will see output like: -``` -...... -[03/12/2024-10:21:08] [TRT] [I] Engine generation completed in 35.9738 seconds. -[03/12/2024-10:21:08] [TRT] [I] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 212 MiB, GPU 775 MiB -[03/12/2024-10:21:08] [TRT] [I] [MemUsageChange] TensorRT-managed allocation in building engine: CPU +0, GPU +775, now: CPU 0, GPU 775 (MiB) -[03/12/2024-10:21:09] [TRT] [I] [MemUsageStats] Peak memory usage during Engine building and serialization: CPU: 6600 MiB -[03/12/2024-10:21:09] [TRT-LLM] [I] Total time of building Unnamed Network 0: 00:00:36 -[03/12/2024-10:21:09] [TRT-LLM] [I] Serializing engine to gpt2/trt_engines/fp16/1-gpu/rank0.engine... -[03/12/2024-10:21:11] [TRT-LLM] [I] Engine serialized. Total time: 00:00:02 -[03/12/2024-10:21:11] [TRT-LLM] [I] Total time of building all engines: 00:00:41 -``` - -#### Fused MultiHead Attention (FMHA) - -`trtllm-build` enables FMHA kernels by default. You may disable it by adding `--context_fmha disable`. - -If you find that the default fp16 accumulation cannot meet the requirement, you can try to enable fp32 accumulation by adding `--enable_context_fmha_fp32_acc` to the inference command (`run.py` or `summarize.py`). However, it is expected to see performance drop. - -Note that the FMHA kernels have to be used together with `gpt_attention_plugin` enabled. - -#### In-flight batching and paged KV cache - -If one wants to use [in-flight batching in C++ runtime](../../docs/in_flight_batching.md), the engine(s) must be built accordingly. In-flight batching in C++ runtime works only with attention plugin, paged KV cache and with packed data. Currently, the `trtllm-build` by default enables `gpt_attention_plugin`, `paged_kv_cache` and `remove_input_padding`, so the built engine(s) can support in-flight batching (unless you explicitly disable one of these options). One can additionally control the size of the block in paged KV cache using `--tokens_per_block=N`. - -### 4. Build TensorRT engine(s) with Random Weights -You can build engine(s) using random weights, which is useful for benchmarking. First, the [`../generate_checkpoint_config.py`](../generate_checkpoint_config.py) script can be used to generate a TensorRT LLM checkpoint config file: - -```bash -# Generate an 8-GPU GPT-175B float16 checkpoint config file. -python3 ../../../generate_checkpoint_config.py --architecture GPTForCausalLM \ - --vocab_size 51200 \ - --hidden_size 12288 \ - --num_hidden_layers 96 \ - --num_attention_heads 96 \ - --dtype float16 \ - --tp_size 8 \ - --output_path gpt_175b/trt_ckpt/fp16/8-gpu/config.json - - -# Generate a 16-GPU GPT-530B float16 checkpoint config file. -python3 ../../../generate_checkpoint_config.py --architecture GPTForCausalLM \ - --vocab_size 51200 \ - --hidden_size 20480 \ - --num_hidden_layers 105 \ - --num_attention_heads 128 \ - --dtype float16 \ - --tp_size 16 \ - --output_path gpt_530b/trt_ckpt/fp16/16-gpu/config.json -``` - -Then, use `trtllm-build` command to build engine(s) with random weights and the model architecture specified by the generated config file. - -```bash -# Build 8-GPU GPT-175B float16 engines using dummy weights, useful for performance tests. -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. -trtllm-build --model_config gpt_175b/trt_ckpt/fp16/8-gpu/config.json \ - --gemm_plugin auto \ - --max_batch_size 256 \ - --output_dir gpt_175b/trt_engines/fp16/8-gpu \ - --workers 8 - -# Build 16-GPU GPT-530B float16 engines using dummy weights, useful for performance tests. -# Enable several TensorRT LLM plugins to increase runtime performance. It also helps with build time. -trtllm-build --model_config gpt_530b/trt_ckpt/fp16/16-gpu/config.json \ - --gemm_plugin auto \ - --max_batch_size 128 \ - --max_input_len 128 \ - --max_seq_len 148 \ - --output_dir gpt_530b/trt_engines/fp16/16-gpu \ - --workers 8 -``` - -### 5. Run inference -#### Single node, single GPU - -The [`run.py`](../../../run.py) script can be used to run inference with the built engine(s). - -```bash -python3 ../../../run.py --engine_dir gpt2/trt_engines/fp16/1-gpu \ - --tokenizer_dir gpt2 \ - --max_output_len 8 -``` - -If the engines are run successfully, you will see output like: -``` -...... -Input [Text 0]: "Born in north-east France, Soyer trained as a" -Output [Text 0 Beam 0]: " chef before moving to London in the early" -``` - -The [`summarize.py`](../../../summarize.py) script can run the built engines to summarize the articles from the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. -For each summary, the script can compute the -[ROUGE](https://en.wikipedia.org/wiki/ROUGE_(metric)) scores and use the `ROUGE-1` score to validate the implementation. -By passing `--test_trt_llm` flag, the script will evaluate TensorRT LLM engines. You may also pass `--test_hf` flag to evaluate the HF model. - -```bash -python3 ../../../summarize.py --engine_dir gpt2/trt_engines/fp16/1-gpu \ - --hf_model_dir gpt2 \ - --test_trt_llm \ - --test_hf -``` -If the engines are run successfully, you will see output like: -``` -...... -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM (total latency: 1.520904541015625 sec) -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM (total output tokens: 0) -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM (tokens per second: 0.0) -[03/13/2024-05:43:18] [TRT-LLM] [I] TensorRT LLM beam 0 result -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge1 : 21.13474087351942 -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge2 : 6.2641616526063775 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeL : 16.693574311238077 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeLsum : 18.477384201634088 -[03/13/2024-05:43:18] [TRT-LLM] [I] Hugging Face (total latency: 8.76440143585205 sec) -[03/13/2024-05:43:18] [TRT-LLM] [I] HF beam 0 result -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge1 : 20.834898522466 -[03/13/2024-05:43:18] [TRT-LLM] [I] rouge2 : 5.6914719275508805 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeL : 16.297064309934132 -[03/13/2024-05:43:18] [TRT-LLM] [I] rougeLsum : 18.018627021792142 -``` - -#### Single node, multiple GPUs - -To run engines using multiple GPUs on a single node, you can use `mpirun` as: - -```bash -mpirun -np 2 \ - python3 ../../../run.py --engine_dir gpt2/trt_engines/fp16/2-gpu \ - --tokenizer_dir gpt2 \ - --max_output_len 8 - -# Note that GPT-175B is built with random weights, so the output will also be random -mpirun -np 8 \ - python3 ../../../run.py --engine_dir gpt_175b/trt_engines/fp16/8-gpu \ - --max_output_len 8 -``` - -#### Multiple nodes, multiple GPUs using [Slurm](https://slurm.schedmd.com) - -To run engines using multiple nodes, you should use a cluster manager like `Slurm`. The following section shows how to configure TensorRT LLM to execute on two nodes using Slurm. - -We start by preparing an `sbatch` script called `tensorrt_llm_run.sub`. That script contains the following code (you must replace the `` strings with your own values): - -```bash -#!/bin/bash -#SBATCH -o logs/tensorrt_llm.out -#SBATCH -e logs/tensorrt_llm.error -#SBATCH -J -#SBATCH -A -#SBATCH -p -#SBATCH --nodes=2 -#SBATCH --ntasks-per-node=8 -#SBATCH --time=00:30:00 - -sudo nvidia-smi -lgc 1410,1410 - -srun --mpi=pmix \ - --container-image \ - --container-mounts : \ - --container-workdir \ - --output logs/tensorrt_llm_%t.out \ - --error logs/tensorrt_llm_%t.error \ - python3 -u ../../../run.py --engine_dir --max_output_len 8 -``` - -Then, submit the job using: - -```bash -sbatch tensorrt_llm_run.sub -``` - -You might have to contact your cluster's administrator to help you customize the above script. - - -## Quantization - -### SmoothQuant - -This section explains how to use SmoothQuant on GPT models with TensorRT-LLM. - -[SmoothQuant](https://arxiv.org/abs/2211.10438) is a post-training quantization -(PTQ) method to quantize LLM models to INT8 for faster inference. As explained -in the article, SmoothQuant modifies a model to enable INT8 quantization -without significantly altering the accuracy. - -#### Model Transformation - -A LLM model is made of multiple matrix-multiplication operations (or GEMMs): `Y -= XW` where `X` of shape `[n, k]`, holds the activation (produced at run-time) -and `W`, of shape `[k, m]` are the learned weights. `Y`, of shape `[n, m]`, is -the matrix product of `X` and `W`. - -SmoothQuant introduces scaling along the `k` dimension by defining a vector of -strictly positive coefficients `s`. `Y = X diag(s)^{-1} diag(s) W`. We now have -`Y = X'W'` where `X' = X diag(s)^{-1}` and `W' = diag(s) W`. This -transformation is introduced so the quantization behaves better. In *normal* -models, `X` tends to be ill-conditioned: it has mostly small-magnitude -coefficients, but also some outliers that makes quantization difficult. -Conversely, the re-scaled `X'` is better suited for INT8 conversion. - -In this example, we only replace Attention's QKV and MLP's FC1 GEMMs to their -Smoothquant'd version since it is sufficient to maintain the accuracy for the -GPT model. During inference, `X'` is computed by fusing the channel-wise -multiplication by `diag(s)^{-1}` with the preceding layernorm's lambda and beta -parameters. `W'` is pre-computed and doesn't need additional modification -during inference. - -#### INT8 inference - -The INT8 quantization scheme used in TensorRT LLM theoretically works on any -GPT model. However, Smoothquant'd models tend to produce more accurate results -with reduced precision. - -INT8 inference modifies GEMMs `Y = XW` so that both `X` and `W` use INT8. The -matrix-multiplication is sped-up because of smaller weight size and fast matrix -products computation thanks to NVIDIA *Tensor Cores* operating on INT8 inputs. - -During inference, X is transformed from its standard floating point (fp) -values: `X_{i8} <- X_{fp} * s_x`. This scaling puts `X` values in the INT8 -range: `[-128, 127]`. Similarly, W is scaled, `W_{i8} <- W_{fp} * s_w` but that -operation is done at model export time, no need for subsequent operations at -run-time. - -The optimized TensorRT LLM GEMM implementation for SmoothQuant does the integer -matrix-multiplication `Y_{i32} <- X_{i8} W_{i8}` and rescales the result to its -original range `Y_{fp} <- Y_{i32} * (s_x)^{-1} * (s_w)^{-1}`. Note that -`Y_{i32}` isn't stored in memory, the re-scaling happens in the GEMM's epilogue -and only `Y_{fp}` gets saved. - -By default `s_x` and `s_w` are single-value coefficients. This is the -*per-tensor* mode. Values for `s_x` and `s_w` are static, estimated at model -export time. - -TensorRT LLM also supports more elaborate modes: - - per-channel: `s_w` is a fixed vector of size `[1, m]`. For that, - TensorRT LLM loads the adequately scaled version of of `W_{i8}` at model - construction time. - - per-token: `s_x` is a vector of size `[n, 1]` determined at run-time, based - on the per-token (a.k.a. per-row) absolute maximum of `X`. -Users can mix-and-match per-channel and per-token options. Both tend to -increase the accuracy of the model at the cost of a slightly increased latency. - -#### Usage -[`convert_checkpoint.py`](./convert_checkpoint.py) features a -`--smoothquant` option. It must be set to a decimal value in `[0, 1]` and -corresponds to the `alpha` parameter in the [SmoothQuant -paper](https://arxiv.org/abs/2211.10438). Setting `--smoothquant` will smooth the model -as explained in [model transformation](#model-transformation) and export the -scaling factors needed for INT8 inference. - -By default, it will run the model in the per-tensor mode, as explained in [INT8 -inference](#int8-inference). You can add any combination of `--per_token` and `--per_channel` to get the corresponding behaviors. - -```bash -# Per-tensor SmoothQuant -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --smoothquant 0.5 \ - --output_dir gpt2/trt_ckpt/int8-sq/1-gpu - -# Per-token per-channel SmoothQuant -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --smoothquant 0.5 \ - --per_token \ - --per_channel \ - --output_dir gpt2/trt_ckpt/int8-sq-ptpc/1-gpu -``` - -Then, use `trtllm-build` to build engine(s). - -```bash -# Per-tensor SmoothQuant -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8-sq/1-gpu \ - --output_dir gpt2/trt_engines/int8-sq/1-gpu - -# Per-token per-channel SmoothQuant -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8-sq-ptpc/1-gpu \ - --output_dir gpt2/trt_engines/int8-sq-ptpc/1-gpu -``` - -Note that GPT attention plugin is required to be enabled for SmoothQuant for now. - -User can also use `ModelOpt` to do INT8 quantization. Especially for gpt variant Starcoder2. -```bash -python3 example/quantization/quantize.py --model_dir starcoder2 \ - --dtype float16 \ - --qformat int8_sq \ - --output_dir starcoder2/trt_ckpt/int8-sq/ -``` -Then, use `trtllm-build` to build engine(s). - -```bash -trtllm-build --checkpoint_dir starcoder2/trt_ckpt/int8-sq/ \ - --output_dir starcoder2/trt_engine/int8-sq/ -``` - - -### INT8 KV Cache - -[`convert_checkpoint.py`](./convert_checkpoint.py) features a -`--int8_kv_cache` option. Setting `--int8_kv_cache` will calibrate the model and export the -scaling factors needed for INT8 KV cache inference. - -```bash -# Int8 KV cache -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --int8_kv_cache \ - --output_dir gpt2/trt_ckpt/int8kv/1-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8kv/1-gpu \ - --output_dir gpt2/trt_engines/int8kv/1-gpu -``` - -INT8 KV cache can be used with or without gpt attention plugin. - -### Weight Only Quantization - -[`convert_checkpoint.py`](./convert_checkpoint.py) features a `--use_weight_only` option that can enable weight-only quantization. You can further set the weight-only precision by passing `int8` or `int4` to the `--weight_only_precision` flag. - -```bash -# Int8 weight-only quantization -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 \ - --output_dir gpt2/trt_ckpt/int8-wo/1-gpu - -# Int4 weight-only quantization -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int4 \ - --output_dir gpt2/trt_ckpt/int4-wo/1-gpu -``` - -Then, use `trtllm-build` to build engine(s). - -```bash -# Int8 weight-only quantization -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int8-wo/1-gpu \ - --output_dir gpt2/trt_engines/int8-wo/1-gpu - -# Int4 weight-only quantization -trtllm-build --checkpoint_dir gpt2/trt_ckpt/int4-wo/1-gpu \ - --output_dir gpt2/trt_engines/int4-wo/1-gpu -``` - -### FP8 Quantization - -[`quantize.py`](../../../quantization/quantize.py) can do FP8 quantization and/or FP8 kv cache quantization, and export TensorRT LLM checkpoint. - -```bash -# FP8 quantization with FP8 kv cache -python3 ../../../quantization/quantize.py --model_dir gpt2 \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir gpt2/trt_ckpt/fp8/1-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp8/1-gpu \ - --output_dir gpt2/trt_engines/fp8/1-gpu -``` - -## Embedding Parallelism -Since the embedding lookup table can be several gigabytes in size. We can distribute this weight across multiple GPUs in order to reduce the memory consumption per GPU. - -### 1. Embedding parallelism -To enable this feature, add the flag `--use_parallel_embedding` to `convert_checkpoint.py`. - -### 2. The sharding dimension for embedding parallelism - -Assume the size of embedding lookup table is (vocab\_size \* hidden\_size), we can shard it along the vocab\_size (`--embedding_sharding_dim 0`) or hidden\_size (`--embedding_sharding_dim 1`) dimension. - -2.1 To shard the embedding lookup table along the hidden\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 1`. Here is an example: - -```bash -# 2-way tensor parallelism with embedding parallelism along hidden dimension -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 1 \ - --output_dir gpt2/trt_ckpt/fp16/2-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/2-gpu \ - --output_dir gpt2/trt_engines/fp16/2-gpu -``` - -2.2 To shard the embedding lookup table along the vocab\_size dimension, set the flag `--use_parallel_embedding --embedding_sharding_dim 0`. In this case, you can optionally enable the lookup plugin when building the engines. - -```bash -# 2-way tensor parallelism with embedding parallelism along vocab dimension -python3 convert_checkpoint.py --model_dir gpt2 \ - --dtype float16 \ - --tp_size 2 \ - --use_parallel_embedding \ - --embedding_sharding_dim 0 \ - --output_dir gpt2/trt_ckpt/fp16/2-gpu - -trtllm-build --checkpoint_dir gpt2/trt_ckpt/fp16/2-gpu \ - --output_dir gpt2/trt_engines/fp16/2-gpu -``` - -## GPT Variant - Granite - -For Granite, the steps are similar to StarCoder. - -```bash -# Download hf granite model -git clone https://huggingface.co/ibm-granite/granite-34b-code-instruct granite - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --model_dir granite \ - --dtype float16 \ - --gpt_variant starcoder \ - --tp_size 4 \ - --output_dir granite/trt_ckpt/fp16/4-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir granite/trt_ckpt/fp16/4-gpu \ - --gemm_plugin auto \ - --output_dir granite/trt_engines/fp16/4-gpu - -# Run inference -mpirun -np 4 \ - python3 ../../../run.py --engine_dir granite/trt_engines/fp16/4-gpu \ - --tokenizer_dir granite \ - --input_text "def print_hello_world():" \ - --max_output_len 20 -``` - -## GPT Variant - SantaCoder - -The SantaCoder extends the existing GPT model with multi-query attention mechanism. The following example shows building a 4-GPU engine and running simple prompt to generate the implementation of `print_hello_world()`. - -```bash -# Download hf santacoder model -git clone https://huggingface.co/bigcode/santacoder - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --model_dir santacoder \ - --dtype float16 \ - --tp_size 4 \ - --output_dir santacoder/trt_ckpt/fp16/4-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir santacoder/trt_ckpt/fp16/4-gpu \ - --gemm_plugin auto \ - --output_dir santacoder/trt_engines/fp16/4-gpu - -# Run inference -mpirun -np 4 \ - python3 ../../../run.py --engine_dir santacoder/trt_engines/fp16/4-gpu \ - --tokenizer_dir santacoder \ - --input_text "def print_hello_world():" \ - --max_output_len 20 -``` - - -## GPT Variant - StarCoder (v1 and v2) - -For StarCoder, the steps are similar to SantaCoder. - -```bash -# Download hf starcoder model -git clone https://huggingface.co/bigcode/starcoder - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --model_dir starcoder \ - --dtype float16 \ - --tp_size 4 \ - --output_dir starcoder/trt_ckpt/fp16/4-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir starcoder/trt_ckpt/fp16/4-gpu \ - --gemm_plugin auto \ - --output_dir starcoder/trt_engines/fp16/4-gpu - -# Run inference -mpirun -np 4 \ - python3 ../../../run.py --engine_dir starcoder/trt_engines/fp16/4-gpu \ - --tokenizer_dir starcoder \ - --input_text "def print_hello_world():" \ - --max_output_len 20 -``` - -For StarCoder2, you can use almost the same steps as shown above. - - Note that StarCoder2 hasn't been merged to the official releases of transformers package yet, so remember using the [main branch of transformers repo](https://github.com/huggingface/transformers). - - Add `--max_attention_window_size 4096` when running with run.py or summarization, which enables the sliding window attention. - - the sliding window size comes from the hf model [config.json](https://huggingface.co/bigcode/starcoder2-15b/blob/main/config.json#L23). - -### Run StarCoder2 with LoRA - -TensorRT LLM supports running StarCoder2 models with FP16/BF16/FP32 LoRA. In this section, we use starcoder2-15b as an example to show how to run an FP8 base model with FP16 LoRA module. - -* download the base model and lora model from HF - -```bash -git-lfs clone https://huggingface.co/bigcode/starcoder2-15b -git-lfs clone https://huggingface.co/KaQyn/peft-lora-starcoder2-15b-unity-copilot -``` - -* Quantize the StarCoder2 model to fp8 from HF -```bash -BASE_STARCODER2_MODEL=./starcoder2-15b -python ../../../quantization/quantize.py --model_dir ${BASE_STARCODER2_MODEL} \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir starcoder2-15b/trt_ckpt/fp8/1-gpu \ - --calib_size 512 -``` - -* Build engine and run inference. -```bash -trtllm-build --checkpoint_dir starcoder2-15b/trt_ckpt/fp8/1-gpu \ - --output_dir starcoder2-15b/trt_engines/fp8_lora/1-gpu \ - --gemm_plugin auto \ - --lora_plugin auto \ - --lora_dir ./peft-lora-starcoder2-15b-unity-copilot - -python ../../../run.py --engine_dir starcoder2-15b/trt_engines/fp8_lora/1-gpu \ - --max_output_len 20 \ - --tokenizer_dir ${BASE_STARCODER2_MODEL} \ - --input_text "def print_hello_world():" \ - --lora_task_uids 0 \ - --no_add_special_tokens \ - --use_py_session -``` - -## GPT-Next - -NVIDIA has released a GPT-like model with some architectural improvements, that you can find here: [https://huggingface.co/nvidia/GPT-2B-001](https://huggingface.co/nvidia/GPT-2B-001). This architecture is also supported by TensorRT-LLM. - -Different from Huggingface's checkpoint, you should specify the NeMo checkpoint path using `--nemo_ckpt_path` for `convert_checkpoint.py`. The script also extracts the tokenizer file from the NeMo checkpoint and saves it to the TensorRT LLM checkpoint folder, which can be used in the inference scripts. - -```bash -# Download NeMo checkpoint -wget https://huggingface.co/nvidia/GPT-2B-001/resolve/main/GPT-2B-001_bf16_tp1.nemo - -# Convert to TensorRT LLM checkpoint -# It also extracts the tokenizer file and saves to the TensorRT LLM checkpoint folder -python3 convert_checkpoint.py --nemo_ckpt_path GPT-2B-001_bf16_tp1.nemo \ - --dtype bfloat16 \ - --output_dir gpt-next-2B/trt_ckpt/bf16/1-gpu - -# Build TensorRT LLM engines -# --gpt_attention_plugin must be set for GPT-Next since Rotary positional embeddings (RoPE) is only supported by the gpt attention plugin at this time. -trtllm-build --checkpoint_dir gpt-next-2B/trt_ckpt/bf16/1-gpu \ - --output_dir gpt-next-2B/trt_engines/bf16/1-gpu - -# Run inference -python3 ../../../run.py --engine_dir gpt-next-2B/trt_engines/bf16/1-gpu \ - --vocab_file gpt-next-2B/trt_ckpt/bf16/1-gpu/tokenizer.model \ - --no_add_special_tokens \ - --max_output_len 8 -``` - -### Prompt-tuning - -For efficient fine-tuning, the NeMo framework allows you to learn virtual tokens to accomplish a downstream task. For more details, please read the -NeMo documentation [here](https://docs.nvidia.com/nemo-framework/user-guide/latest/overview.html). - -TensorRT LLM supports inference with those virtual tokens. To enable it, pass the prompt embedding table's maximum size at build time with `--max_prompt_embedding_table_size N`. For example: - -```bash -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --nemo_ckpt_path megatron_converted_8b_tp4_pp1.nemo \ - --dtype float16 \ - --output_dir gpt-next-8B/trt_ckpt/fp16/1-gpu - -# Build TensorRT LLM engines with prompt-tuning enabled -trtllm-build --checkpoint_dir gpt-next-8B/trt_ckpt/fp16/1-gpu \ - --max_prompt_embedding_table_size 100 \ - --output_dir gpt-next-8B/trt_engines/fp16/1-gpu -``` - -You can now export the learned embedding table with: -```bash -python3 nemo_prompt_convert.py -i email_composition.nemo -o email_composition.npy -``` -It'll give you a summary of the different tasks in the table, that you can specify at runtime. - -Finally, you can run inference on pre-defined tokens: -```bash -python3 ../../../run.py --engine_dir gpt-next-8B/trt_engines/fp16/1-gpu \ - --vocab_file gpt-next-8B/trt_ckpt/fp16/1-gpu/tokenizer.model \ - --no_add_special_tokens \ - --prompt_table_path email_composition.npy \ - --prompt_tasks 0 \ - --max_output_len 8 -``` - - -### MultiLoRA with the Nemo checkpoint - -```bash -# Download NeMo checkpoint -wget https://huggingface.co/nvidia/GPT-2B-001/resolve/main/GPT-2B-001_bf16_tp1.nemo - -# Convert to TensorRT LLM checkpoint -python3 convert_checkpoint.py --nemo_ckpt_path GPT-2B-001_bf16_tp1.nemo \ - --dtype float16 \ - --output_dir gpt-next-2B/trt_ckpt/fp16/1-gpu - -# Build TensorRT LLM engines -trtllm-build --checkpoint_dir gpt-next-2B/trt_ckpt/fp16/1-gpu \ - --lora_plugin auto \ - --lora_dir gpt2b_lora-900.nemo gpt2b_lora-stories.nemo \ - --lora_ckpt_source "nemo" \ - --lora_target_modules attn_qkv \ - --max_batch_size 4 \ - --max_beam_width 2 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --output_dir gpt-next-2B/trt_engines/fp16/1-gpu - -# Run inference directly from NeMo LoRA checkpoint -# --lora_task_ids correspond to the index of the models given with --lora_dir. -1 means no LoRA -python3 ../../../run.py --engine_dir gpt-next-2B/trt_engines/fp16/1-gpu \ - --vocab_file gpt-next-2B/trt_ckpt/fp16/1-gpu/tokenizer.model \ - --no_add_special_tokens \ - --max_output_len 20 \ - --use_py_session \ - --lora_task_uids 0 -1 1 \ - --input_text "After Washington had returned to Williamsburg, Dinwiddie ordered him to lead a larger force to assist Trent in his work. While en route, Washington learned of Trent's retreat. Since Tanaghrisson had promised support to the British, Washington continued toward Fort Duquesne and met with the Mingo leader. Learning of a French scouting party in the area, Washington, with Tanaghrisson and his party, surprised the Canadians on May 28 in what became known as the Battle of Jumonville Glen. They killed many of the Canadians, including their commanding officer, Joseph Coulon de Jumonville, whose head was reportedly split open by Tanaghrisson with a tomahawk. The historian Fred Anderson suggests that Tanaghrisson was acting to gain the support of the British and regain authority over his own people. They had been inclined to support the French, with whom they had long trading relationships. One of Tanaghrisson's men told Contrecoeur that Jumonville had been killed by British musket fire. Question: Upon learning of a French scounting party in the area, what did Washington do? Answer:" "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" -``` - -The output would look like (Note that in this case the adapters have only been trained for a few epochs, so the result quality is poor): - -``` -...... -Input [Text 0]: "After Washington had returned to Williamsburg, Dinwiddie ordered him to lead a larger force to assist Trent in his work. While en route, Washington learned of Trent's retreat. Since Tanaghrisson had promised support to the British, Washington continued toward Fort Duquesne and met with the Mingo leader. Learning of a French scouting party in the area, Washington, with Tanaghrisson and his party, surprised the Canadians on May 28 in what became known as the Battle of Jumonville Glen. They killed many of the Canadians, including their commanding officer, Joseph Coulon de Jumonville, whose head was reportedly split open by Tanaghrisson with a tomahawk. The historian Fred Anderson suggests that Tanaghrisson was acting to gain the support of the British and regain authority over his own people. They had been inclined to support the French, with whom they had long trading relationships. One of Tanaghrisson's men told Contrecoeur that Jumonville had been killed by British musket fire. Question: Upon learning of a French scounting party in the area, what did Washington do? Answer:" -Output [Text 0 Beam 0]: "He surprised the Canadians on May 28 in what became known as the Battle of Jumonville" -Input [Text 1]: "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" -Output [Text 1 Beam 0]: ". - -The game is played with a deck of cards, and the player who has the most" -Input [Text 2]: "You hold the job title in the Wizarding World of Harry Potter where you say random words looking for spells" -Output [Text 2 Beam 0]: ". - -You are a wizard who is a wizard. - -You are a wizard who is" -``` diff --git a/examples/models/core/gpt/convert_checkpoint.py b/examples/models/core/gpt/convert_checkpoint.py deleted file mode 100644 index c8bbebe39925..000000000000 --- a/examples/models/core/gpt/convert_checkpoint.py +++ /dev/null @@ -1,343 +0,0 @@ -import argparse -import json -import os -import shutil -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from pathlib import Path - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import release_gc -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models import GPTForCausalLM -from tensorrt_llm.models.gpt.convert import (UnpackedNemoCheckpointDir, - copy_tokenizer_files, load_hf_gpt, - unpack_nemo_ckpt, - update_tokenizer_paths) -from tensorrt_llm.models.modeling_utils import QuantConfig -from tensorrt_llm.quantization import QuantAlgo - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--nemo_ckpt_path', type=str, default=None) - parser.add_argument( - '--gpt_variant', - default=None, - choices=[ - None, 'gpt2', 'santacoder', 'starcoder', 'starcoder2', 'persimmon', - 'kosmos-2', 'nemotron' - ], - help= - "By default the script will try to infer the gpt_variant from model_dir. " - "Or users may overwrite gpt_variant by explicitly passing the variant.") - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument( - '--dtype', - type=str, - default='auto', - choices=['auto', 'float16', 'bfloat16', 'float32'], - help= - "The data type for the model weights and activations if not quantized. " - "If 'auto', the data type is automatically inferred from the source model; " - "however, if the source dtype is float32, it is converted to float16.") - parser.add_argument("--load_model_on_cpu", action="store_true") - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - - parser.add_argument( - '--calib_dataset', - type=str, - default='lambada', - help= - "The huggingface dataset name or the local directory of the dataset for calibration." - ) - parser.add_argument( - '--int8_kv_cache', - default=False, - action="store_true", - help= - 'By default, we use dtype for KV cache. int8_kv_cache chooses int8 quantization for KV' - ) - parser.add_argument( - '--per_channel', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor for the GEMM\'s result. ' - 'per_channel instead uses a different static scaling factor for each channel. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - '--per_token', - default=False, - action="store_true", - help= - 'By default, we use a single static scaling factor to scale activations in the int8 range. ' - 'per_token chooses at run time, and for each token, a custom scaling factor. ' - 'The latter is usually more accurate, but a little slower.') - parser.add_argument( - "--smoothquant", - "-sq", - type=float, - default=None, - help="Set the α parameter (see https://arxiv.org/pdf/2211.10438.pdf)" - " to Smoothquant the model, and output int8 weights." - " A good first try is 0.5. Must be in [0, 1]") - parser.add_argument("--dataset_cache_dir", - type=str, - default=None, - help="cache dir to load the hugging face dataset") - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - parser.add_argument( - '--nemo_rename_key', - type=str, - nargs='+', - default=[], - help= - "Change a layer name when loading a NeMo checkpoint. Should follow :" - ) - - args = parser.parse_args() - - tensorrt_llm.logger.set_level(args.log_level) - return args - - -def args_to_quant_config(args: argparse.Namespace) -> QuantConfig: - '''return config dict with quantization info based on the command line args - ''' - quant_config = QuantConfig() - if args.use_weight_only: - if args.weight_only_precision == 'int8': - quant_config.quant_algo = QuantAlgo.W8A16 - elif args.weight_only_precision == 'int4': - quant_config.quant_algo = QuantAlgo.W4A16 - elif args.smoothquant: - quant_config.smoothquant_val = args.smoothquant - if args.per_channel: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_CHANNEL_PER_TENSOR_PLUGIN - else: - if args.per_token: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PER_TOKEN_PLUGIN - else: - quant_config.quant_algo = QuantAlgo.W8A8_SQ_PER_TENSOR_PLUGIN - - if args.int8_kv_cache: - quant_config.kv_cache_quant_algo = QuantAlgo.INT8 - - # Check if model ckpt is pre-quantized to fp8. - hf_quant_config_path = f"{args.model_dir}/hf_quant_config.json" - if os.path.exists(hf_quant_config_path): - with open(hf_quant_config_path, 'r') as f: - hf_quant_config = json.load(f) - if hf_quant_config.get("producer", {}).get("name") == "modelopt": - modelopt_quant_config = hf_quant_config.get("quantization", {}) - if modelopt_quant_config.get("quant_algo", None) == QuantAlgo.FP8: - quant_config.quant_algo = QuantAlgo.FP8 - if modelopt_quant_config.get("kv_cache_quant_algo", - None) == QuantAlgo.FP8: - quant_config.kv_cache_quant_algo = QuantAlgo.FP8 - - return quant_config - - -def convert_and_save_hf(args): - model_dir = args.model_dir - load_model_on_cpu = args.load_model_on_cpu - world_size = args.tp_size * args.pp_size - - override_fields = { - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'gpt_variant': args.gpt_variant, - } - - quant_config = args_to_quant_config(args) - is_prequantized_to_fp8 = quant_config.quant_algo == QuantAlgo.FP8 - if is_prequantized_to_fp8: - override_fields.update({'is_prequantized_to_fp8': True}) - - if args.smoothquant is not None or args.int8_kv_cache: - mapping = Mapping(world_size=world_size, - tp_size=args.tp_size, - pp_size=args.pp_size) - GPTForCausalLM.quantize( - args.model_dir, - args.output_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - device='cpu' if args.load_model_on_cpu else 'cuda', - calib_dataset=args.calib_dataset, - **override_fields) - else: - # Defer weight loading if checkpoint is prequantized to fp8. - if is_prequantized_to_fp8: - hf_model_or_dir = model_dir - else: - hf_model_or_dir = load_hf_gpt(model_dir, load_model_on_cpu) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - model = GPTForCausalLM.from_hugging_face(hf_model_or_dir, - args.dtype, - mapping=mapping, - quant_config=quant_config, - **override_fields) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - -def execute(workers, func, args): - if workers == 1: - for rank, f in enumerate(func): - f(args, rank) - else: - with ThreadPoolExecutor(max_workers=workers) as p: - futures = [p.submit(f, args, rank) for rank, f in enumerate(func)] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - -def convert_and_save_nemo(args): - world_size = args.tp_size * args.pp_size - quant_config = args_to_quant_config(args) - - override_fields = { - 'use_parallel_embedding': True, - 'embedding_sharding_dim': 0, - } - - nemo_ckpt_dir = os.path.join(args.output_dir, "unpacked") - nemo_ckpt_dir = unpack_nemo_ckpt(args.nemo_ckpt_path, nemo_ckpt_dir) - - def convert_and_save_rank(args, rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - model = GPTForCausalLM.from_nemo( - nemo_ckpt_dir, - dtype=args.dtype, - mapping=mapping, - quant_config=quant_config, - load_model_on_cpu=args.load_model_on_cpu, - nemo_rename_key=args.nemo_rename_key, - **override_fields) - model.save_checkpoint(args.output_dir, save_config=(rank == 0)) - del model - - execute(args.workers, [convert_and_save_rank] * world_size, args) - release_gc() - - # Copy tokenizer files - unpacked_checkpoints_dir = UnpackedNemoCheckpointDir( - nemo_ckpt_dir, load_checkpoints_to_cpu=args.load_model_on_cpu) - nemo_model_config = unpacked_checkpoints_dir.model_config - tokenizer_config = update_tokenizer_paths( - nemo_model_config["tokenizer"], - unpacked_checkpoints_dir.get_all_tokenizer_file_paths()) - copy_tokenizer_files(tokenizer_config, Path(args.output_dir)) - - # Clean up unpacked nemo checkpoint - shutil.rmtree(nemo_ckpt_dir) - - -def main(): - emit_engine_arch_deprecation("convert_checkpoint.py") - # TODO(qijun): Currently, the convert script depends on a torch op: - # torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix, - # which is included in tensorrt_llm Python package. Otherwise, the convert - # script does not need to import tensorrt_llm. Will remove it after reimplementing - # the op with PyTorch. - print(tensorrt_llm.__version__) - args = parse_arguments() - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - if args.model_dir is not None: - convert_and_save_hf(args) - elif args.nemo_ckpt_path is not None: - convert_and_save_nemo(args) - else: - raise NotImplementedError("No source model path specified!") - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') - - -if __name__ == '__main__': - main() diff --git a/examples/models/core/gpt/merge_ptuning_tables.py b/examples/models/core/gpt/merge_ptuning_tables.py deleted file mode 100644 index 7431541f8425..000000000000 --- a/examples/models/core/gpt/merge_ptuning_tables.py +++ /dev/null @@ -1,44 +0,0 @@ -#! /usr/bin/env python3 -# 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. -import argparse -from pathlib import Path -from typing import List - -import numpy as np - - -def combine_tables(input_tables: List[Path], output_table: Path): - tables = [np.load(table) for table in input_tables] - max_vocab_size = max(table.shape[1] for table in tables) - padded_tables = [ - np.pad(table, [(0, 0), (0, max_vocab_size - table.shape[1]), (0, 0)]) - for table in tables - ] - merged_table = np.concatenate(padded_tables, axis=0) - np.save(output_table, merged_table) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument("input_tables", - type=Path, - nargs="+", - help="paths to tables to merge.") - parser.add_argument("output_table", - type=Path, - help="path where to save the combined table") - args = parser.parse_args() - combine_tables(args.input_tables, args.output_table) diff --git a/examples/models/core/gpt/nemo_lora_convert.py b/examples/models/core/gpt/nemo_lora_convert.py deleted file mode 100644 index 59e675402ec2..000000000000 --- a/examples/models/core/gpt/nemo_lora_convert.py +++ /dev/null @@ -1,206 +0,0 @@ -#! /usr/bin/env python3 -# 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. -import argparse -import datetime -import logging -import tempfile -from pathlib import Path - -import numpy as np -import torch -import yaml - -from tensorrt_llm._utils import str_dtype_to_torch, to_json_file, torch_to_numpy -from tensorrt_llm.lora_manager import LoraManager, get_all_nemo_lora_weights -from tensorrt_llm.models.gpt.convert import cpu_map_location, unpack_nemo_ckpt - -log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s" -logging.basicConfig(format=log_format) -LOGGER = logging.getLogger(__name__) - - -def get_lora_keys(layer_idx): - in_key = f'model.language_model.encoder.layers.{layer_idx}.self_attention.adapter_layer.lora_kqv_adapter.linear_in.weight' - out_key = f'model.language_model.encoder.layers.{layer_idx}.self_attention.adapter_layer.lora_kqv_adapter.linear_out.weight' - return in_key, out_key - - -def save_val(val, dir, key, tp_num=None, write_npy=False): - ext = "npy" if write_npy else "bin" - suffix = ext if tp_num is None else f"{tp_num}.{ext}" - if write_npy: - np.save(dir / f"model.{key}.{suffix}", val) - else: - val.tofile(dir / f"model.{key}.{suffix}") - - -def lora_convert(out_dir, lora_config, lora_weights, customization_id, - precision): - saved_dir = Path(out_dir) - saved_dir.mkdir(parents=True, exist_ok=True) - num_layers = int(lora_config["num_layers"]) - config = {"lora_config": {"lora_kqv_adapter": {}}} - config['lora_config']['precision'] = precision - layer_weights = get_all_nemo_lora_weights(lora_weights) - for layer_idx in range(num_layers): - linear_in_weight = layer_weights[layer_idx]['in'] - linear_out_weight = layer_weights[layer_idx]['out'] - config["lora_config"]["lora_kqv_adapter"]["0"] = { - "key": f"{customization_id}", - "low_rank": f"{linear_in_weight.shape[0]}", - } - - # do something else here. just choose some key instead of basing it on the nemo key - in_key, out_key = get_lora_keys(layer_idx) - - save_val( - torch_to_numpy( - linear_in_weight.transpose( - 1, 0).contiguous().to(dtype=str_dtype_to_torch(precision))), - saved_dir, - in_key.replace("lora_kqv_adapter", f"lora_kqv_adapter.{0}")) - save_val( - torch_to_numpy( - linear_out_weight.transpose( - 1, 0).contiguous().to(dtype=str_dtype_to_torch(precision))), - saved_dir, - out_key.replace("lora_kqv_adapter", f"lora_kqv_adapter.{0}")) - - to_json_file(config, saved_dir / "lora_weights.json") - - -def lora_convert_cpp_runtime(out_dir, - lora_config, - lora_weights, - precision='float16'): - saved_dir = Path(out_dir) - saved_dir.mkdir(parents=True, exist_ok=True) - num_layers = int(lora_config["num_layers"]) - weights = [] - weight_config = [] - layer_weights = get_all_nemo_lora_weights(lora_weights) - for layer_idx in range(num_layers): - in_weights = layer_weights[layer_idx]['in'] - out_weights = layer_weights[layer_idx]['out'] - LOGGER.debug(f"layer {layer_idx} in_weights: {in_weights.shape}") - LOGGER.debug(f"layer {layer_idx} out_weights: {out_weights.shape}") - in_out_weights = [] - adapter_size = 0 - for w, inout in ((in_weights, "in"), (out_weights, "out")): - assert len(w.shape) == 2 - # assume that the hidden dim is the larger of the 2 - dim0 = w.shape[0] - dim1 = w.shape[1] - adapter_size = min(dim0, dim1) - # in_weights should have shape [adaper_size, hidden] - if dim1 < dim0 and inout == "in": - adapter_size = dim1 - w = w.transpose(1, 0) - # out_weights should have shape [hidden, adapter_size] - elif dim0 < dim1 and inout == "out": - adapter_size = dim0 - w = w.transpose(1, 0) - - w = w.contiguous().flatten().to(dtype=str_dtype_to_torch(precision)) - in_out_weights.append(w) - in_out_weights = torch.concatenate(in_out_weights).flatten().numpy() - weights.append(in_out_weights) - weight_config.append( - np.array([ - LoraManager.LORA_MODULE_IDS["attn_qkv"], layer_idx, adapter_size - ], - dtype=np.int32)) - all_weights = np.expand_dims(np.stack(weights), 0) - all_configs = np.expand_dims(np.stack(weight_config), 0) - - save_val(all_weights, - saved_dir, - "lora_weights", - tp_num=None, - write_npy=True) - save_val(all_configs, saved_dir, "lora_config", tp_num=None, write_npy=True) - - -def main(args): - start_time = datetime.datetime.now() - with tempfile.TemporaryDirectory() as prompt_out_dir: - prompt_out_dir = Path(prompt_out_dir) - unpack_nemo_ckpt(args.in_file, prompt_out_dir) - LOGGER.info("Spent %s (h:m:s) to unpack NeMo prompt archive", - datetime.datetime.now() - start_time) - - model_weights_ckpt = "model_weights.ckpt" - with open(prompt_out_dir / "model_config.yaml") as f: - prompt_config = yaml.full_load(f) - LOGGER.debug(prompt_config) - - start_time = datetime.datetime.now() - weight_path = prompt_out_dir / model_weights_ckpt - - prompt_weights = torch.load( - weight_path, - map_location=cpu_map_location, - ) - if args.write_cpp_runtime_tensors: - lora_convert_cpp_runtime(args.out_dir, - prompt_config, - prompt_weights, - precision=args.storage_type) - else: - lora_convert(args.out_dir, - prompt_config, - prompt_weights, - args.customization_id, - precision=args.storage_type) - - LOGGER.info("Spent %s (h:m:s) to convert the prompt model", - datetime.datetime.now() - start_time) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - '--out-dir', - '-o', - type=Path, - help='path to output embedding table file in the .npy format', - required=True) - parser.add_argument('--in-file', - '-i', - type=Path, - help='path to input prompt-tuning checkpoint file', - required=True) - parser.add_argument("--verbose", - action="store_true", - help="Provide verbose messages") - parser.add_argument("--customization-id", type=str, default="lora") - parser.add_argument("--write-cpp-runtime-tensors", - action="store_true", - default=False) - parser.add_argument("--storage-type", - type=str, - default="float16", - choices=["float32", "float16", "bfloat16"]) - args = parser.parse_args() - - LOGGER.setLevel(logging.DEBUG if args.verbose else logging.INFO) - - print("\n=============== Argument ===============") - for key in vars(args): - print(f"{key}: {vars(args)[key]}") - print("========================================") - - main(args) diff --git a/examples/models/core/gpt/nemo_prompt_convert.py b/examples/models/core/gpt/nemo_prompt_convert.py deleted file mode 100755 index c55131759c9d..000000000000 --- a/examples/models/core/gpt/nemo_prompt_convert.py +++ /dev/null @@ -1,145 +0,0 @@ -#! /usr/bin/env python3 -# 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. -import argparse -import datetime -import logging -import tempfile -from pathlib import Path - -import numpy as np -import torch -import yaml - -from tensorrt_llm._utils import torch_to_numpy -from tensorrt_llm.models.gpt.convert import cpu_map_location, unpack_nemo_ckpt - -log_format = "%(asctime)s %(name)s [%(levelname)s] %(message)s" -logging.basicConfig(format=log_format) -LOGGER = logging.getLogger(__name__) - - -def prompt_convert(out_file, prompt_config, prompt_weights): - nemo_type = "peft_tuning" if "peft" in prompt_config else "prompt_learning" - - vtokens_embeddings = [] - vtokens_len = [] - - if nemo_type == "peft_tuning": - ptuning_key = "model.embedding.adapter_layer.ptuning_adapter.inference_table" - if ptuning_key not in prompt_weights: - key_match = "adapter_layer.ptuning_adapter" - for k in prompt_weights.keys(): - if key_match in k: - ptuning_key = k - break - else: - raise ValueError( - "Could not find a suitable ptuning key in Nemo dict." - f" Tried {ptuning_key} or any key matching *{key_match}*") - prompt_task_weights = prompt_weights[ptuning_key] - - if 'hidden_size' in prompt_config: - assert prompt_config['hidden_size'] == prompt_task_weights.shape[ - 1], "P-Tuning hidden size does not match the model's." - - vtokens_embeddings.append(prompt_task_weights) - vtokens_len.append(prompt_task_weights.shape[0]) - else: - prompt_templates = prompt_config["task_templates"] - actual_task_id = 0 - for task_name_id, prompt_task in enumerate(prompt_templates): - prompt_task_name = prompt_task["taskname"] - LOGGER.info(f"Task {actual_task_id}: {prompt_task['taskname']}") - prompt_task_weights = prompt_weights["prompt_table"].get( - f"prompt_table.{prompt_task_name}.prompt_embeddings.weight") - if prompt_task_weights is None: - continue - vtokens_embeddings.append(prompt_task_weights) - vtokens_len.append(prompt_task_weights.shape[0]) - actual_task_id += 1 - - max_vtoken_len = max(vtokens_len) - embedding_dim = vtokens_embeddings[0].shape[1] - - # pad tasks to longest task embedding table - for i, vtoken_emb_table in enumerate(vtokens_embeddings): - padded_table = torch.zeros((max_vtoken_len, embedding_dim)) - padded_table[:vtoken_emb_table.shape[0], :] = vtoken_emb_table - vtokens_embeddings[i] = padded_table - - vtokens_embeddings = torch.stack(vtokens_embeddings) - np.save(out_file, torch_to_numpy(vtokens_embeddings)) - - -def main(args): - start_time = datetime.datetime.now() - with tempfile.TemporaryDirectory() as prompt_out_dir: - prompt_out_dir = Path(prompt_out_dir) - unpack_nemo_ckpt(args.in_file, prompt_out_dir) - LOGGER.info("Spent %s (h:m:s) to unpack NeMo prompt archive", - datetime.datetime.now() - start_time) - - model_weights_ckpt = "model_weights.ckpt" - with open(prompt_out_dir / "model_config.yaml") as f: - prompt_config = yaml.full_load(f) - LOGGER.debug(prompt_config) - - start_time = datetime.datetime.now() - weight_path = prompt_out_dir / model_weights_ckpt - if not weight_path.exists(): - weight_path = prompt_out_dir / "mp_rank_00" / model_weights_ckpt - - prompt_weights = torch.load( - weight_path, - map_location=cpu_map_location, - ) - prompt_convert(args.out_file, prompt_config, prompt_weights) - - LOGGER.info("Spent %s (h:m:s) to convert the prompt model", - datetime.datetime.now() - start_time) - - -if __name__ == "__main__": - parser = argparse.ArgumentParser() - parser.add_argument( - '--out-file', - '-o', - type=Path, - help='path to output embedding table file in the .npy format', - required=True) - parser.add_argument('--in-file', - '-i', - type=Path, - help='path to input prompt-tuning checkpoint file', - required=True) - parser.add_argument("--storage-type", - "-t", - type=str, - default="fp32", - choices=["fp32", "fp16"]) - parser.add_argument("--verbose", - action="store_true", - help="Provide verbose messages") - args = parser.parse_args() - - LOGGER.setLevel(logging.DEBUG if args.verbose else logging.INFO) - - print("\n=============== Argument ===============") - for key in vars(args): - print(f"{key}: {vars(args)[key]}") - print("========================================") - - main(args) diff --git a/examples/models/core/gpt/requirements.txt b/examples/models/core/gpt/requirements.txt deleted file mode 100644 index 592e01e5ba6d..000000000000 --- a/examples/models/core/gpt/requirements.txt +++ /dev/null @@ -1,6 +0,0 @@ --c ../../../constraints.txt -tensorrt_llm>=0.0.0.dev0 -datasets==3.1.0 -evaluate -rouge_score -SentencePiece>=0.1.99 diff --git a/examples/models/core/gpt/run_hf.py b/examples/models/core/gpt/run_hf.py deleted file mode 100644 index 850d44b04f03..000000000000 --- a/examples/models/core/gpt/run_hf.py +++ /dev/null @@ -1,162 +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. -import argparse -import csv -from pathlib import Path - -import numpy as np -import torch -from transformers import AutoModelForCausalLM, AutoTokenizer, T5Tokenizer - -import tensorrt_llm - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--max_output_len', type=int, required=True) - parser.add_argument('--log_level', type=str, default='warning') - parser.add_argument('--model_dir', type=str, default='gpt2') - parser.add_argument('--data_type', - type=str, - choices=['fp32', 'fp16'], - default='fp32') - parser.add_argument('--input_text', - type=str, - default='Born in north-east France, Soyer trained as a') - parser.add_argument( - '--input_tokens', - dest='input_file', - type=str, - help= - 'CSV or Numpy file containing tokenized input. Alternative to text input.', - default=None) - parser.add_argument('--output_csv', - type=str, - help='CSV file where the tokenized output is stored.', - default=None) - parser.add_argument('--output_npy', - type=str, - help='Numpy file where the tokenized output is stored.', - default=None) - parser.add_argument('--tokenizer', - dest='tokenizer_path', - help="HF tokenizer config path", - default='gpt2') - parser.add_argument('--vocab_file', - help="Used for sentencepiece tokenizers") - return parser.parse_args() - - -def generate( - max_output_len: int, - log_level: str = 'error', - model_dir: str = 'gpt2', - data_type: str = 'fp32', - input_text: str = 'Born in north-east France, Soyer trained as a', - input_file: str = None, - output_csv: str = None, - output_npy: str = None, - tokenizer_path='gpt2', - vocab_file=None, -): - tensorrt_llm.logger.set_level(log_level) - - model = AutoModelForCausalLM.from_pretrained(model_dir, - trust_remote_code=True) - model.cuda() - if data_type == 'fp16': - model.half() - - if vocab_file is not None: - tokenizer = T5Tokenizer(vocab_file=vocab_file) - END_ID = 50256 - else: - tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) - END_ID = tokenizer.eos_token_id - - input_tokens = [] - if input_file is None: - input_tokens.append( - tokenizer.encode(input_text, add_special_tokens=False)) - else: - if input_file.endswith('.csv'): - with open(input_file, 'r') as csv_file: - csv_reader = csv.reader(csv_file, delimiter=',') - for line in csv_reader: - input_tokens.append(np.array(line, dtype='int32')) - elif input_file.endswith('.npy'): - inputs = np.load(input_file) - for row in inputs: - row = row[row != END_ID] - input_tokens.append(row) - else: - print('Input file format not supported.') - raise SystemExit - - input_ids = None - input_lengths = None - if input_file is None: - input_ids = torch.tensor(input_tokens, dtype=torch.int32, device='cuda') - input_lengths = torch.tensor([input_ids.size(1)], - dtype=torch.int32, - device='cuda') - max_input_length = torch.max(input_lengths).item() - else: - input_lengths = torch.tensor([len(x) for x in input_tokens], - dtype=torch.int32, - device='cuda') - max_input_length = torch.max(input_lengths).item() - input_ids = np.full((len(input_lengths), max_input_length), END_ID) - for i in range(len(input_lengths)): - input_ids[i][-len(input_tokens[i]):] = input_tokens[i] - input_ids = torch.tensor(input_ids, dtype=torch.int32, device='cuda') - - top_k = 1 - temperature = 1 - output_ids = model.generate(input_ids, - max_length=max_input_length + max_output_len, - top_k=top_k, - temperature=temperature, - eos_token_id=END_ID, - pad_token_id=END_ID) - torch.cuda.synchronize() - - if output_csv is None and output_npy is None: - for b in range(input_lengths.size(0)): - inputs = input_tokens[b] - input_text = tokenizer.decode(inputs) - print(f'Input: {input_text}') - outputs = output_ids[b][max_input_length:].tolist() - output_text = tokenizer.decode(outputs) - print(f'Output: {output_text}') - - if output_csv is not None: - output_file = Path(output_csv) - output_file.parent.mkdir(exist_ok=True, parents=True) - outputs = output_ids[:, max_input_length:].tolist() - with open(output_file, 'w') as csv_file: - writer = csv.writer(csv_file, delimiter=',') - writer.writerows(outputs) - - if output_npy is not None: - output_file = Path(output_npy) - output_file.parent.mkdir(exist_ok=True, parents=True) - outputs = output_ids[:, max_input_length:].tolist() - np.save(output_file, np.array(outputs, dtype='int32')) - - -if __name__ == '__main__': - args = parse_arguments() - generate(**vars(args)) diff --git a/examples/models/core/internlm2/.gitignore b/examples/models/core/internlm2/.gitignore deleted file mode 100644 index 7ce339719a3e..000000000000 --- a/examples/models/core/internlm2/.gitignore +++ /dev/null @@ -1,2 +0,0 @@ -internlm* -tokenizer.model diff --git a/examples/models/core/internlm2/README.md b/examples/models/core/internlm2/README.md deleted file mode 100644 index 0c1f8de5b908..000000000000 --- a/examples/models/core/internlm2/README.md +++ /dev/null @@ -1,207 +0,0 @@ -# InternLM2 - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run InternLM2 7B / 20B models in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -## Overview - -The TensorRT LLM InternLM2 implementation is based on the LLaMA model. The implementation can -be found in [model.py](../../../../tensorrt_llm/models/llama/model.py). -The TensorRT LLM InternLM2 example code lies in [`examples/models/core/internlm2`](./): - -* [`convert_checkpoint.py`](./convert_checkpoint.py) converts the Huggingface Model of InternLM2 into TensorRT LLM checkpoint. - - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * FP16 / BF16 - * INT8 & INT4 Weight-Only - * Tensor Parallel - -## Usage - -The TensorRT LLM InternLM2 example code locates at [examples/models/core/internlm2](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first to make sure the example uses matched `tensorrt_llm` version: - -```bash -pip install -r requirements.txt -``` - -TensorRT LLM InternLM2 builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -InternLM2 has released several checkpoints of different size or capabilities under https://huggingface.co/internlm. Users can pick any one repository and follow instructions to prepare the checkpoint. - -Below examples use [internlm2-chat-7b](https://huggingface.co/internlm/internlm2-chat-7b) and [internlm2-chat-20b](https://huggingface.co/internlm/internlm2-chat-20b) and assume these repositories are cloned or linked under this directory, for example `./internlm2-chat-7b`. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `--workers` feature only supports single node. - -Here're some examples: - -```bash -# Build a single-GPU float16 engine from HF weights. -# gpt_attention_plugin is necessary in InternLM2. -# Try use_gemm_plugin to prevent accuracy issue. -cd examples/models/core/internlm2 - -# Convert the InternLM2 7B model using a single GPU and FP16. -python convert_checkpoint.py --model_dir ./internlm2-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm2-chat-7b/trt_engines/fp16/1-gpu/ -# Note: setting `--dtype bfloat16` to use bfloat16 precision. - -# BUild the InternLM2 7B model using a single GPU -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/fp16/1-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Convert the InternLM2 7B model using a single GPU and apply INT8 weight-only quantization.. -python convert_checkpoint.py --model_dir ./internlm2-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm2-chat-7b/trt_engines/int8/1-gpu/ \ - --use_weight_only \ - --weight_only_precision int8 - -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/int8/1-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Note: setting `--weight_only_precision int4` to use INT4 weight-only quantization - -# Build InternLM2 7B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./internlm2-chat-7b/ \ - --dtype float16 \ - --output_dir ./internlm2-chat-7b/trt_engines/fp16/2-gpu/ \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/fp16/2-gpu/ \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 - -# Build InternLM2 20B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./internlm2-chat-20b/ \ - --dtype bfloat16 \ - --output_dir ./internlm2-chat-20b/trt_engines/bf16/2-gpu/ \ - --tp_size 2 --workers 2 - -trtllm-build --checkpoint_dir ./internlm2-chat-7b/trt_engines/bf16/2-gpu/ \ - --output_dir ./engine_outputs \ - --gpt_attention_plugin bfloat16 \ - --gemm_plugin bfloat16 -``` - -#### INT8 weight only - -Examples: - -```bash -cd examples/models/core/internlm2 - -# For 7B models -python convert_checkpoint.py --model_dir ./internlm2-chat-7b \ - --output_dir ./internlm2-chat-7b/w8a16/ \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 - -# Build 7B model with both INT8 weight-only -trtllm-build --checkpoint_dir ./internlm2-chat-7b/w8a16 \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 -``` - - -```bash -cd examples/models/core/internlm2 - -# For 20B models -python convert_checkpoint.py --model_dir ./internlm2-chat-20b \ - --output_dir ./internlm2-chat-20b/w8a16 \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 - -# Build 20B model with both INT8 weight-only -trtllm-build --checkpoint_dir ./internlm2-chat-20b/w8a16 \ - --output_dir ./engine_outputs \ - --gemm_plugin float16 \ -``` - -### Run - -To run a TensorRT LLM InternLM2 model using the engines generated by `trtllm-build` - -```bash -# InternLM2 7B with fp16 -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/fp16/1-gpu/ - -# InternLM2 7B with bf16 -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/bf16/1-gpu/ - -# InternLM2 7B with int8 weight only quantization -python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/weight_only/1-gpu/ - -# InternLM2 7B with fp16 and tensor parallelism -mpirun -n 2 --allow-run-as-root \ - python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/fp16/2-gpu/ - -# InternLM2 20B with fp16 and tensor parallelism and pipeline parallelism -mpirun -n 4 --allow-run-as-root \ - python ../../../run.py --max_output_len=120 \ - --input_text 'Tell me about yourself.' \ - --tokenizer_dir ./internlm2-chat-7b/ \ - --engine_dir=./internlm2-chat-7b/trt_engines/bf16/4-gpu/ -``` - -### Summarization using the InternLM2 model - -```bash -# Run summarization using the InternLM2 7B model in FP16. -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./engine_outputs - -# Run summarization using the InternLM2 7B model quantized to w8a16. -python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./engine_outputs - -# Run summarization using the InternLM2 7B model in FP16 using two GPUs. -mpirun -n 2 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-7b/ \ - --data_type fp16 \ - --engine_dir ./internlm2-chat-7b/trt_engines/fp16/2-gpu/ - -# Run summarization using the InternLM2 20B model in BF16 using 4 GPUs. -mpirun -n 4 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm --test_hf \ - --hf_model_dir ./internlm2-chat-20b/ \ - --data_type bf16 \ - --engine_dir ./internlm2-chat-20b/trt_engines/bf16/4-gpu/ -``` diff --git a/examples/models/core/internlm2/convert_checkpoint.py b/examples/models/core/internlm2/convert_checkpoint.py deleted file mode 100644 index 44c80d6d51ff..000000000000 --- a/examples/models/core/internlm2/convert_checkpoint.py +++ /dev/null @@ -1,549 +0,0 @@ -import argparse -import json -import os -import time -import traceback -from concurrent.futures import ThreadPoolExecutor, as_completed -from typing import Dict, Optional - -import numpy as np -import safetensors -import torch -from einops import rearrange -from transformers import AutoConfig, AutoModelForCausalLM - -import tensorrt_llm -from tensorrt_llm._deprecation import emit_engine_arch_deprecation -from tensorrt_llm._utils import get_hf_rope_theta, release_gc -from tensorrt_llm.mapping import Mapping -from tensorrt_llm.models.llama import convert - - -def parse_arguments(): - parser = argparse.ArgumentParser() - parser.add_argument('--model_dir', type=str, default=None) - parser.add_argument('--tp_size', - type=int, - default=1, - help='N-way tensor parallelism size') - parser.add_argument('--pp_size', - type=int, - default=1, - help='N-way pipeline parallelism size') - parser.add_argument('--dtype', - type=str, - default='float16', - choices=['float32', 'bfloat16', 'float16']) - parser.add_argument( - '--use_parallel_embedding', - action="store_true", - default=False, - help= - 'By default embedding parallelism is disabled. By setting this flag, embedding parallelism is enabled' - ) - parser.add_argument( - '--embedding_sharding_dim', - type=int, - default=0, - choices=[0, 1], - help= - 'By default the embedding lookup table is sharded along vocab dimension (embedding_sharding_dim=0). ' - 'To shard it along hidden dimension, set embedding_sharding_dim=1' - 'Note: embedding sharing is only enabled when embedding_sharding_dim = 0' - ) - - parser.add_argument( - '--use_weight_only', - default=False, - action="store_true", - help='Quantize weights for the various GEMMs to INT4/INT8.' - 'See --weight_only_precision to set the precision') - parser.add_argument( - '--weight_only_precision', - const='int8', - type=str, - nargs='?', - default='int8', - choices=['int8', 'int4'], - help= - 'Define the precision for the weights when using weight-only quantization.' - 'You must also use --use_weight_only for that argument to have an impact.' - ) - parser.add_argument('--output_dir', - type=str, - default='tllm_checkpoint', - help='The path to save the TensorRT LLM checkpoint') - parser.add_argument( - '--workers', - type=int, - default=1, - help='The number of workers for converting checkpoint in parallel') - parser.add_argument('--log_level', type=str, default='info') - args = parser.parse_args() - - tensorrt_llm.logger.set_level(args.log_level) - return args - - -def get_qkv_weight(weight: torch.Tensor, - hidden_size: int, - num_heads: int, - tp_size: int, - tp_rank: int, - is_bias: bool, - num_kv_heads: Optional[int] = None) -> torch.Tensor: - """ Splits the QKV matrix according to tensor parallelism """ - head_size = hidden_size // num_heads - num_kv_groups = num_heads // num_kv_heads - mha_mode = num_kv_heads == num_heads - weight = rearrange(weight, - '(h gs d) dim -> h gs d dim', - gs=2 + num_kv_groups, - d=head_size) - q_w, k_w, v_w = torch.split(weight, [num_kv_groups, 1, 1], dim=1) - if is_bias: - q_w = q_w.ravel() - k_w = k_w.ravel() - v_w = v_w.ravel() - qkv_w = torch.cat((q_w, k_w, v_w)) - qkv_w = convert.split_qkv_bias_tp(qkv_w, num_heads, hidden_size, - tp_size, tp_rank) - else: - q_w = rearrange(q_w, 'h gs d dim -> (h gs d) dim') - k_w = rearrange(k_w, 'h gs d dim -> (h gs d) dim') - v_w = rearrange(v_w, 'h gs d dim -> (h gs d) dim') - if not mha_mode: - if num_kv_heads < tp_size: - k_w = convert.dup_kv_weight(k_w, num_kv_heads, tp_size) - v_w = convert.dup_kv_weight(v_w, num_kv_heads, tp_size) - assert (k_w.shape[0] % (tp_size * head_size)) == 0 - assert (v_w.shape[0] % (tp_size * head_size)) == 0 - wq = convert.split(q_w, tp_size, tp_rank) - wk = convert.split(k_w, tp_size, tp_rank) - wv = convert.split(v_w, tp_size, tp_rank) - qkv_w = torch.concat((wq, wk, wv)) - - else: - qkv_w = torch.cat([q_w, k_w, v_w], dim=0) - - qkv_w = convert.split_qkv_tp(qkv_w, num_heads, hidden_size, tp_size, - tp_rank) - return qkv_w - - -def get_tllm_linear_weight( - weight: torch.Tensor, - prefix: str, - bias: Optional[torch.Tensor] = None, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8 -) -> Dict[str, torch.Tensor]: - results = {} - if use_weight_only: - v = weight.t().contiguous() - processed_torch_weights, torch_weight_scales = \ - torch.ops.trtllm.symmetric_quantize_last_axis_of_batched_matrix( - v, plugin_weight_only_quant_type) - results[f'{prefix}.weight'] = processed_torch_weights - results[f'{prefix}.per_channel_scale'] = torch_weight_scales - else: - results[f'{prefix}.weight'] = weight.contiguous() - - if bias is not None: - results[f'{prefix}.bias'] = bias - - return results - - -def convert_from_hf(hf_model, - hf_config, - mapping: Mapping, - dtype: str = 'float32', - use_parallel_embedding: bool = False, - sharding_dim: int = 0, - use_weight_only: bool = False, - plugin_weight_only_quant_type: torch.dtype = torch.int8): - weights = {} - tik = time.time() - - model_params = dict(hf_model.named_parameters()) - #This is for InternVL2 - if hf_config.architectures[0] == 'InternLM2ForCausalLM': - keys_to_rename = [ - key for key in model_params.keys() if 'language_model.' in key - ] - keys_to_delete = [ - key for key in model_params.keys() if 'vision_model.' in key - ] - for key in keys_to_rename: - keys_rename = key.replace('language_model.', '') - model_params[keys_rename] = model_params[key] - del model_params[key] - for key in keys_to_delete: - del model_params[key] - - dtype = getattr(torch, dtype) - num_attention_heads = hf_config.num_attention_heads - hidden_size = hf_config.hidden_size - vocab_size = hf_config.vocab_size - num_kv_heads = hf_config.num_key_value_heads - num_hidden_layers = hf_config.num_hidden_layers - layers_range = mapping.pp_layers(num_hidden_layers) - is_xcomposer2 = ( - hf_config.architectures[0] == 'InternLMXComposer2ForCausalLM') - lora_weights = {} - for l in layers_range: - prefix = f'model.layers.{l}' - tllm_prex = f'transformer.layers.{l - layers_range[0]}' - - qkv_weight = convert.get_weight(model_params, - f'{prefix}.attention.wqkv', dtype) - qkv_w = get_qkv_weight(qkv_weight, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=False, - num_kv_heads=num_kv_heads) - - if is_xcomposer2: - lora_prefix = f'base_model.model.model.layers.{l}' - assert num_attention_heads % num_kv_heads == 0 - num_key_value_groups = num_attention_heads // num_kv_heads - - qkv_loraA = convert.get_weight(model_params, - prefix + '.attention.wqkv.Plora_A', - dtype) - qkv_loraB = convert.get_weight(model_params, - prefix + '.attention.wqkv.Plora_B', - dtype) - q_lora_name = f"{lora_prefix}.self_attn.q_proj" - k_lora_name = f"{lora_prefix}.self_attn.k_proj" - v_lora_name = f"{lora_prefix}.self_attn.v_proj" - - #save qkv_loraA to be (q/k/v)_loraA - lora_weights[f"{q_lora_name}.lora_A.weight"] = qkv_loraA - lora_weights[f"{k_lora_name}.lora_A.weight"] = qkv_loraA - lora_weights[f"{v_lora_name}.lora_A.weight"] = qkv_loraA - - qkv_lora_rank = qkv_loraB.shape[-1] - head_size = hidden_size // num_attention_heads - - qkv_loraB = qkv_loraB.reshape(-1, num_key_value_groups + 2, - head_size, qkv_lora_rank) - q_loraB = qkv_loraB[:, :num_key_value_groups, :, :].reshape( - -1, qkv_lora_rank).contiguous() - k_loraB = qkv_loraB[:, - -2, :, :].reshape(-1, - qkv_lora_rank).contiguous() - v_loraB = qkv_loraB[:, - -1, :, :].reshape(-1, - qkv_lora_rank).contiguous() - - #save (q/k/v)_loraB - lora_weights[f"{q_lora_name}.lora_B.weight"] = q_loraB - lora_weights[f"{k_lora_name}.lora_B.weight"] = k_loraB - lora_weights[f"{v_lora_name}.lora_B.weight"] = v_loraB - - wo_loraA = convert.get_weight(model_params, - prefix + '.attention.wo.Plora_A', - dtype) - wo_loraB = convert.get_weight(model_params, - prefix + '.attention.wo.Plora_B', - dtype) - - lora_weights[ - f"{lora_prefix}.self_attn.o_proj.lora_A.weight"] = wo_loraA - lora_weights[ - f"{lora_prefix}.self_attn.o_proj.lora_B.weight"] = wo_loraB - - mlp_gate_loraA = convert.get_weight( - model_params, prefix + '.feed_forward.w3.Plora_A', dtype) - mlp_gate_loraB = convert.get_weight( - model_params, prefix + '.feed_forward.w3.Plora_B', dtype) - lora_weights[ - f"{lora_prefix}.mlp.up_proj.lora_A.weight"] = mlp_gate_loraA - lora_weights[ - f"{lora_prefix}.mlp.up_proj.lora_B.weight"] = mlp_gate_loraB - - mlp_fc_loraA = convert.get_weight( - model_params, prefix + '.feed_forward.w1.Plora_A', dtype) - mlp_fc_loraB = convert.get_weight( - model_params, prefix + '.feed_forward.w1.Plora_B', dtype) - lora_weights[ - f"{lora_prefix}.mlp.gate_proj.lora_A.weight"] = mlp_fc_loraA - lora_weights[ - f"{lora_prefix}.mlp.gate_proj.lora_B.weight"] = mlp_fc_loraB - - mlp_proj_loraA = convert.get_weight( - model_params, prefix + '.feed_forward.w2.Plora_A', dtype) - mlp_proj_loraB = convert.get_weight( - model_params, prefix + '.feed_forward.w2.Plora_B', dtype) - lora_weights[ - f"{lora_prefix}.mlp.down_proj.lora_A.weight"] = mlp_proj_loraA - lora_weights[ - f"{lora_prefix}.mlp.down_proj.lora_B.weight"] = mlp_proj_loraB - - qkv_bias = None - if f'{prefix}.attention.wqkv.bias' in model_params: - qkv_bias = convert.get_bias(model_params, - f'{prefix}.attention.wqkv', dtype) - if qkv_bias is None: - qkv_b = None - else: - qkv_b = get_qkv_weight(qkv_bias, - hidden_size, - num_attention_heads, - mapping.tp_size, - mapping.tp_rank, - is_bias=True, - num_kv_heads=num_kv_heads) - weights.update( - get_tllm_linear_weight( - qkv_w, - f'{tllm_prex}.attention.qkv', - qkv_b, - use_weight_only, - plugin_weight_only_quant_type, - )) - - attn_dense_weight = convert.get_weight(model_params, - f'{prefix}.attention.wo', dtype) - attn_dense_w = convert.split_matrix_tp(attn_dense_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - attn_dense_bias = None - if f'{prefix}.attention.wo.bias' in model_params: - attn_dense_bias = convert.get_bias(model_params, - f'{prefix}.attention.wo', dtype) - - weights.update( - get_tllm_linear_weight( - attn_dense_w, - f'{tllm_prex}.attention.dense', - attn_dense_bias, - use_weight_only, - plugin_weight_only_quant_type, - )) - - mlp_fc_weight = convert.get_weight(model_params, - f'{prefix}.feed_forward.w1', dtype) - mlp_fc_w = convert.split_matrix_tp(mlp_fc_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - mlp_fc_b = None - weights.update( - get_tllm_linear_weight(mlp_fc_w, f'{tllm_prex}.mlp.fc', mlp_fc_b, - use_weight_only, - plugin_weight_only_quant_type)) - - mlp_proj_weight = convert.get_weight(model_params, - f'{prefix}.feed_forward.w2', dtype) - mlp_proj_w = convert.split_matrix_tp(mlp_proj_weight, - mapping.tp_size, - mapping.tp_rank, - dim=1) - mlp_proj_bias = None - weights.update( - get_tllm_linear_weight(mlp_proj_w, f'{tllm_prex}.mlp.proj', - mlp_proj_bias, use_weight_only, - plugin_weight_only_quant_type)) - - mlp_gate_weight = convert.get_weight(model_params, - f'{prefix}.feed_forward.w3', dtype) - mlp_gate_w = convert.split_matrix_tp(mlp_gate_weight, - mapping.tp_size, - mapping.tp_rank, - dim=0) - mlp_gate_bias = None - weights.update( - get_tllm_linear_weight(mlp_gate_w, f'{tllm_prex}.mlp.gate', - mlp_gate_bias, use_weight_only, - plugin_weight_only_quant_type)) - - # Layer norms do not use tensor parallelism - input_ln_weight = convert.get_weight(model_params, - f'{prefix}.attention_norm', dtype) - weights[f'{tllm_prex}.input_layernorm.weight'] = input_ln_weight - - post_ln_weight = convert.get_weight(model_params, f'{prefix}.ffn_norm', - dtype) - weights[f'{tllm_prex}.post_layernorm.weight'] = post_ln_weight - - release_gc() - - if is_xcomposer2: - torch.save(lora_weights, 'adapter_model.bin') - adapter_config = { - "base_model_name_or_path": - "Internlm-xcomposer-2-7b-vl-hf", - "bias": - "none", - "enable_lora": - None, - "fan_in_fan_out": - False, - "inference_mode": - True, - "lora_alpha": - 256.0, - "lora_dropout": - 0.05, - "merge_weights": - False, - "modules_to_save": - None, - "peft_type": - "LORA", - "r": - 256, - "target_modules": [ - "q_proj", "v_proj", "k_proj", "o_proj", "gate_proj", - "down_proj", "up_proj" - ], - "task_type": - "CAUSAL_LM" - } - with open(os.path.join('adapter_config.json'), 'w') as f: - json.dump(adapter_config, f, indent=4) - - embed_w = convert.get_weight(model_params, 'model.tok_embeddings', dtype) - if use_parallel_embedding: - embed_w = convert.split_matrix_tp(embed_w, - mapping.tp_size, - mapping.tp_rank, - dim=sharding_dim) - if mapping.is_first_pp_rank(): - weights['transformer.vocab_embedding.weight'] = embed_w - lm_head_weights = convert.get_weight(model_params, 'output', dtype) - if mapping.is_last_pp_rank(): - if vocab_size % mapping.tp_size != 0: - # padding - vocab_size_padded = convert.pad_vocab_size(vocab_size, - mapping.tp_size) - pad_width = vocab_size_padded - vocab_size - - lm_head_weights = torch.from_numpy( - np.pad(lm_head_weights.detach().cpu().numpy(), - ((0, pad_width), (0, 0)), - 'constant', - constant_values=0)) - weights['lm_head.weight'] = convert.split_matrix_tp(lm_head_weights, - mapping.tp_size, - mapping.tp_rank, - dim=0) - ln_f_w = convert.get_weight(model_params, 'model.norm', dtype) - weights['transformer.ln_f.weight'] = ln_f_w - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Weights loaded. Total time: {t}') - return weights - - -if __name__ == '__main__': - emit_engine_arch_deprecation("convert_checkpoint.py") - args = parse_arguments() - world_size = args.tp_size * args.pp_size - - tik = time.time() - - if not os.path.exists(args.output_dir): - os.makedirs(args.output_dir) - - quant_algo = None - plugin_weight_only_quant_type = None - if args.use_weight_only and args.weight_only_precision == 'int8': - plugin_weight_only_quant_type = torch.int8 - quant_algo = 'W8A16' - elif args.use_weight_only and args.weight_only_precision == 'int4': - plugin_weight_only_quant_type = torch.quint4x2 - quant_algo = 'W4A16' - - hf_config = AutoConfig.from_pretrained(args.model_dir, - trust_remote_code=True) - #This is for InternVL2 - if hasattr(hf_config, 'llm_config'): - hf_config = hf_config.llm_config - - config = { - 'architecture': hf_config.architectures[0], - 'dtype': args.dtype, - 'logits_dtype': 'float32', - 'num_hidden_layers': hf_config.num_hidden_layers, - 'num_attention_heads': hf_config.num_attention_heads, - 'num_key_value_heads': hf_config.num_key_value_heads, - 'hidden_size': hf_config.hidden_size, - 'intermediate_size': hf_config.intermediate_size, - 'norm_epsilon': hf_config.rms_norm_eps, - 'vocab_size': hf_config.vocab_size, - 'position_embedding_type': 'rope_gpt_neox', - 'rotary_base': get_hf_rope_theta(hf_config, 10000.0), - 'max_position_embeddings': hf_config.max_position_embeddings, - 'hidden_act': hf_config.hidden_act, - 'use_parallel_embedding': args.use_parallel_embedding, - 'embedding_sharding_dim': args.embedding_sharding_dim, - 'quantization': { - 'quant_algo': quant_algo, - }, - 'mapping': { - 'world_size': world_size, - 'tp_size': args.tp_size, - 'pp_size': args.pp_size, - }, - 'has_partial_lora_mask': - hf_config.architectures[0] == 'InternLMXComposer2ForCausalLM', - 'attn_bias': getattr(hf_config, 'bias', False), - 'rotary_scaling': getattr(hf_config, "rope_scaling", None) - } - - with open(os.path.join(args.output_dir, 'config.json'), 'w') as f: - json.dump(config, f, indent=4) - - def covert_and_save(rank): - mapping = Mapping(world_size=world_size, - rank=rank, - tp_size=args.tp_size, - pp_size=args.pp_size) - - hf_model = AutoModelForCausalLM.from_pretrained(args.model_dir, - trust_remote_code=True, - dtype="auto") - weights = convert_from_hf( - hf_model, - hf_config, - mapping, - dtype=args.dtype, - use_parallel_embedding=args.use_parallel_embedding, - sharding_dim=args.embedding_sharding_dim, - use_weight_only=args.use_weight_only, - plugin_weight_only_quant_type=plugin_weight_only_quant_type) - del hf_model - save_file = os.path.join(args.output_dir, f'rank{rank}.safetensors') - print(f'Saving to {save_file}') - safetensors.torch.save_file(weights, save_file) - - if args.workers == 1: - for rank in range(world_size): - covert_and_save(rank) - else: - with ThreadPoolExecutor(max_workers=args.workers) as p: - futures = [ - p.submit(covert_and_save, rank) for rank in range(world_size) - ] - exceptions = [] - for future in as_completed(futures): - try: - future.result() - except Exception as e: - traceback.print_exc() - exceptions.append(e) - assert len( - exceptions - ) == 0, "Checkpoint conversion failed, please check error log." - - tok = time.time() - t = time.strftime('%H:%M:%S', time.gmtime(tok - tik)) - print(f'Total time of converting checkpoints: {t}') diff --git a/examples/models/core/internlm2/requirements.txt b/examples/models/core/internlm2/requirements.txt deleted file mode 100644 index d27fa26c6812..000000000000 --- a/examples/models/core/internlm2/requirements.txt +++ /dev/null @@ -1 +0,0 @@ -einops diff --git a/examples/models/core/llama/.gitignore b/examples/models/core/llama/.gitignore deleted file mode 100644 index 02915ec1a252..000000000000 --- a/examples/models/core/llama/.gitignore +++ /dev/null @@ -1,5 +0,0 @@ -llama* -tokenizer.model -*output* -*.safetensors -*.json diff --git a/examples/models/core/llama/README.md b/examples/models/core/llama/README.md deleted file mode 100644 index f427053284ce..000000000000 --- a/examples/models/core/llama/README.md +++ /dev/null @@ -1,1592 +0,0 @@ -# LLaMA - -> [!WARNING] -> The `convert_checkpoint.py` / `trtllm-build` / `run.py` workflow described -> below is **legacy** and will not receive new features. New projects should use -> [`trtllm-serve`](https://nvidia.github.io/TensorRT-LLM/quick-start-guide.html) -> or the [LLM Python API](https://nvidia.github.io/TensorRT-LLM/llm-api/index.html) instead. - -This document shows how to build and run a LLaMA model in TensorRT LLM on both single GPU, single node multi-GPU and multi-node multi-GPU. - -- [LLaMA](#llama) - - [Overview](#overview) - - [Support Matrix](#support-matrix) - - [Usage](#usage) - - [Build TensorRT engine(s)](#build-tensorrt-engines) - - [LLaMA v2 Updates](#llama-v2-updates) - - [LLaMA v3 Updates](#llama-v3-updates) - - [Long context length](#long-context-length) - - [Long context evaluation](#long-context-evaluation) - - [1M long context test case](#1m-long-context-test-case) - - [INT8 KV cache](#int8-kv-cache) - - [SmoothQuant](#smoothquant) - - [FP8 Post-Training Quantization](#fp8-post-training-quantization) - - [Groupwise quantization (AWQ/GPTQ)](#groupwise-quantization-awqgptq) - - [AWQ](#awq) - - [GPTQ](#gptq) - - [w4aINT8 quantization (QServe)](#w4aint8-quantization-qserve) - - [NVFP4 quantization](#nvfp4-quantization) - - [Run](#run) - - [Multi-GPU multi-node (MGMN) support](#multi-gpu-multi-node-mgmn-support) - - [Summarization using the LLaMA model](#summarization-using-the-llama-model) - - [Mistral v0.1](#mistral-v01) - - [Mistral Nemo](#mistral-nemo) - - [Running CodeLlama](#running-codellama) - - [Build](#build) - - [Run](#run-1) - - [Run models with LoRA](#run-models-with-lora) - - [Run LLaMa with several lora checkpoints](#run-llama-with-several-lora-checkpoints) - - [Run FP8 Mistral v0.1 with FP16 lora checkpoint](#run-fp8-mistral-v01-with-fp16-lora-checkpoint) - - [Run INT4-AWQ LLaMa with several FP16 lora checkpoints](#run-int4-awq-llama-with-several-fp16-lora-checkpoints) - - [Run LLaMa with StreamingLLM](#run-llama-with-streamingllm) - - [Run LLaMA-3.1 405B Model](#run-llama-31-405b-model) - - [Convert Checkpoint to TensorRT LLM Unified Checkpoint](#convert-checkpoint-to-tensorrt-llm-unified-checkpoint) - - [Build Engine](#build-engine) - - [Run Inference](#run-inference) - - [Run LLaMa-3.3 70B Model on PyTorch Backend](#run-llama-33-70b-model-on-pytorch-backend) - - [Prepare TensorRT LLM extra configs](#prepare-tensorrt-llm-extra-configs) - - [Launch trtllm-serve OpenAI-compatible API server](#launch-trtllm-serve-openai-compatible-api-server) - - [Run performance benchmarks](#run-performance-benchmarks) - -## Overview - -The TensorRT LLM LLaMA implementation can be found in [tensorrt_llm/models/llama/model.py](../../../../tensorrt_llm/models/llama/model.py). The TensorRT LLM LLaMA example code is located in [`examples/models/core/llama`](./). There is one main file: - -* [`convert_checkpoint.py`](./convert_checkpoint.py) to convert the LLaMA model into TensorRT LLM checkpoint format. - -In addition, there are two shared files in the parent folder [`examples`](../../../) for inference and evaluation: - -* [`run.py`](../../../run.py) to run the inference on an input text; -* [`summarize.py`](../../../summarize.py) to summarize the articles in the [cnn_dailymail](https://huggingface.co/datasets/abisee/cnn_dailymail) dataset. - -## Support Matrix - * BF16/FP16 - * FP8 - * INT8 & INT4 Weight-Only - * SmoothQuant - * Groupwise quantization (AWQ/GPTQ) - * w4aINT8 quantization (QServe) - * FP8 KV CACHE - * INT8 KV CACHE (+ AWQ/per-channel weight-only) - * Tensor Parallel + Pipeline Parallel, Tensor Parallel + Context Parallel - * STRONGLY TYPED - -## Usage - -The TensorRT LLM LLaMA example code locates at [examples/models/core/llama](./). It takes HF weights as input, and builds the corresponding TensorRT engines. The number of TensorRT engines depends on the number of GPUs used to run inference. - -### Build TensorRT engine(s) - -Please install required packages first to make sure the example uses matched `tensorrt_llm` version: - -```bash -pip install --upgrade -r requirements.txt -``` - -Need to prepare the HF LLaMA checkpoint by following the guides here https://huggingface.co/docs/transformers/main/en/model_doc/llama. - -The `trtllm-build` command builds TensorRT engine(s) from HF checkpoint. If no checkpoint directory is specified, TensorRT LLM will build engine(s) with dummy weights. - -`trtllm-build` command has a variety of options. In particular, the plugin-related options have two categories: -* Plugin options that requires a data type (e.g., `gpt_attention_plugin`), you can - * explicitly specify `float16`/`bfloat16`/`float32`, so that the plugins are enabled with the specified precision; - * implicitly specify `auto`, so that the plugins are enabled with the precision automatically inferred from model dtype (i.e., the dtype specified in weight conversion); or - * disable the plugin by `disable`. -* Other features that requires a boolean (e.g., `context_fmha`, `paged_kv_cache`, `remove_input_padding`), you can - * enable/disable the feature by specifying `enable`/`disable`. - -The defaults have been carefully tuned for better performance. For example, `gpt_attention_plugin`, `context_fmha`, `paged_kv_cache` and `remove_input_padding` are enabled by default. See more details by `trtllm-build --help`. - -Normally `trtllm-build` only requires single GPU, but if you've already got all the GPUs needed for inference, you could enable parallel building to make the engine building process faster by adding `--workers` argument. Please note that currently `workers` feature only supports single node. - -`--use_fused_mlp=enable` enables GEMM horizontal fusion in gated MLP layer, which reduces input traffic and potentially improves performance. For FP8 PTQ, the downside is slight reduction of accuracy because one of the quantization scaling factors are discarded (accuracy 0.45734 vs 0.45755 for LLaMA-v2 7B using modelopt/examples/hf/instruct_eval/mmlu.py). - -`--use_fused_mlp=enable --gemm_swiglu_plugin ` fuses 2 GEMMs without biases and SwiGLU into one kernel. This is a preview feature and is only supported for dtype `fp8`. The supported architecture is SM90. - -Here're some examples: - -```bash -# Build a single-GPU float16 engine from HF weights. -# Try use_gemm_plugin to prevent accuracy issue. - -# Build the LLaMA 7B model using a single GPU and FP16. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_fp16 \ - --dtype float16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp16 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/1-gpu \ - --gemm_plugin auto - -# Build the LLaMA 7B model using a single GPU and BF16. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_bf16 \ - --dtype bfloat16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_bf16 \ - --output_dir ./tmp/llama/7B/trt_engines/bf16/1-gpu \ - --gemm_plugin auto - -# Build the LLaMA 7B model using a single GPU and apply INT8 weight-only quantization. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_fp16_wq \ - --dtype float16 \ - --use_weight_only \ - --weight_only_precision int8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp16_wq \ - --output_dir ./tmp/llama/7B/trt_engines/weight_only/1-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 7B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_2gpu_tp2 \ - --dtype float16 \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_tp2 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/2-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 7B using 2-way tensor parallelism and 2-way pipeline parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_4gpu_tp2_pp2 \ - --dtype float16 \ - --tp_size 2 \ - --pp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_tp2_pp2 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 7B using 2-way tensor parallelism and 2-way context parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_4gpu_tp2_cp2 \ - --dtype float16 \ - --tp_size 2 \ - --cp_size 2 -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_tp2_cp2 \ - --output_dir ./tmp/llama/7B/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 30B using 2-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/30B/hf/ \ - --output_dir ./tllm_checkpoint_2gpu_tp2 \ - --dtype float16 \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_tp2 \ - --output_dir ./tmp/llama/30B/trt_engines/fp16/2-gpu/ \ - --gemm_plugin auto -``` - -#### LLaMA v2 Updates -The LLaMA v2 models with 7B and 13B are compatible with the LLaMA v1 implementation. The above -commands still work. - - -For LLaMA v2 70B, there is a restriction on tensor parallelism that the number of KV heads -must be **divisible by the number of GPUs**. For example, since the 70B model has 8 KV heads, you can run it with -2, 4 or 8 GPUs (1 GPU as well for FP8). - - -```bash -# Build LLaMA 70B using 8-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 70B using 4-way tensor parallelism and 2-way pipeline parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --dtype float16 \ - --tp_size 4 \ - --pp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA 70B TP=8 using Meta checkpoints directly. -python convert_checkpoint.py --meta_ckpt_dir ./tmp/llama/70B/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto -``` - -Same instructions can be applied to fine-tuned versions of the LLaMA v2 models (e.g. 7Bf or llama-2-7b-chat). - -#### LLaMA v3 Updates -The LLaMA 3.0 models with 8B and 70b are compatible with the LLaMA v2 implementation. The above -commands still work. - -Note that the `rope_theta` and `vocab_size` are larger in LLaMA v3 models and these values are now inferred -or pickup up from the `params.json` when using the `meta_ckpt_dir`. - -LLaMA 3.2 models are also supported now. For text only model like [Llama-3.2-1B](https://huggingface.co/meta-llama/Llama-3.2-1B), the steps are same to v3.0. For vision model like [Llama-3.2-11B-Vision](https://huggingface.co/meta-llama/Llama-3.2-11B-Vision), please refer to the [examples/models/core/mllama/README.md](../mllama/README.md) - -```bash -# Build LLaMA v3 8B TP=1 using HF checkpoints directly. -python convert_checkpoint.py --model_dir ./tmp/llama/8B/hf/ \ - --output_dir ./tllm_checkpoint_1gpu_tp1 \ - --dtype float16 \ - --tp_size 1 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_tp1 \ - --output_dir ./tmp/llama/8B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 8B TP=1 using Meta checkpoints directly. -python convert_checkpoint.py --meta_ckpt_dir ./tmp/llama/8B/ \ - --output_dir ./tllm_checkpoint_1gpu_tp1 \ - --dtype float16 \ - --tp_size 1 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_tp1 \ - --output_dir ./tmp/llama/8B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 70B using 8-way tensor parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 70B using 4-way tensor parallelism and 2-way pipeline parallelism. -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --dtype float16 \ - --tp_size 4 \ - --pp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp4_pp2 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Build LLaMA v3 70B TP=8 using Meta checkpoints directly. -python convert_checkpoint.py --meta_ckpt_dir ./tmp/llama/70B/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto -``` - -Same instructions can be applied to fine-tuned versions of the LLaMA v2 models (e.g. 7Bf or llama-2-7b-chat). - -### Long context length -With long context lengths, multi_block_mode is turned on by default to enable faster decoding in multi-head attention. To disable this feature, add `--multi_block_mode=False` to the runtime command. - - -A few LLaMA models are fine-tuned for long context length that TRT-LLM can support today. For example https://huggingface.co/Yukang/LongAlpaca-70B employs rotary scaling plus fine-tuning to support up to 32K context length. The following show the steps for running LongAlpaca-70B in TRT-LLM: - - -```bash -# Build 8-GPU engine with long context LLaMA model -python convert_checkpoint.py --model_dir ./tmp/LongAlpaca-70B/ \ - --output_dir ./tllm_checkpoint_8gpu_tp8 \ - --dtype float16 \ - --tp_size 8 \ - -trtllm-build --checkpoint_dir ./tllm_checkpoint_8gpu_tp8 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --gemm_plugin auto - -# Get the long text data from Gutenberg Project -wget https://www.gutenberg.org/cache/epub/64317/pg64317.txt - -# Replace the line breaks with special character '\n' and append "Summarize this story:" at end of text -awk '{printf "%s\\n", $0} END {printf "\\nSummarize this story:"}' pg64317.txt > pg64317_sanitized.txt - -# Run with 8 GPUs -# Notice, `--max_input_length ` is a convenience option to limit the input length for the data. -# It should be set to the maximum context length the model supports. Here the limit is set to 32K. -mpirun -n 8 --allow-run-as-root \ - python ../../../run.py \ - --max_output_len 128 \ - --max_input_length 32768 \ - --input_file pg64317_sanitized.txt \ - --engine_dir ./tmp/llama/70B/trt_engines/fp16/8-gpu/ \ - --tokenizer_dir ./tmp/LongAlpaca-70B/ -``` - -Note that if engine is built with contiguous KV cache (i.e., without the flag `--paged_kv_cache`), you may need to reduce the max batch size (`--max_batch_size`) to fit the whole model and the KV cache in the GPU memory. The ballpark estimate for runtime memory consumption is given by - -``` -Total memory = (Model size + KV cache size + Activation memory) / Parallelism -``` - -where -- The model size is `the number of parameters * the size of data type`. -- The KV cache size is `the total number of tokens * the size of KV cache data type * the number of layers * the KV hidden dimension` -- The activation memory is determined by TRT engine, which can be a few GBs regardless of the degree of parallelism used - -For LLaMA v2 70B FP16 weights + FP8 KV cache, the model size is 70B parameters * 2 bytes = 140GB. The KV cache size is 32K tokens * 1 bytes * 80 layers * 2048 KV hidden dimension = 5GB per 32K tokens. We have 145GB spread across 8 GPUs. The end result is ~18GB per GPU plus some GBs of flat scratch/activation memory allocated by TRT engine and the TRT-LLM runtime. - -Note that the KV hidden dimension is derived by the number of KV heads times hidden dimension of each head. LLaMA v2 70B has hidden dimension of 8192, and uses grouped-query attention where 8 key heads and 8 value heads are associated with 64 query heads. Each head has hidden dimension of 8192/64 = 128. So the hidden dimension for KV in total is 128 * 8 * 2 = 2048. - -The total number of tokens is determined by beam width, batch size, and maximum sequence length. - -#### Long context evaluation - -* Download dataset and model - -```bash -git-lfs clone https://huggingface.co/datasets/DKYoon/SlimPajama-6B -git-lfs clone https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k/ -``` - -* Run examples with max_input_len 16384 - -To evaluate the PPL of very long context, we need to enable `use_paged_context_fmha` and setup `max_num_tokens` to enable the chunked context inference, reducing the activation memory requirement. Also, we need to enable `gather_context_logits` to return the logits to compute the PPL. - -```bash -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --dtype float16 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --gather_context_logits \ - --max_num_tokens 4096 \ - --max_input_len 16384 \ - --max_seq_len 16394 \ - --use_paged_context_fmha enable - -python ./examples/summarize.py --test_trt_llm \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --data_type fp16 \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --eval_task eval_context_ppl \ - --max_input_len 16384 \ - --use_py_session \ - --dataset_dir ./SlimPajama-6B/ -``` - -* Run evaluation on passkey task - -To evaluate the accuracy of very long context on `needle in haystack`, we need to enable `use_paged_context_fmha` and setup `max_num_tokens` to enable the chunked context inference, reducing the activation memory requirement. To save memory, we don't enable the `gather_context_logits` here because we don't need logits. - -```bash -python3 examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 4 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_input_len 131072 \ - --max_seq_len 131082 \ - --use_paged_context_fmha enable - -python examples/eval_long_context.py --task passkey \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --stop_idx 10 \ - --max_input_length 131072 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 131136 -``` - -* Run evaluation on kv_retrieval - -`kv_retrieval` is harder than `passkey` and is helpful to distinguish the model capability. - -To run the kv_retrieval, we need a third-party repo to prepare the keys. - -```bash -git clone git@github.com:nelson-liu/lost-in-the-middle.git -pip install -r lost-in-the-middle/requirements.txt -python -u lost-in-the-middle/scripts/make_kv_retrieval_data.py --num-keys 3000 --num-examples 500 --output-path kv-retrieval-3000_keys.jsonl.gz -gzip -d kv-retrieval-3000_keys.jsonl.gz -``` - -Prepare input data and run evaluation. - -```bash -python examples/infinitebench/construct_synthetic_dataset.py --test_case build_kv_retrieval --test_level 0 - -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --dtype float16 \ - --tp_size 1 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_input_len 131072 \ - --max_seq_len 131082 \ - --use_paged_context_fmha enable - -python examples/eval_long_context.py --task kv_retrieval \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --stop_idx 10 \ - --max_input_length 131072 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 131136 \ - --tensorrt_llm_accuracy_threshold 0.6 -``` - -expected results: - -```bash -[05/28/2024-03:31:43] [TRT-LLM] [I] ==== Evaluation ==== -[05/28/2024-03:31:43] [TRT-LLM] [I] # examples: 500 -[05/28/2024-03:31:43] [TRT-LLM] [I] Start index: 0 -[05/28/2024-03:31:43] [TRT-LLM] [I] Stop index: 10 -[05/28/2024-03:31:43] [TRT-LLM] [I] Max tokens: 50 -[05/28/2024-03:34:50] [TRT-LLM] [I] Compute the score -10it [00:00, 131072.00it/s] -[05/28/2024-03:34:51] [TRT-LLM] [I] Evaluation takes: 187.19733428955078 sec. -[05/28/2024-03:34:51] [TRT-LLM] [I] accuracy of 10 examples: 0.6 -``` - -#### 1M long context test case - -- Prepare 1M needle-in-a-haystack datasets - -```bash -python examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 7 -``` - -- Llama-3-8B example - -```bash -git-lfs clone https://huggingface.co/gradientai/Llama-3-8B-Instruct-Gradient-1048k/ - -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --dtype float16 \ - --tp_size 4 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-8B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-8B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_batch_size 1 \ - --max_seq_len 1048576 \ - --use_paged_context_fmha enable \ - --workers 4 - -mpirun -n 4 --allow-run-as-root python examples/eval_long_context.py --task passkey \ - --engine_dir /tmp/llama-3-8B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-8B-Instruct-Gradient-1048k/ \ - --stop_idx 1 \ - --max_input_length 1048566 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 1100000 -``` - -- Llama-3-70B example - -For the 70B model, at least 8 A100 80GB GPUs are required. - -```bash -git-lfs clone https://huggingface.co/gradientai/Llama-3-70B-Instruct-Gradient-1048k/ - -python examples/models/core/llama/convert_checkpoint.py --model_dir ./Llama-3-70B-Instruct-Gradient-1048k/ \ - --output_dir /tmp/llama-3-70B-1048k/trt_ckpts \ - --dtype float16 \ - --tp_size 8 - -python -m tensorrt_llm.commands.build --checkpoint_dir /tmp/llama-3-70B-1048k/trt_ckpts \ - --output_dir /tmp/llama-3-70B-1048k/trt_engines \ - --gemm_plugin float16 \ - --max_num_tokens 4096 \ - --max_batch_size 1 \ - --max_seq_len 1048576 \ - --use_paged_context_fmha enable \ - --workers 8 - -mpirun -n 8 --allow-run-as-root python examples/eval_long_context.py --task passkey \ - --engine_dir /tmp/llama-3-70B-1048k/trt_engines \ - --tokenizer_dir ./Llama-3-70B-Instruct-Gradient-1048k/ \ - --stop_idx 1 \ - --max_input_length 1048566 \ - --enable_chunked_context \ - --max_tokens_in_paged_kv_cache 1100000 -``` - -expected result: - -```bash -[05/27/2024-10:30:45] [TRT-LLM] [I] Compute the score -1it [00:00, 4215.38it/s] -[05/27/2024-10:30:45] [TRT-LLM] [I] accuracy of 1 examples: 1.0 -``` - -### INT8 KV cache -INT8 KV cache could be enabled to reduce memory footprint. It will bring more performance gains when batch size gets larger. - -For INT8 KV cache, [`convert_checkpoint.py`](./convert_checkpoint.py) features a -`--int8_kv_cache` option. Setting `--int8_kv_cache` will calibrate the model, -and then export the scaling factors needed for INT8 KV cache inference. - -Example: - -```bash -python convert_checkpoint.py --model_dir ./llama-models/llama-7b-hf \ - --output_dir ./llama-models/llama-7b-hf/int8_kv_cache/ \ - --dtype float16 \ - --int8_kv_cache -``` - -[`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 KV cache. - - -**INT8 KV cache + per-channel weight-only quantization** - -INT8 KV cache could be combined with per-channel weight-only quantization, as follows: - -Examples of INT8 weight-only quantization + INT8 KV cache - -```bash -# Build model with both INT8 weight-only and INT8 KV cache enabled -python convert_checkpoint.py --model_dir ./llama-models/llama-7b-hf \ - --output_dir ./tllm_checkpoint_1gpu_int8_kv_wq \ - --dtype float16 \ - --int8_kv_cache \ - --use_weight_only \ - --weight_only_precision int8 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_int8_kv_wq \ - --output_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_weight_only/1-gpu \ - --gemm_plugin auto -``` - -Test with `summarize.py`: - -```bash -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./llama-models/llama-7b-hf \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_weight_only/1-gpu \ - --test_hf -``` - -**INT8 KV cache + AWQ** - -In addition, you can enable INT8 KV cache together with AWQ (per-group INT4 weight-only quantization)like the following command. - -```bash -python ../../../quantization/quantize.py --model_dir /tmp/llama-7b-hf \ - --output_dir ./tllm_checkpoint_1gpu_awq_int8_kv_cache \ - --dtype float16 \ - --qformat int4_awq \ - --awq_block_size 128 \ - --kv_cache_dtype int8 \ - --calib_size 32 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_awq_int8_kv_cache \ - --output_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_int4_AWQ/1-gpu/ \ - --gemm_plugin auto \ -``` - -Test with `summarize.py`: - -```bash -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir /tmp/llama-7b-hf \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/int8_kv_cache_int4_AWQ/1-gpu \ - --test_hf -``` - -### SmoothQuant - -The smoothquant supports both LLaMA v1 and LLaMA v2. Unlike the FP16 build where the HF weights are processed and loaded into the TensorRT LLM directly, the SmoothQuant needs to load INT8 weights which should be pre-processed before building an engine. - -Example: -```bash -python3 convert_checkpoint.py --model_dir /llama-models/llama-7b-hf --output_dir /tmp/tllm_checkpoint_1gpu_sq --dtype float16 --smoothquant 0.5 -trtllm-build --checkpoint_dir /tmp/tllm_checkpoint_1gpu_sq \ - --output_dir ./engine_outputs \ - --gemm_plugin auto -``` - -[`convert_checkpoint.py`](./convert_checkpoint.py) add new options for the support of INT8 inference of SmoothQuant models. - -`--smoothquant` is the starting point of INT8 inference. By default, it -will run the model in the _per-tensor_ mode. - -Then, you can add any combination of `--per-token` and `--per-channel` to get the corresponding behaviors. - -Examples of build invocations: - -```bash -# Build model for SmoothQuant in the _per_token_ + _per_channel_ mode -python3 convert_checkpoint.py --model_dir /llama-models/llama-7b-hf \ - --output_dir /tmp/tllm_checkpoint_1gpu_sq \ - --dtype float16 \ - --smoothquant 0.5 \ - --per_token \ - --per_channel - -trtllm-build --checkpoint_dir /tmp/tllm_checkpoint_1gpu_sq \ - --output_dir ./engine_outputs \ - --gemm_plugin auto -``` - -### FP8 Post-Training Quantization - -The examples below uses the NVIDIA Modelopt (AlgorithMic Model Optimization) toolkit for the model quantization process. - -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - - -```bash -# Quantize HF LLaMA 70B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir ./tmp/llama/70B \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_2gpu_fp8 \ - --calib_size 512 \ - --tp_size 2 - -# Build trtllm engines from the trtllm checkpoint -# Enable fp8 context fmha to get further acceleration by setting `--use_fp8_context_fmha enable` -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_fp8 \ - --output_dir ./engine_outputs \ - --gemm_plugin auto \ - --workers 2 -``` - -**Note**: A LLaMA 70B model with BF16 is about 140GB, a LLaMA 70B model with FP8 is about 70GB. -The peak GPU memory consumption when doing FP8 quantizaton is more than 210GB (there is also some activation memory occupation when doing calibration). -So you need a node with at least 4 H100(A100) to run the quantization command. After quantization, 2 GPUs are okay to for building and run. - -Note: use FP8 GEMV to optimize performance in FP8 small-batch-size cases. - -```bash -# Quantize HF LLaMA 7B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir /tmp/llama-7b-hf \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_fp8 \ - --calib_size 512 - -# Build trtllm engines from the trtllm checkpoint -# Enable fp8 gemm plugin to get acceleration in small-batch-size cases -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp8 \ - --output_dir ./engine_outputs \ - --gemm_plugin fp8 -``` - -**Note**: FP8 gemv plugin uses CUDA cores to compute, by contrast to Tensor Core gemm kernel within cuBLAS. Over last year, as cuBLAS have improved their performance by a lot under small M case for Hopper(sm90), FP8 gemv kernel may or may not surpass cuBLAS, depending on specific gemm problem shape. Nonetheless, we still strongly recommend FP8 gemv kernel for Ada (sm89) as cuBLAS still falls behind gemv on it. - -### Groupwise quantization (AWQ/GPTQ) -One can enable AWQ/GPTQ INT4 weight only quantization with these options when building engine with `trtllm-build`: - -- `--use_weight_only` enables weight only GEMMs in the network. -- `--per_group` enable groupwise weight only quantization, for GPT-J example, we support AWQ with the group size default as 128. -- `--weight_only_precision` should specify the weight only quantization format. Supported formats are `int4_awq` or `int4_gptq`. -- `--quant_ckpt_path` passes the quantized checkpoint to build the engine. - -AWQ/GPTQ examples below involves 2 steps: -1. Weight quantization -2. Build TRT-LLM engine - -#### AWQ -1. Weight quantization: - - NVIDIA Modelopt toolkit is used for AWQ weight quantization. Please see [examples/quantization/README.md](/examples/quantization/README.md#preparation) for Modelopt installation instructions. - - ```bash - # Quantize HF LLaMA 7B checkpoint into INT4 AWQ format - python ../../../quantization/quantize.py --model_dir ./tmp/llama-7b-hf \ - --dtype float16 \ - --qformat int4_awq \ - --awq_block_size 128 \ - --output_dir ./quantized_int4-awq \ - --calib_size 32 - ``` - HF checkpoints generated with [AutoAWQ](https://github.com/casper-hansen/AutoAWQ) are also supported through the following conversion script: - - ```bash - # Convert AutoAWQ HF checkpoints into TRT-LLM checkpoint - python convert_checkpoint.py --model_dir ./tmp/Llama-2-7B-AWQ \ - --output_dir ./quantized_int4-awq - ``` - -2. Build TRT-LLM engine: - - ```bash - trtllm-build --checkpoint_dir ./quantized_int4-awq \ - --output_dir ./tmp/llama/7B/trt_engines/int4_AWQ/1-gpu/ \ - --gemm_plugin auto - ``` - -#### GPTQ -To run the GPTQ LLaMa example, the following steps are required: - -1. Weight quantization: - - Quantized weights for GPTQ are generated using [AutoGPTQ](https://github.com/AutoGPTQ/AutoGPTQ) as follow: - - ```bash - git clone https://github.com/AutoGPTQ/AutoGPTQ - cd AutoGPTQ - pip install . - - # Download the quant_autogptq script - wget https://gist.githubusercontent.com/TheBloke/b47c50a70dd4fe653f64a12928286682/raw/ebcee019d90a178ee2e6a8107fdd7602c8f1192a/quant_autogptq.py - - # Quantize weights into INT4 and save as safetensors - # Quantized weight with parameter "--act-order" is not supported in TRT-LLM - python quant_autogptq.py ./tmp/llama/7B ./llama-7b-4bit-gs128.safetensors wikitext --bits 4 --group_size 128 --desc_act 0 --damp 0.1 --dtype float16 --seqlen 4096 --num_samples 3 --use_fast - ``` - Then we can convert the saved `./llama-7b-4bit-gs128.safetensors` into TRT-LLM checkpoints by: - ```bash - # Build the LLaMA 7B model using 2-way tensor parallelism and apply INT4 GPTQ quantization. - # Compressed checkpoint safetensors are generated separately from GPTQ. - python convert_checkpoint.py --model_dir /tmp/llama-7b-hf \ - --output_dir ./tllm_checkpoint_2gpu_gptq \ - --dtype float16 \ - --quant_ckpt_path ./llama-7b-4bit-gs128.safetensors \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --per_group \ - --tp_size 2 - ``` - HF checkpoints generated with AutoGPTQ are also supported through the following conversion script: - - ```bash - # Convert AutoGPTQ HF checkpoints into 2-way tensor parallelism TRT-LLM checkpoint - python convert_checkpoint.py --model_dir ./tmp/Llama-2-7B-GPTQ \ - --output_dir ./tllm_checkpoint_2gpu_gptq \ - --tp_size 2 - ``` - -2. Build TRT-LLM engine: - - ```bash - # Build the LLaMA 7B model using 2-way tensor parallelism and apply INT4 GPTQ quantization. - # Compressed checkpoint safetensors are generated separately from GPTQ. - python convert_checkpoint.py --model_dir /tmp/llama-7b-hf \ - --output_dir ./tllm_checkpoint_2gpu_gptq \ - --dtype float16 \ - --quant_ckpt_path ./llama-7b-4bit-gs128.safetensors \ - --use_weight_only \ - --weight_only_precision int4_gptq \ - --per_group \ - --tp_size 2 - - trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu_gptq \ - --output_dir ./tmp/llama/7B/trt_engines/int4_GPTQ/2-gpu/ \ - --gemm_plugin auto - ``` - -### w4aINT8 quantization (QServe) - -TensorRT LLM integrates the quantized GEMM from [QServe](https://arxiv.org/abs/2405.04532), which employs 4-bit quantization for weights and 8-bit quantization for activations. This technique offers versatile performance benefits across different scenarios. When the GEMM's m dimension is small, as in small batch-size decoding, it achieves performance comparable to w4a16 by reducing the memory bandwidth required for weight access. Conversely, for larger m dimensions, such as during prefilling or large batch-size decoding, it matches the performance of w8a8 by leveraging INT8 Tensor Cores. - -Please follow the steps to run the model using QServe w4aINT8: - -1. Weight quantization: - - Currently we rely on the 3rd-party repo [deepcompressor](https://github.com/mit-han-lab/deepcompressor) to prepare the fake-quantized checkpoint. Follow the [instructions](https://github.com/mit-han-lab/deepcompressor/blob/main/examples/llm/README.md#usage) to quantize the model. Please use the `configs/qoq-g128.yaml` for per-group quantization, and `configs/qoq-gchn.yaml` for the per-channel quantization. Do not forget to add the flag `--save-model path/to/deepcompressor/ckpt` so that the quantized model is dumped to the disk. - - After quantization, the weights in the original Hugging Face checkpoint (assume under `path/to/huggingface/ckpt/`) will be quantized and the following files are obtained under your `path/to/deepcompressor/ckpt`: - - - `model.pt`: fake-quantized fp16 weights. - - `scale.pt`: quantization scales and zeros. - -2. Checkpoint conversion: - - Convert the DeepCompressor checkpoint into TensorRT LLM checkpoint, potentially with tensor parallelism: - - ```bash - export TRTLLM_DISABLE_UNIFIED_CONVERTER=1 # The current checkpoint conversion code requires legacy path - python convert_checkpoint.py --model_dir path/to/huggingface/ckpt/ \ - --output_dir path/to/trtllm/ckpt/ \ - --dtype float16 \ - --quant_ckpt_path path/to/deepcompressor/ckpt \ - --use_qserve \ - --per_group \ # Add this option if using per-group quantization - --tp_size 2 - ``` - -3. Build engine: - - ```bash - trtllm-build --checkpoint_dir path/to/trtllm/ckpt/ \ - --output_dir path/to/trtllm/engine \ - --gemm_plugin auto - ``` - -### NVFP4 quantization - -TRTLLM supports NVFP4 precision with blocksize=16 for both activations and GEMM weights. - -Please follow the steps to run the model using: - -1. Weight quantization and activation calibration using modelopt: - - ```bash - python example/quantization/quantize.py --model_dir path/to/huggingface/ckpt/ \ - --output_dir path/to/trtllm/ckpt/ \ - --dtype float16 \ - --qformat nvfp4 \ - --kv_cache_dtype fp8 \ - --tp_size 1 - ``` - -2. Build engine: - - ```bash - trtllm-build --checkpoint_dir path/to/trtllm/ckpt/ \ - --output_dir path/to/trtllm/engine - - # with FP8 paged context FMHA for better performance - trtllm-build --checkpoint_dir path/to/trtllm/ckpt/ \ - --output_dir path/to/trtllm/engine \ - --use_paged_context_fmha enable \ - --use_fp8_context_fmha enable - ``` - -### Run - -To run a TensorRT LLM LLaMA model using the engines generated by `trtllm-build` - -```bash -# With fp16 inference -python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/7B/ \ - --engine_dir=./tmp/llama/7B/trt_engines/fp16/1-gpu/ - -# With bf16 inference -python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/7B/ \ - --engine_dir=./tmp/llama/7B/trt_engines/bf16/1-gpu/ -``` - -### Multi-GPU multi-node (MGMN) support - -In MGMN case, you can still convert and build engines on a single node and then run the model on a multi-node environment, such as [Slurm](https://slurm.schedmd.com/documentation.html). - -For example, to build LLaMA 70B for 2 nodes with 8 GPUs per node, we can use 8-way tensor parallelism and 2-way pipeline parallelism: - -```bash -python convert_checkpoint.py --model_dir ./tmp/llama/70B/hf/ \ - --output_dir ./tllm_checkpoint_16gpu_tp8_pp2 \ - --dtype float16 \ - --tp_size 8 \ - --pp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_16gpu_tp8_pp2 \ - --output_dir ./tmp/llama/70B/trt_engines/fp16/16-gpu/ \ - --workers 8 \ - --gemm_plugin auto -``` - -Note that `–-workers` is still set to 8 to build all engines within a single node. - -To run the LLaMA 70B model on 2 nodes via Slurm, you need to prepare a Slurm script to submit the task, the script contains the following lines: - -```bash -#SBATCH -N 2 -#SBATCH --ntasks-per-node=8 -#SBATCH -p -# more sbatch options here... - -srun --container-image= \ - --mpi=pmix \ - ... \ # more srun options here - python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/70B/hf/ \ - --engine_dir=./tmp/llama/70B/trt_engines/fp16/16-gpu/ -``` - -Finally, you can submit the task with `sbatch .sh`. - -Considering the Slurm or other cluster management systems may be highly customized and the task-submit command may be variant, the forementioned example is for reference only. The key point is to submit the Python script with the MPI runtime, and TensorRT LLM will take care of the rest. - -### Summarization using the LLaMA model - -```bash -# Run summarization using the LLaMA 7B model in FP16. -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./tmp/llama/7B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/fp16/1-gpu/ - -# Run summarization using the LLaMA 7B model quantized to INT8. -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./tmp/llama/7B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/weight_only/1-gpu/ - -# Run summarization using the LLaMA 7B model in FP16 using two GPUs. -mpirun -n 2 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./tmp/llama/7B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/7B/trt_engines/fp16/2-gpu/ - -# Run summarization using the LLaMA 30B model in FP16 using two GPUs. -mpirun -n 2 --allow-run-as-root \ - python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./tmp/llama/30B/ \ - --data_type fp16 \ - --engine_dir ./tmp/llama/30B/trt_engines/fp16/2-gpu/ -``` - -#### Mistral v0.1 -Mistral v0.1 is compatible with LLaMA interface and can be built and run using the same instructions. -Setting `--max_input_len`, corresponding to the `max_position_embeddings` in the original Mistral config explicitly regulates context size. -The `--max_attention_window_size` parameter is set to the `sliding_window` value in the config and regulates both sliding window attention in the context phase and rolling buffer cache in the generation phase. - -```bash -# Build Mistral 7B with max input length 32256 -python convert_checkpoint.py --model_dir ./mistral-7b-v0.1 \ - --output_dir ./tllm_checkpoint_1gpu_mistral \ - --dtype float16 -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_mistral \ - --output_dir ./tmp/mistral/7B/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto \ - --max_input_len 32256 - -# Run Mistral 7B fp16 inference with sliding window/cache size 4096 -python ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./mistral-7b-v0.1 \ - --engine_dir=./tmp/mistral/7B/trt_engines/fp16/1-gpu/ \ - --max_attention_window_size=4096 -``` - -Note that if you are comparing TRT-LLM with Huggingface, -you should install `transformers` with version >= 4.34.1 in order to have Mistral model supported. -And upgrade `flash-attn` package by `pip install --upgrade flash-attn` or you may see wrong results generated by the huggingface implementation. - -#### Mistral Nemo -[Mistral Nemo](https://mistral.ai/news/mistral-nemo/) is compatible with LLaMA interface and can be built and run using the same instructions. -Please upgrade the [transformers](https://pypi.org/project/transformers/) to 4.43.0.dev0 or higher release for running this model. - -```bash -# Build Mistral Nemo with max input length 10240 -python convert_checkpoint.py --model_dir ./Mistral-Nemo-Instruct-2407 \ - --output_dir ./tllm_checkpoint_1gpu_mistral_nemo \ - --dtype bfloat16 \ - --smoothquant 0.5 \ - --per_channel \ - --per_token - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_mistral_nemo \ - --output_dir ./tmp/mistral_nemo/trt_engines/bf16/1-gpu/ \ - --gemm_plugin bfloat16 \ - --max_input_len 10240 - -# Run summarization using the Mistral Nemo model quantized to INT8. -python ../../../summarize.py --test_trt_llm \ - --hf_model_dir ./Mistral-Nemo-Instruct-2407 \ - --data_type bf16 \ - --engine_dir ./tmp/mistral_nemo/trt_engines/bf16/1-gpu// -``` - -## Running CodeLlama -Those examples can be used to build and run the CodeLlama models. All 7b, 13b, and 34b sizes and variants are supported. - -There are a couple of differences in CodeLlama in comparison to LLaMA v1/v2 models: rotary_base (`theta=1000000.0f`) and vocabulary size (`32016` (1)). - -_(1): Only applicable to 7b and 13b model sizes_. 34b model variants use `32000`. - -### Build -Use the following command to build `CodeLlama-7b-Instruct`: -```bash -python convert_checkpoint.py --model_dir /tmp/CodeLlama-7b-Instruct-hf \ - --output_dir ./tllm_checkpoint_1gpu_codellama \ - --dtype float16 - - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_codellama \ - --output_dir ./tmp/codellama/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto -``` -The example below uses the NVIDIA ModelOpt (AlgorithMic Model Optimization) toolkit for the model quantization process. -First make sure Modelopt toolkit is installed (see [examples/quantization/README.md](/examples/quantization/README.md#preparation)) - -```bash -# Quantize HF CodeLlama 7B into FP8 and export trtllm checkpoint -python ../../../quantization/quantize.py --model_dir /tmp/CodeLlama-7b-Instruct-hf \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_fp8 \ - --calib_size 512 - -# Build trtllm engines from the trtllm checkpoint -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp8 \ - --output_dir ./engine_outputs \ - --gemm_plugin auto -``` - -Use the following command to build `CodeLlama-34b-Instruct` for 4 GPUs (TP=4): -```bash -python convert_checkpoint.py --model_dir /tmp/CodeLlama-34b-Instruct-hf \ - --output_dir ./tllm_checkpoint_4gpu_codellama \ - --dtype float16 \ - --tp_size 4 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_codellama \ - --output_dir ./tmp/codellama/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto -``` - -NOTE: CodeLlama uses the `max_position_embeddings` of 16K. -To build the engine for running similarly long input/output, you need to specify that during build. - -Use `--max_input_len` (which defaults to `1024`) and `--max_seq_len` (which by default is deduced from `max_position_embeddings`) according to your use case, e.g.: -```bash -python convert_checkpoint.py --model_dir /tmp/CodeLlama-34b-Instruct-hf \ - --output_dir ./tllm_checkpoint_4gpu_codellama \ - --dtype float16 \ - --tp_size 8 \ - --use_parallel_embedding - -trtllm-build --checkpoint_dir ./tllm_checkpoint_4gpu_codellama \ - --output_dir ./tmp/codellama/trt_engines/fp16/4-gpu/ \ - --gemm_plugin auto \ - --max_input_len 15360 \ - --max_seq_len 16384 \ - --max_batch_size 4 -``` - -### Run -Use the following command to run the 7b engine from above: -```bash -python ../run.py --max_output_len=40 --tokenizer_dir . --engine_dir codellama_7b --input_text "In Bash, how do I list all text files?" -``` -Use the following command to run the 34b engine with long input/output from above: -```bash -mpirun -n 8 --allow-run-as-root \ - python ../../../run.py --max_output_len=160 --tokenizer_dir ./CodeLlama-34b-Instruct \ - --engine_dir codellama_34b --input_text "In python, write a function for binary searching an element in an integer array." -``` - -## Run models with LoRA - -Download the base model and lora model from HF: - -```bash -git-lfs clone https://huggingface.co/meta-llama/Llama-2-13b-hf -git-lfs clone https://huggingface.co/hfl/chinese-llama-2-lora-13b -``` - -Build engine, setting `--lora_plugin` and `--lora_dir`. If lora has separate lm_head and embedding, they will replace lm_head and embedding of base model. - -```bash -python convert_checkpoint.py --model_dir Llama-2-13b-hf \ - --output_dir ./tllm_checkpoint_2gpu \ - --dtype float16 \ - --tp_size 2 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_2gpu \ - --output_dir /tmp/new_lora_13b/trt_engines/fp16/2-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 1 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir chinese-llama-2-lora-13b -``` - -Run inference. Remember to use lora tokenizer because lora model has larger vocab size. - -```bash -mpirun -n 2 python ../../../run.py --engine_dir "/tmp/new_lora_13b/trt_engines/fp16/2-gpu/" \ - --max_output_len 50 \ - --tokenizer_dir "chinese-llama-2-lora-13b/" \ - --input_text "今天天气很好，我到公园的时候，" \ - --lora_task_uids 0 \ - --no_add_special_tokens \ - --use_py_session - - Input: "今天天气很好，我到公园的时候，" -Output: "发现公园里到处都是人，有的在跑步，有的在打羽毛球，还有的在跳绳，我和妈妈一起在公园里散步，我和妈妈在公园里散步的时候，看见了一位老爷爷在打羽毛球" -``` - -Users who want to skip LoRA module may pass uid -1 with `--lora_task_uids -1`. -In that case, the model will not run the LoRA module and the results will be -different. Since the LoRA tokenizer, embedding and LM head are still used, -the results will also be different with vanilla LLaMA and significantly degrade compared with `--lora_task_uids 0`. - -```bash -mpirun -n 2 python ../../../run.py --engine_dir "/tmp/new_lora_13b/trt_engines/fp16/2-gpu/" \ - --max_output_len 50 \ - --tokenizer_dir "chinese-llama-2-lora-13b/" \ - --input_text "今天天气很好，我到公园的时候，" \ - --lora_task_uids -1 \ - --no_add_special_tokens \ - --use_py_session - - Input: "今天天气很好，我到公园的时候，" -Output: "看见好多人们都看书，看书书看书书，看书书看书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书书" -``` - -### Run LLaMa with several lora checkpoints - -In this section, we show how to run a model with multiple LoRA modules at the same time. Note that if one of the LoRA module has a -fine-tuned embedding table or logit GEMM, users should guarantee that all the instances of the model can use the same fine-tuned -embedding table or logit GEMM. -Here, we use two LoRA checkpoints as examples. These two LoRA checkponits add LoRA modules to `q_proj` and `v_proj`. Because we only -support adding lora modules on `q`, `k` and `v` at the same time, we need to add `--lora_target_modules "attn_q" "attn_k" "attn_v"`. -In this case, we assign null pointers for the `k` LoRA module in TensorRT LLM and skip the computation at runtime. - -As the rank of the LoRA modules of both checkpoints is 8, we can set `--max_lora_rank 8` to reduce the memory requirement for the LoRA plugin. - -In this example, we use a LoRA checkpoint fine-tuned on the Chinese dataset `luotuo-lora-7b-0.1` and a LoRA checkpoint fine-tuned on -the Japanese dataset `Japanese-Alpaca-LoRA-7b-v0`. For the `lora_manager` to load several checkpoints, we pass several directories -of LoRA checkpoints at the same time: `--lora_dir "luotuo-lora-7b-0.1/" "Japanese-Alpaca-LoRA-7b-v0/"`. -Then, `lora_manager` will assign `lora_task_uids` to these checkpoints. `lora_task_uids -1` is a predefined value, which corresponds to -the base model. If we pass `lora_task_uids 0 1`, this means we want to use the first LoRA checkpoint on first sentence and use the second LoRA checkpoint on the second sentence. - -To verify the correctness, we pass the same Chinese input `美国的首都在哪里? \n答案:` three times as well as the same Japanese input `アメリカ合衆国の首都はどこですか? \n答え:` three times (both inputs mean `Where is the capital of America? \nAnswer`). We run on base model, `luotuo-lora-7b-0.1` and `Japanese-Alpaca-LoRA-7b-v0/`. - -```bash -git-lfs clone https://huggingface.co/qychen/luotuo-lora-7b-0.1 -git-lfs clone https://huggingface.co/kunishou/Japanese-Alpaca-LoRA-7b-v0 -BASE_LLAMA_MODEL=llama-7b-hf/ - -python convert_checkpoint.py --model_dir ${BASE_LLAMA_MODEL} \ - --output_dir ./tllm_checkpoint_1gpu \ - --dtype float16 -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu \ - --output_dir /tmp/llama_7b_with_lora_qkv/trt_engines/fp16/1-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 8 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir "luotuo-lora-7b-0.1/" "Japanese-Alpaca-LoRA-7b-v0/" \ - --max_lora_rank 8 \ - --lora_target_modules attn_q attn_k attn_v - -python ../../../run.py --engine_dir "/tmp/llama_7b_with_lora_qkv/trt_engines/fp16/1-gpu/" \ - --max_output_len 10 \ - --tokenizer_dir ${BASE_LLAMA_MODEL} \ - --input_text "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" \ - --lora_task_uids -1 0 1 -1 0 1 \ - --use_py_session --top_p 0.5 --top_k 0 -``` - -The results would be like - -```bash -Input [Text 0]: " 美国的首都在哪里? \n答案:" -Output [Text 0 Beam 0]: "Washington, D.C. -What is the" - -Input [Text 1]: " 美国的首都在哪里? \n答案:" -Output [Text 1 Beam 0]: "华盛顿。 -" - -Input [Text 2]: " 美国的首都在哪里? \n答案:" -Output [Text 2 Beam 0]: "Washington D.C.�����" - -Input [Text 3]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 3 Beam 0]: "Washington, D.C. -Which of" - -Input [Text 4]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 4 Beam 0]: "华盛顿。 -" - -Input [Text 5]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 5 Beam 0]: "ワシントン D.C." -``` - -We can observe that `luotuo-lora-7b-0.1` produces correct answers on the first sentence and the fifth sentence (in Chinese), `Japanese-Alpaca-LoRA-7b-v0` produces correct answers on the sixth sentence (in Japanese). - -### Run FP8 Mistral v0.1 with FP16 lora checkpoint - -In this section, we use Mistral v0.1 as an example show how to run an FP8 base model with FP16 LoRA module. - -* download the base model and lora model from HF - -```bash -git-lfs clone https://huggingface.co/davidkim205/komt-mistral-7b-v1 -git-lfs clone https://huggingface.co/davidkim205/komt-mistral-7b-v1-lora -``` - -* Quantize the Mistral v0.1 model to fp8 from HF -```bash -BASE_MISTRAL_MODEL=komt-mistral-7b-v1/ -python ../../../quantization/quantize.py --model_dir ${BASE_MISTRAL_MODEL} \ - --dtype float16 \ - --qformat fp8 \ - --kv_cache_dtype fp8 \ - --output_dir ./tllm_checkpoint_1gpu_fp8 \ - --calib_size 512 -``` - -* Build engine and run inference with sliding window/cache size 4096. -```bash -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_fp8 \ - --output_dir /tmp/mistral_komt_lora/7B/trt_engines/fp8/1-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 8 \ - --max_input_len 32256 \ - --max_seq_len 33280 \ - --lora_dir ./komt-mistral-7b-v1-lora - -python ../../../run.py --max_output_len=1024 \ - --tokenizer_dir ./komt-mistral-7b-v1 \ - --engine_dir=/tmp/mistral_komt_lora/7B/trt_engines/fp8/1-gpu/ \ - --input_text "[INST]오늘은 날씨가 아주 좋다 내가 공원에 갔을 때 [/INST]" \ - --max_attention_window_size=4096 \ - --lora_task_uids 0 \ - --use_py_session \ - --temperature 0.8 \ - --top_p 0.8 \ - --top_k 100 -``` - -The results would be like - -```bash -Input [Text 0]: " [INST]오늘은 날씨가 아주 좋다 내가 공원에 갔을 때 [/INST]" -Output [Text 0 Beam 0]: "날씨가 아주 좋은 날에 공원에 갔을 때는 산책이나 운동을 즐기는 것이 좋습니다. 공원에서 걷거나 조깅을 하면서 신선한 공기를 마시고 자연 속에서 휴식을 취할 수 있습니다. 또한, 공원에서 가족이나 친구와 함께 피크닉을 즐기거나 야외 스포츠를 즐길 수도 있습니다. 날씨가 좋을 때 공원에 가는 것은 건강과 웰빙에 좋은 방법입니다." -``` - - -### Run INT4-AWQ LLaMa with several FP16 lora checkpoints - -TensorRT LLM can also support Quantized base model + FP16/BF16 LoRA. We can first quantize the base model and build engine with the quantized checkpoint and different LoRA adapters. In this section, we show how to run an INT4-AWQ llama model with multiple FP16 LoRA modules. - -* Quantize the llama model to INT4-AWQ from HF -```bash -BASE_LLAMA_MODEL=llama-7b-hf/ -python ../../../quantization/quantize.py --model_dir ${BASE_LLAMA_MODEL} \ - --output_dir ./tllm_checkpoint_1gpu_awq \ - --dtype float16 \ - --qformat int4_awq \ - --awq_block_size 128 \ - --calib_size 32 -``` - -* Download the lora model, build engine, and run inference. -```bash -git-lfs clone https://huggingface.co/qychen/luotuo-lora-7b-0.1 -git-lfs clone https://huggingface.co/kunishou/Japanese-Alpaca-LoRA-7b-v0 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_awq \ - --output_dir /tmp/llama_7b_with_lora_qkv/trt_engines/int4_AWQ/1-gpu/ \ - --gemm_plugin auto \ - --lora_plugin auto \ - --max_batch_size 8 \ - --max_input_len 512 \ - --max_seq_len 562 \ - --lora_dir "luotuo-lora-7b-0.1/" "Japanese-Alpaca-LoRA-7b-v0/" \ - --max_lora_rank 8 \ - --lora_target_modules attn_q attn_k attn_v - -python ../../../run.py --engine_dir "/tmp/llama_7b_with_lora_qkv/trt_engines/int4_AWQ/1-gpu/" \ - --max_output_len 10 \ - --tokenizer_dir ${BASE_LLAMA_MODEL} \ - --input_text "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "美国的首都在哪里? \n答案:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" "アメリカ合衆国の首都はどこですか? \n答え:" \ - --lora_task_uids -1 0 1 -1 0 1 \ - --use_py_session --top_p 0.5 --top_k 0 -``` - -The results would be like - -```bash -Input [Text 0]: " 美国的首都在哪里? \n答案:" -Output [Text 0 Beam 0]: "Washington, D.C. -What is the" - -Input [Text 1]: " 美国的首都在哪里? \n答案:" -Output [Text 1 Beam 0]: "华盛顿。 -" - -Input [Text 2]: " 美国的首都在哪里? \n答案:" -Output [Text 2 Beam 0]: "洛爱睿" - -Input [Text 3]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 3 Beam 0]: "Washington, D.C. -Copyright " - -Input [Text 4]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 4 Beam 0]: "华盛顿。 -" - -Input [Text 5]: " アメリカ合衆国の首都はどこですか? \n答え:" -Output [Text 5 Beam 0]: "ワシントン、D.C" -``` - -## Run LLaMa with StreamingLLM - -* Build engine. Set `--streamingllm enable` to enable StreamingLLM. - -```bash -# Build the LLaMA 7B model with StreamingLLM feature using a single GPU and FP16. -python convert_checkpoint.py --model_dir ./tmp/llama/7B/ \ - --output_dir ./tllm_checkpoint_1gpu_streamlingllm \ - --dtype float16 - -trtllm-build --checkpoint_dir ./tllm_checkpoint_1gpu_streamlingllm \ - --output_dir ./tmp/llama/7B/trt_engines/fp16_StreamingLLM/1-gpu/ \ - --gemm_plugin auto \ - --streamingllm enable - -``` - -* Run inference. Use `--sink_token_length` to set the number of sink tokens, and use `--max_attention_window_size` to set the `sliding_window` value. - -```bash -# Run LLaMA 7B fp16 inference with sliding window/cache size 2048 and sink token length 4. -python3 ../../../run.py --max_output_len=50 \ - --tokenizer_dir ./tmp/llama/7B/ \ - --engine_dir=./tmp/llama/7B/trt_engines/fp16_StreamingLLM/1-gpu/ \ - --max_attention_window_size=2048 \ - --sink_token_length=4 -``` - -Note that the sink tokens is included in the sliding attention tokens, and there are at most `max_attention_window_size` tokens are stored in the KV cache. - -## Run LLaMA-3.1 405B Model - -Currently, TensorRT LLM supports Meta checkpoint and Huggingface checkpoint for LLaMA-3.1. In this section, we demonstrate how to run the LLaMA-3.1 405B model via TensorRT-LLM. Here, we assume users have downloaded the checkpoints and placed them at `llama_3.1_405B_meta_model/` (Meta BF16 checkpoint), `llama_3.1_405B_HF_model/` (HF BF16 checkpoint) and `llama_3.1_405B_HF_FP8_model/` (HF FP8 checkpoint). Before converting the checkpoints to TensorRT LLM unified checkpoints, **please check that `{"rope_scaling": {"rope_type": "llama3"}}` is set in the configuration file**. With this flag, TensorRT LLM will enable the rope scaling of LLaMA-3.1. If not, please add it to the config file. - -Users can run the LLaMA-3.1 model with higher precision (bf16/fp16) or fp8. Here, to prevent accuracy drop, we perform per-channel per-token fp8 quantization (leveraged from https://github.com/pytorch/FBGEMM) on MLP layers, keeping other layers at higher precision. Note that per-channel per-token fp8 quantization is only supported on Huggingface checkpoint now. We will support it on Meta checkpoint soon. Note that this feature only supports SM90. - -### Convert Checkpoint to TensorRT LLM Unified Checkpoint - -To use the fp8 quantization, please add the `--use_fp8_rowwise` flag during the checkpoint conversion. In this demonstration, we convert the Meta checkpoint to bfloat16 with TP8-PP2 and the HF checkpoint to FP8 with TP8. - -Note: You may need to update your transformers installation via `pip install --upgrade transformers`. -Note: For 405B HF model cloned before 09 Aug 2024, there are duplicated kv head weights (The `num_key_value_heads` in `config.json` is 16). Users could use `--remove_duplicated_kv_heads` to remove them. The new checkpoint without duplicated kv heads is uploaded on 09 Aug 2024 and the `num_key_value_heads` is 8 now. For the new checkpoint, adding the flag `--remove_duplicated_kv_heads` would lead to error. - -```bash -# Run BF16 model by BF16 -python examples/models/core/llama/convert_checkpoint.py --meta_ckpt_dir llama_3.1_405B_meta_model/ \ - --output_dir llama_3.1_405B_meta_model/trt_ckpts/tp8-pp2/ \ - --dtype bfloat16 \ - --tp_size 8 \ - --pp_size 2 \ - --load_by_shard \ - --workers 2 - -# Run BF16 model by FP8 -python examples/models/core/llama/convert_checkpoint.py --model_dir llama_3.1_405B_HF_model/ \ - --output_dir llama_3.1_405B_HF_model/trt_ckpts/tp8-pp1/ \ - --dtype bfloat16 \ - --use_fp8_rowwise \ - --tp_size 8 \ - --pp_size 1 \ - --load_by_shard \ - --workers 8 \ - -# Run FP8 model by FP8 -# The source HF model is in FP8 format, so --use_fp8_rowwise is enabled automatically -# Optionally enable --use_meta_fp8_rowwise_recipe to strictly follow the original Meta's LLaMA 3.1 recipe: -# (1) Skip quantization for the first and last Transformer layers -# (2) Skip quantization for the Attention layers -python examples/models/core/llama/convert_checkpoint.py --model_dir llama_3.1_405B_HF_FP8_model/ \ - --output_dir llama_3.1_405B_HF_FP8_model/trt_ckpts/tp8-pp1/ \ - --dtype bfloat16 \ - --tp_size 8 \ - --pp_size 1 \ - --use_meta_fp8_rowwise_recipe \ - --load_by_shard \ - --workers 8 -``` - -### Build Engine - -```bash -trtllm-build --checkpoint_dir llama_3.1_405B_meta_model/trt_ckpts/tp8-pp2/ \ - --output_dir llama_3.1_405B_meta_model/trt_engines/tp8-pp2/ \ - --max_num_tokens 4096 \ - --max_input_len 255000 \ - --max_seq_len 256000 \ - --use_paged_context_fmha enable \ - --workers 8 - -trtllm-build --checkpoint_dir llama_3.1_405B_HF_model/trt_ckpts/tp8-pp1/ \ - --output_dir llama_3.1_405B_HF_model/trt_engines/tp8-pp1/ \ - --max_num_tokens 4096 \ - --max_input_len 64000 \ - --max_seq_len 65000 \ - --use_paged_context_fmha enable \ - --workers 8 - -trtllm-build --checkpoint_dir llama_3.1_405B_HF_FP8_model/trt_ckpts/tp8-pp1/ \ - --output_dir llama_3.1_405B_HF_FP8_model/trt_engines/tp8-pp1/ \ - --max_num_tokens 4096 \ - --max_input_len 64000 \ - --max_seq_len 65000 \ - --use_paged_context_fmha enable \ - --workers 8 -``` - -### Run Inference - -To run inference on the 405B model, we often need to use multi-node to accommodate the entire model. Here, we use slurm to launch the job on multiple nodes. - -Notes: -* For convenience, we use the Huggingface tokenizer for tokenization. - -The following script shows how to run evaluation on long context: - -```bash -# Long context test for 128k -python ./examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 4; mkdir -p 128k_context; mv passkey.jsonl 128k_context; - -srun --mpi pmi2 -N 2 -n 16 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/eval_long_context.py --task passkey \ - --engine_dir llama_3.1_405B_meta_model/trt_engines/tp8-pp2/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --stop_idx 6 \ - --max_input_length 255000 \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 256064 \ - --data_dir 128k_context \ - --output_dir 128k_context_tp8' - -# output would be like -# 6it [00:00, 15354.38it/s] -# [07/22/2024-19:09:54] [TRT-LLM] [I] Evaluation takes: 372.16685914993286 sec. -# [07/22/2024-19:09:54] [TRT-LLM] [I] accuracy of 6 examples: 1.0 - -# Long context test for 64k -python ./examples/infinitebench/construct_synthetic_dataset.py --test_case build_passkey --test_level 3; mkdir -p 64k_context; mv passkey.jsonl 64k_context; - -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/eval_long_context.py --task passkey \ - --engine_dir llama_3.1_405B_HF_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --stop_idx 6 \ - --max_input_length 64000 \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064 \ - --data_dir 64k_context \ - --output_dir 64k_context_tp8' - -# Long context test for 64k -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/eval_long_context.py --task passkey \ - --engine_dir llama_3.1_405B_HF_FP8_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_FP8_model/ \ - --stop_idx 6 \ - --max_input_length 64000 \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064 \ - --data_dir 64k_context \ - --output_dir 64k_context_tp8' -``` - -The following script shows how to run evaluation on MMLU tasks: - -```bash -srun --mpi pmi2 -N 2 -n 16 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/mmlu.py --test_trt_llm \ - --engine_dir llama_3.1_405B_meta_model/trt_engines/tp8-pp2/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 256064' - -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/mmlu.py --test_trt_llm \ - --engine_dir llama_3.1_405B_HF_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_model/ \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064' - -srun --mpi pmi2 -N 1 -n 8 --ntasks-per-node 8 --container-image \ ---container-mounts \ ---container-name llama-3.1-405b \ ---container-workdir \ -bash -c 'python ./examples/mmlu.py --test_trt_llm \ - --engine_dir llama_3.1_405B_HF_FP8_model/trt_engines/tp8-pp1/ \ - --tokenizer_dir llama_3.1_405B_HF_FP8_model/ \ - --enable_chunked_context \ - --kv_cache_free_gpu_memory_fraction 0.999 \ - --max_tokens_in_paged_kv_cache 65064' -``` - -## Run LLaMa-3.3 70B Model on PyTorch Backend -This section provides the steps to run LLaMa-3.3 70B model FP8 precision on PyTorch backend by launching TensorRT LLM server and run performance benchmarks. - -### Prepare TensorRT LLM extra configs -```bash -cat >./config.yml <