github-featured/nvidia--tensorrt
2
0
Fork 0
Code Issues Pull Requests Deployments Models Agent Notes Packages Projects Releases Wiki Activity

TensorRT Release 10.14 (#4633)

Browse Source 
Signed-off-by: Asfiya Baig <asfiyab@nvidia.com>
This commit is contained in:
asfiyab-nvidia
2025-11-07 16:11:33 -08:00
committed by GitHub
parent a833f79454
commit 3b4ddc1d45
308 changed files with 9475 additions and 3565 deletions
Download Patch File Download Diff File Expand all files Collapse all files
CHANGELOG.md
+17 -2
View File
@@ -1,11 +1,27 @@
# TensorRT OSS Release Changelog
## 10.14 GA - 2025-11-7
- Sample changes
- Replace all pycuda usages with cuda-python APIs
- Removed the efficientnet samples
- Deprecated tensorflow_object_detection and efficientdet samples
- Samples will no longer be released with the packages. The TensorRT GitHub repository will be the single source.
- Parsers:
- Added support for the `Attention` operator
- Improved refit for `ConstantOfShape` nodes
- Demos
- demoDiffusion:
- Added support for the Cosmos-Predict2 text2image and video2world pipelines
## 10.13.3 GA - 2025-9-8
- Added support for TensorRT API Capture and Replay feature, see the [developer guide](https://docs.nvidia.com/deeplearning/tensorrt/latest/inference-library/advanced.html) for more information.
- Demo changes
- Added support for Flux Kontext pipeline.
## 10.13.2 GA - 2025-8-18
- Added support for CUDA 13.0, dropped support for CUDA 11.X
- Dropped support for Ubuntu 20.04
@@ -24,7 +40,6 @@
- Added `loadModelProto`, `loadInitializer` and `refitModelProto` APIs for IParserRefitter. These APIs are meant to be used to load user initializers when refitting ONNX models.
- Deprecated `IParser::parseWithWeightDescriptors`.
## 10.12.0 GA - 2025-6-10
- Plugin changes
- Migrated `IPluginV2`-descendent version 1 of `cropAndResizeDynamic`, to version 2, which implements `IPluginV3`.
CMakeLists.txt
+16 -8
View File
@@ -67,6 +67,9 @@ endif()
set(CMAKE_SKIP_BUILD_RPATH True)
# CUDA targets
set(DEFAULT_CUDA_VERSION 13.0.0)
set_ifndef(CUDA_VERSION ${DEFAULT_CUDA_VERSION})
message(STATUS "CUDA version set to ${CUDA_VERSION}")
if (DEFINED GPU_ARCHS)
message(STATUS "GPU_ARCHS defined as ${GPU_ARCHS}. Generating CUDA code for SM ${GPU_ARCHS}")
@@ -81,6 +84,9 @@ else()
list(APPEND CMAKE_CUDA_ARCHITECTURES 100 120)
endif()
if(CUDA_VERSION VERSION_GREATER_EQUAL 13.0)
list(APPEND CMAKE_CUDA_ARCHITECTURES 110)
endif()
message(STATUS "GPU_ARCHS is not defined. Generating CUDA code for default SMs: ${CMAKE_CUDA_ARCHITECTURES}")
endif()
set(BERT_GENCODES)
@@ -105,6 +111,10 @@ project(TensorRT
DESCRIPTION "TensorRT is a C++ library that facilitates high-performance inference on NVIDIA GPUs and deep learning accelerators."
HOMEPAGE_URL "https://github.com/NVIDIA/TensorRT")
if (WIN32)
enable_language(C)
endif()
if(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
set(CMAKE_INSTALL_PREFIX ${TRT_LIB_DIR}/../ CACHE PATH "TensorRT installation" FORCE)
endif(CMAKE_INSTALL_PREFIX_INITIALIZED_TO_DEFAULT)
@@ -113,7 +123,7 @@ option(BUILD_PLUGINS "Build TensorRT plugin" ON)
option(BUILD_PARSERS "Build TensorRT parsers" ON)
option(BUILD_SAMPLES "Build TensorRT samples" ON)
# C++14
# C++17
set(CMAKE_CXX_STANDARD 17)
set(CMAKE_CXX_STANDARD_REQUIRED ON)
set(CMAKE_CXX_EXTENSIONS OFF)
@@ -141,14 +151,10 @@ endif()
############################################################################################
# Dependencies
set(DEFAULT_CUDA_VERSION 12.2.0)
set(DEFAULT_CUDNN_VERSION 8.9)
set(DEFAULT_PROTOBUF_VERSION 3.20.1)
set(DEFAULT_PROTOBUF_VERSION 3.20.3)
# Dependency Version Resolution
set_ifndef(CUDA_VERSION ${DEFAULT_CUDA_VERSION})
message(STATUS "CUDA version set to ${CUDA_VERSION}")
set_ifndef(CUDNN_VERSION ${DEFAULT_CUDNN_VERSION})
message(STATUS "cuDNN version set to ${CUDNN_VERSION}")
set_ifndef(PROTOBUF_VERSION ${DEFAULT_PROTOBUF_VERSION})
@@ -221,16 +227,18 @@ endif()
############################################################################################
# TensorRT
set(HINT_PATHS "${TRT_OUT_DIR}" "${TRT_LIB_DIR}")
if(BUILD_PLUGINS)
add_subdirectory(plugin)
else()
find_library_create_target(nvinfer_plugin ${nvinfer_plugin_lib_name} SHARED "${TRT_OUT_DIR}" "${TRT_LIB_DIR}")
find_library_create_target(${nvinfer_plugin_lib_name} ${nvinfer_plugin_lib_name} SHARED "${HINT_PATHS}")
endif()
if(BUILD_PARSERS)
add_subdirectory(parsers)
else()
find_library_create_target(nvonnxparser ${nvonnxparser_lib_name} SHARED "${TRT_OUT_DIR}" "${TRT_LIB_DIR}")
find_library_create_target(${nvonnxparser_lib_name} ${nvonnxparser_lib_name} SHARED "${HINT_PATHS}")
endif()
if(BUILD_SAMPLES)
README.md
+27 -18
View File
@@ -32,7 +32,7 @@ To build the TensorRT-OSS components, you will first need the following software
**TensorRT GA build**
- TensorRT v10.13.3.9
- TensorRT v10.14.1.48
- Available from direct download links listed below
**System Packages**
@@ -86,24 +86,24 @@ To build the TensorRT-OSS components, you will first need the following software
Else download and extract the TensorRT GA build from [NVIDIA Developer Zone](https://developer.nvidia.com) with the direct links below:
- [TensorRT 10.13.3.9 for CUDA 13.0, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz)
- [TensorRT 10.13.3.9 for CUDA 12.9, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz)
- [TensorRT 10.13.3.9 for CUDA 13.0, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/zip/TensorRT-10.13.3.9.Windows.win10.cuda-13.0.zip)
- [TensorRT 10.13.3.9 for CUDA 12.9, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/zip/TensorRT-10.13.3.9.Windows.win10.cuda-12.9.zip)
- [TensorRT 10.14.1.48 for CUDA 13.0, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz)
- [TensorRT 10.14.1.48 for CUDA 12.9, Linux x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz)
- [TensorRT 10.14.1.48 for CUDA 13.0, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/zip/TensorRT-10.14.1.48.Windows.win10.cuda-13.0.zip)
- [TensorRT 10.14.1.48 for CUDA 12.9, Windows x86_64](https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/zip/TensorRT-10.14.1.48.Windows.win10.cuda-12.9.zip)
**Example: Ubuntu 22.04 on x86-64 with cuda-13.0**
```bash
cd ~/Downloads
tar -xvzf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz
export TRT_LIBPATH=`pwd`/TensorRT-10.13.3.9
tar -xvzf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz
export TRT_LIBPATH=`pwd`/TensorRT-10.14.1.48
```
**Example: Windows on x86-64 with cuda-12.9**
```powershell
Expand-Archive -Path TensorRT-10.13.3.9.Windows.win10.cuda-12.9.zip
$env:TRT_LIBPATH="$pwd\TensorRT-10.13.3.9\lib"
Expand-Archive -Path TensorRT-10.14.1.48.Windows.win10.cuda-12.9.zip
$env:TRT_LIBPATH="$pwd\TensorRT-10.14.1.48\lib"
```
## Setting Up The Build Environment
@@ -112,16 +112,16 @@ For Linux platforms, we recommend that you generate a docker container for build
1. #### Generate the TensorRT-OSS build container.
**Example: Ubuntu 22.04 on x86-64 with cuda-13.0 (default)**
**Example: Ubuntu 24.04 on x86-64 with cuda-13.0 (default)**
```bash
./docker/build.sh --file docker/ubuntu-22.04.Dockerfile --tag tensorrt-ubuntu22.04-cuda13.0
./docker/build.sh --file docker/ubuntu-24.04.Dockerfile --tag tensorrt-ubuntu24.04-cuda13.0
```
**Example: Rockylinux8 on x86-64 with cuda-12.9**
**Example: Rockylinux8 on x86-64 with cuda-13.0**
```bash
./docker/build.sh --file docker/rockylinux8.Dockerfile --tag tensorrt-rockylinux8-cuda12.9
./docker/build.sh --file docker/rockylinux8.Dockerfile --tag tensorrt-rockylinux8-cuda13.0
```
**Example: Ubuntu 24.04 cross-compile for Jetson (aarch64) with cuda-13.0 (JetPack SDK)**
@@ -137,9 +137,9 @@ For Linux platforms, we recommend that you generate a docker container for build
```
2. #### Launch the TensorRT-OSS build container.
**Example: Ubuntu 22.04 build container**
**Example: Ubuntu 24.04 build container**
```bash
./docker/launch.sh --tag tensorrt-ubuntu22.04-cuda13.0 --gpus all
./docker/launch.sh --tag tensorrt-ubuntu24.04-cuda13.0 --gpus all
```
> NOTE:
> <br> 1. Use the `--tag` corresponding to build container generated in Step 1.
@@ -175,7 +175,7 @@ For Linux platforms, we recommend that you generate a docker container for build
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DTRT_PLATFORM_ID=aarch64 -DGPU_ARCHS=110
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DTRT_OUT_DIR=`pwd`/out -DTRT_PLATFORM_ID=aarch64
CC=/usr/bin/gcc make -j$(nproc)
```
@@ -186,7 +186,16 @@ For Linux platforms, we recommend that you generate a docker container for build
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_cross.toolchain -DGPU_ARCHS=110
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_cross.toolchain
make -j$(nproc)
```
**Example: Ubuntu 24.04 Cross-Compile for DriveOS (aarch64) with cuda-13.0**
```bash
cd $TRT_OSSPATH
mkdir -p build && cd build
cmake .. -DTRT_LIB_DIR=$TRT_LIBPATH -DCMAKE_TOOLCHAIN_FILE=$TRT_OSSPATH/cmake/toolchains/cmake_aarch64_dos_cross.toolchain
make -j$(nproc)
```
@@ -196,7 +205,7 @@ For Linux platforms, we recommend that you generate a docker container for build
cd $TRT_OSSPATH
mkdir -p build
cd -p build
cmake .. -DTRT_LIB_DIR="$env:TRT_LIBPATH" -DCUDNN_ROOT_DIR="$env:CUDNN_PATH" -DTRT_OUT_DIR="$pwd\\out"
cmake .. -DTRT_LIB_DIR="$env:TRT_LIBPATH" -DTRT_OUT_DIR="$pwd\\out"
msbuild TensorRT.sln /property:Configuration=Release -m:$env:NUMBER_OF_PROCESSORS
```
VERSION
+1 -1
View File
@@ -1 +1 @@
10.13.3.9
10.14.1.48
cmake/modules/BundleLibraries.cmake
+324
View File
@@ -0,0 +1,324 @@
# 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.
# This file handles logic relating to building "fat" static libraries, that is, ones which contain other static libraries.
# Idiomatically, CMake would prefer we ship all static libraries as part of our release and ship a file which specifies the linkages.
# However, for various reasons, this is both inadvisable and annoying. Instead, we would prefer to ship one static library containing all dependencies.
#
# To do that, we need rules to bundle static libraries into other static libraries. Hence, this class.
define_property(TARGET
PROPERTY BUNDLED_LIBRARY_TEMPLATE_PATH
BRIEF_DOCS "File path to the template script which will be written to when calling target_bundle_libraries."
)
define_property(TARGET
PROPERTY BUNDLE_LIBRARIES
BRIEF_DOCS "The list of libraries that have been bundled into this target by target_bundle_libraries."
)
define_property(TARGET
PROPERTY BUNDLE_LIBRARY_KNOWN_TYPE
BRIEF_DOCS "Fallback type used when a target's TYPE is UNKNOWN_LIBRARY."
)
# Internal helper to prefix all messages with "[target_bundle_libraries]: ".
#
# \param mode The message mode to be passed to message(...)
# \param argn The message contents.
macro(__bundleMessage mode)
message(${mode} "[target_bundle_libraries]: " ${ARGN})
endmacro()
# Illegal genex magic™
# Given a string containing a generator expression (var), edits the variable in-place to escape any generator expression literals.
# The escaped generator expression is then suitable for use within a LIST:TRANSFORM replacement block.
# This more or less allows for mapping lists to generator expressions, which can be recursively evaluated.
function(escape_generator_expression var)
string(REPLACE ">" "__ANGLE_R__" ${var} "${${var}}")
string(REPLACE "$" "$<1:$>" ${var} "${${var}}")
string(REPLACE "," "$<COMMA>" ${var} "${${var}}")
string(REPLACE "__ANGLE_R__" "$<ANGLE-R>" ${var} "${${var}}")
return(PROPAGATE ${var})
endfunction()
# Recursively unwraps alias targets until finding the real target. Non-targets are returned verbatim.
# We need to ensure that the generated .mri script contains a deduplicated list of bundled targets
# so we need to resolve aliases, otherwise we may end up with multiple entries for the same target.
#
# \param target_name The name of the target to unwrap.
# \param result_var The variable to store the unwrapped target name in.
function(unwrapAlias target_name result_var)
# First, try to unwrap common generator expressions that may wrap the target name.
string(REGEX MATCH "\\$<LINK_LIBRARY:WHOLE_ARCHIVE,([a-zA-Z0-9_.:]+)>" _ ${target_name})
if(TARGET ${CMAKE_MATCH_1})
set(target_name ${CMAKE_MATCH_1})
endif()
string(REGEX MATCH "\\$<LINK_ONLY:([a-zA-Z0-9_.:]+)>" _ ${target_name})
if(TARGET ${CMAKE_MATCH_1})
set(target_name ${CMAKE_MATCH_1})
endif()
if(TARGET ${target_name})
get_target_property(aliased_target ${target_name} ALIASED_TARGET)
if(aliased_target)
# Recursively unwrap in case there are multiple levels
unwrapAlias(${aliased_target} unwrapped)
set(${result_var} ${unwrapped} PARENT_SCOPE)
else()
# Not an alias, return the original name
set(${result_var} ${target_name} PARENT_SCOPE)
endif()
else()
# Not a target at all, return the original name
set(${result_var} ${target_name} PARENT_SCOPE)
endif()
endfunction()
# Internal function to retrieve the type of library for a given target.
# This function will fallback to the value of BUNDLE_LIBRARY_KNOWN_TYPE if it encounters an UNKNOWN_LIBRARY.
#
# \param lib A target to evaluate the type for.
# \param outVar The output variable to store the type name in.
function(__get_lib_type lib outVar)
get_target_property(libType ${lib} TYPE)
if (${libType} STREQUAL UNKNOWN_LIBRARY)
get_target_property(knownType ${lib} BUNDLE_LIBRARY_KNOWN_TYPE)
if (NOT ${knownType} STREQUAL "knownType-NOTFOUND")
set(libType ${knownType})
endif()
__bundleMessage(DEBUG "Using known type of unknown library ${lib}: ${knownType}")
endif()
set(${outVar} ${libType} PARENT_SCOPE)
endfunction()
# This is an internal-only function called the first time that target_bundle_libraries is called.
# It "registers" a target for bundling by creating the base template file and populating the target property BUNDLED_LIBRARY_TEMPLATE_PATH.
# Additionally, it registers the file generation logic for making the "final" script, as well as the custom command for running it after the build.
#
# \param lib The mainLib from target_bundle_libraries.
# \param templatePath The file path to the template file that will be created.
function(__registerTargetForBundling lib templatePath)
if(MSVC)
set(scriptPath $<TARGET_FILE_DIR:${lib}>/archive-${lib}.bat)
set(template "/OUT:\"$<TARGET_FILE:${lib}>\" \"$<TARGET_FILE:${lib}>\"\n")
# Windows-syntax version of the same logic from the linux build below.
# Main differences is that windows does not have `addlib`, and uses `\n` instead of `\n`. Additionally, the values need to be quoted to account for spaces in file paths.
set(replaceExpr "\"$<IF:$<TARGET_EXISTS:\\1>,$<TARGET_FILE:\\1>,\\1>\"")
escape_generator_expression(replaceExpr)
string(APPEND template "$<TARGET_GENEX_EVAL:${lib},$<JOIN:$<LIST:TRANSFORM,$<TARGET_PROPERTY:${lib},BUNDLE_LIBRARIES>,REPLACE,(.+),${replaceExpr}>,\n>>\n")
file(WRITE ${templatePath} ${template})
file(GENERATE
OUTPUT ${scriptPath}
INPUT ${templatePath}
)
add_custom_command(TARGET ${lib} POST_BUILD
COMMAND ${CMAKE_AR} /NOLOGO @\"${scriptPath}\"
COMMAND ${CMAKE_COMMAND} -E echo "Bundled $<LIST:LENGTH,$<TARGET_PROPERTY:${lib},BUNDLE_LIBRARIES>> static libraries into target ${lib}. Script: ${scriptPath}"
WORKING_DIRECTORY $<TARGET_FILE_DIR:${lib}>
)
else()
set(scriptPath $<TARGET_FILE_DIR:${lib}>/archive-${lib}.mri)
set(template "create $<TARGET_FILE:${lib}>\n")
string(APPEND template "addlib $<TARGET_FILE:${lib}>\n")
# Expand BUNDLE_LIBRARIES into the appropriate chain of addlib commands needed.
# BUNDLE_LIBRARIES will contain either (a) target names or (b) absolute file paths to libraries to include.
# This first part will disambiguate between (a) and (b) by evaluating (a) to `addlib $<TARGET_FILE:lib>` and (b) to `addlib [[filepath]]`
set(replaceExpr "addlib $<IF:$<TARGET_EXISTS:\\1>,$<TARGET_FILE:\\1>,\\1>")
escape_generator_expression(replaceExpr)
# The second part maps every element in BUNDLE_LIBRARIES to `replaceExpr` and evaluates the resulting replacement, which produces the final .mri file.
string(APPEND template "$<TARGET_GENEX_EVAL:${lib},$<JOIN:$<LIST:TRANSFORM,$<TARGET_PROPERTY:${lib},BUNDLE_LIBRARIES>,REPLACE,(.+),${replaceExpr}>,\n>>\n")
string(APPEND template "save\n")
string(APPEND template "end\n")
file(WRITE ${templatePath} ${template})
file(GENERATE
OUTPUT ${scriptPath}
INPUT ${templatePath}
)
add_custom_command(TARGET ${lib} POST_BUILD
COMMAND ${CMAKE_AR} -M < ${scriptPath}
COMMAND ${CMAKE_RANLIB} $<TARGET_FILE:${lib}>
COMMAND ${CMAKE_COMMAND} -E echo "Bundled $<LIST:LENGTH,$<TARGET_PROPERTY:${lib},BUNDLE_LIBRARIES>> static libraries into target ${lib}. Script: ${scriptPath}"
WORKING_DIRECTORY $<TARGET_FILE_DIR:${lib}>
)
endif()
set_target_properties(${lib}
PROPERTIES BUNDLED_LIBRARY_TEMPLATE_PATH ${templatePath}
)
endfunction()
# Subcomponent of target_bundle_libraries which is responsible for walking the provided depLibs and recursively calling target_bundle_libraries.
# This macro must only be used within target_bundle_libraries.
#
# \param mainLib The current main library from target_bundle_libraries
# \param linkVis The link visibility from target_bundle_libraries
# \param argn The dependencies to walk. Usually the INTERFACE_LINK_LIBRARIES of a target currently being bundled.
function(__bundleRecursiveDeps mainLib linkVis)
if(ARGN)
get_target_property(bundledLibs ${mainLib} BUNDLE_LIBRARIES)
foreach(dep IN LISTS ARGN)
unwrapAlias(${dep} dep)
if(${dep} IN_LIST bundledLibs)
continue() # Skip bundling of already-bundled libs to avoid many invocations of the same warnings.
endif()
if(TARGET ${dep})
__get_lib_type(${dep} depType)
if (${depType} STREQUAL STATIC_LIBRARY)
target_bundle_libraries(${mainLib} ${linkVis} ${dep})
elseif(${depType} STREQUAL INTERFACE_LIBRARY)
# For interface libraries, we want to add all of the static libraries they may be pointing to, without the library itself (since it is not a static).
get_target_property(interfaceLibs ${dep} INTERFACE_LINK_LIBRARIES)
__bundleRecursiveDeps(${mainLib} ${linkVis} ${interfaceLibs})
elseif(${depType} STREQUAL SHARED_LIBRARY)
# Skip SO's, since static libraries at the SO-boundary should not be bundled into the target mainLib.
else()
__bundleMessage(DEBUG "Skipping unhandled dependency ${dep} (a dependency of ${bundledLib}) with type ${depType} in ${mainLib}")
endif()
else()
__bundleMessage(DEBUG "Failed to recursively bundle dependency ${dep} (a dependency of ${bundledLib}) in ${mainLib}")
endif()
endforEach()
endif()
endfunction()
# This function acts as a replacement target_link_libraries, instead linking
# one or more "bundled" static libraries into a "main" static library.
# The bundled libs are embedded inside the main lib during the post-build step.
# CMake interface properties for definitions, etc. are propagated using the specified visibility.
#
# This function is recursive. Any static dependencies of any bundled library will also be bundled into the mainLib.
#
# \param mainLib The main library that other libraries are to be bundled into.
# \param linkVis The link visiblity for interface properties. One of PRIVATE or PUBLIC.
# Using PRIVATE hides all interface properties of bundled libraries from consumers of this library.
# Using PUBLIC will share interface properties with dependents.
# \param argn Variadic arguments - The list of targets to be linked into the mainLib.
function(target_bundle_libraries mainLib linkVis)
if (NOT ${linkVis} STREQUAL PUBLIC AND NOT ${linkVis} STREQUAL PRIVATE)
__bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries with unknown visibility ${linkVis}. Must be either \"PUBLIC\" or \"PRIVATE\".")
endif()
# TODO: Allow direct insertion of absolute file paths when CMAKE_LINK_LIBRARIES_ONLY_TARGETS is off, instead of always forcing targets.
if (NOT TARGET ${mainLib})
__bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries for target ${mainLib}, but no target with that name is known.")
endif()
__get_lib_type(${mainLib} mainLibType)
if (NOT mainLibType STREQUAL STATIC_LIBRARY)
__bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries for target ${mainLib}, but it is not a static library.")
endif()
get_target_property(isMainLibImported ${mainLib} IMPORTED)
if(isMainLibImported)
__bundleMessage(FATAL_ERROR "Error: Called target_bundle_libraries for target ${mainLib}, but it is an imported target.")
endif()
if(NOT ARGN)
__bundleMessage(WARNING "Called target_bundle_libraries with no libraries for target ${mainLib}")
endif()
__bundleMessage(DEBUG "Bundling ${ARGN} into ${mainLib}")
get_target_property(templatePath ${mainLib} BUNDLED_LIBRARY_TEMPLATE_PATH)
if(NOT EXISTS ${templatePath})
if(MSVC)
set(templatePath ${PROJECT_BINARY_DIR}/archive-${mainLib}.bat.template)
else()
set(templatePath ${PROJECT_BINARY_DIR}/archive-${mainLib}.mri.template)
endif()
__bundleMessage(STATUS "Registering default script path ${templatePath} for target ${mainLib}")
__registerTargetForBundling(${mainLib} ${templatePath})
endif()
get_target_property(bundledLibs ${mainLib} BUNDLE_LIBRARIES)
if (NOT bundledLibs)
set(bundledLibs "")
endif()
foreach(bundledLib IN LISTS ARGN)
unwrapAlias(${bundledLib} bundledLib)
__get_lib_type(${bundledLib} bundledLibType)
if(${bundledLibType} STREQUAL INTERFACE_LIBRARY)
# Interface libraries are not bundled, since they do not contain any static libraries.
# Their dependencies will be bundled recursively in the next step.
continue()
endif()
if (NOT ${bundledLibType} STREQUAL STATIC_LIBRARY AND NOT ${bundledLibType} STREQUAL OBJECT_LIBRARY)
__bundleMessage(FATAL_ERROR "Attempted to bundle ${bundledLibType} library ${bundledLib} into target ${mainLib} (only static and object libraries may be bundled)")
endif()
if (${bundledLibType} STREQUAL STATIC_LIBRARY)
list(APPEND bundledLibs ${bundledLib})
else()
# Exclude object libs from the BUNDLE_LIBRARIES property as they get added into the static lib using normal means.
endif()
endforeach()
list(REMOVE_DUPLICATES bundledLibs)
set_target_properties(${mainLib}
PROPERTIES BUNDLE_LIBRARIES "${bundledLibs}"
)
foreach(bundledLib IN LISTS ARGN)
unwrapAlias(${bundledLib} bundledLib)
__get_lib_type(${bundledLib} bundledLibType)
# Recursively bundle all static dependencies of each lib to be bundled.
get_target_property(depLibs ${bundledLib} LINK_LIBRARIES)
__bundleRecursiveDeps(${mainLib} ${linkVis} ${depLibs})
# Since we want the main library to be linkable standalone, we need both the LINK_LIBRARIES and INTERFACE_LINK_LIBRARIES bundled in.
# Otherwise, private static dependencies may be lost.
get_target_property(depLibs ${bundledLib} INTERFACE_LINK_LIBRARIES)
__bundleRecursiveDeps(${mainLib} ${linkVis} ${depLibs})
# Use `target_link_libraries` to propagate INTERFACE definitions from bundled libs
# BUILD_LOCAL_INTERFACE prevents clients from seeing this internal link relationship
# COMPILE_ONLY prevents the bundled lib from appearing in the link command redundantly
if (${bundledLibType} STREQUAL STATIC_LIBRARY OR ${bundledLibType} STREQUAL INTERFACE_LIBRARY)
target_link_libraries(${mainLib} ${linkVis}
$<BUILD_LOCAL_INTERFACE:$<COMPILE_ONLY:${bundledLib}>>
)
else()
# Include Object Libraries as full libraries, since they do not get added by the bundling stage.
# To do this without breaking the link dependency logic, we need to steal the target objects and link the lib as local + compile only.
target_link_libraries(${mainLib} ${linkVis}
$<BUILD_LOCAL_INTERFACE:$<COMPILE_ONLY:${bundledLib}>>
)
target_sources(${mainLib} PRIVATE $<TARGET_OBJECTS:${bundledLib}>)
endif()
# require that bundled libs are built before we try to bundle them
add_dependencies(${mainLib} ${bundledLib})
endforeach()
if(NOT EXISTS ${templatePath})
__bundleMessage(FATAL_ERROR "Template file ${templatePath} for target ${mainLib} does not exist.")
endif()
endfunction()
cmake/modules/ImportDL.cmake
+33
View File
@@ -0,0 +1,33 @@
# 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_guard()
if(NOT TARGET dl)
# libdl is included in the system library on Windows and QNX.
if(CMAKE_SYSTEM_NAME STREQUAL "Linux")
find_library(DL_LIB_PATH
NAMES ${CMAKE_DL_LIBS}
REQUIRED
)
message(STATUS "Creating imported target 'dl' for ${DL_LIB_PATH}")
add_library(dl SHARED IMPORTED)
set_target_properties(dl PROPERTIES IMPORTED_LOCATION "${DL_LIB_PATH}")
else()
message(STATUS "Creating no-op target 'dl' since libdl is not available on this platform.")
add_library(dl INTERFACE) # Add a fake dl target so we can still call target_link_libraries without error, even though it's a no-op.
endif()
endif()
cmake/modules/InstallUtils.cmake
+12 -1
View File
@@ -13,6 +13,7 @@
# See the License for the specific language governing permissions and
# limitations under the License.
include_guard()
include(GNUInstallDirs)
@@ -26,7 +27,7 @@ include(GNUInstallDirs)
function(installLibraries)
cmake_parse_arguments(
ARG # Prefix for parsed args
"OPTIONAL" # Options (flags)
"OPTIONAL;RUNTIME_ONLY" # Options (flags)
"COMPONENT" # Single value args
"TARGETS;CONFIGURATIONS" # Multi-value args
${ARGN}
@@ -50,12 +51,22 @@ function(installLibraries)
set(optional_arg OPTIONAL)
endif()
# When RUNTIME_ONLY is passed, we only want to install .dll files.
# Instead of also installing the import library (.lib) files.
# This is only relevant on Windows since Linux doesn't have this distinction.
if(ARG_RUNTIME_ONLY AND WIN32)
set(runtime_only_arg
RUNTIME DESTINATION ${CMAKE_INSTALL_BINDIR}
)
endif()
# Install the libraries
install(
TARGETS ${ARG_TARGETS}
${optional_arg}
${component_arg}
${config_arg}
${runtime_only_arg}
)
# Install PDB files for MSVC builds
cmake/modules/Platforms.cmake
+54
View File
@@ -0,0 +1,54 @@
# 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.
# Contains constants for the various platform names TRT supports.
set(TRT_PLATFORM_X86
"x86_64"
CACHE INTERNAL "Linux")
set(TRT_PLATFORM_AARCH64
"aarch64"
CACHE INTERNAL "ARM Linux")
set(TRT_PLATFORM_QNX
"qnx"
CACHE INTERNAL "QNX")
set(TRT_PLATFORM_QNX_SAFE
"qnx-safe"
CACHE INTERNAL "QNX Safe")
set(TRT_PLATFORM_WIN10
"win10"
CACHE INTERNAL "Windows 10")
# Checks if the current build platform matches any of the passed (ARGN) platforms.
#
# \param outVar The output variable name.
# \param argn The list of platforms to check against.
# \returns TRUE if TRT_BUILD_PLATFORM matches any of the platforms, FALSE otherwise.
function(checkPlatform outVar)
if(NOT DEFINED TRT_BUILD_PLATFORM)
message(FATAL_ERROR "checkPlatform was called before TRT_BUILD_PLATFORM was defined!")
endif()
set(isPlatform FALSE)
foreach(platform IN LISTS ARGN)
if(${platform} STREQUAL ${TRT_BUILD_PLATFORM})
set(isPlatform TRUE)
endif()
endforeach()
set(${outVar} ${isPlatform} PARENT_SCOPE)
endfunction()
cmake/modules/ShouldCompileKernel.cmake
+37
View File
@@ -35,6 +35,43 @@ function(get_all_numeric_sms OUT_VAR)
set(${OUT_VAR} ${ALL_NUMERIC_SMS} PARENT_SCOPE)
endfunction()
# \brief Converts the list returned by get_all_numeric_sms into a list of arch values.
# \returns the list in the name specified by OUT_VAR for native platform and OUT_VAR_CROSS for cross OS support. e.g. ptx, sm75, sm80, sm86, sm89, sm100, sm120.
function(get_all_fatbin_archs OUT_VAR OUT_VAR_CROSS)
# Use get_all_numeric_sms to get SM values and convert them to sm-prefixed format
set(ARCH_LIST "")
set(ARCH_LIST_CROSS "")
get_all_numeric_sms(NUMERIC_SMS)
foreach(SM IN LISTS NUMERIC_SMS)
list(APPEND ARCH_LIST "sm${SM}")
endforeach()
# Note: sm89 it is missing in NUMERIC_SMS since TRT treats sm89 as sm86.
# We should add sm89 to the list to generate the builder resource for sm89.
# If only sm86 is in the list, it means this build only supports sm86,
# so no need to add sm89.
list(FIND ARCH_LIST "sm86" SM86_INDEX)
list(FIND ARCH_LIST "sm89" SM89_INDEX)
list(LENGTH ARCH_LIST ARCH_LIST_COUNT)
if(${SM86_INDEX} GREATER_EQUAL 0 AND ${SM89_INDEX} EQUAL -1 AND ${ARCH_LIST_COUNT} GREATER 1)
list(APPEND ARCH_LIST "sm89")
endif()
# There is also a klib which only contains PTX code.
list(APPEND ARCH_LIST "ptx")
set(ARCH_LIST_CROSS ${ARCH_LIST})
# Cask5 does not include sm100 cubins. Exclude sm100 for both
# cross-OS support and native Windows build.
list(FILTER ARCH_LIST_CROSS EXCLUDE REGEX "sm100")
if(${TRT_BUILD_PLATFORM} STREQUAL ${TRT_PLATFORM_WIN10})
list(FILTER ARCH_LIST EXCLUDE REGEX "sm100")
endif()
set(${OUT_VAR} ${ARCH_LIST} PARENT_SCOPE)
set(${OUT_VAR_CROSS} ${ARCH_LIST_CROSS} PARENT_SCOPE)
endfunction()
# Certain cubins are binary compatible between different SM versions, so they are reused.
# This function checks if a SM-named file should be compiled based on current SM enablement.
# Specifically, the SM80 files are compiled if either 80, 86, or 89 are enabled.
cmake/modules/find_library_create_target.cmake
+5 -1
View File
@@ -31,6 +31,10 @@ macro(find_library_create_target target_name lib libtype hints)
find_library(${lib}_LIB_PATH ${lib})
message(STATUS "Library that was found ${${lib}_LIB_PATH}")
add_library(${target_name} ${libtype} IMPORTED)
set_property(TARGET ${target_name} PROPERTY IMPORTED_LOCATION ${${lib}_LIB_PATH})
if(MSVC)
set_property(TARGET ${target_name} PROPERTY IMPORTED_IMPLIB ${${lib}_LIB_PATH})
else()
set_property(TARGET ${target_name} PROPERTY IMPORTED_LOCATION ${${lib}_LIB_PATH})
endif()
message(STATUS "==========================================================================================")
endmacro()
cmake/toolchains/cmake_aarch64_dos_cross.toolchain
+60
View File
@@ -0,0 +1,60 @@
#
# 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.
#
set(CMAKE_SYSTEM_NAME Linux)
set(CMAKE_SYSTEM_PROCESSOR aarch64)
set(CMAKE_FIND_ROOT_PATH /usr/aarch64-linux-gnu/)
set(TRT_PLATFORM_ID "aarch64")
set(CUDA_PLATFORM_ID "aarch64-linux")
set(CMAKE_C_COMPILER /usr/bin/aarch64-linux-gnu-gcc)
set(CMAKE_CXX_COMPILER /usr/bin/aarch64-linux-gnu-g++)
set(CMAKE_C_COMPILER_ID "GNU")
set(CMAKE_CXX_COMPILER_ID "GNU")
set(CMAKE_CXX_COMPILE_FEATURES cxx_std_17)
set(CMAKE_C_FLAGS "" CACHE STRING "" FORCE)
set(CMAKE_CXX_FLAGS "" CACHE STRING "" FORCE)
set(CMAKE_C_COMPILER_TARGET aarch64-linux-gnu)
set(CMAKE_CXX_COMPILER_TARGET aarch64-linux-gnu)
set(CMAKE_C_COMPILER_FORCED TRUE)
set(CMAKE_CXX_COMPILER_FORCED TRUE)
set(CUDA_ROOT /usr/local/cuda/targets/${CUDA_PLATFORM_ID} CACHE STRING "CUDA ROOT dir")
set(CUDNN_LIB /usr/lib/aarch64-linux-gnu/libcudnn.so)
set(BUILD_LIBRARY_ONLY 1)
set(CUDA_TOOLKIT_ROOT_DIR ${CUDA_ROOT})
set(CUDA_INCLUDE_DIRS ${CUDA_ROOT}/include)
set(CMAKE_CUDA_COMPILER /usr/local/cuda/bin/nvcc)
set(CMAKE_CUDA_HOST_COMPILER ${CMAKE_CXX_COMPILER} CACHE STRING "" FORCE)
set(CMAKE_CUDA_FLAGS "-I${CUDA_INCLUDE_DIRS} -Xcompiler=\"-fPIC ${CMAKE_CXX_FLAGS}\"" CACHE STRING "" FORCE)
set(CMAKE_CUDA_COMPILER_FORCED TRUE)
set(CUDA_LIBS -L${CUDA_ROOT}/lib)
set(ADDITIONAL_PLATFORM_LIB_FLAGS ${CUDA_LIBS} -lcublas -lcudart -lstdc++ -lm)
link_directories(${CUDA_ROOT}/lib)
demo/BERT/README.md
+2 -2
View File
@@ -86,7 +86,7 @@ This demo BERT application can be run within the TensorRT OSS build container. I
- [NGC CLI](https://ngc.nvidia.com/setup/installers/cli) - for downloading BERT checkpoints from NGC.
- PyPI Packages:
- [pycuda](https://pypi.org/project/pycuda/) (tested v2019.1.2)
- [cuda-python](https://pypi.org/project/cuda-python/) (tested v13.0.1)
- [onnx](https://pypi.org/project/onnx) (tested v1.12.0)
- [tensorflow](https://pypi.org/project/tensorflow/) (tested v2.9.1)
- [torch](https://pypi.org/project/torch/) (tested v1.11.0)
@@ -162,7 +162,7 @@ Completing these steps should resolve the error you encountered and allow the co
6. Run inference with CUDA Graph support.
A separate python `inference_c.py` script is provided to run inference with CUDA Graph support. This is necessary since CUDA Graph is only supported through CUDA C/C++ APIs, not pyCUDA. The `inference_c.py` script uses pybind11 to interface with C/C++ for CUDA graph capturing and launching. The cmdline interface is the same as `inference.py` except for an extra `--enable-graph` option.
A separate python `inference_c.py` script is provided to run inference with CUDA Graph support. This is necessary since CUDA Graph is only supported through CUDA C/C++ APIs. The `inference_c.py` script uses pybind11 to interface with C/C++ for CUDA graph capturing and launching. The cmdline interface is the same as `inference.py` except for an extra `--enable-graph` option.
```bash
mkdir -p build; pushd build
demo/BERT/builder.py
+2 -3
View File
@@ -26,8 +26,7 @@ import re
import sys
import time
import onnx
import pycuda.autoinit
from helpers.cuda_utils import getComputeCapacity
# TensorRT
import tensorrt as trt
from helpers.calibrator import BertCalibrator as BertCalibrator
@@ -669,7 +668,7 @@ def main():
args.batch_size = args.batch_size or [1]
args.sequence_length = args.sequence_length or [128]
cc = pycuda.autoinit.device.compute_capability()
cc = getComputeCapacity()
if cc[0] * 10 + cc[1] < 75 and args.force_int8_multihead:
raise RuntimeError("--force-int8-multihead option is only supported on Turing+ GPU.")
if cc[0] * 10 + cc[1] < 72 and args.force_int8_skipln:
demo/BERT/builder_varseqlen.py
+2 -4
View File
@@ -26,8 +26,7 @@ import re
import sys
import time
import onnx
import pycuda.autoinit
from helpers.cuda_utils import getComputeCapacity
# TensorRT
import tensorrt as trt
from builder_utils import load_tf_weights, load_pytorch_weights_and_quant, load_onnx_weights_and_quant, load_megatron_pickle_weights
@@ -352,7 +351,6 @@ def bert_model(config, init_dict, network, input_tensor, residual, mask_idx, cu_
squad_logits_out.name = "logits_out"
network.mark_output(squad_logits_out)
def squad_output(prefix, config, init_dict, network, input_tensor):
"""
Create the squad output
@@ -645,7 +643,7 @@ def main():
if args.verbose:
TRT_LOGGER.min_severity = TRT_LOGGER.VERBOSE
cc = pycuda.autoinit.device.compute_capability()
cc = getComputeCapacity()
if cc[0] * 10 + cc[1] < 72:
raise RuntimeError("This variable-length BERT demo only support Xavier+ GPU.")
demo/BERT/helpers/calibrator.py
+8 -7
View File
@@ -19,8 +19,8 @@
import tensorrt as trt
import os
import pycuda.driver as cuda
import pycuda.autoinit
from helpers.cuda_utils import cuda_call, memcpy_host_to_device
from cuda.bindings import driver as cuda, runtime as cudart
import numpy as np
import helpers.tokenization as tokenization
import helpers.data_processing as dp
@@ -44,11 +44,12 @@ class BertCalibrator(trt.IInt8LegacyCalibrator):
self.max_query_length = 64
# Allocate enough memory for a whole batch.
self.device_inputs = [cuda.mem_alloc(self.max_seq_length * trt.int32.itemsize * self.batch_size) for binding in range(3)]
self.device_inputs = [cuda_call(cudart.cudaMalloc(self.max_seq_length * trt.int32.itemsize * self.batch_size)) for binding in range(3)]
def free(self):
for dinput in self.device_inputs:
dinput.free()
# dinput is a device pointer (int) returned by cudaMalloc
cuda_call(cudart.cudaFree(dinput))
def get_batch_size(self):
return self.batch_size
@@ -80,9 +81,9 @@ class BertCalibrator(trt.IInt8LegacyCalibrator):
segment_ids = features[0].segment_ids
input_mask = features[0].input_mask
cuda.memcpy_htod(self.device_inputs[0], input_ids.ravel())
cuda.memcpy_htod(self.device_inputs[1], segment_ids.ravel())
cuda.memcpy_htod(self.device_inputs[2], input_mask.ravel())
memcpy_host_to_device(self.device_inputs[0], input_ids.ravel())
memcpy_host_to_device(self.device_inputs[1], segment_ids.ravel())
memcpy_host_to_device(self.device_inputs[2], input_mask.ravel())
self.current_index += self.batch_size
return self.device_inputs
demo/BERT/helpers/cuda_utils.py
+120
View File
@@ -0,0 +1,120 @@
#!/usr/bin/env python3
#
# 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.
#
import numpy as np
import logging
from cuda.bindings import driver as cuda, runtime as cudart
class CudaStreamContext:
"""CUDA stream lifecycle management with context manager support"""
def __init__(self):
"""Initialize CUDA stream"""
self._stream = cuda_call(cudart.cudaStreamCreate())
def __enter__(self):
"""Create CUDA stream when entering context (if not already created)"""
if self._stream is None:
self._stream = cuda_call(cudart.cudaStreamCreate())
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Destroy CUDA stream when exiting context"""
self.free()
@property
def stream(self) -> cudart.cudaStream_t:
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
return self._stream
def synchronize(self):
"""Synchronize the stream"""
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
cuda_call(cudart.cudaStreamSynchronize(self._stream))
def free(self):
"""Explicitly free the CUDA stream"""
if self._stream is not None:
try:
cuda_call(cudart.cudaStreamDestroy(self._stream))
self._stream = None
except Exception as e:
logging.warning(f"Failed to destroy CUDA stream: {e}")
def __del__(self):
"""Cleanup stream on destruction"""
if hasattr(self, '_stream') and self._stream is not None:
self.free()
def __str__(self):
return f"CudaStreamContext: {self._stream}"
def __repr__(self):
return self.__str__()
def cuda_call(call):
"""Helper function to make CUDA calls and check for errors"""
def _cudaGetErrorEnum(error):
if isinstance(error, cuda.CUresult):
err, name = cuda.cuGetErrorName(error)
return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
elif isinstance(error, cudart.cudaError_t):
return cudart.cudaGetErrorName(error)[1]
else:
raise RuntimeError("Unknown error type: {}".format(error))
err, res = call[0], call[1:]
if err.value:
raise RuntimeError(
"CUDA error code={}({})".format(
err.value, _cudaGetErrorEnum(err)
)
)
if len(res) == 1:
return res[0]
elif len(res) == 0:
return None
else:
return res
def getComputeCapacity(devID=0):
major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
return (major, minor)
def memcpy_host_to_device_async(device_ptr: int, host_arr: np.ndarray, stream):
"""Wrapper for async host-to-device memory copy"""
cuda_call(cudart.cudaMemcpyAsync(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream))
def memcpy_device_to_host_async(host_arr: np.ndarray, device_ptr: int, stream):
"""Wrapper for async device-to-host memory copy"""
cuda_call(cudart.cudaMemcpyAsync(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream))
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
"""Wrapper for synchronous host-to-device memory copy"""
cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
"""Wrapper for synchronous device-to-host memory copy"""
cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))
# Initialize CUDA
cuda_call(cudart.cudaFree(0))
demo/BERT/inference.ipynb
+46 -46
View File
@@ -221,8 +221,8 @@
"metadata": {},
"outputs": [],
"source": [
"import pycuda.driver as cuda\n",
"import pycuda.autoinit\n",
"from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host\n",
"from cuda.bindings import driver as cuda, runtime as cudart\n",
"import collections\n",
"import numpy as np\n",
"import time\n",
@@ -238,9 +238,7 @@
" input_nbytes = trt.volume(input_shape) * trt.int32.itemsize\n",
" \n",
" # Allocate device memory for inputs.\n",
" d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(3)]\n",
" # Create a stream in which to copy inputs/outputs and run inference.\n",
" stream = cuda.Stream()\n",
" d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]\n",
"\n",
" # Specify input shapes. These must be within the min/max bounds of the active profile (0th profile in this case)\n",
" # Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.\n",
@@ -250,57 +248,59 @@
" assert context.all_binding_shapes_specified\n",
"\n",
" # Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
" h_output = cuda.pagelocked_empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n",
" d_output = cuda.mem_alloc(h_output.nbytes)\n",
" h_output = np.empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n",
" cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))\n",
" d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))\n",
"\n",
" print(\"\\nRunning Inference...\")\n",
" with CudaStreamContext() as stream:\n",
" print(\"\\nRunning Inference...\")\n",
"\n",
" _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name\n",
" \"NetworkOutput\",\n",
" [\"start_logits\", \"end_logits\", \"feature_index\"])\n",
" networkOutputs = []\n",
" _NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name\n",
" \"NetworkOutput\",\n",
" [\"start_logits\", \"end_logits\", \"feature_index\"])\n",
" networkOutputs = []\n",
"\n",
" eval_time_elapsed = 0\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" input_ids = cuda.register_host_memory(np.ascontiguousarray(feature.input_ids.ravel()))\n",
" segment_ids = cuda.register_host_memory(np.ascontiguousarray(feature.segment_ids.ravel()))\n",
" input_mask = cuda.register_host_memory(np.ascontiguousarray(feature.input_mask.ravel()))\n",
" eval_time_elapsed = 0\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" input_ids = np.ascontiguousarray(feature.input_ids.ravel())\n",
" segment_ids = np.ascontiguousarray(feature.segment_ids.ravel())\n",
" input_mask = np.ascontiguousarray(feature.input_mask.ravel())\n",
"\n",
" eval_start_time = time.time()\n",
" cuda.memcpy_htod_async(d_inputs[0], input_ids, stream)\n",
" cuda.memcpy_htod_async(d_inputs[1], segment_ids, stream)\n",
" cuda.memcpy_htod_async(d_inputs[2], input_mask, stream)\n",
" eval_start_time = time.time()\n",
" memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)\n",
"\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
"\n",
" # Run inference\n",
" context.execute_async_v3(stream_handle=stream.handle)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
" # Run inference\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
"\n",
" # Transfer predictions back from GPU\n",
" cuda.memcpy_dtoh_async(h_output, d_output, stream)\n",
" stream.synchronize()\n",
" # Transfer predictions back from GPU\n",
" memcpy_device_to_host_async(h_output, d_output, stream.stream)\n",
" stream.synchronize()\n",
"\n",
" for index, batch in enumerate(h_output):\n",
" # Data Post-processing\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(batch.squeeze()[:, 0]),\n",
" end_logits = np.array(batch.squeeze()[:, 1]),\n",
" feature_index = feature_index\n",
" ))\n",
" for index, batch in enumerate(h_output):\n",
" # Data Post-processing\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(batch.squeeze()[:, 0]),\n",
" end_logits = np.array(batch.squeeze()[:, 1]),\n",
" feature_index = feature_index\n",
" ))\n",
"\n",
" eval_time_elapsed /= len(features)\n",
" \n",
" print(\"-----------------------------\")\n",
" print(\"Running Inference at {:.3f} Sentences/Sec\".format(1.0/eval_time_elapsed))\n",
" print(\"-----------------------------\")\n"
" eval_time_elapsed /= len(features)\n",
" \n",
" print(\"-----------------------------\")\n",
" print(\"Running Inference at {:.3f} Sentences/Sec\".format(1.0/eval_time_elapsed))\n",
" print(\"-----------------------------\")\n"
]
},
{
demo/BERT/inference.py
+103 -96
View File
@@ -30,8 +30,8 @@ import argparse
import collections
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async
from cuda.bindings import driver as cuda, runtime as cudart
import helpers.tokenization as tokenization
import helpers.data_processing as dp
@@ -143,126 +143,133 @@ if __name__ == '__main__':
raise RuntimeError("Could not find any profile that can run batch size {}.".format(args.batch_size))
# Create a stream in which to copy inputs/outputs and run inference.
stream = cuda.Stream()
with CudaStreamContext() as stream:
context.set_optimization_profile_async(selected_profile, stream.stream)
binding_idx_offset = selected_profile * engine.num_io_tensors
context.set_optimization_profile_async(selected_profile, stream.handle)
binding_idx_offset = selected_profile * engine.num_io_tensors
# Specify input shapes. These must be within the min/max bounds of the active profile
# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.
input_shape = (args.batch_size, max_seq_length)
input_nbytes = trt.volume(input_shape) * trt.int32.itemsize
for name in ["input_ids", "segment_ids", "input_mask"]:
context.set_input_shape(name, input_shape)
assert len(context.infer_shapes()) == 0
# Specify input shapes. These must be within the min/max bounds of the active profile
# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.
input_shape = (args.batch_size, max_seq_length)
input_nbytes = trt.volume(input_shape) * trt.int32.itemsize
for name in ["input_ids", "segment_ids", "input_mask"]:
context.set_input_shape(name, input_shape)
assert len(context.infer_shapes()) == 0
# Allocate device memory for inputs.
d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]
# Allocate device memory for inputs.
d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(3)]
# Allocate output buffer by querying the size from the context. This may be different for different input shapes.
h_output = np.empty(tuple(context.get_tensor_shape("logits_out")), dtype=np.float32)
cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))
# Pin the memory for faster transfers
d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))
# Allocate output buffer by querying the size from the context. This may be different for different input shapes.
h_output = cuda.pagelocked_empty(tuple(context.get_tensor_shape("logits_out")), dtype=np.float32)
d_output = cuda.mem_alloc(h_output.nbytes)
def inference(features, tokens):
global h_output
def inference(features, tokens):
global h_output
_NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
"NetworkOutput",
["start_logits", "end_logits", "feature_index"])
networkOutputs = []
_NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
"NetworkOutput",
["start_logits", "end_logits", "feature_index"])
networkOutputs = []
eval_time_elapsed = 0
for feature_index, feature in enumerate(features):
# Copy inputs
input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), args.batch_size, axis=0)
segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), args.batch_size, axis=0)
input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), args.batch_size, axis=0)
eval_time_elapsed = 0
for feature_index, feature in enumerate(features):
# Copy inputs
input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), args.batch_size, axis=0)
segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), args.batch_size, axis=0)
input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), args.batch_size, axis=0)
input_ids = np.ascontiguousarray(input_ids_batch.ravel())
segment_ids = np.ascontiguousarray(segment_ids_batch.ravel())
input_mask = np.ascontiguousarray(input_mask_batch.ravel())
input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids_batch.ravel()))
segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids_batch.ravel()))
input_mask = cuda.register_host_memory(np.ascontiguousarray(input_mask_batch.ravel()))
eval_start_time = time.time()
memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)
eval_start_time = time.time()
cuda.memcpy_htod_async(d_inputs[0], input_ids, stream)
cuda.memcpy_htod_async(d_inputs[1], segment_ids, stream)
cuda.memcpy_htod_async(d_inputs[2], input_mask, stream)
bindings = [0 for _ in range(binding_idx_offset)] + [int(d_inp) for d_inp in d_inputs] + [int(d_output)]
bindings = [0 for _ in range(binding_idx_offset)] + [int(d_inp) for d_inp in d_inputs] + [int(d_output)]
# allocate address for IO tensor
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])
# allocate address for IO tensor
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])
# Run inference
context.execute_async_v3(stream_handle=stream.handle)
# Synchronize the stream
stream.synchronize()
eval_time_elapsed += (time.time() - eval_start_time)
# Run inference
context.execute_async_v3(stream_handle=stream.stream)
# Synchronize the stream
stream.synchronize()
eval_time_elapsed += (time.time() - eval_start_time)
# Transfer predictions back from GPU
cuda.memcpy_dtoh_async(h_output, d_output, stream)
stream.synchronize()
# Transfer predictions back from GPU
memcpy_device_to_host_async(h_output, d_output, stream.stream)
stream.synchronize()
# Only retrieve and post-process the first batch
batch = h_output[0]
networkOutputs.append(_NetworkOutput(
start_logits = np.array(batch.squeeze()[:, 0]),
end_logits = np.array(batch.squeeze()[:, 1]),
feature_index = feature_index
))
# Only retrieve and post-process the first batch
batch = h_output[0]
networkOutputs.append(_NetworkOutput(
start_logits = np.array(batch.squeeze()[:, 0]),
end_logits = np.array(batch.squeeze()[:, 1]),
feature_index = feature_index
))
eval_time_elapsed /= len(features)
eval_time_elapsed /= len(features)
# Total number of n-best predictions to generate in the nbest_predictions.json output file
n_best_size = 20
# Total number of n-best predictions to generate in the nbest_predictions.json output file
n_best_size = 20
# The maximum length of an answer that can be generated. This is needed
# because the start and end predictions are not conditioned on one another
max_answer_length = 30
# The maximum length of an answer that can be generated. This is needed
# because the start and end predictions are not conditioned on one another
max_answer_length = 30
prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
networkOutputs, args.n_best_size, args.max_answer_length)
prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
networkOutputs, args.n_best_size, args.max_answer_length)
return eval_time_elapsed, prediction, nbest_json
return eval_time_elapsed, prediction, nbest_json
def print_single_query(eval_time_elapsed, prediction, nbest_json):
print("------------------------")
print("Running inference in {:.3f} Sentences/Sec".format(args.batch_size/eval_time_elapsed))
print("------------------------")
def print_single_query(eval_time_elapsed, prediction, nbest_json):
print("------------------------")
print("Running inference in {:.3f} Sentences/Sec".format(args.batch_size/eval_time_elapsed))
print("------------------------")
print("Answer: '{}'".format(prediction))
print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0))
print("Answer: '{}'".format(prediction))
print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0))
if squad_examples:
all_predictions = collections.OrderedDict()
if squad_examples:
all_predictions = collections.OrderedDict()
for example in squad_examples:
features = question_features(example.doc_tokens, example.question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens)
all_predictions[example.id] = prediction
with open(output_prediction_file, "w") as f:
f.write(json.dumps(all_predictions, indent=4))
print("\nOutput dump to {}".format(output_prediction_file))
else:
# Extract tokecs from the paragraph
doc_tokens = dp.convert_doc_tokens(paragraph_text)
if question_text:
print("\nPassage: {}".format(paragraph_text))
print("\nQuestion: {}".format(question_text))
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
for example in squad_examples:
features = question_features(example.doc_tokens, example.question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens)
all_predictions[example.id] = prediction
with open(output_prediction_file, "w") as f:
f.write(json.dumps(all_predictions, indent=4))
print("\nOutput dump to {}".format(output_prediction_file))
else:
# If no question text is provided, loop until the question is 'exit'
EXIT_CMDS = ["exit", "quit"]
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
# Extract tokecs from the paragraph
doc_tokens = dp.convert_doc_tokens(paragraph_text)
if question_text:
print("\nPassage: {}".format(paragraph_text))
print("\nQuestion: {}".format(question_text))
while question_text.strip() not in EXIT_CMDS:
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
else:
# If no question text is provided, loop until the question is 'exit'
EXIT_CMDS = ["exit", "quit"]
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
while question_text.strip() not in EXIT_CMDS:
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
# free allocated memory
for d_input in d_inputs:
cuda_call(cudart.cudaFree(d_input))
cuda_call(cudart.cudaFree(d_output))
cuda_call(cudart.cudaHostUnregister(h_output))
demo/BERT/inference_varseqlen.py
+102 -95
View File
@@ -30,8 +30,9 @@ import argparse
import collections
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async
from cuda.bindings import driver as cuda, runtime as cudart
import helpers.tokenization as tokenization
import helpers.data_processing as dp
@@ -131,127 +132,133 @@ if __name__ == '__main__':
with open(args.engine, 'rb') as f, trt.Runtime(TRT_LOGGER) as runtime, \
runtime.deserialize_cuda_engine(f.read()) as engine, engine.create_execution_context() as context:
# Create a stream in which to copy inputs/outputs and run inference.
stream = cuda.Stream()
with CudaStreamContext() as stream:
# select engine profile
context.set_optimization_profile_async(0, stream.stream)
# select engine profile
context.set_optimization_profile_async(0, stream.handle)
input_nbytes = max_seq_length * trt.int32.itemsize
input_nbytes = max_seq_length * trt.int32.itemsize
# Allocate device memory for inputs.
d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(4)]
# Allocate device memory for inputs.
d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(4)]
# Allocate output buffer by querying the size from the context. This may be different for different input shapes.
h_output = cuda.pagelocked_empty((2 * max_seq_length), dtype=np.float32)
d_output = cuda.mem_alloc(h_output.nbytes)
# Allocate output buffer by querying the size from the context. This may be different for different input shapes.
h_output = np.empty((2 * max_seq_length), dtype=np.float32)
cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))
d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))
def inference(features, tokens):
global h_output
def inference(features, tokens):
global h_output
_NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
"NetworkOutput",
["start_logits", "end_logits", "feature_index"])
networkOutputs = []
_NetworkOutput = collections.namedtuple( # pylint: disable=invalid-name
"NetworkOutput",
["start_logits", "end_logits", "feature_index"])
networkOutputs = []
eval_time_elapsed = 0
for feature_index, feature in enumerate(features):
# Copy inputs
B = 1
S = np.sum(feature.input_mask)
input_ids = feature.input_ids[0:S]
segment_ids = feature.segment_ids[0:S]
cu_seq_lens = np.array([0, S], dtype=np.int32);
eval_time_elapsed = 0
for feature_index, feature in enumerate(features):
# Copy inputs
B = 1
S = np.sum(feature.input_mask)
input_ids = feature.input_ids[0:S]
segment_ids = feature.segment_ids[0:S]
cu_seq_lens = np.array([0, S], dtype=np.int32);
input_dim0_shape = {"input_ids": S, "segment_ids": S, "cu_seqlens": 2, "max_seqlen": S}
for name, val in input_dim0_shape.items():
if context.get_tensor_shape(name)[0] != val:
context.set_input_shape(name, (val,))
input_dim0_shape = {"input_ids": S, "segment_ids": S, "cu_seqlens": 2, "max_seqlen": S}
for name, val in input_dim0_shape.items():
if context.get_tensor_shape(name)[0] != val:
context.set_input_shape(name, (val,))
h_input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids.ravel()))
h_segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids.ravel()))
h_cu_seq_lens = cuda.register_host_memory(np.ascontiguousarray(cu_seq_lens.ravel()))
h_input_ids = np.ascontiguousarray(input_ids.ravel())
h_segment_ids = np.ascontiguousarray(segment_ids.ravel())
h_cu_seq_lens = np.ascontiguousarray(cu_seq_lens.ravel())
eval_start_time = time.time()
cuda.memcpy_htod_async(d_inputs[0], h_input_ids, stream)
cuda.memcpy_htod_async(d_inputs[1], h_segment_ids, stream)
cuda.memcpy_htod_async(d_inputs[2], h_cu_seq_lens, stream)
eval_start_time = time.time()
memcpy_host_to_device_async(d_inputs[0], h_input_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[1], h_segment_ids, stream.stream)
memcpy_host_to_device_async(d_inputs[2], h_cu_seq_lens, stream.stream)
# Setup tensor address
bindings = [int(d_inputs[i]) for i in range(4)] + [int(d_output)]
# Setup tensor address
bindings = [int(d_inputs[i]) for i in range(4)] + [int(d_output)]
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
# Run inference
context.execute_async_v3(stream_handle=stream.handle)
# Synchronize the stream
stream.synchronize()
eval_time_elapsed += (time.time() - eval_start_time)
# Run inference
context.execute_async_v3(stream_handle=stream.stream)
# Synchronize the stream
stream.synchronize()
eval_time_elapsed += (time.time() - eval_start_time)
# Transfer predictions back from GPU
cuda.memcpy_dtoh_async(h_output, d_output, stream)
stream.synchronize()
# Transfer predictions back from GPU
memcpy_device_to_host_async(h_output, d_output, stream.stream)
stream.synchronize()
# Only retrieve and post-process the first batch
networkOutputs.append(_NetworkOutput(
start_logits = np.array(h_output[0:S]),
end_logits = np.array(h_output[S:S*2]),
feature_index = feature_index
))
# Only retrieve and post-process the first batch
networkOutputs.append(_NetworkOutput(
start_logits = np.array(h_output[0:S]),
end_logits = np.array(h_output[S:S*2]),
feature_index = feature_index
))
eval_time_elapsed /= len(features)
eval_time_elapsed /= len(features)
# Total number of n-best predictions to generate in the nbest_predictions.json output file
n_best_size = 20
# Total number of n-best predictions to generate in the nbest_predictions.json output file
n_best_size = 20
# The maximum length of an answer that can be generated. This is needed
# because the start and end predictions are not conditioned on one another
max_answer_length = 30
# The maximum length of an answer that can be generated. This is needed
# because the start and end predictions are not conditioned on one another
max_answer_length = 30
prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
networkOutputs, args.n_best_size, args.max_answer_length)
prediction, nbest_json, scores_diff_json = dp.get_predictions(tokens, features,
networkOutputs, args.n_best_size, args.max_answer_length)
return eval_time_elapsed, prediction, nbest_json
return eval_time_elapsed, prediction, nbest_json
def print_single_query(eval_time_elapsed, prediction, nbest_json):
print("------------------------")
print("Running inference in {:.3f} Sentences/Sec".format(1/eval_time_elapsed))
print("------------------------")
def print_single_query(eval_time_elapsed, prediction, nbest_json):
print("------------------------")
print("Running inference in {:.3f} Sentences/Sec".format(1/eval_time_elapsed))
print("------------------------")
print("Answer: '{}'".format(prediction))
print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0))
print("Answer: '{}'".format(prediction))
print("With probability: {:.3f}".format(nbest_json[0]['probability'] * 100.0))
if squad_examples:
all_predictions = collections.OrderedDict()
if squad_examples:
all_predictions = collections.OrderedDict()
for example in squad_examples:
features = question_features(example.doc_tokens, example.question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens)
all_predictions[example.id] = prediction
with open(output_prediction_file, "w") as f:
f.write(json.dumps(all_predictions, indent=4))
print("\nOutput dump to {}".format(output_prediction_file))
else:
# Extract tokecs from the paragraph
doc_tokens = dp.convert_doc_tokens(paragraph_text)
if question_text:
print("\nPassage: {}".format(paragraph_text))
print("\nQuestion: {}".format(question_text))
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
for example in squad_examples:
features = question_features(example.doc_tokens, example.question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, example.doc_tokens)
all_predictions[example.id] = prediction
with open(output_prediction_file, "w") as f:
f.write(json.dumps(all_predictions, indent=4))
print("\nOutput dump to {}".format(output_prediction_file))
else:
# If no question text is provided, loop until the question is 'exit'
EXIT_CMDS = ["exit", "quit"]
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
# Extract tokecs from the paragraph
doc_tokens = dp.convert_doc_tokens(paragraph_text)
if question_text:
print("\nPassage: {}".format(paragraph_text))
print("\nQuestion: {}".format(question_text))
while question_text.strip() not in EXIT_CMDS:
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
else:
# If no question text is provided, loop until the question is 'exit'
EXIT_CMDS = ["exit", "quit"]
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
while question_text.strip() not in EXIT_CMDS:
features = question_features(doc_tokens, question_text)
eval_time_elapsed, prediction, nbest_json = inference(features, doc_tokens)
print_single_query(eval_time_elapsed, prediction, nbest_json)
question_text = input("Question (to exit, type one of {:}): ".format(EXIT_CMDS))
# free allocated memory
for d_input in d_inputs:
cuda_call(cudart.cudaFree(d_input))
cuda_call(cudart.cudaFree(d_output))
cuda_call(cudart.cudaHostUnregister(h_output))
demo/BERT/notebooks/Q-and-A.ipynb
+116 -99
View File
@@ -58,6 +58,8 @@
"\n",
"import sys\n",
"sys.path.append('/workspace/TensorRT/demo/BERT')\n",
"import os\n",
"sys.path.append(os.path.abspath('..'))\n",
"\n",
"import ipywidgets as widgets\n",
"import tensorrt as trt;\n",
@@ -71,8 +73,8 @@
"import collections\n",
"import numpy as np\n",
"import tensorrt as trt\n",
"import pycuda.driver as cuda\n",
"import pycuda.autoinit\n",
"from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device_async, memcpy_device_to_host_async\n",
"from cuda.bindings import driver as cuda, runtime as cudart\n",
"\n",
"from helpers import tokenization as tokenization\n",
"from helpers import data_processing as dp\n",
@@ -109,44 +111,45 @@
" networkOutputs = []\n",
"\n",
" eval_time_elapsed = 0\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), 1, axis=0)\n",
" segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), 1, axis=0)\n",
" input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), 1, axis=0)\n",
" with CudaStreamContext() as stream:\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" input_ids_batch = np.repeat(np.expand_dims(feature.input_ids, 0), 1, axis=0)\n",
" segment_ids_batch = np.repeat(np.expand_dims(feature.segment_ids, 0), 1, axis=0)\n",
" input_mask_batch = np.repeat(np.expand_dims(feature.input_mask, 0), 1, axis=0)\n",
"\n",
" input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids_batch.ravel()))\n",
" segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids_batch.ravel()))\n",
" input_mask = cuda.register_host_memory(np.ascontiguousarray(input_mask_batch.ravel()))\n",
" input_ids = np.ascontiguousarray(input_ids_batch.ravel())\n",
" segment_ids = np.ascontiguousarray(segment_ids_batch.ravel())\n",
" input_mask = np.ascontiguousarray(input_mask_batch.ravel())\n",
"\n",
" eval_start_time = time.time()\n",
" cuda.memcpy_htod_async(d_inputs[0], input_ids, stream)\n",
" cuda.memcpy_htod_async(d_inputs[1], segment_ids, stream)\n",
" cuda.memcpy_htod_async(d_inputs[2], input_mask, stream)\n",
" eval_start_time = time.time()\n",
" memcpy_host_to_device_async(d_inputs[0], input_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[1], segment_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[2], input_mask, stream.stream)\n",
"\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
"\n",
" # Run inference\n",
" trt_context.execute_async_v3(stream_handle=stream.handle)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
" # Run inference\n",
" trt_context.execute_async_v3(stream_handle=stream.stream)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
"\n",
" # Transfer predictions back from GPU\n",
" cuda.memcpy_dtoh_async(h_output, d_output, stream)\n",
" stream.synchronize()\n",
" # Transfer predictions back from GPU\n",
" memcpy_device_to_host_async(h_output, d_output, stream.stream)\n",
" stream.synchronize()\n",
"\n",
" for index, batch in enumerate(h_output):\n",
" # Data Post-processing\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(batch.squeeze()[:, 0]),\n",
" end_logits = np.array(batch.squeeze()[:, 1]),\n",
" feature_index = feature_index\n",
" ))\n",
" for index, batch in enumerate(h_output):\n",
" # Data Post-processing\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(batch.squeeze()[:, 0]),\n",
" end_logits = np.array(batch.squeeze()[:, 1]),\n",
" feature_index = feature_index\n",
" ))\n",
"\n",
" eval_time_elapsed /= len(features)\n",
"\n",
@@ -172,59 +175,60 @@
" networkOutputs = []\n",
"\n",
" eval_time_elapsed = 0\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" B = 1\n",
" S = np.sum(feature.input_mask)\n",
" input_ids = feature.input_ids[0:S]\n",
" segment_ids = feature.segment_ids[0:S]\n",
" cu_seq_lens = np.array([0, S], dtype=np.int32)\n",
" with CudaStreamContext() as stream:\n",
" for feature_index, feature in enumerate(features):\n",
" # Copy inputs\n",
" B = 1\n",
" S = np.sum(feature.input_mask)\n",
" input_ids = feature.input_ids[0:S]\n",
" segment_ids = feature.segment_ids[0:S]\n",
" cu_seq_lens = np.array([0, S], dtype=np.int32)\n",
"\n",
" first_tensor_name = engine.get_tensor_name(0)\n",
" second_tensor_name = engine.get_tensor_name(1)\n",
" third_tensor_name = engine.get_tensor_name(2)\n",
" forth_tensor_name = engine.get_tensor_name(3)\n",
" first_tensor_name = engine.get_tensor_name(0)\n",
" second_tensor_name = engine.get_tensor_name(1)\n",
" third_tensor_name = engine.get_tensor_name(2)\n",
" forth_tensor_name = engine.get_tensor_name(3)\n",
"\n",
" if context.get_tensor_shape(first_tensor_name)[0] != S:\n",
" context.set_input_shape(first_tensor_name, (S,))\n",
" if context.get_tensor_shape(second_tensor_name)[0] != S:\n",
" context.set_input_shape(second_tensor_name, (S,))\n",
" if context.get_tensor_shape(third_tensor_name)[0] != 2:\n",
" context.set_input_shape(third_tensor_name, (2,))\n",
" if context.get_tensor_shape(forth_tensor_name)[0] != S:\n",
" context.set_input_shapee(forth_tensor_name, (S,))\n",
" if context.get_tensor_shape(first_tensor_name)[0] != S:\n",
" context.set_input_shape(first_tensor_name, (S,))\n",
" if context.get_tensor_shape(second_tensor_name)[0] != S:\n",
" context.set_input_shape(second_tensor_name, (S,))\n",
" if context.get_tensor_shape(third_tensor_name)[0] != 2:\n",
" context.set_input_shape(third_tensor_name, (2,))\n",
" if context.get_tensor_shape(forth_tensor_name)[0] != S:\n",
" context.set_input_shape(forth_tensor_name, (S,))\n",
"\n",
" h_input_ids = cuda.register_host_memory(np.ascontiguousarray(input_ids.ravel()))\n",
" h_segment_ids = cuda.register_host_memory(np.ascontiguousarray(segment_ids.ravel()))\n",
" h_cu_seq_lens = cuda.register_host_memory(np.ascontiguousarray(cu_seq_lens.ravel()))\n",
" h_input_ids = np.ascontiguousarray(input_ids.ravel())\n",
" h_segment_ids = np.ascontiguousarray(segment_ids.ravel())\n",
" h_cu_seq_lens = np.ascontiguousarray(cu_seq_lens.ravel())\n",
"\n",
" eval_start_time = time.time()\n",
" cuda.memcpy_htod_async(d_inputs[0], h_input_ids, INT8_stream)\n",
" cuda.memcpy_htod_async(d_inputs[1], h_segment_ids, INT8_stream)\n",
" cuda.memcpy_htod_async(d_inputs[2], h_cu_seq_lens, INT8_stream)\n",
" eval_start_time = time.time()\n",
" memcpy_host_to_device_async(d_inputs[0], h_input_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[1], h_segment_ids, stream.stream)\n",
" memcpy_host_to_device_async(d_inputs[2], h_cu_seq_lens, stream.stream)\n",
"\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
" # Setup tensor address\n",
" bindings = [int(d_inputs[i]) for i in range(3)] + [int(d_output)]\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i])\n",
"\n",
" # Run inference\n",
" trt_context.execute_async_v3(stream_handle=INT8_stream.handle)\n",
" # Synchronize the stream\n",
" INT8_stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
" # Run inference\n",
" trt_context.execute_async_v3(stream_handle=stream.stream)\n",
" # Synchronize the stream\n",
" stream.synchronize()\n",
" eval_time_elapsed += (time.time() - eval_start_time)\n",
"\n",
" # Transfer predictions back from GPU\n",
" cuda.memcpy_dtoh_async(h_output, d_output, INT8_stream)\n",
" INT8_stream.synchronize()\n",
" # Transfer predictions back from GPU\n",
" memcpy_device_to_host_async(h_output, d_output, stream.stream)\n",
" stream.synchronize()\n",
"\n",
" # Only retrieve and post-process the first batch\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(h_output[0:S]),\n",
" end_logits = np.array(h_output[S:S*2]),\n",
" feature_index = feature_index\n",
" ))\n",
" # Only retrieve and post-process the first batch\n",
" networkOutputs.append(_NetworkOutput(\n",
" start_logits = np.array(h_output[0:S]),\n",
" end_logits = np.array(h_output[S:S*2]),\n",
" feature_index = feature_index\n",
" ))\n",
"\n",
" eval_time_elapsed /= len(features)\n",
"\n",
@@ -268,7 +272,7 @@
"input_nbytes = trt.volume(input_shape) * trt.int32.itemsize\n",
"\n",
"# Allocate device memory for inputs.\n",
"d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(3)]\n",
"d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(3)]\n",
"\n",
"# Specify input shapes. These must be within the min/max bounds of the active profile (0th profile in this case)\n",
"# Note that input shapes can be specified on a per-inference basis, but in this case, we only have a single shape.\n",
@@ -278,34 +282,35 @@
"assert context.all_binding_shapes_specified\n",
"\n",
"# Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
"h_output = cuda.pagelocked_empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n",
"d_output = cuda.mem_alloc(h_output.nbytes)\n",
"h_output = np.empty(tuple(context.get_tensor_shape(engine.get_tensor_name(3))), dtype=np.float32)\n",
"cuda_call(cudart.cudaHostRegister(h_output, h_output.nbytes, 0))\n",
"d_output = cuda_call(cudart.cudaMalloc(h_output.nbytes))\n",
"\n",
"# Create a stream in which to copy inputs/outputs and run inference.\n",
"stream = cuda.Stream()\n",
"with CudaStreamContext() as stream:\n",
" ### INT8 TRT model\n",
" engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n",
" max_seq_length = 384\n",
"\n",
"### INT8 TRT model\n",
"engine_path = \"engines_%s/megatron_large_seqlen384_int8qat_sparse.engine\"%TRT_VERSION\n",
"max_seq_length = 384\n",
" INT8_runtime = trt.Runtime(TRT_LOGGER)\n",
" INT8_engine = INT8_runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n",
" INT8_context = INT8_engine.create_execution_context()\n",
"\n",
"INT8_runtime = trt.Runtime(TRT_LOGGER)\n",
"INT8_engine = INT8_runtime.deserialize_cuda_engine(open(engine_path, 'rb') .read()) \n",
"INT8_context = INT8_engine.create_execution_context()\n",
" # select engine profile\n",
" INT8_context.set_optimization_profile_async(0, stream.stream)\n",
"\n",
"# select engine profile\n",
"INT8_context.set_optimization_profile_async(0, stream.handle)\n",
" input_nbytes = max_seq_length * trt.int32.itemsize\n",
"\n",
"input_nbytes = max_seq_length * trt.int32.itemsize\n",
" # Allocate device memory for inputs.\n",
" INT8_d_inputs = [cuda_call(cudart.cudaMalloc(input_nbytes)) for binding in range(4)]\n",
"\n",
"# Allocate device memory for inputs.\n",
"INT8_d_inputs = [cuda.mem_alloc(input_nbytes) for binding in range(4)]\n",
" # Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
" INT8_h_output = np.empty((2 * max_seq_length), dtype=np.float32)\n",
" cuda_call(cudart.cudaHostRegister(INT8_h_output, INT8_h_output.nbytes, 0))\n",
" INT8_d_output = cuda_call(cudart.cudaMalloc(INT8_h_output.nbytes))\n",
"\n",
"# Allocate output buffer by querying the size from the context. This may be different for different input shapes.\n",
"INT8_h_output = cuda.pagelocked_empty((2 * max_seq_length), dtype=np.float32)\n",
"INT8_d_output = cuda.mem_alloc(INT8_h_output.nbytes)\n",
"\n",
"# Create a stream in which to copy inputs/outputs and run inference.\n",
"INT8_stream = cuda.Stream()\n"
" # No separate INT8 stream; using context manager stream\n",
" INT8_stream = stream.stream\n"
]
},
{
@@ -439,7 +444,19 @@
"id": "musical-right",
"metadata": {},
"outputs": [],
"source": []
"source": [
"\n",
"# free allocated memory\n",
"for d_input in d_inputs:\n",
" cuda_call(cudart.cudaFree(d_input))\n",
"cuda_call(cudart.cudaFree(d_output))\n",
"cuda_call(cudart.cudaHostUnregister(h_output))\n",
"\n",
"for INT8_d_input in INT8_d_inputs:\n",
" cuda_call(cudart.cudaFree(INT8_d_input))\n",
"cuda_call(cudart.cudaFree(INT8_d_output))\n",
"cuda_call(cudart.cudaHostUnregister(INT8_h_output))"
]
}
],
"metadata": {
demo/BERT/notebooks/benchmark.ipynb
+91 -88
View File
@@ -41,7 +41,7 @@
},
{
"cell_type": "code",
"execution_count": 5,
"execution_count": null,
"id": "c89592ab-e50b-47c9-9b5b-6b5067d0b22a",
"metadata": {
"jupyter": {
@@ -56,6 +56,8 @@
"\n",
"import sys\n",
"sys.path.append('/workspace/TensorRT/demo/BERT')\n",
"import os\n",
"sys.path.append(os.path.abspath('..'))\n",
"\n",
"import tensorrt as trt;\n",
"TRT_VERSION = trt.__version__\n",
@@ -65,8 +67,8 @@
"import ctypes\n",
"import numpy as np\n",
"import tensorrt as trt\n",
"import pycuda.driver as cuda\n",
"import pycuda.autoinit\n",
"from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host\n",
"from cuda.bindings import driver as cuda, runtime as cudart\n",
"import numpy as np\n",
"from tqdm import tqdm\n",
"\n",
@@ -80,15 +82,16 @@
"\n",
"TRT_LOGGER = trt.Logger(trt.Logger.ERROR)\n",
"\n",
"\n",
"class DeviceBuffer(object):\n",
" def __init__(self, shape, dtype=trt.int32):\n",
" self.buf = cuda.mem_alloc(trt.volume(shape) * dtype.itemsize)\n",
" self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize))\n",
"\n",
" def binding(self):\n",
" return int(self.buf)\n",
"\n",
" def free(self):\n",
" self.buf.free()\n",
" cuda_call(cudart.cudaFree(self.buf))\n",
"\n",
"doc_stride = 128\n",
"max_query_length = 64\n",
@@ -138,56 +141,57 @@
" test_cu_seq_lens = np.arange(0, args.sequence_length * max(args.batch_size) + 1, args.sequence_length, dtype=np.int32)\n",
"\n",
" # Copy input h2d\n",
" cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())\n",
" cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())\n",
" cuda.memcpy_htod(buffers[2].buf, test_cu_seq_lens.ravel())\n",
" memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())\n",
" memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())\n",
" memcpy_host_to_device(buffers[2].buf, test_cu_seq_lens.ravel())\n",
"\n",
" bench_times = {}\n",
" stream = cuda.Stream()\n",
" with CudaStreamContext() as stream:\n",
"\n",
" tensor_name = engine.get_tensor_name(engine.num_io_tensors-1)\n",
" for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
" for idx in range(engine.num_optimization_profiles):\n",
" profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n",
" if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
" context.set_optimization_profile_async(idx, stream.handle)\n",
" binding_idx_offset = idx * engine.num_io_tensors\n",
" break\n",
" tensor_name = engine.get_tensor_name(engine.num_io_tensors-1)\n",
" for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
" for idx in range(engine.num_optimization_profiles):\n",
" profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n",
" if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
" context.set_optimization_profile_async(idx, stream.stream)\n",
" binding_idx_offset = idx * engine.num_io_tensors\n",
" break\n",
"\n",
" # Each profile has unique bindings\n",
" bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
" input_shape = (batch_size, args.sequence_length)\n",
" for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
" assert context.all_binding_shapes_specified\n",
" # Each profile has unique bindings\n",
" bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
" input_shape = (batch_size, args.sequence_length)\n",
" for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
" assert context.all_binding_shapes_specified\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
"\n",
" # Inference\n",
" total_time = 0\n",
" start = cuda.Event()\n",
" end = cuda.Event()\n",
" # Inference\n",
" total_time = 0\n",
" start = cuda_call(cudart.cudaEventCreate())\n",
" end = cuda_call(cudart.cudaEventCreate())\n",
"\n",
" # Warmup\n",
" for _ in range(args.warm_up_runs):\n",
" context.execute_async_v3(stream_handle=stream.handle)\n",
" stream.synchronize()\n",
" # Warmup\n",
" for _ in range(args.warm_up_runs):\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" stream.synchronize()\n",
"\n",
" # Timing loop\n",
" times = []\n",
" progress_bar.value = 0\n",
" for _ in range(iteration_selector.value):\n",
" start.record(stream)\n",
" context.execute_async_v3(stream_handle=stream.handle)\n",
" end.record(stream)\n",
" stream.synchronize()\n",
" times.append(end.time_since(start))\n",
" progress_bar.value +=1\n",
" # Timing loop\n",
" times = []\n",
" progress_bar.value = 0\n",
" for _ in range(iteration_selector.value):\n",
" cuda_call(cudart.cudaEventRecord(start, stream.stream))\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" cuda_call(cudart.cudaEventRecord(end, stream.stream))\n",
" stream.synchronize()\n",
" elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end))\n",
" times.append(elapsed_time)\n",
" progress_bar.value +=1\n",
"\n",
" # Compute average time, 95th percentile time and 99th percentile time.\n",
" bench_times[batch_size] = times\n",
" # Compute average time, 95th percentile time and 99th percentile time.\n",
" bench_times[batch_size] = times\n",
"\n",
" [b.free() for b in buffers]\n",
"\n",
@@ -227,57 +231,56 @@
" test_input_mask = np.ones((max(args.batch_size), args.sequence_length), dtype=np.int32)\n",
"\n",
" # Copy input h2d\n",
" cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())\n",
" cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())\n",
" cuda.memcpy_htod(buffers[2].buf, test_input_mask.ravel())\n",
" memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())\n",
" memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())\n",
" memcpy_host_to_device(buffers[2].buf, test_input_mask.ravel())\n",
"\n",
" bench_times = {}\n",
" stream = cuda.Stream()\n",
" with CudaStreamContext() as stream:\n",
"\n",
" tensor_name = engine.get_tensor_name(engine.num_io_tensors - 1)\n",
" for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
" for idx in range(engine.num_optimization_profiles):\n",
" profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n",
" if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
" context.set_optimization_profile_async(idx, stream.handle)\n",
" binding_idx_offset = idx * engine.num_io_tensors\n",
" break\n",
" tensor_name = engine.get_tensor_name(engine.num_io_tensors - 1)\n",
" for idx, batch_size in enumerate(sorted(args.batch_size)):\n",
" for idx in range(engine.num_optimization_profiles):\n",
" profile_shape = engine.get_tensor_profile_shape(name = tensor_name, profile_index = idx)\n",
" if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size:\n",
" context.set_optimization_profile_async(idx, stream.stream)\n",
" binding_idx_offset = idx * engine.num_io_tensors\n",
" break\n",
"\n",
" # Each profile has unique bindings\n",
" bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
" input_shape = (batch_size, args.sequence_length)\n",
" for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
" assert context.all_binding_shapes_specified\n",
" # Each profile has unique bindings\n",
" bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]\n",
" input_shape = (batch_size, args.sequence_length)\n",
" for binding in range(3):\n",
" tensor_name = engine.get_tensor_name(binding)\n",
" context.set_input_shape(tensor_name, input_shape)\n",
" assert context.all_binding_shapes_specified\n",
"\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
" for i in range(engine.num_io_tensors):\n",
" context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])\n",
"\n",
" # Inference\n",
" total_time = 0\n",
" start = cuda.Event()\n",
" end = cuda.Event()\n",
" stream = cuda.Stream()\n",
" # Inference\n",
" total_time = 0\n",
" start = cuda_call(cudart.cudaEventCreate())\n",
" end = cuda_call(cudart.cudaEventCreate())\n",
"\n",
" # Warmup\n",
" for _ in range(args.warm_up_runs):\n",
" context.execute_async_v3(stream_handle=stream.handle)\n",
" stream.synchronize()\n",
" # Warmup\n",
" for _ in range(args.warm_up_runs):\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" stream.synchronize()\n",
"\n",
" # Timing loop\n",
" times = []\n",
" progress_bar.value = 0\n",
" for _ in range(iteration_selector.value):\n",
" start.record(stream)\n",
" context.execute_async_v3(stream_handle=stream.handle)\n",
" end.record(stream)\n",
" stream.synchronize()\n",
" times.append(end.time_since(start))\n",
" progress_bar.value +=1\n",
" # Timing loop\n",
" times = []\n",
" progress_bar.value = 0\n",
" for _ in range(iteration_selector.value):\n",
" start.record(stream.stream)\n",
" context.execute_async_v3(stream_handle=stream.stream)\n",
" end.record(stream.stream)\n",
" stream.synchronize()\n",
" times.append(end.time_since(start))\n",
" progress_bar.value +=1\n",
"\n",
" # Compute average time, 95th percentile time and 99th percentile time.\n",
" bench_times[batch_size] = times\n",
" # Compute average time, 95th percentile time and 99th percentile time.\n",
" bench_times[batch_size] = times\n",
"\n",
" [b.free() for b in buffers]\n",
"\n",
demo/BERT/perf.py
+51 -52
View File
@@ -18,24 +18,22 @@
import argparse
import ctypes
import time
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host
from cuda.bindings import driver as cuda, runtime as cudart
import numpy as np
TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
class DeviceBuffer(object):
def __init__(self, shape, dtype=trt.int32):
self.buf = cuda.mem_alloc(trt.volume(shape) * dtype.itemsize)
self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize))
def binding(self):
return int(self.buf)
def free(self):
self.buf.free()
cuda_call(cudart.cudaFree(self.buf))
def main():
@@ -73,63 +71,64 @@ def main():
test_input_mask = np.ones((max(args.batch_size), args.sequence_length), dtype=np.int32)
# Copy input h2d
cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())
cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())
cuda.memcpy_htod(buffers[2].buf, test_input_mask.ravel())
memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())
memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())
memcpy_host_to_device(buffers[2].buf, test_input_mask.ravel())
bench_times = {}
stream = cuda.Stream()
for batch_size in sorted(args.batch_size):
# Select engine profile
selected_profile = -1
for idx in range(engine.num_optimization_profiles):
profile_shape = engine.get_tensor_profile_shape(name = "input_ids", profile_index = idx)
if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size and profile_shape[0][1] <= args.sequence_length and profile_shape[2][1] >= args.sequence_length:
selected_profile = idx
break
if selected_profile == -1:
raise RuntimeError("None of the dynamic shape profiles meets the requirement batch = {} and sequence = {}.".format(batch_size, args.sequence_length))
context.set_optimization_profile_async(selected_profile, stream.handle)
with CudaStreamContext() as stream:
for batch_size in sorted(args.batch_size):
# Select engine profile
selected_profile = -1
for idx in range(engine.num_optimization_profiles):
profile_shape = engine.get_tensor_profile_shape(name = "input_ids", profile_index = idx)
if profile_shape[0][0] <= batch_size and profile_shape[2][0] >= batch_size and profile_shape[0][1] <= args.sequence_length and profile_shape[2][1] >= args.sequence_length:
selected_profile = idx
break
if selected_profile == -1:
raise RuntimeError("None of the dynamic shape profiles meets the requirement batch = {} and sequence = {}.".format(batch_size, args.sequence_length))
context.set_optimization_profile_async(selected_profile, stream.stream)
# Each profile has unique bindings
binding_idx_offset = selected_profile * engine.num_io_tensors
bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]
# Each profile has unique bindings
binding_idx_offset = selected_profile * engine.num_io_tensors
bindings = [0] * binding_idx_offset + [buf.binding() for buf in buffers]
input_shape = (batch_size, args.sequence_length)
for name in ["input_ids", "segment_ids", "input_mask"]:
context.set_input_shape(name, input_shape)
assert len(context.infer_shapes()) == 0
input_shape = (batch_size, args.sequence_length)
for name in ["input_ids", "segment_ids", "input_mask"]:
context.set_input_shape(name, input_shape)
assert len(context.infer_shapes()) == 0
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i + binding_idx_offset])
# Inference
total_time = 0
start = cuda.Event()
end = cuda.Event()
# Inference
total_time = 0
start = cuda_call(cudart.cudaEventCreate())
end = cuda_call(cudart.cudaEventCreate())
# Warmup
for _ in range(args.warm_up_runs):
context.execute_async_v3(stream_handle=stream.handle)
stream.synchronize()
# Warmup
for _ in range(args.warm_up_runs):
context.execute_async_v3(stream_handle=stream.stream)
stream.synchronize()
# Timing loop
times = []
actual_iterations = 0
start_time = time.time()
while actual_iterations < args.iterations or (time.time() - start_time) < args.duration:
start.record(stream)
context.execute_async_v3(stream_handle=stream.handle)
end.record(stream)
stream.synchronize()
times.append(end.time_since(start))
actual_iterations += 1
# Timing loop
times = []
actual_iterations = 0
start_time = time.time()
while actual_iterations < args.iterations or (time.time() - start_time) < args.duration:
cuda_call(cudart.cudaEventRecord(start, stream.stream))
context.execute_async_v3(stream_handle=stream.stream)
cuda_call(cudart.cudaEventRecord(end, stream.stream))
stream.synchronize()
elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end))
times.append(elapsed_time)
actual_iterations += 1
# Compute average time, 95th percentile time and 99th percentile time.
bench_times[batch_size] = times
# Compute average time, 95th percentile time and 99th percentile time.
bench_times[batch_size] = times
[b.free() for b in buffers]
[b.free() for b in buffers]
for batch_size, times in bench_times.items():
total_time = sum(times)
demo/BERT/perf_varseqlen.py
+44 -43
View File
@@ -20,8 +20,8 @@ import ctypes
import time
import numpy as np
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
from helpers.cuda_utils import cuda_call, CudaStreamContext, memcpy_host_to_device, memcpy_device_to_host
from cuda.bindings import driver as cuda, runtime as cudart
import numpy as np
@@ -29,13 +29,13 @@ TRT_LOGGER = trt.Logger(trt.Logger.ERROR)
class DeviceBuffer(object):
def __init__(self, shape, dtype=trt.int32):
self.buf = cuda.mem_alloc(trt.volume(shape) * dtype.itemsize)
self.buf = cuda_call(cudart.cudaMalloc(trt.volume(shape) * dtype.itemsize))
def binding(self):
return int(self.buf)
def free(self):
self.buf.free()
cuda_call(cudart.cudaFree(self.buf))
def main():
@@ -74,57 +74,58 @@ def main():
test_cu_seq_lens = np.arange(0, args.sequence_length * max(args.batch_size) + 1, args.sequence_length, dtype=np.int32)
# Copy input h2d
cuda.memcpy_htod(buffers[0].buf, test_word_ids.ravel())
cuda.memcpy_htod(buffers[1].buf, test_segment_ids.ravel())
cuda.memcpy_htod(buffers[2].buf, test_cu_seq_lens.ravel())
memcpy_host_to_device(buffers[0].buf, test_word_ids.ravel())
memcpy_host_to_device(buffers[1].buf, test_segment_ids.ravel())
memcpy_host_to_device(buffers[2].buf, test_cu_seq_lens.ravel())
bench_times = {}
for idx, batch_size in enumerate(sorted(args.batch_size)):
stream = cuda.Stream()
context.set_optimization_profile_async(0, stream.handle)
with CudaStreamContext() as stream:
context.set_optimization_profile_async(0, stream.stream)
# Each profile has unique bindings
bindings = [buf.binding() for buf in buffers]
# Each profile has unique bindings
bindings = [buf.binding() for buf in buffers]
shapes = {
"input_ids": (args.sequence_length * batch_size, ),
"segment_ids": (args.sequence_length * batch_size, ),
"cu_seqlens": (batch_size + 1, ),
"max_seqlen": (args.sequence_length, ),
}
shapes = {
"input_ids": (args.sequence_length * batch_size, ),
"segment_ids": (args.sequence_length * batch_size, ),
"cu_seqlens": (batch_size + 1, ),
"max_seqlen": (args.sequence_length, ),
}
for binding, shape in shapes.items():
context.set_input_shape(binding, shape)
assert len(context.infer_shapes()) == 0
for binding, shape in shapes.items():
context.set_input_shape(binding, shape)
assert len(context.infer_shapes()) == 0
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
for i in range(engine.num_io_tensors):
context.set_tensor_address(engine.get_tensor_name(i), bindings[i])
# Inference
total_time = 0
start = cuda.Event()
end = cuda.Event()
# Inference
total_time = 0
start = cuda_call(cudart.cudaEventCreate())
end = cuda_call(cudart.cudaEventCreate())
# Warmup
for _ in range(args.warm_up_runs):
context.execute_async_v3(stream_handle=stream.handle)
stream.synchronize()
# Warmup
for _ in range(args.warm_up_runs):
context.execute_async_v3(stream_handle=stream.stream)
stream.synchronize()
# Timing loop
times = []
actual_iterations = 0
start_time = time.time()
while actual_iterations < args.iterations or (time.time() - start_time) < args.duration:
start.record(stream)
context.execute_async_v3(stream_handle=stream.handle)
end.record(stream)
stream.synchronize()
times.append(end.time_since(start))
actual_iterations += 1
# Timing loop
times = []
actual_iterations = 0
start_time = time.time()
while actual_iterations < args.iterations or (time.time() - start_time) < args.duration:
cuda_call(cudart.cudaEventRecord(start, stream.stream))
context.execute_async_v3(stream_handle=stream.stream)
cuda_call(cudart.cudaEventRecord(end, stream.stream))
stream.synchronize()
elapsed_time = cuda_call(cudart.cudaEventElapsedTime(start, end))
times.append(elapsed_time)
actual_iterations += 1
# Compute average time, 95th percentile time and 99th percentile time.
bench_times[batch_size] = times
# Compute average time, 95th percentile time and 99th percentile time.
bench_times[batch_size] = times
[b.free() for b in buffers]
demo/DeBERTa/cuda_utils.py
+128
View File
@@ -0,0 +1,128 @@
#!/usr/bin/env python3
#
# 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.
#
import numpy as np
import logging
from cuda.bindings import driver as cuda, runtime as cudart
class CudaStreamContext:
"""CUDA stream lifecycle management with context manager support"""
def __init__(self):
"""Initialize CUDA stream"""
self._stream = cuda_call(cudart.cudaStreamCreate())
def __enter__(self):
"""Create CUDA stream when entering context (if not already created)"""
if self._stream is None:
self._stream = cuda_call(cudart.cudaStreamCreate())
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Destroy CUDA stream when exiting context"""
self.free()
@property
def stream(self) -> cudart.cudaStream_t:
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
return self._stream
def synchronize(self):
"""Synchronize the stream"""
if self._stream is None:
raise RuntimeError("Stream not created. Use 'with' statement.")
cuda_call(cudart.cudaStreamSynchronize(self._stream))
def free(self):
"""Explicitly free the CUDA stream"""
if self._stream is not None:
try:
cuda_call(cudart.cudaStreamDestroy(self._stream))
self._stream = None
except Exception as e:
logging.warning(f"Failed to destroy CUDA stream: {e}")
def __del__(self):
"""Cleanup stream on destruction"""
if hasattr(self, '_stream') and self._stream is not None:
self.free()
def __str__(self):
return f"CudaStreamContext: {self._stream}"
def __repr__(self):
return self.__str__()
def cuda_call(call):
"""Helper function to make CUDA calls and check for errors"""
def _cudaGetErrorEnum(error):
if isinstance(error, cuda.CUresult):
err, name = cuda.cuGetErrorName(error)
return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
elif isinstance(error, cudart.cudaError_t):
return cudart.cudaGetErrorName(error)[1]
else:
raise RuntimeError("Unknown error type: {}".format(error))
err, res = call[0], call[1:]
if err.value:
raise RuntimeError(
"CUDA error code={}({})".format(
err.value, _cudaGetErrorEnum(err)
)
)
if len(res) == 1:
return res[0]
elif len(res) == 0:
return None
else:
return res
def getComputeCapacity(devID=0):
major = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
minor = cuda_call(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
return (major, minor)
def memcpy_host_to_device_async(device_ptr: int, host_arr: np.ndarray, stream):
"""Wrapper for async host-to-device memory copy"""
cuda_call(cudart.cudaMemcpyAsync(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice, stream))
def memcpy_device_to_host_async(host_arr: np.ndarray, device_ptr: int, stream):
"""Wrapper for async device-to-host memory copy"""
cuda_call(cudart.cudaMemcpyAsync(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost, stream))
def memcpy_device_to_device_async(dst_device_ptr: int, src_device_ptr: int, nbytes: int, stream):
"""Wrapper for async device-to-device memory copy"""
cuda_call(cudart.cudaMemcpyAsync(dst_device_ptr, src_device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice, stream))
def memcpy_host_to_device(device_ptr: int, host_arr: np.ndarray):
"""Wrapper for synchronous host-to-device memory copy"""
cuda_call(cudart.cudaMemcpy(device_ptr, host_arr.ctypes.data, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice))
def memcpy_device_to_host(host_arr: np.ndarray, device_ptr: int):
"""Wrapper for synchronous device-to-host memory copy"""
cuda_call(cudart.cudaMemcpy(host_arr.ctypes.data, device_ptr, host_arr.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost))
def memcpy_device_to_device(dst_device_ptr: int, src_device_ptr: int, nbytes: int):
"""Wrapper for synchronous device-to-device memory copy"""
cuda_call(cudart.cudaMemcpy(dst_device_ptr, src_device_ptr, nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToDevice))
# Initialize CUDA
cuda_call(cudart.cudaFree(0))
demo/DeBERTa/deberta_tensorrt_inference.py
+22 -13
View File
@@ -40,9 +40,17 @@ import torch
import tensorrt as trt
import os, sys, argparse
import numpy as np
import pycuda.driver as cuda
import pycuda.autoinit # without this, "LogicError: explicit_context_dependent failed: invalid device context - no currently active context?"
from time import time
from cuda.bindings import driver as cuda, runtime as cudart
from cuda_utils import (
cuda_call,
CudaStreamContext,
memcpy_host_to_device_async,
memcpy_device_to_host_async,
memcpy_host_to_device,
memcpy_device_to_device_async,
memcpy_device_to_device,
)
TRT_VERSION = int(trt.__version__[:3].replace('.','')) # e.g., version 8.4.1.5 becomes 84
@@ -167,7 +175,7 @@ class TRTModel:
inputs = []
outputs = []
bindings = []
stream = cuda.Stream()
stream = CudaStreamContext()
for i in range(engine.num_io_tensors):
tensor_name = engine.get_tensor_name(i)
@@ -175,10 +183,10 @@ class TRTModel:
dtype = trt.nptype(engine.get_tensor_dtype(tensor_name))
# Allocate host and device buffers
host_mem = cuda.pagelocked_empty(size, dtype) # page-locked memory buffer (won't swapped to disk)
device_mem = cuda.mem_alloc(host_mem.nbytes)
host_mem = np.empty(size, dtype)
device_mem = cuda_call(cudart.cudaMalloc(host_mem.nbytes))
# Append the device buffer address to device bindings. When cast to int, it's a linear index into the context's memory (like memory address). See https://documen.tician.de/pycuda/driver.html#pycuda.driver.DeviceAllocation
# Append the device buffer address to device bindings. When cast to int, it's a linear index into the context's memory (like memory address).
bindings.append(int(device_mem))
# Append to the appropriate input/output list.
@@ -226,18 +234,19 @@ class TRTModel:
# fill host memory with flattened input data
np.copyto(self.inputs[i].host, model_input.ravel())
elif TORCH:
nbytes = model_input.element_size() * model_input.nelement()
if timing:
cuda.memcpy_dtod(self.inputs[i].device, model_input.data_ptr(), model_input.element_size() * model_input.nelement())
memcpy_device_to_device(self.inputs[i].device, model_input.data_ptr(), nbytes)
else:
# for Torch GPU tensor it's easier, can just do Device to Device copy
cuda.memcpy_dtod_async(self.inputs[i].device, model_input.data_ptr(), model_input.element_size() * model_input.nelement(), self.stream) # dtod need size in bytes
memcpy_device_to_device_async(self.inputs[i].device, model_input.data_ptr(), nbytes, self.stream.stream)
if NUMPY:
if timing:
[cuda.memcpy_htod(inp.device, inp.host) for inp in self.inputs]
[memcpy_host_to_device(inp.device, inp.host) for inp in self.inputs]
else:
# input, Host to Device
[cuda.memcpy_htod_async(inp.device, inp.host, self.stream) for inp in self.inputs]
[memcpy_host_to_device_async(inp.device, inp.host, self.stream.stream) for inp in self.inputs]
for i in range(self.engine.num_io_tensors):
self.context.set_tensor_address(self.engine.get_tensor_name(i), self.bindings[i])
@@ -250,13 +259,13 @@ class TRTModel:
duration = end_time - start_time
else:
# run inference
self.context.execute_async_v3(stream_handle=self.stream.handle)
self.context.execute_async_v3(stream_handle=self.stream.stream)
if timing:
[cuda.memcpy_dtoh(out.host, out.device) for out in self.outputs]
[cuda_call(cudart.cudaMemcpy(out.host.ctypes.data, out.device, out.host.nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)) for out in self.outputs]
else:
# output, Device to Host
[cuda.memcpy_dtoh_async(out.host, out.device, self.stream) for out in self.outputs]
[memcpy_device_to_host_async(out.host, out.device, self.stream.stream) for out in self.outputs]
if not timing:
# synchronize to ensure completion of async calls
demo/Diffusion/README.md
+50 -5
View File
@@ -7,7 +7,7 @@ This demo application ("demoDiffusion") showcases the acceleration of Stable Dif
### Clone the TensorRT OSS repository
```bash
git clone git@github.com:NVIDIA/TensorRT.git -b release/10.13.3 --single-branch
git clone git@github.com:NVIDIA/TensorRT.git -b release/10.14 --single-branch
cd TensorRT
```
@@ -23,10 +23,18 @@ NOTE: The demo supports CUDA>=12.0
### Install the required packages
To install dependencies for modern pipelines (SDXL and later):
```bash
source setup.sh
```
To install dependencies for legacy pipelines (SD 1.5/2.1):
```bash
REQUIREMENTS_FILE=requirements_legacy.txt source setup.sh
```
Check your installed version using:
`python3 -c 'import tensorrt;print(tensorrt.__version__)'`
@@ -41,7 +49,7 @@ onnx 1.18.0
onnx-graphsurgeon 0.5.2
onnxruntime 1.19.2
polygraphy 0.49.22
tensorrt 10.13.3.9
tensorrt 10.14.1.48
tokenizers 0.13.3
torch 2.8.0a0+5228986c39.nv25.6
transformers 4.52.4
@@ -225,12 +233,18 @@ Note that a denosing-percentage is applied to the number of denoising-steps when
### Generate an image with Stable Diffusion v3.5-large with ControlNet guided by an image and a text prompt
```bash
# Depth
# Depth BF16
python3 demo_controlnet_sd35.py "a photo of a man" --controlnet-type depth --hf-token=$HF_TOKEN --denoising-steps 40 --guidance-scale 4.5 --bf16 --download-onnx-models
# Canny
# Depth FP8
python3 demo_controlnet_sd35.py "a photo of a man" --version=3.5-large --fp8 --controlnet-type depth --download-onnx-models --denoising-steps=40 --guidance-scale 4.5 --hf-token=$HF_TOKEN
# Canny BF16
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models
# Canny FP8
python3 demo_controlnet_sd35.py "A Night time photo taken by Leica M11, portrait of a Japanese woman in a kimono, looking at the camera, Cherry blossoms" --version=3.5-large --fp8 --controlnet-type canny --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --download-onnx-models
# Blur
python3 demo_controlnet_sd35.py "generated ai art, a tiny, lost rubber ducky in an action shot close-up, surfing the humongous waves, inside the tube, in the style of Kelly Slater" --controlnet-type blur --hf-token=$HF_TOKEN --denoising-steps 60 --guidance-scale 3.5 --bf16 --download-onnx-models
```
@@ -455,7 +469,37 @@ Memory usage captured below excludes the ONNX export step, and assumes use of th
NOTE: The FP8 and FP4 Pipelines are supported on Hopper/Ada/Blackwell devices only. The FP4 pipeline is most performant on Blackwell devices.
### Specify Custom Paths for ONNX models and TensorRT engines (FLUX only)
### Run Cosmos2 World Foundation Models
Select the prompts and export them as below
```bash
export PROMPT="A close-up shot captures a vibrant yellow scrubber vigorously working on a grimy plate, its bristles moving in circular motions to lift stubborn grease and food residue. The dish, once covered in remnants of a hearty meal, gradually reveals its original glossy surface. Suds form and bubble around the scrubber, creating a satisfying visual of cleanliness in progress. The sound of scrubbing fills the air, accompanied by the gentle clinking of the dish against the sink. As the scrubber continues its task, the dish transforms, gleaming under the bright kitchen lights, symbolizing the triumph of cleanliness over mess."
export NEGATIVE_PROMPT="The video captures a series of frames showing ugly scenes, static with no motion, motion blur, over-saturation, shaky footage, low resolution, grainy texture, pixelated images, poorly lit areas, underexposed and overexposed scenes, poor color balance, washed out colors, choppy sequences, jerky movements, low frame rate, artifacting, color banding, unnatural transitions, outdated special effects, fake elements, unconvincing visuals, poorly edited content, jump cuts, visual noise, and flickering. Overall, the video is of poor quality."
```
#### 1. Generate an Image from a Text Prompt
##### Run Cosmos-Predict2-2B-Text2Image
```bash
# BF16
python3 demo_txt2image_cosmos.py "$PROMPT" --negative-prompt="$NEGATIVE_PROMPT" --hf-token=$HF_TOKEN
```
#### 2. Generate a Video guided by an Initial Video Conditioning and a Text Prompt
##### Run Cosmos-Predict2-2B-Video2World (only PyTorch backend enabled)
```bash
# BF16
python3 demo_vid2world_cosmos.py "$PROMPT" --negative-prompt="$NEGATIVE_PROMPT" --hf-token=$HF_TOKEN
```
### Specify Custom Paths for ONNX models and TensorRT engines (FLUX, Stable Diffusion 3.5 and Cosmos only)
Custom override paths to pre-exported ONNX model files can be provided using `--custom-onnx-paths`. These ONNX models are directly used to build TRT engines without further optimization on the ONNX graphs. Paths should be a comma-separated list of <model_name>:<path> pairs. For example: `--custom-onnx-paths=transformer:/path/to/transformer.onnx,vae:/path/to/vae.onnx`. Call <PipelineClass>.get_model_names(...) for the list of supported model names.
@@ -467,3 +511,4 @@ Custom override paths to pre-built engine files can be provided using `--custom-
- To accelerate engine building time use `--timing-cache <path to cache file>`. The cache file will be created if it does not already exist. Note that performance may degrade if cache files are used across multiple GPU targets. It is recommended to use timing caches only during development. To achieve the best perfromance in deployment, please build engines without timing cache.
- Specify new directories for storing onnx and engine files when switching between versions, LoRAs, ControlNets, etc. This can be done using `--onnx-dir <new onnx dir>` and `--engine-dir <new engine dir>`.
- Inference performance can be improved by enabling [CUDA graphs](https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#cuda-graphs) using `--use-cuda-graph`. Enabling CUDA graphs requires fixed input shapes, so this flag must be combined with `--build-static-batch` and cannot be combined with `--build-dynamic-shape`.
demo/Diffusion/demo_diffusion/dd_argparse.py
+17 -6
View File
@@ -79,6 +79,10 @@ def add_arguments(parser):
"flux.1-dev-canny",
"flux.1-dev-depth",
"flux.1-kontext-dev",
"cosmos-predict2-2b-text2image",
"cosmos-predict2-14b-text2image",
"cosmos-predict2-2b-video2world",
"cosmos-predict2-14b-video2world",
),
help="Version of Stable Diffusion",
)
@@ -141,7 +145,7 @@ def add_arguments(parser):
"--custom-onnx-paths",
type=parse_key_value_pairs,
help=(
"[FLUX only] Custom override paths to pre-exported ONNX model files. These ONNX models are directly used to "
"[FLUX, Stable Diffusion 3.5-large, Cosmos only] Custom override paths to pre-exported ONNX model files. These ONNX models are directly used to "
"build TRT engines without further optimization on the ONNX graphs. Paths should be a comma-separated list "
"of <model_name>:<path> pairs. For example: "
"--custom-onnx-paths=transformer:/path/to/transformer.onnx,vae:/path/to/vae.onnx. Call "
@@ -278,6 +282,7 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic
is_flux = args.version.startswith("flux")
is_sd35 = args.version.startswith("3.5")
is_cosmos = args.version.startswith("cosmos")
if args.height % 8 != 0 or args.width % 8 != 0:
raise ValueError(
@@ -322,9 +327,11 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic
)
# Check controlnet compatibility
if hasattr(args, "controlnet_type") and args.version != "xl-1.0":
raise ValueError("fp8 controlnet quantization is only supported for SDXL.")
if getattr(args, "controlnet_type", None) is not None:
if args.version not in ("xl-1.0", "3.5-large"):
raise ValueError("fp8 controlnet quantization is only supported for SDXL and SD3.5-large.")
if args.version == "3.5-large" and args.controlnet_type == "blur":
raise ValueError("Blur controlnet type is not supported for SD3.5.")
# Check for conflicting quantization
if args.int8:
raise ValueError("Cannot apply both int8 and fp8 quantization, please choose only one.")
@@ -386,7 +393,9 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic
# Torch-fallback and Torch-inference
if args.torch_fallback and not args.torch_inference:
assert is_flux or is_sd35, "PyTorch Fallback is only supported for Flux and Stable Diffusion 3.5 pipelines."
assert (
is_flux or is_sd35 or is_cosmos
), "PyTorch Fallback is only supported for Flux, Stable Diffusion 3.5 and Cosmos pipelines."
args.torch_fallback = args.torch_fallback.split(",")
if args.torch_fallback and args.torch_inference:
@@ -397,7 +406,9 @@ def process_pipeline_args(args: argparse.Namespace) -> Tuple[Dict[str, Any], Dic
# low-vram
if args.low_vram:
assert is_flux or is_sd35, "low-vram mode is only supported for Flux and Stable Diffusion 3.5 pipelines."
assert (
is_flux or is_sd35 or is_cosmos
), "low-vram mode is only supported for Flux, Stable Diffusion 3.5 and Cosmos pipelines."
# Pack arguments
kwargs_init_pipeline = {
demo/Diffusion/demo_diffusion/model/__init__.py
+7 -1
View File
@@ -26,12 +26,13 @@ from demo_diffusion.model.clip import (
SD3_T5XXLModel,
get_clip_embedding_dim,
)
from demo_diffusion.model.controlnet import SD3ControlNet
from demo_diffusion.model.diffusion_transformer import (
CosmosTransformerModel,
FluxTransformerModel,
SD3_MMDiTModel,
SD3TransformerModel,
)
from demo_diffusion.model.controlnet import SD3ControlNet
from demo_diffusion.model.gan import VQGANModel
from demo_diffusion.model.load import unload_torch_model
from demo_diffusion.model.lora import FLUXLoraLoader, SDLoraLoader, merge_loras
@@ -47,6 +48,8 @@ from demo_diffusion.model.unet import (
UNetXLModelControlNet,
)
from demo_diffusion.model.vae import (
AutoencoderKLWanEncoderModel,
AutoencoderKLWanModel,
SD3_VAEDecoderModel,
SD3_VAEEncoderModel,
TorchVAEEncoder,
@@ -67,6 +70,7 @@ __all__ = [
"SD3_T5XXLModel",
"CLIPVisionWithProjModel",
"CLIPImageProcessorModel",
"CosmosTransformerModel",
# diffusion_transformer
"SD3_MMDiTModel",
"FluxTransformerModel",
@@ -98,6 +102,8 @@ __all__ = [
"TorchVAEEncoder",
"VAEEncoderModel",
"SD3_VAEEncoderModel",
"AutoencoderKLWanModel",
"AutoencoderKLWanEncoderModel",
# load
"unload_torch_model",
]
demo/Diffusion/demo_diffusion/model/base_model.py
+13 -5
View File
@@ -19,13 +19,13 @@ import json
import os
import numpy as np
import onnx
import torch
from diffusers import DiffusionPipeline
from onnx import numpy_helper
import onnx
from demo_diffusion.model import load, optimizer
from demo_diffusion.model.lora import merge_loras
from onnx import numpy_helper
class BaseModel:
@@ -50,6 +50,7 @@ class BaseModel:
):
self.name = self.__class__.__name__
self.pipeline_type = pipeline
self.pipeline = pipeline.name
self.version = version
self.path = load.get_path(version, pipeline)
@@ -70,7 +71,7 @@ class BaseModel:
self.min_batch = 1
self.max_batch = max_batch_size
self.min_image_shape = 256 # min image resolution: 256x256
self.max_image_shape = 1344 # max image resolution: 1344x1344
self.max_image_shape = 1360 # max image resolution: 1360x1360
self.min_latent_shape = self.min_image_shape // self.compression_factor
self.max_latent_shape = self.max_image_shape // self.compression_factor
@@ -123,6 +124,7 @@ class BaseModel:
enable_lora_merge=False,
static_shape=False,
lora_loader=None,
dynamo=False,
):
onnx_opt_graph = None
# Export optimized ONNX model (if missing)
@@ -135,6 +137,11 @@ class BaseModel:
assert lora_loader is not None
model = merge_loras(model, lora_loader)
export_kwargs = {}
if dynamo:
export_kwargs["dynamic_shapes"] = self.get_dynamic_axes()
else:
export_kwargs["dynamic_axes"] = self.get_dynamic_axes()
inputs = self.get_sample_input(1, opt_image_height, opt_image_width, static_shape)
torch.onnx.export(
model,
@@ -145,8 +152,9 @@ class BaseModel:
do_constant_folding=self.do_constant_folding,
input_names=self.get_input_names(),
output_names=self.get_output_names(),
dynamic_axes=self.get_dynamic_axes(),
verbose=False,
dynamo=dynamo,
**export_kwargs,
)
if custom_model:
@@ -223,7 +231,7 @@ class BaseModel:
print(f"[I] Found cached weights map: {weights_map_path} ")
def optimize(self, onnx_graph, return_onnx=True, **kwargs):
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose)
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose, version=self.version)
opt.info(self.name + ": original")
opt.cleanup()
opt.info(self.name + ": cleanup")
demo/Diffusion/demo_diffusion/model/clip.py
+2 -2
View File
@@ -174,7 +174,7 @@ class CLIPModel(base_model.BaseModel):
return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)
def optimize(self, onnx_graph):
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose)
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose, version=self.version)
opt.info(self.name + ": original")
keep_outputs = [0, 1] if self.keep_pooled_output else [0]
opt.select_outputs(keep_outputs)
@@ -356,7 +356,7 @@ class SD3_CLIPGModel(CLIPModel):
return output
def optimize(self, onnx_graph):
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose)
opt = optimizer.Optimizer(onnx_graph, verbose=self.verbose, version=self.version)
opt.info(self.name + ": original")
opt.select_outputs([0, 1])
opt.cleanup()
demo/Diffusion/demo_diffusion/model/diffusion_transformer.py
+225 -23
View File
@@ -27,7 +27,7 @@ from demo_diffusion.utils_sd3.other_impls import load_into
from demo_diffusion.utils_sd3.sd3_impls import BaseModel as BaseModelSD3
# List of models to import from diffusers.models
models_to_import = ["FluxTransformer2DModel", "SD3Transformer2DModel"]
models_to_import = ["FluxTransformer2DModel", "SD3Transformer2DModel", "CosmosTransformer3DModel"]
for model in models_to_import:
globals()[model] = import_from_diffusers(model, "diffusers.models")
@@ -475,9 +475,8 @@ class SD3TransformerModel(base_model.BaseModel):
"encoder_hidden_states",
"pooled_projections",
"timestep",
"block_controlnet_hidden_states"
]
if not self.fp8:
input_names.append("block_controlnet_hidden_states")
return input_names
def get_output_names(self):
@@ -491,9 +490,8 @@ class SD3TransformerModel(base_model.BaseModel):
"pooled_projections": {0: xB},
"timestep": {0: xB},
"latent": {0: xB, 2: "H", 3: "W"},
"block_controlnet_hidden_states": {1: xB, 2: "latent_dim"}
}
if not self.fp8:
dynamic_axes["block_controlnet_hidden_states"] = {1: xB, 2: "latent_dim"}
return dynamic_axes
def get_input_profile(
@@ -535,9 +533,7 @@ class SD3TransformerModel(base_model.BaseModel):
(self.xB * max_batch, self.config["pooled_projection_dim"]),
],
"timestep": [(self.xB * min_batch,), (self.xB * batch_size,), (self.xB * max_batch,)],
}
if not self.fp8:
input_profile["block_controlnet_hidden_states"] = [
"block_controlnet_hidden_states": [
(
self.num_controlnet_layers,
self.xB * min_batch,
@@ -557,6 +553,7 @@ class SD3TransformerModel(base_model.BaseModel):
self.config["num_attention_heads"] * self.config["attention_head_dim"],
),
]
}
return input_profile
@@ -568,14 +565,13 @@ class SD3TransformerModel(base_model.BaseModel):
"pooled_projections": (self.xB * batch_size, self.config["pooled_projection_dim"]),
"timestep": (self.xB * batch_size,),
"latent": (self.xB * batch_size, self.out_channels, latent_height, latent_width),
}
if not self.fp8:
shape_dict["block_controlnet_hidden_states"] = (
"block_controlnet_hidden_states": (
self.num_controlnet_layers,
self.xB * batch_size,
latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"],
self.config["num_attention_heads"] * self.config["attention_head_dim"],
)
}
return shape_dict
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
@@ -600,19 +596,225 @@ class SD3TransformerModel(base_model.BaseModel):
),
torch.randn(self.xB * batch_size, self.config["pooled_projection_dim"], dtype=dtype, device=self.device),
torch.randn(self.xB * batch_size, dtype=torch.float32, device=self.device),
{
"block_controlnet_hidden_states": torch.randn(
self.num_controlnet_layers,
self.xB * batch_size,
latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"],
self.config["num_attention_heads"] * self.config["attention_head_dim"],
dtype=dtype,
device=self.device,
),
}
)
if not self.fp8:
sample_input += (
{
"block_controlnet_hidden_states": torch.randn(
self.num_controlnet_layers,
self.xB * batch_size,
latent_height // self.config["patch_size"] * latent_width // self.config["patch_size"],
self.config["num_attention_heads"] * self.config["attention_head_dim"],
dtype=dtype,
device=self.device,
),
}
return sample_input
class CosmosTransformerModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
tf32=False,
int8=False,
fp8=False,
bf16=False,
max_batch_size=16,
text_maxlen=77,
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
):
super(CosmosTransformerModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
int8=int8,
fp8=fp8,
bf16=bf16,
max_batch_size=max_batch_size,
text_maxlen=text_maxlen,
)
self.subfolder = "transformer"
self.transformer_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(self.transformer_model_dir):
self.config = CosmosTransformer3DModel.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load CosmosTransformer3DModel config from: {self.transformer_model_dir}")
self.config = CosmosTransformer3DModel.load_config(self.transformer_model_dir)
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
if not load.is_model_cached(self.transformer_model_dir, model_opts, self.hf_safetensor):
model = CosmosTransformer3DModel.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(self.transformer_model_dir, **model_opts)
else:
print(f"[I] Load CosmosTransformer3DModel model from: {self.transformer_model_dir}")
model = CosmosTransformer3DModel.from_pretrained(self.transformer_model_dir, **model_opts).to(self.device)
if torch_inference:
model.to(memory_format=torch.channels_last)
if self.fp16:
model.transformer_blocks[6].attn1.norm_q.float().to(self.device)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model.to(self.device)
def get_input_names(self):
input_names = [
"hidden_states",
"timestep",
"encoder_hidden_states",
"padding_mask",
]
if self.pipeline_type.is_video2world():
input_names.append("fps")
input_names.append("condition_mask")
return input_names
def get_output_names(self):
return ["latent"]
def get_dynamic_axes(self):
dynamic_axes = {
"hidden_states": {0: "B", 2: "latent_frames", 3: "latent_H", 4: "latent_W"},
"timestep": {0: "B"},
"encoder_hidden_states": {0: "B"},
"padding_mask": {0: "B", 2: "H", 3: "W"},
}
if self.pipeline_type.is_video2world():
dynamic_axes["fps"] = {0: "B"}
dynamic_axes["condition_mask"] = {0: "B", 2: "latent_frames", 3: "latent_H", 4: "latent_W"}
return dynamic_axes
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
(
min_batch,
max_batch,
min_image_height,
max_image_height,
min_image_width,
max_image_width,
min_latent_height,
max_latent_height,
min_latent_width,
max_latent_width,
) = self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
latent_frames = 24 if self.pipeline_type.is_video2world() else 1
latent_channels = (
self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"]
)
input_profile = {
"hidden_states": [
(min_batch, latent_channels, latent_frames, min_latent_height, min_latent_width),
(batch_size, latent_channels, latent_frames, latent_height, latent_width),
(max_batch, latent_channels, latent_frames, max_latent_height, max_latent_width),
],
"timestep": [(min_batch,), (batch_size,), (max_batch,)],
"encoder_hidden_states": [
(min_batch, self.text_maxlen, self.config["text_embed_dim"]),
(batch_size, self.text_maxlen, self.config["text_embed_dim"]),
(max_batch, self.text_maxlen, self.config["text_embed_dim"]),
],
"padding_mask": [
(1, 1, min_image_height, min_image_width),
(1, 1, image_height, image_width),
(1, 1, max_image_height, max_image_width),
],
}
if self.pipeline_type.is_video2world():
input_profile["fps"] = [(min_batch,), (batch_size,), (max_batch,)]
input_profile["condition_mask"] = [
(min_batch, 1, latent_frames, min_latent_height, min_latent_width),
(batch_size, 1, latent_frames, latent_height, latent_width),
(max_batch, 1, latent_frames, max_latent_height, max_latent_width),
]
return input_profile
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
# TODO: get latent_frames from infer call
latent_frames = 24 if self.pipeline_type.is_video2world() else 1
latent_channels = (
self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"]
)
shape_dict = {
"hidden_states": (batch_size, latent_channels, latent_frames, latent_height, latent_width),
"timestep": (batch_size,),
"encoder_hidden_states": (batch_size, self.text_maxlen, self.config["text_embed_dim"]),
"padding_mask": (1, 1, image_height, image_width),
"latent": (batch_size, self.config["in_channels"], latent_frames, latent_height, latent_width),
}
if self.pipeline_type.is_video2world():
shape_dict["fps"] = (batch_size,)
shape_dict["condition_mask"] = (batch_size, 1, latent_frames, latent_height, latent_width)
return shape_dict
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float32
assert not (self.fp16 and self.bf16), "fp16 and bf16 cannot be enabled simultaneously"
tensor_dtype = torch.bfloat16 if self.bf16 else (torch.float16 if self.fp16 else torch.float32)
latent_frames = 1
latent_channels = (
self.config["in_channels"] - 1 if self.pipeline_type.is_video2world() else self.config["in_channels"]
)
sample_input = (
{
"hidden_states": torch.randn(
batch_size,
latent_channels,
latent_frames,
latent_height,
latent_width,
dtype=tensor_dtype,
device=self.device,
),
"timestep": torch.tensor([1.0] * batch_size, dtype=tensor_dtype, device=self.device),
"encoder_hidden_states": torch.randn(
batch_size, self.text_maxlen, self.config["text_embed_dim"], dtype=tensor_dtype, device=self.device
),
"padding_mask": torch.ones(
batch_size, 1, image_height, image_width, dtype=tensor_dtype, device=self.device
),
},
)
if self.pipeline_type.is_video2world():
sample_input[-1]["fps"] = torch.tensor([30] * batch_size, dtype=dtype, device=self.device)
sample_input[-1]["condition_mask"] = torch.randn(
batch_size,
1,
latent_frames,
latent_height,
latent_width,
dtype=tensor_dtype,
device=self.device,
)
return sample_input
demo/Diffusion/demo_diffusion/model/load.py
+8
View File
@@ -97,6 +97,14 @@ def get_path(version: str, pipeline: "pipeline.DiffusionPipeline", controlnets:
return "black-forest-labs/FLUX.1-Depth-dev"
elif version == "flux.1-kontext-dev":
return "black-forest-labs/FLUX.1-Kontext-dev"
elif version == "cosmos-predict2-2b-text2image":
return "nvidia/Cosmos-Predict2-2B-Text2Image"
elif version == "cosmos-predict2-14b-text2image":
return "nvidia/Cosmos-Predict2-14B-Text2Image"
elif version == "cosmos-predict2-2b-video2world":
return "nvidia/Cosmos-Predict2-2B-Video2World"
elif version == "cosmos-predict2-14b-video2world":
return "nvidia/Cosmos-Predict2-14B-Video2World"
else:
raise ValueError(f"Unsupported version {version} + pipeline {pipeline.name}")
demo/Diffusion/demo_diffusion/model/optimizer.py
+13 -1
View File
@@ -23,6 +23,7 @@ import onnx
import onnx_graphsurgeon as gs
import torch
from onnx import shape_inference
from onnxconverter_common.float16 import convert_float_to_float16
from polygraphy.backend.onnx.loader import fold_constants
from demo_diffusion.model import load
@@ -46,9 +47,11 @@ def optimize_checkpoint(model, torch_inference: str):
class Optimizer:
def __init__(self, onnx_graph, verbose=False):
def __init__(self, onnx_graph, verbose=False, version=None):
self.graph = gs.import_onnx(onnx_graph)
self.verbose = verbose
self.version = version
def info(self, prefix):
if self.verbose:
@@ -186,6 +189,15 @@ class Optimizer:
# Convert INT8 Zero to FP8.
onnx_graph = convert_zp_fp8(onnx_graph)
# WAR for legacy SD pipelines
legacy_versions = (
"1.4",
"1.5",
"2.1",
)
if any(self.version.startswith(prefix) for prefix in legacy_versions):
onnx_graph = convert_float_to_float16(onnx_graph, keep_io_types=False, disable_shape_infer=True)
self.graph = gs.import_onnx(onnx_graph)
# Add cast nodes to Resize I/O.
cast_resize_io(self.graph)
demo/Diffusion/demo_diffusion/model/t5.py
+21 -2
View File
@@ -27,6 +27,7 @@ from demo_diffusion.model import base_model, load, optimizer
class T5Model(base_model.BaseModel):
def __init__(
self,
version,
@@ -44,6 +45,7 @@ class T5Model(base_model.BaseModel):
build_strongly_typed=False,
weight_streaming=False,
weight_streaming_budget_percentage=None,
use_attention_mask=False,
):
super(T5Model, self).__init__(
version,
@@ -70,6 +72,7 @@ class T5Model(base_model.BaseModel):
self.build_strongly_typed = build_strongly_typed
self.weight_streaming = weight_streaming
self.weight_streaming_budget_percentage = weight_streaming_budget_percentage
self.use_attention_mask = use_attention_mask
def get_model(self, torch_inference=""):
model_opts = (
@@ -91,12 +94,16 @@ class T5Model(base_model.BaseModel):
return model
def get_input_names(self):
if self.use_attention_mask:
return ["input_ids", "attention_mask"]
return ["input_ids"]
def get_output_names(self):
return ["text_embeddings"]
def get_dynamic_axes(self):
if self.use_attention_mask:
return {"input_ids": {0: "B"}, "attention_mask": {0: "B"}, "text_embeddings": {0: "B"}}
return {"input_ids": {0: "B"}, "text_embeddings": {0: "B"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
@@ -104,9 +111,16 @@ class T5Model(base_model.BaseModel):
min_batch, max_batch, _, _, _, _, _, _, _, _ = self.get_minmax_dims(
batch_size, image_height, image_width, static_batch, static_shape
)
return {
profile = {
"input_ids": [(min_batch, self.text_maxlen), (batch_size, self.text_maxlen), (max_batch, self.text_maxlen)]
}
if self.use_attention_mask:
profile["attention_mask"] = [
(min_batch, self.text_maxlen),
(batch_size, self.text_maxlen),
(max_batch, self.text_maxlen),
]
return profile
def get_shape_dict(self, batch_size, image_height, image_width):
self.check_dims(batch_size, image_height, image_width)
@@ -114,8 +128,13 @@ class T5Model(base_model.BaseModel):
"input_ids": (batch_size, self.text_maxlen),
"text_embeddings": (batch_size, self.text_maxlen, self.config.d_model),
}
if self.use_attention_mask:
output["attention_mask"] = (batch_size, self.text_maxlen)
return output
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
self.check_dims(batch_size, image_height, image_width)
return torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)
inputs = {"input_ids": torch.zeros(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)}
if self.use_attention_mask:
inputs["attention_mask"] = torch.ones(batch_size, self.text_maxlen, dtype=torch.int32, device=self.device)
return inputs
demo/Diffusion/demo_diffusion/model/vae.py
+161 -4
View File
@@ -27,10 +27,7 @@ from demo_diffusion.utils_sd3.other_impls import load_into
from demo_diffusion.utils_sd3.sd3_impls import SDVAE
# List of models to import from diffusers.models
models_to_import = [
"AutoencoderKL",
"AutoencoderKLTemporalDecoder",
]
models_to_import = ["AutoencoderKL", "AutoencoderKLTemporalDecoder", "AutoencoderKLWan"]
for model in models_to_import:
globals()[model] = import_from_diffusers(model, "diffusers.models")
@@ -507,3 +504,163 @@ class SD3_VAEEncoderModel(base_model.BaseModel):
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
dtype = torch.float16 if self.fp16 else torch.float32
return torch.randn(batch_size, 3, image_height, image_width, dtype=dtype, device=self.device)
class AutoencoderKLWanModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
tf32=False,
bf16=False,
max_batch_size=16,
):
super(AutoencoderKLWanModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
bf16=bf16,
max_batch_size=max_batch_size,
)
self.subfolder = "vae"
self.vae_decoder_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(self.vae_decoder_model_dir):
self.config = AutoencoderKLWan.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load AutoencoderKLWan (decoder) config from: {self.vae_decoder_model_dir}")
self.config = AutoencoderKLWan.load_config(self.vae_decoder_model_dir)
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
if not load.is_model_cached(self.vae_decoder_model_dir, model_opts, self.hf_safetensor):
model = AutoencoderKLWan.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(self.vae_decoder_model_dir, **model_opts)
else:
print(f"[I] Load AutoencoderKLWan (decoder) model from: {self.vae_decoder_model_dir}")
model = AutoencoderKLWan.from_pretrained(self.vae_decoder_model_dir, **model_opts).to(self.device)
model.forward = model.decode
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
def get_input_names(self):
return ["latent"]
def get_output_names(self):
return ["images"]
def get_dynamic_axes(self):
return {"latent": {0: "B", 3: "H", 4: "W"}, "images": {0: "B", 3: "8H", 4: "8W"}}
def get_input_profile(self, batch_size, image_height, image_width, static_batch, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
min_batch, max_batch, _, _, _, _, min_latent_height, max_latent_height, min_latent_width, max_latent_width = (
self.get_minmax_dims(batch_size, image_height, image_width, static_batch, static_shape)
)
return {
"latent": [
(min_batch, self.config["z_dim"], 1, min_latent_height, min_latent_width),
(batch_size, self.config["z_dim"], 1, latent_height, latent_width),
(max_batch, self.config["z_dim"], 1, max_latent_height, max_latent_width),
]
}
def get_shape_dict(self, batch_size, image_height, image_width):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
return {
"latent": (batch_size, self.config["z_dim"], 1, latent_height, latent_width),
"images": (batch_size, 3, 1, image_height, image_width),
}
def get_sample_input(self, batch_size, image_height, image_width, static_shape):
latent_height, latent_width = self.check_dims(batch_size, image_height, image_width)
dtype = torch.float16 if self.fp16 else torch.bfloat16 if self.bf16 else torch.float32
return torch.randn(
batch_size, self.config["z_dim"], 1, latent_height, latent_width, dtype=dtype, device=self.device
)
class AutoencoderKLWanEncoderModelWrapper(torch.nn.Module):
def __init__(self, model):
super().__init__()
self.model = model
def forward(self, x):
return self.model.encode(x).latent_dist.sample()
class AutoencoderKLWanEncoderModel(base_model.BaseModel):
def __init__(
self,
version,
pipeline,
device,
hf_token,
verbose,
framework_model_dir,
fp16=False,
tf32=False,
bf16=False,
max_batch_size=16,
):
super(AutoencoderKLWanEncoderModel, self).__init__(
version,
pipeline,
device=device,
hf_token=hf_token,
verbose=verbose,
framework_model_dir=framework_model_dir,
fp16=fp16,
tf32=tf32,
bf16=bf16,
max_batch_size=max_batch_size,
)
self.subfolder = "vae"
self.vae_encoder_model_dir = load.get_checkpoint_dir(
self.framework_model_dir, self.version, self.pipeline, self.subfolder
)
if not os.path.exists(self.vae_encoder_model_dir):
self.config = AutoencoderKLWan.load_config(self.path, subfolder=self.subfolder, token=self.hf_token)
else:
print(f"[I] Load AutoencoderKLWan (encoder) config from: {self.vae_encoder_model_dir}")
self.config = AutoencoderKLWan.load_config(self.vae_encoder_model_dir)
def get_model(self, torch_inference=""):
model_opts = (
{"torch_dtype": torch.float16} if self.fp16 else {"torch_dtype": torch.bfloat16} if self.bf16 else {}
)
if not load.is_model_cached(self.vae_encoder_model_dir, model_opts, self.hf_safetensor):
model = AutoencoderKLWan.from_pretrained(
self.path,
subfolder=self.subfolder,
use_safetensors=self.hf_safetensor,
token=self.hf_token,
**model_opts,
).to(self.device)
model.save_pretrained(self.vae_encoder_model_dir, **model_opts)
else:
print(f"[I] Load AutoencoderKLWan (encoder) model from: {self.vae_encoder_model_dir}")
model = AutoencoderKLWan.from_pretrained(self.vae_encoder_model_dir, **model_opts).to(self.device)
model = AutoencoderKLWanEncoderModelWrapper(model)
model = optimizer.optimize_checkpoint(model, torch_inference)
return model
demo/Diffusion/demo_diffusion/pipeline/__init__.py
+71 -12
View File
@@ -14,19 +14,11 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
import importlib
from typing import TYPE_CHECKING
from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline
from demo_diffusion.pipeline.flux_pipeline import FluxKontextPipeline, FluxPipeline
from demo_diffusion.pipeline.stable_cascade_pipeline import StableCascadePipeline
from demo_diffusion.pipeline.stable_diffusion_3_pipeline import StableDiffusion3Pipeline
from demo_diffusion.pipeline.stable_diffusion_35_pipeline import (
StableDiffusion35Pipeline,
)
from demo_diffusion.pipeline.stable_diffusion_pipeline import StableDiffusionPipeline
from demo_diffusion.pipeline.stable_video_diffusion_pipeline import (
StableVideoDiffusionPipeline,
)
from demo_diffusion.pipeline.type import PIPELINE_TYPE
# Expose public API while avoiding importing optional dependencies at module import time.
# Each attribute is imported on first access via __getattr__.
__all__ = [
"DiffusionPipeline",
@@ -36,6 +28,73 @@ __all__ = [
"StableDiffusion3Pipeline",
"StableDiffusion35Pipeline",
"StableDiffusionPipeline",
"CosmosPipeline",
"StableVideoDiffusionPipeline",
"PIPELINE_TYPE",
]
_LAZY_ATTRS = {
# Core/base
"DiffusionPipeline": ("demo_diffusion.pipeline.diffusion_pipeline", "DiffusionPipeline"),
"PIPELINE_TYPE": ("demo_diffusion.pipeline.type", "PIPELINE_TYPE"),
# Stable Diffusion family
"StableDiffusionPipeline": ("demo_diffusion.pipeline.stable_diffusion_pipeline", "StableDiffusionPipeline"),
"StableDiffusion3Pipeline": ("demo_diffusion.pipeline.stable_diffusion_3_pipeline", "StableDiffusion3Pipeline"),
"StableDiffusion35Pipeline": ("demo_diffusion.pipeline.stable_diffusion_35_pipeline", "StableDiffusion35Pipeline"),
# Stable Cascade
"StableCascadePipeline": ("demo_diffusion.pipeline.stable_cascade_pipeline", "StableCascadePipeline"),
# Stable Video Diffusion
"StableVideoDiffusionPipeline": ("demo_diffusion.pipeline.stable_video_diffusion_pipeline", "StableVideoDiffusionPipeline"),
# Flux family (optional dependency: `flux`)
"FluxPipeline": ("demo_diffusion.pipeline.flux_pipeline", "FluxPipeline"),
"FluxKontextPipeline": ("demo_diffusion.pipeline.flux_pipeline", "FluxKontextPipeline"),
# Cosmos (optional dependency: `flux`)
"CosmosPipeline": ("demo_diffusion.pipeline.cosmos_pipeline", "CosmosPipeline"),
}
def __getattr__(name):
if name not in _LAZY_ATTRS:
raise AttributeError(f"module 'demo_diffusion.pipeline' has no attribute {name!r}")
module_path, attr_name = _LAZY_ATTRS[name]
try:
module = importlib.import_module(module_path)
except ModuleNotFoundError as e:
missing_pkg = e.name or "<unknown>"
raise ModuleNotFoundError(
f"Optional dependency '{missing_pkg}' is required for '{name}'. "
"Install the appropriate extras/requirements for the selected pipeline "
"(e.g., use the non-legacy requirements for Flux/Cosmos), or install the missing package."
) from e
try:
return getattr(module, attr_name)
except AttributeError as e:
raise AttributeError(
f"'{module_path}' does not export attribute '{attr_name}' (while resolving '{name}')."
) from e
if TYPE_CHECKING:
from demo_diffusion.pipeline.diffusion_pipeline import DiffusionPipeline as DiffusionPipeline
from demo_diffusion.pipeline.type import PIPELINE_TYPE as PIPELINE_TYPE
from demo_diffusion.pipeline.stable_diffusion_pipeline import (
StableDiffusionPipeline as StableDiffusionPipeline,
)
from demo_diffusion.pipeline.stable_diffusion_3_pipeline import (
StableDiffusion3Pipeline as StableDiffusion3Pipeline,
)
from demo_diffusion.pipeline.stable_diffusion_35_pipeline import (
StableDiffusion35Pipeline as StableDiffusion35Pipeline,
)
from demo_diffusion.pipeline.stable_cascade_pipeline import (
StableCascadePipeline as StableCascadePipeline,
)
from demo_diffusion.pipeline.stable_video_diffusion_pipeline import (
StableVideoDiffusionPipeline as StableVideoDiffusionPipeline,
)
# Optional pipelines
from demo_diffusion.pipeline.flux_pipeline import (
FluxPipeline as FluxPipeline,
FluxKontextPipeline as FluxKontextPipeline,
)
from demo_diffusion.pipeline.cosmos_pipeline import CosmosPipeline as CosmosPipeline
demo/Diffusion/demo_diffusion/pipeline/cosmos_pipeline.py
+1009
View File
File diff suppressed because it is too large Load Diff
demo/Diffusion/demo_diffusion/pipeline/diffusion_pipeline.py
+14 -2
View File
@@ -100,6 +100,10 @@ class DiffusionPipeline(ABC):
"flux.1-dev-depth",
"flux.1-schnell",
"flux.1-kontext-dev",
"cosmos-predict2-2b-text2image",
"cosmos-predict2-14b-text2image",
"cosmos-predict2-2b-video2world",
"cosmos-predict2-14b-video2world",
)
SCHEDULER_DEFAULTS = {
"1.4": "PNDM",
@@ -120,6 +124,10 @@ class DiffusionPipeline(ABC):
"flux.1-dev-depth": "FlowMatchEuler",
"flux.1-schnell": "FlowMatchEuler",
"flux.1-kontext-dev": "FlowMatchEuler",
"cosmos-predict2-2b-text2image": "FlowMatchEuler",
"cosmos-predict2-14b-text2image": "FlowMatchEuler",
"cosmos-predict2-2b-video2world": "FlowMatchEuler",
"cosmos-predict2-14b-video2world": "FlowMatchEuler",
}
def __init__(
@@ -385,9 +393,10 @@ class DiffusionPipeline(ABC):
), "fp8 quantization only supported for SDXL, SD1.5, SD2.1, SD3.5-large and FLUX pipelines"
if (
(self.pipeline_type.is_sd_xl() and model_name == "unetxl")
or (self.version.startswith("flux.1") and model_name == "transformer")
or (
(self.version.startswith("flux.1") or self.version.startswith("3.5-large"))
and model_name == "transformer"
self.version.startswith("3.5-large")
and ("transformer" in model_name or "controlnet" in model_name)
)
or (model_name == "unet")
):
@@ -622,6 +631,7 @@ class DiffusionPipeline(ABC):
else:
self.is_native_export_supported(model_config)
dynamo = True if self.pipeline_type.is_video2world() and model_name == "transformer" else False
if do_export_onnx or do_export_weights_map:
if not model_config['use_int8'] and not model_config['use_fp8']:
obj.export_onnx(
@@ -633,6 +643,7 @@ class DiffusionPipeline(ABC):
enable_lora_merge=model_config["do_lora_merge"],
static_shape=static_shape,
lora_loader=self.lora_loader,
dynamo=dynamo,
)
else:
print(f"[I] Generating quantized ONNX model: {model_config['onnx_path']}")
@@ -656,6 +667,7 @@ class DiffusionPipeline(ABC):
opt_image_width,
custom_model=quantized_model,
static_shape=static_shape,
dynamo=dynamo,
)
# FIXME do_export_weights_map needs ONNX graph
demo/Diffusion/demo_diffusion/pipeline/flux_pipeline.py
+1 -1
View File
@@ -794,7 +794,7 @@ class FluxPipeline(DiffusionPipeline):
print("|-----------------|--------------|")
print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms))
print("|-----------------|--------------|")
print("Throughput: {:.2f} image/s".format(batch_size * 1000.0 / walltime_ms))
print("Throughput: {:.5f} image/s".format(batch_size * 1000.0 / walltime_ms))
def _check_integrity(self, images):
integrity_checker = PixtralContentFilter(self.device)
demo/Diffusion/demo_diffusion/pipeline/stable_cascade_pipeline.py
+18 -15
View File
@@ -93,7 +93,7 @@ class StableCascadePipeline(StableDiffusionPipeline):
self.models['vqgan'] = VQGANModel(**models_args, fp16=self.fp16, bf16=self.bf16, latent_dim_scale = self.latent_dim_scale)
def encode_prompt(self, prompt, negative_prompt, encoder='clip', pooled_outputs=False, output_hidden_states=False):
self.profile_start('clip', color='green')
self.profile_start(encoder, color='green')
tokenizer = self.tokenizer
@@ -140,7 +140,7 @@ class StableCascadePipeline(StableDiffusionPipeline):
if output_hidden_states:
text_embeddings = torch.cat([text_hidden_states, uncond_hidden_states]) if self.do_classifier_free_guidance else text_hidden_states
self.profile_stop('clip')
self.profile_stop(encoder)
if pooled_outputs:
return text_embeddings, pooled_output
return text_embeddings
@@ -157,7 +157,7 @@ class StableCascadePipeline(StableDiffusionPipeline):
do_autocast = False
with torch.autocast('cuda', enabled=do_autocast):
self.profile_start('denoise', color='blue')
self.profile_start(denoiser, color='blue')
for step_index, timestep in enumerate(timesteps):
# ratio input required for stable cascade prior
timestep_ratio = timestep.expand(latents.size(0)).to(latents.dtype)
@@ -197,31 +197,34 @@ class StableCascadePipeline(StableDiffusionPipeline):
latents = latents.to(dtype=torch.bfloat16 if self.bf16 else torch.float32)
self.profile_stop('denoise')
self.profile_stop(denoiser)
return latents
def decode_latent(self, latents):
self.profile_start('vqgan', color='red')
latents = self.models['vqgan'].scale_factor * latents
def decode_latent(self, latents, model_name='vqgan'):
self.profile_start(model_name, color='red')
latents = self.models[model_name].scale_factor * latents
if self.torch_inference:
images = self.torch_models['vqgan'](latents)['sample']
images = self.torch_models[model_name](latents)['sample']
else:
images = self.runEngine('vqgan', {'latent': latents})['images']
self.profile_stop('vqgan')
images = self.runEngine(model_name, {'latent': latents})['images']
self.profile_stop(model_name)
return images
def print_summary(self, denoising_steps, walltime_ms, batch_size):
print('|-----------------|--------------|')
print('| {:^15} | {:^12} |'.format('Module', 'Latency'))
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('CLIP', cudart.cudaEventElapsedTime(self.events['clip'][0], self.events['clip'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('UNet'+' x '+str(denoising_steps), cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1]))
if 'vqgan' in self.stages:
print('| {:^15} | {:>9.2f} ms |'.format('VQGAN', cudart.cudaEventElapsedTime(self.events['vqgan'][0], self.events['vqgan'][1])[1]))
for stage in self.stages:
stage_name = stage + ' x ' + str(denoising_steps) if stage == 'unet' else stage
print(
"| {:^15} | {:>9.2f} ms |".format(
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
)
)
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
print('|-----------------|--------------|')
print('Throughput: {:.2f} image/s'.format(batch_size*1000./walltime_ms))
print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms))
def infer(
self,
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_35_pipeline.py
+19 -34
View File
@@ -178,6 +178,12 @@ class StableDiffusion35Pipeline(DiffusionPipeline):
elif "controlnet" in model_name:
hf_download_path_cnet = hf_download_path.replace("large", "controlnets")
dirname = f"controlnet_{self.controlnet}"
if "blur" in model_name:
pass
elif model_config["use_fp8"]:
dirname = os.path.join(dirname, "fp8")
elif self.bf16:
dirname = os.path.join(dirname, "bf16")
elif model_name in self.stages:
dirname = model_name
else:
@@ -241,7 +247,7 @@ class StableDiffusion35Pipeline(DiffusionPipeline):
self.bf16 = True if int8 or fp8 or fp4 else self.bf16
self.fp16 = True if not self.bf16 else False
self.tf32=True
self.tf32 = True
self.fp8 = fp8
self.int8 = int8
self.fp4 = fp4
@@ -330,44 +336,24 @@ class StableDiffusion35Pipeline(DiffusionPipeline):
print("|-----------------|--------------|")
print("| {:^15} | {:^12} |".format("Module", "Latency"))
print("|-----------------|--------------|")
if "vae_encoder" in self.stages:
for stage in self.stages:
# controlnet is profiled in the denoising step
if "controlnet" in stage:
continue
stage_name = stage
if "transformer" in stage:
if f"controlnet_{self.controlnet}" in self.stages:
stage_name += '+cnet'
stage_name += ' x ' + str(denoising_steps)
print(
"| {:^15} | {:>9.2f} ms |".format(
"VAE Encoder",
cudart.cudaEventElapsedTime(self.events["vae_encoder"][0], self.events["vae_encoder"][1])[1],
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
)
)
print(
"| {:^15} | {:>9.2f} ms |".format(
"CLIP-G", cudart.cudaEventElapsedTime(self.events["clip_g"][0], self.events["clip_g"][1])[1]
)
)
print(
"| {:^15} | {:>9.2f} ms |".format(
"CLIP-L", cudart.cudaEventElapsedTime(self.events["clip_l"][0], self.events["clip_l"][1])[1]
)
)
print(
"| {:^15} | {:>9.2f} ms |".format(
"T5", cudart.cudaEventElapsedTime(self.events["t5"][0], self.events["t5"][1])[1]
)
)
print(
"| {:^15} | {:>9.2f} ms |".format(
"MMDiT" + " x " + str(denoising_steps),
cudart.cudaEventElapsedTime(self.events["transformer"][0], self.events["transformer"][1])[1],
)
)
print(
"| {:^15} | {:>9.2f} ms |".format(
"VAE Decoder",
cudart.cudaEventElapsedTime(self.events["vae"][0], self.events["vae"][1])[1],
)
)
print("|-----------------|--------------|")
print("| {:^15} | {:>9.2f} ms |".format("Pipeline", walltime_ms))
print("|-----------------|--------------|")
print("Throughput: {:.2f} image/s".format(self.batch_size * 1000.0 / walltime_ms))
print("Throughput: {:.5f} image/s".format(self.batch_size * 1000.0 / walltime_ms))
@staticmethod
def _tokenize(
@@ -711,9 +697,8 @@ class StableDiffusion35Pipeline(DiffusionPipeline):
"timestep": timestep_inp,
"encoder_hidden_states": prompt_embeds,
"pooled_projections": pooled_prompt_embeds,
"block_controlnet_hidden_states": control_block_samples,
}
if not self.fp8:
params["block_controlnet_hidden_states"] = control_block_samples
# Predict the noise residual
if self.torch_inference or self.torch_fallback[denoiser]:
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_3_pipeline.py
+30 -28
View File
@@ -126,9 +126,9 @@ class StableDiffusion3Pipeline:
# Pipeline type
self.pipeline_type = pipeline_type
self.stages = ['clip_g', 'clip_l', 't5xxl', 'mmdit', 'vae_decoder']
self.stages = ['clip_g', 'clip_l', 't5xxl', 'transformer', 'vae_decoder']
if input_image is not None:
self.stages += ['vae_encoder']
self.stages = ['vae_encoder'] + self.stages
self.config = {}
self.config['clip_hidden_states'] = True
@@ -161,7 +161,7 @@ class StableDiffusion3Pipeline:
self.generator = torch.Generator(device="cuda").manual_seed(seed)
# Create CUDA events and stream
for stage in ['clip_g', 'clip_l', 't5xxl', 'denoise', 'vae_encode', 'vae_decode']:
for stage in self.stages:
self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
self.stream = cudart.cudaStreamCreate()[1]
@@ -261,9 +261,9 @@ class StableDiffusion3Pipeline:
if 't5xxl' in self.stages:
self.models['t5xxl'] = SD3_T5XXLModel(**models_args, fp16=True, embedding_dim=get_clip_embedding_dim(self.version, self.pipeline_type))
# Load MMDiT model
if 'mmdit' in self.stages:
self.models['mmdit'] = SD3_MMDiTModel(**models_args, fp16=True, shift=self.shift)
# Load Transformer model
if 'transformer' in self.stages:
self.models['transformer'] = SD3_MMDiTModel(**models_args, fp16=True, shift=self.shift)
# Load VAE Encoder model
if 'vae_encoder' in self.stages:
@@ -320,7 +320,7 @@ class StableDiffusion3Pipeline:
# Load torch models
for model_name, obj in self.models.items():
if self.torch_fallback[model_name] or model_name == 'mmdit':
if self.torch_fallback[model_name] or model_name == 'transformer':
self.torch_models[model_name] = obj.get_model(torch_inference=self.torch_inference)
def calculateMaxDeviceMemory(self):
@@ -361,17 +361,19 @@ class StableDiffusion3Pipeline:
print('|-----------------|--------------|')
print('| {:^15} | {:^12} |'.format('Module', 'Latency'))
print('|-----------------|--------------|')
if 'vae_encoder' in self.stages:
print('| {:^15} | {:>9.2f} ms |'.format('VAE Encoder', cudart.cudaEventElapsedTime(self.events['vae_encode'][0], self.events['vae_encode'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('CLIP-G', cudart.cudaEventElapsedTime(self.events['clip_g'][0], self.events['clip_g'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('CLIP-L', cudart.cudaEventElapsedTime(self.events['clip_l'][0], self.events['clip_l'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('T5XXL', cudart.cudaEventElapsedTime(self.events['t5xxl'][0], self.events['t5xxl'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('MMDiT'+' x '+str(denoising_steps), cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('VAE Decoder', cudart.cudaEventElapsedTime(self.events['vae_decode'][0], self.events['vae_decode'][1])[1]))
for stage in self.stages:
stage_name = stage
if "transformer" in stage:
stage_name += ' x ' + str(denoising_steps)
print(
"| {:^15} | {:>9.2f} ms |".format(
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
)
)
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
print('|-----------------|--------------|')
print('Throughput: {:.2f} image/s'.format(batch_size*1000./walltime_ms))
print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms))
def save_image(self, images, pipeline, prompt, seed):
# Save image
@@ -414,7 +416,7 @@ class StableDiffusion3Pipeline:
neg_conditioning = tokenize(negative_prompt[0])
return conditioning, neg_conditioning
def denoise_latent(self, latent, conditioning, neg_conditioning, model_name='mmdit'):
def denoise_latent(self, latent, conditioning, neg_conditioning, model_name='transformer'):
def get_noise(latent):
return torch.randn(latent.size(), dtype=torch.float32, layout=latent.layout, generator=self.generator, device="cuda").to(latent.dtype)
@@ -456,7 +458,7 @@ class StableDiffusion3Pipeline:
scaled = neg_out + (pos_out - neg_out) * cond_scale
return scaled
self.profile_start('denoise', color='blue')
self.profile_start(model_name, color='blue')
latent = latent.half().cuda()
noise = get_noise(latent).cuda()
@@ -470,32 +472,32 @@ class StableDiffusion3Pipeline:
latent = sample_euler(cfg_denoiser, noise_scaled, sigmas, extra_args=extra_args)
latent = SD3LatentFormat().process_out(latent)
self.profile_stop('denoise')
self.profile_stop(model_name)
return latent
def encode_image(self):
def encode_image(self, model_name='vae_encoder'):
self.input_image = self.input_image.to(self.device)
self.profile_start('vae_encode', color='orange')
self.profile_start(model_name, color='orange')
if self.torch_inference:
with torch.autocast("cuda", dtype=torch.float16):
latent = self.torch_models['vae_encoder'](self.input_image)
latent = self.torch_models[model_name](self.input_image)
else:
latent = self.runEngine('vae_encoder', {'images': self.input_image})['latent']
latent = self.runEngine(model_name, {'images': self.input_image})['latent']
latent = SD3LatentFormat().process_in(latent)
self.profile_stop('vae_encode')
self.profile_stop(model_name)
return latent
def decode_latent(self, latent):
self.profile_start('vae_decode', color='red')
def decode_latent(self, latent, model_name='vae_decoder'):
self.profile_start(model_name, color='red')
if self.torch_inference:
with torch.autocast("cuda", dtype=torch.float16):
image = self.torch_models['vae_decoder'](latent)
image = self.torch_models[model_name](latent)
else:
image = self.runEngine('vae_decoder', {'latent': latent})['images']
image = self.runEngine(model_name, {'latent': latent})['images']
image = image.float()
self.profile_stop('vae_decode')
self.profile_stop(model_name)
return image
def infer(
demo/Diffusion/demo_diffusion/pipeline/stable_diffusion_pipeline.py
+18 -12
View File
@@ -49,10 +49,10 @@ from demo_diffusion.model import (
CLIPModel,
CLIPWithProjModel,
SDLoraLoader,
UNet2DConditionControlNetModel,
UNetModel,
UNetXLModel,
UNetXLModelControlNet,
UNet2DConditionControlNetModel,
VAEEncoderModel,
VAEModel,
get_clip_embedding_dim,
@@ -265,7 +265,7 @@ class StableDiffusionPipeline:
self.generator = torch.Generator(device="cuda").manual_seed(seed)
# Create CUDA events and stream
for stage in ['clip', 'denoise', 'vae', 'vae_encoder', 'vqgan']:
for stage in self.stages:
self.events[stage] = [cudart.cudaEventCreate()[1], cudart.cudaEventCreate()[1]]
self.stream = cudart.cudaStreamCreate()[1]
@@ -713,7 +713,7 @@ class StableDiffusionPipeline:
return images
def encode_prompt(self, prompt, negative_prompt, encoder='clip', pooled_outputs=False, output_hidden_states=False):
self.profile_start('clip', color='green')
self.profile_start(encoder, color='green')
tokenizer = self.tokenizer2 if encoder == 'clip2' else self.tokenizer
@@ -756,7 +756,7 @@ class StableDiffusionPipeline:
if output_hidden_states:
text_embeddings = torch.cat([uncond_hidden_states, text_hidden_states]).to(dtype=torch.float16) if self.do_classifier_free_guidance else text_hidden_states
self.profile_stop('clip')
self.profile_stop(encoder)
if pooled_outputs:
return text_embeddings, pooled_output
return text_embeddings
@@ -818,7 +818,7 @@ class StableDiffusionPipeline:
do_autocast = self.torch_inference != '' and self.models[denoiser].fp16
with torch.autocast('cuda', enabled=do_autocast):
self.profile_start('denoise', color='blue')
self.profile_start(denoiser, color='blue')
for step_index, timestep in enumerate(timesteps):
# Expand the latents if we are doing classifier free guidance
latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents
@@ -870,7 +870,7 @@ class StableDiffusionPipeline:
latents = 1. / self.vae_scaling_factor * latents
latents = latents.to(dtype=torch.float32)
self.profile_stop('denoise')
self.profile_stop(denoiser)
return latents
def encode_image(self, input_image):
@@ -901,15 +901,21 @@ class StableDiffusionPipeline:
print('|-----------------|--------------|')
print('| {:^15} | {:^12} |'.format('Module', 'Latency'))
print('|-----------------|--------------|')
if 'vae_encoder' in self.stages:
print('| {:^15} | {:>9.2f} ms |'.format('VAE-Enc', cudart.cudaEventElapsedTime(self.events['vae_encoder'][0], self.events['vae_encoder'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('CLIP', cudart.cudaEventElapsedTime(self.events['clip'][0], self.events['clip'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('UNet'+('+CNet' if self.pipeline_type.is_controlnet() else '')+' x '+str(denoising_steps), cudart.cudaEventElapsedTime(self.events['denoise'][0], self.events['denoise'][1])[1]))
print('| {:^15} | {:>9.2f} ms |'.format('VAE-Dec', cudart.cudaEventElapsedTime(self.events['vae'][0], self.events['vae'][1])[1]))
for stage in self.stages:
stage_name = stage
if "unet" in stage:
if self.pipeline_type.is_controlnet():
stage_name += '+cnet'
stage_name += ' x ' + str(denoising_steps)
print(
"| {:^15} | {:>9.2f} ms |".format(
stage_name, cudart.cudaEventElapsedTime(self.events[stage][0], self.events[stage][1])[1],
)
)
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
print('|-----------------|--------------|')
print('Throughput: {:.2f} image/s'.format(batch_size*1000./walltime_ms))
print('Throughput: {:.5f} image/s'.format(batch_size*1000./walltime_ms))
def save_image(self, images, pipeline, prompt, seed):
# Save image
demo/Diffusion/demo_diffusion/pipeline/stable_video_diffusion_pipeline.py
+6 -2
View File
@@ -118,7 +118,11 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
# TODO user configurable cuda_device_id
cuda_device_id = 0
vram_size = cudart.cudaGetDeviceProperties(cuda_device_id)[1].totalGlobalMem
properties = cudart.cudaGetDeviceProperties(cuda_device_id)
if properties[0] != 0:
total_device_count = cudart.cudaGetDeviceCount()[1]
raise ValueError(f"Failed to get device properties for device {cuda_device_id}, total device count: {total_device_count}")
vram_size = properties[1].totalGlobalMem
self.low_vram = vram_size < _GiB(40)
if self.low_vram:
print(f"[W] WARNING low VRAM ({vram_size/_GiB(1):.2f} GB) mode selected. Certain optimizations may be skipped.")
@@ -397,7 +401,7 @@ class StableVideoDiffusionPipeline(StableDiffusionPipeline):
print('|-----------------|--------------|')
print('| {:^15} | {:>9.2f} ms |'.format('Pipeline', walltime_ms))
print('|-----------------|--------------|')
print('Throughput: {:.2f} videos/min ({} frames)'.format(batch_size*60000./walltime_ms, num_frames))
print('Throughput: {:.5f} videos/min ({} frames)'.format(batch_size*60000./walltime_ms, num_frames))
def save_video(self, frames, pipeline, seed):
video_name_prefix = '-'.join([pipeline, 'fp16', str(seed), str(random.randint(1000,9999))])
demo/Diffusion/demo_diffusion/pipeline/type.py
+4
View File
@@ -29,6 +29,7 @@ class PIPELINE_TYPE(enum.Enum):
XL_REFINER = enum.auto()
CASCADE_PRIOR = enum.auto()
CASCADE_DECODER = enum.auto()
VIDEO2WORLD = enum.auto()
def is_txt2img(self):
return self in (self.TXT2IMG, self.CONTROLNET)
@@ -62,3 +63,6 @@ class PIPELINE_TYPE(enum.Enum):
def is_cascade(self):
return self.is_cascade_prior() or self.is_cascade_decoder()
def is_video2world(self):
return self == self.VIDEO2WORLD
demo/Diffusion/demo_diffusion/utils_modelopt.py
+2 -1
View File
@@ -23,7 +23,6 @@ from typing import Set
import modelopt.torch.quantization as mtq
import numpy as np
import onnx
import onnx_graphsurgeon as gs
import torch
import torch.nn.functional as F
@@ -38,6 +37,8 @@ from modelopt.torch.quantization.calib.max import MaxCalibrator
from PIL import Image
from torch.utils.data import Dataset, Sampler
import onnx
USE_PEFT = True
try:
from peft.tuners.lora.layer import Conv2d as PEFTLoRAConv2d
demo/Diffusion/demo_txt2image_cosmos.py
+147
View File
@@ -0,0 +1,147 @@
#
# 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.
#
import argparse
from cuda import cudart
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parse_args():
parser = argparse.ArgumentParser(description="Options for Cosmos text2image Demo", conflict_handler="resolve")
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="cosmos-predict2-2b-text2image",
choices=("cosmos-predict2-2b-text2image", "cosmos-predict2-14b-text2image"),
help="Version of Cosmos",
)
parser.add_argument(
"--height",
type=int,
default=768,
help="Height of image to generate (must be multiple of 8)",
)
parser.add_argument(
"--width",
type=int,
default=1360,
help="Width of image to generate (must be multiple of 8)",
)
parser.add_argument("--denoising-steps", type=int, default=35, help="Number of denoising steps")
parser.add_argument(
"--guidance-scale",
type=float,
default=7.0,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument(
"--num-images-per-prompt", type=int, default=1, help="The number of images to generate per prompt."
)
parser.add_argument(
"--max_sequence_length",
type=int,
default=512,
help="Maximum sequence length to use with the prompt.",
)
parser.add_argument(
"--t5-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ",
)
parser.add_argument(
"--transformer-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights.",
)
parser.add_argument(
"--bf16",
action="store_true",
default=True,
help="Use bfloat16 precision by default.",
)
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
negative_prompt = args.negative_prompt
# Process input args
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`negative_prompt` must be of type `str` list, but is {type(negative_prompt)}")
negative_prompt = negative_prompt * batch_size
kwargs_run_demo = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"height": args.height,
"width": args.width,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
"num_images_per_prompt": args.num_images_per_prompt,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Cosmos text2image demo using TensorRT")
args = parse_args()
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
demo = pipeline_module.CosmosPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.TXT2IMG)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if args.onnx_export_only:
print("[I] ONNX export completed. Exiting...")
demo.teardown()
exit(0)
# In low-vram mode we allocate the required device memory individually before each model is run.
if demo.low_vram:
demo.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
# Run inference
images = demo.run(**kwargs_run_demo)
demo.teardown()
# save images
demo.save_images(kwargs_run_demo["prompt"], images, check_integrity=True)
demo/Diffusion/demo_vid2world_cosmos.py
+170
View File
@@ -0,0 +1,170 @@
#
# 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.
#
import argparse
from cuda import cudart
from diffusers.utils import load_image, load_video
from demo_diffusion import dd_argparse
from demo_diffusion import pipeline as pipeline_module
def parse_args():
parser = argparse.ArgumentParser(description="Options for Cosmos video2world Demo", conflict_handler="resolve")
parser = dd_argparse.add_arguments(parser)
parser.add_argument(
"--version",
type=str,
default="cosmos-predict2-2b-video2world",
choices=("cosmos-predict2-2b-video2world", "cosmos-predict2-14b-video2world"),
help="Version of Cosmos",
)
parser.add_argument('--input-image', type=str, default=None, help="Path to the input image")
parser.add_argument('--input-video', type=str, default=None, help="Path to the input video")
parser.add_argument(
"--height",
type=int,
default=704,
help="Height of image to generate (must be multiple of 8)",
)
parser.add_argument(
"--width",
type=int,
default=1280,
help="Width of image to generate (must be multiple of 8)",
)
parser.add_argument("--denoising-steps", type=int, default=35, help="Number of denoising steps")
parser.add_argument(
"--guidance-scale",
type=float,
default=7.0,
help="Value of classifier-free guidance scale (must be greater than 1)",
)
parser.add_argument("--num-frames", type=int, default=93, help="The number of frames in the generated video.")
parser.add_argument("--fps", type=int, default=16, help="The frames per second of the generated video.")
parser.add_argument("--num-videos-per-prompt", type=int, default=1, help="The number of videos to generate per prompt.")
parser.add_argument(
"--max_sequence_length",
type=int,
default=512,
help="Maximum sequence length to use with the prompt.",
)
parser.add_argument(
"--t5-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the T5 model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights. ",
)
parser.add_argument(
"--transformer-ws-percentage",
type=int,
default=None,
help="Set runtime weight streaming budget as the percentage of the size of streamable weights for the transformer model. This argument only takes effect when --ws is set. 0 streams the most weights and 100 or None streams no weights.",
)
parser.add_argument(
"--bf16",
action="store_true",
default=True,
help="Use bfloat16 precision by default.",
)
return parser.parse_args()
def process_demo_args(args):
batch_size = args.batch_size
prompt = args.prompt
negative_prompt = args.negative_prompt
# Process input args
if not isinstance(prompt, list):
raise ValueError(f"`prompt` must be of type `str` list, but is {type(prompt)}")
prompt = prompt * batch_size
if not isinstance(negative_prompt, list):
raise ValueError(f"`negative_prompt` must be of type `str` list, but is {type(negative_prompt)}")
negative_prompt = negative_prompt * batch_size
# process input image and input video
if args.input_image and args.input_video:
raise ValueError("Only one of --input-image or --input-video can be provided")
if args.input_image:
args.input_image = load_image(args.input_image)
elif args.input_video:
args.input_video = load_video(args.input_video)
else:
# load default image
args.input_image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/yellow-scrubber.png")
kwargs_run_demo = {
"prompt": prompt,
"negative_prompt": negative_prompt,
"height": args.height,
"width": args.width,
"batch_count": args.batch_count,
"num_warmup_runs": args.num_warmup_runs,
"use_cuda_graph": args.use_cuda_graph,
"num_frames": args.num_frames,
"fps": args.fps,
"input_image": args.input_image,
"input_video": args.input_video,
"num_videos_per_prompt": args.num_videos_per_prompt,
}
return kwargs_run_demo
if __name__ == "__main__":
print("[I] Initializing Cosmos video2world demo using TensorRT")
args = parse_args()
# Enforce torch-inference is enabled
if not args.torch_inference:
print("[W] The video2world demo only supports the PyTorch backend. Enabling torch-inference with 'eager' mode.")
args.torch_inference = "eager"
_, kwargs_load_engine, _ = dd_argparse.process_pipeline_args(args)
kwargs_run_demo = process_demo_args(args)
# Initialize demo
demo = pipeline_module.CosmosPipeline.FromArgs(args, pipeline_type=pipeline_module.PIPELINE_TYPE.VIDEO2WORLD)
# Load TensorRT engines and pytorch modules
demo.load_engines(
framework_model_dir=args.framework_model_dir,
**kwargs_load_engine,
)
if args.onnx_export_only:
print("[I] ONNX export completed. Exiting...")
demo.teardown()
exit(0)
# In low-vram mode we allocate the required device memory individually before each model is run.
if demo.low_vram:
demo.device_memory_sizes = demo.get_device_memory_sizes()
else:
_, shared_device_memory = cudart.cudaMalloc(demo.calculate_max_device_memory())
demo.activate_engines(shared_device_memory)
demo.load_resources(args.height, args.width, args.batch_size, args.seed)
# Run inference
videos = demo.run(**kwargs_run_demo)
demo.teardown()
# save video
demo.save_video(kwargs_run_demo["prompt"], videos, check_integrity=True)
demo/Diffusion/docs/support_matrix.md
+29 -27
View File
@@ -6,32 +6,34 @@ This demo supports Diffusion models that are popular in the Generative AI commun
## Pipeline Support Matrix
| Pipeline | Version | Task | Supported Precisions | Additional features | Hub | Restrictions |
|------------------------|---------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------|------------------------|--------------------------|---------------------------------|------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| Stable Diffusion | 1.4 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) |
| Stable Diffusion | [1.5](../README.md#generate-an-image-guided-by-a-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [KiwiXR/stable-diffusion-v1-5](https://huggingface.co/KiwiXR/stable-diffusion-v1-5) |
| Stable Diffusion | 1.4, [1.5](../README.md#generate-an-inpainted-image-guided-by-an-image-mask-and-a-text-prompt) | <ul><li>Inpainting</li></ul> | FP16 | N/A | [benjamin-paine/stable-diffusion-v1-5-inpainting](https://huggingface.co/benjamin-paine/stable-diffusion-v1-5-inpainting) |
| Stable Diffusion | dreamshaper-7 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [Lykon/dreamshaper-7](https://huggingface.co/Lykon/dreamshaper-7) |
| Stable Diffusion | 2.0-base | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) |
| Stable Diffusion | 2.0 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) |
| Stable Diffusion | 2.0, 2.0-base | <ul><li>Inpainting</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-2-inpainting](https://huggingface.co/stabilityai/stable-diffusion-2-inpainting) |
| Stable Diffusion | 2.1-base | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) |
| Stable Diffusion | 2.1 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) |
| Stable Diffusion | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | LoRA (FP16, BF16, FP8) | [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) |
| Stable Diffusion | [XL 1.0-refiner](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-xl-refiner-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0) |
| Stable Diffusion | [XL-Turbo](../README.md#faster-text-to-image-using-sdxl-turbo) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/sdxl-turbo](https://huggingface.co/stabilityai/sdxl-turbo) |
| Stable Diffusion | [3](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) |
| Stable Diffusion | [3.5-medium](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16, BF16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) |
| Stable Diffusion | [3.5-large](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16, BF16, FP8 | N/A | [stabilityai/stable-diffusion-3-large](https://huggingface.co/stabilityai/stable-diffusion-3-large) |
| ControlNet | [1.5](../README.md#generate-an-image-with-controlnet-guided-by-images-and-text-prompts) | <ul><li>Image-to-image</li></ul> | FP16 | N/A | <ul><li>[lllyasviel/sd-controlnet-canny](https://huggingface.co/lllyasviel/sd-controlnet-canny)</li><li>[lllyasviel/sd-controlnet-depth](https://huggingface.co/lllyasviel/sd-controlnet-depth)</li><li>[lllyasviel/sd-controlnet-hed](https://huggingface.co/lllyasviel/sd-controlnet-hed)</li><li>[lllyasviel/sd-controlnet-mlsd](https://huggingface.co/lllyasviel/sd-controlnet-mlsd)</li><li>[lllyasviel/sd-controlnet-normal](https://huggingface.co/lllyasviel/sd-controlnet-normal)</li><li>[lllyasviel/sd-controlnet_openpose](https://huggingface.co/lllyasviel/sd-controlnet-openpose)</li><li>[lllyasviel/sd-controlnet_scribble](https://huggingface.co/lllyasviel/sd-controlnet-scribble)</li><li>[lllyasviel/sd-controlnet_seg](https://huggingface.co/lllyasviel/sd-controlnet-seg)</li></ul> |
| ControlNet | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Image-to-image</li></ul> | FP16, FP8 | N/A | [stabilityai/controlnet-canny-sdxl-1.0](https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0) |
| ControlNet | [3.5-large](../README.md#generate-an-image-with-stable-diffusion-v35-large-with-controlnet-guided-by-an-image-and-a-text-prompt) | <ul><li>Image-to-image</li></ul> | FP16, BF16 | N/A | <ul><li>[stabilityai/stable-diffusion-3.5-large-controlnet-canny](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-canny)</li><li>[stabilityai/stable-diffusion-3.5-large-controlnet-depth](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-depth)</li><li>[stabilityai/stable-diffusion-3.5-large-controlnet-blur](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-blur)</li></ul> |
| Stable Video Diffusion | [XT-1.1](../README.md#generate-a-video-guided-by-an-initial-image-using-stable-video-diffusion) | <ul><li>Text-to-video</li></ul> | FP16, FP8 | N/A | [stabilityai/stable-video-diffusion-img2vid-xt-1-1](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt-1-1) |
| Stable Cascade | [N/A](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-cascade) | <ul><li>Text-to-image</li></ul> | BF16 | N/A | <ul><li>[stabilityai/stable-cascade-prior](https://huggingface.co/stabilityai/stable-cascade-prior)</li><li>[stabilityai/stable-cascade](https://huggingface.co/stabilityai/stable-cascade)</li></ul> |
| Flux | [1-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) |
| Flux | [1-Schnell](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) |
| Flux | [1-Canny-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Canny-dev](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev) |
| Flux | [1-Depth-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev) |
| Flux | [1-Kontext-Dev](../README.md#5-edit-an-image-using-flux-kontext) | <ul><li>Image-to-image</li></ul> | BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) |
| Pipeline | Version | Task | Supported Precisions | Additional features | Hub | Restrictions |
|------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------|-------------------------------------------------------------|-----------------------------------------------|--------------------------|--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
| Stable Diffusion | 1.4 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [CompVis/stable-diffusion-v1-4](https://huggingface.co/CompVis/stable-diffusion-v1-4) |
| Stable Diffusion | [1.5](../README.md#generate-an-image-guided-by-a-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [KiwiXR/stable-diffusion-v1-5](https://huggingface.co/KiwiXR/stable-diffusion-v1-5) |
| Stable Diffusion | 1.4, [1.5](../README.md#generate-an-inpainted-image-guided-by-an-image-mask-and-a-text-prompt) | <ul><li>Inpainting</li></ul> | FP16 | N/A | [benjamin-paine/stable-diffusion-v1-5-inpainting](https://huggingface.co/benjamin-paine/stable-diffusion-v1-5-inpainting) |
| Stable Diffusion | dreamshaper-7 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [Lykon/dreamshaper-7](https://huggingface.co/Lykon/dreamshaper-7) |
| Stable Diffusion | 2.0-base | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-2-base](https://huggingface.co/stabilityai/stable-diffusion-2-base) |
| Stable Diffusion | 2.0 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-2](https://huggingface.co/stabilityai/stable-diffusion-2) |
| Stable Diffusion | 2.0, 2.0-base | <ul><li>Inpainting</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-2-inpainting](https://huggingface.co/stabilityai/stable-diffusion-2-inpainting) |
| Stable Diffusion | 2.1-base | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1-base](https://huggingface.co/stabilityai/stable-diffusion-2-1-base) |
| Stable Diffusion | 2.1 | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | N/A | [stabilityai/stable-diffusion-2-1](https://huggingface.co/stabilityai/stable-diffusion-2-1) |
| Stable Diffusion | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, FP8, INT8 * | LoRA (FP16, BF16, FP8) | [stabilityai/stable-diffusion-xl-base-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-base-1.0) |
| Stable Diffusion | [XL 1.0-refiner](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-xl-refiner-1.0](https://huggingface.co/stabilityai/stable-diffusion-xl-refiner-1.0) |
| Stable Diffusion | [XL-Turbo](../README.md#faster-text-to-image-using-sdxl-turbo) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16 | N/A | [stabilityai/sdxl-turbo](https://huggingface.co/stabilityai/sdxl-turbo) |
| Stable Diffusion | [3](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) |
| Stable Diffusion | [3.5-medium](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16, BF16 | N/A | [stabilityai/stable-diffusion-3-medium](https://huggingface.co/stabilityai/stable-diffusion-3-medium) |
| Stable Diffusion | [3.5-large](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-diffusion-3) | <ul><li>Text-to-image</li></ul> | FP16, BF16, FP8 | N/A | [stabilityai/stable-diffusion-3-large](https://huggingface.co/stabilityai/stable-diffusion-3-large) |
| ControlNet | [1.5](../README.md#generate-an-image-with-controlnet-guided-by-images-and-text-prompts) | <ul><li>Image-to-image</li></ul> | FP16 | N/A | <ul><li>[lllyasviel/sd-controlnet-canny](https://huggingface.co/lllyasviel/sd-controlnet-canny)</li><li>[lllyasviel/sd-controlnet-depth](https://huggingface.co/lllyasviel/sd-controlnet-depth)</li><li>[lllyasviel/sd-controlnet-hed](https://huggingface.co/lllyasviel/sd-controlnet-hed)</li><li>[lllyasviel/sd-controlnet-mlsd](https://huggingface.co/lllyasviel/sd-controlnet-mlsd)</li><li>[lllyasviel/sd-controlnet-normal](https://huggingface.co/lllyasviel/sd-controlnet-normal)</li><li>[lllyasviel/sd-controlnet_openpose](https://huggingface.co/lllyasviel/sd-controlnet-openpose)</li><li>[lllyasviel/sd-controlnet_scribble](https://huggingface.co/lllyasviel/sd-controlnet-scribble)</li><li>[lllyasviel/sd-controlnet_seg](https://huggingface.co/lllyasviel/sd-controlnet-seg)</li></ul> |
| ControlNet | [XL 1.0-base](../README.md#generate-an-image-with-stable-diffusion-xl-guided-by-a-single-text-prompt) | <ul><li>Image-to-image</li></ul> | FP16, FP8 | N/A | [stabilityai/controlnet-canny-sdxl-1.0](https://huggingface.co/diffusers/controlnet-canny-sdxl-1.0) |
| ControlNet | [3.5-large](../README.md#generate-an-image-with-stable-diffusion-v35-large-with-controlnet-guided-by-an-image-and-a-text-prompt) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8 (canny and depth only) | N/A | <ul><li>[stabilityai/stable-diffusion-3.5-large-controlnet-canny](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-canny)</li><li>[stabilityai/stable-diffusion-3.5-large-controlnet-depth](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-depth)</li><li>[stabilityai/stable-diffusion-3.5-large-controlnet-blur](https://huggingface.co/stabilityai/stable-diffusion-3.5-large-controlnet-blur)</li></ul> |
| Stable Video Diffusion | [XT-1.1](../README.md#generate-a-video-guided-by-an-initial-image-using-stable-video-diffusion) | <ul><li>Text-to-video</li></ul> | FP16, FP8 | N/A | [stabilityai/stable-video-diffusion-img2vid-xt-1-1](https://huggingface.co/stabilityai/stable-video-diffusion-img2vid-xt-1-1) |
| Stable Cascade | [N/A](../README.md#generate-an-image-guided-by-a-text-prompt-using-stable-cascade) | <ul><li>Text-to-image</li></ul> | BF16 | N/A | <ul><li>[stabilityai/stable-cascade-prior](https://huggingface.co/stabilityai/stable-cascade-prior)</li><li>[stabilityai/stable-cascade](https://huggingface.co/stabilityai/stable-cascade)</li></ul> |
| Flux | [1-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) |
| Flux | [1-Schnell](../README.md#generate-an-image-guided-by-a-text-prompt-using-flux) | <ul><li>Text-to-image</li><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 * | LoRA (FP16, BF16, FP8) | [black-forest-labs/FLUX.1-schnell](https://huggingface.co/black-forest-labs/FLUX.1-schnell) |
| Flux | [1-Canny-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Canny-dev](https://huggingface.co/black-forest-labs/FLUX.1-Canny-dev) |
| Flux | [1-Depth-Dev](../README.md#generate-an-image-guided-by-a-text-prompt-and-a-control-image-using-flux-controlnet) | <ul><li>Image-to-image</li></ul> | FP16, BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Depth-dev](https://huggingface.co/black-forest-labs/FLUX.1-Depth-dev) |
| Flux | [1-Kontext-Dev](../README.md#5-edit-an-image-using-flux-kontext) | <ul><li>Image-to-image</li></ul> | BF16, FP8, FP4 | N/A | [black-forest-labs/FLUX.1-Kontext-dev](https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev) |
| Cosmos | [cosmos-predict2-2b-text2image](../README.md#1-generate-an-image-from-a-text-prompt-1), cosmos-predict2-14b-text2image | <ul><li>Text-to-image</li></ul> | BF16 | N/A | <ul><li>[nvidia/Cosmos-Predict2-2B-Text2Image](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Text2Image)</li><li>[nvidia/Cosmos-Predict2-14B-Text2Image](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Text2Image) |
| Cosmos | [cosmos-predict2-2b-video2world](../README.md#2-generate-a-video-guided-by-an-initial-video-conditioning-and-a-text-prompt), cosmos-predict2-14b-video2world | <ul><li>Video-to-World</li></ul> | BF16 | N/A | <ul><li>[nvidia/Cosmos-Predict2-2B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-2B-Video2World)</li><li>[nvidia/Cosmos-Predict2-14B-Video2World](https://huggingface.co/nvidia/Cosmos-Predict2-14B-Video2World) |
*Note: Only the text2image pipelines support FP4/FP8/INT8 quantization. The image2image pipelines don't support quantization.
demo/Diffusion/requirements.txt
+2 -1
View File
@@ -9,7 +9,6 @@ ftfy
matplotlib
nvtx
onnx==1.18.0
onnxruntime==1.19.2
onnxscript==0.3.2
opencv-python-headless==4.8.0.74
scipy
@@ -21,3 +20,5 @@ peft==0.17.0
polygraphy==0.49.22
sentencepiece
numpy==1.26.4
imageio-ffmpeg
demo/Diffusion/requirements_legacy.txt
+24
View File
@@ -0,0 +1,24 @@
apex==0.9.10dev
accelerate==1.2.1
colored==2.3.1
controlnet_aux==0.0.6
cuda-python==13.0.2
diffusers==0.35.0
ftfy==6.3.1
matplotlib==3.10.7
nvtx==0.2.13
onnx==1.17.0
onnxscript==0.3.2
onnxconverter-common==1.14.0
opencv-python-headless==4.8.0.74
scipy==1.15.3
transformers==4.52.4
--extra-index-url https://pypi.nvidia.com
nvidia-modelopt[torch,onnx]==0.29.0
onnx-graphsurgeon==0.5.2
peft==0.17.0
polygraphy==0.49.22
sentencepiece==0.2.1
numpy==1.26.4
imageio-ffmpeg==0.6.0
demo/Diffusion/setup.sh
+6 -1
View File
@@ -1,10 +1,15 @@
#!/bin/bash
# Minimal and backward-compatible: default to modern requirements, allow override via env.
REQ_FILE="${REQUIREMENTS_FILE:-requirements.txt}"
echo "[I] Using requirements file: ${REQ_FILE} (override with REQUIREMENTS_FILE=<file>)"
# Upgrade pip and install TensorRT
python3 -m pip install --upgrade pip
pip3 install --pre tensorrt-cu12
# Install the required packages
PIP_CONSTRAINT= pip3 install -r requirements.txt
PIP_CONSTRAINT= pip3 install -r "${REQ_FILE}"
# Install libgl1
# Check if apt-get is available
docker/rockylinux8.Dockerfile
+9 -9
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-rockylinux8
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.13.3.9
ENV TRT_VERSION 10.14.1.48
SHELL ["/bin/bash", "-c"]
# Setup user account
@@ -51,15 +51,15 @@ RUN dnf install -y python38 python38-devel &&\
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp38-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp38-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp38-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp38-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
docker/rockylinux9.Dockerfile
+9 -9
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-rockylinux9
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.13.3.9
ENV TRT_VERSION 10.14.1.48
SHELL ["/bin/bash", "-c"]
# Setup user account
@@ -56,15 +56,15 @@ RUN dnf -y install \
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp39-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp39-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp39-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp39-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
docker/ubuntu-22.04.Dockerfile
+9 -9
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu22.04
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.13.3.9
ENV TRT_VERSION 10.14.1.48
SHELL ["/bin/bash", "-c"]
# Setup user account
@@ -70,15 +70,15 @@ RUN apt-get install -y --no-install-recommends \
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp310-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp310-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp310-none-linux_x86_64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib64 \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp310-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
docker/ubuntu-24.04-aarch64.Dockerfile
+11 -10
View File
@@ -20,7 +20,7 @@ ARG CUDA_VERSION=13.0.0
# Multi-arch container support available in non-cudnn containers.
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04
ENV TRT_VERSION 10.13.3.9
ENV TRT_VERSION 10.14.1.48
SHELL ["/bin/bash", "-c"]
# Setup user account and edit default account
@@ -60,7 +60,8 @@ RUN apt-get update && apt-get install -y --no-install-recommends \
lintian \
fakeroot \
dh-make \
build-essential
build-essential \
libffi-dev
# Install python3
RUN apt-get install -y --no-install-recommends \
@@ -79,15 +80,15 @@ ENV PATH="/opt/venv/bin:$PATH"
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp312-none-linux_aarch64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_aarch64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-10.13.3.9/python/tensorrt-10.13.3.9-cp312-none-linux_aarch64.whl ;\
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/aarch64-linux-gnu/ \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_aarch64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
docker/ubuntu-24.04.Dockerfile
+121
View File
@@ -0,0 +1,121 @@
#
# 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.
#
ARG CUDA_VERSION=13.0.0
# Multi-arch container support available in non-cudnn containers.
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu24.04
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION=10.14.1.48
SHELL ["/bin/bash", "-c"]
# Setup user account and edit default account
RUN if id "ubuntu" &>/dev/null; then \
usermod -u 1234 ubuntu && \
groupmod -g 1234 ubuntu; \
fi
ARG uid=1000
ARG gid=1000
RUN groupadd -r -f -g ${gid} trtuser && useradd -o -r -l -u ${uid} -g ${gid} -ms /bin/bash trtuser
RUN usermod -aG sudo trtuser
RUN echo 'trtuser:nvidia' | chpasswd
RUN mkdir -p /workspace && chown trtuser /workspace
ENV DEBIAN_FRONTEND=noninteractive
# Update CUDA signing key
RUN apt-key adv --fetch-keys https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2404/x86_64/3bf863cc.pub
# Install requried libraries
RUN apt-get update && apt-get install -y software-properties-common
RUN add-apt-repository ppa:ubuntu-toolchain-r/test
RUN apt-get update && apt-get install -y --no-install-recommends \
libcurl4-openssl-dev \
wget \
git \
pkg-config \
sudo \
ssh \
libssl-dev \
pbzip2 \
pv \
bzip2 \
unzip \
devscripts \
lintian \
fakeroot \
dh-make \
build-essential
# Install python3
RUN apt-get install -y --no-install-recommends \
python3 \
python3-pip \
python3-dev \
python3-wheel \
python3-venv &&\
cd /usr/local/bin &&\
ln -s /usr/bin/python3 python &&\
ln -s /usr/bin/pip3 pip;
# Create python3 virtualenv
RUN python3 -m venv /opt/venv
ENV PATH="/opt/venv/bin:$PATH"
# Install TensorRT
RUN if [ "${CUDA_VERSION:0:2}" = "13" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-13.0.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/x86_64-linux-gnu/ \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_x86_64.whl ;\
elif [ "${CUDA_VERSION:0:2}" = "12" ]; then \
wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& tar -xf TensorRT-10.14.1.48.Linux.x86_64-gnu.cuda-12.9.tar.gz \
&& cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/x86_64-linux-gnu/ \
&& pip install TensorRT-10.14.1.48/python/tensorrt-10.14.1.48-cp312-none-linux_x86_64.whl ;\
else \
echo "Invalid CUDA_VERSION"; \
exit 1; \
fi
# Install PyPI packages
RUN pip3 install --upgrade pip
RUN pip3 install setuptools>=41.0.0
RUN pip3 install jupyter jupyterlab
# Workaround to remove numpy installed with tensorflow
RUN pip3 install --upgrade numpy
# Install Cmake
RUN cd /tmp && \
wget https://github.com/Kitware/CMake/releases/download/v3.27.9/cmake-3.27.9-Linux-x86_64.sh && \
chmod +x cmake-3.27.9-Linux-x86_64.sh && \
./cmake-3.27.9-Linux-x86_64.sh --prefix=/usr/local --exclude-subdir --skip-license && \
rm ./cmake-3.27.9-Linux-x86_64.sh
# Download NGC client
RUN cd /usr/local/bin && wget https://ngc.nvidia.com/downloads/ngccli_cat_linux.zip && unzip ngccli_cat_linux.zip && chmod u+x ngc-cli/ngc && rm ngccli_cat_linux.zip ngc-cli.md5 && echo "no-apikey\nascii\n" | ngc-cli/ngc config set
# Set environment and working directory
ENV TRT_LIBPATH=/usr/lib/x86_64-linux-gnu/
ENV TRT_OSSPATH=/workspace/TensorRT
ENV PATH="/workspace/TensorRT/build/out:${PATH}:/usr/local/bin/ngc-cli"
ENV LD_LIBRARY_PATH="${LD_LIBRARY_PATH}:${TRT_OSSPATH}/build/out:${TRT_LIBPATH}"
WORKDIR /workspace
USER trtuser
RUN ["/bin/bash"]
docker/ubuntu-cross-aarch64.Dockerfile
+4 -4
View File
@@ -21,7 +21,7 @@ ARG OS_VERSION=24.04
FROM nvidia/cuda:${CUDA_VERSION}-devel-ubuntu${OS_VERSION}
LABEL maintainer="NVIDIA CORPORATION"
ENV TRT_VERSION 10.13.3.9
ENV TRT_VERSION 10.14.1.48
ENV DEBIAN_FRONTEND=noninteractive
# Setup user account and edit default account
@@ -84,9 +84,9 @@ RUN wget https://developer.download.nvidia.com/compute/cuda/13.0.0/local_install
# Unpack libnvinfer.
RUN wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.13.3/tars/TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz && \
tar -xf TensorRT-10.13.3.9.Linux.aarch64-gnu.cuda-13.0.tar.gz && \
cp -a TensorRT-10.13.3.9/lib/*.so* /usr/lib/aarch64-linux-gnu
RUN wget https://developer.nvidia.com/downloads/compute/machine-learning/tensorrt/10.14.1/tars/TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz && \
tar -xf TensorRT-10.14.1.48.Linux.aarch64-gnu.cuda-13.0.tar.gz && \
cp -a TensorRT-10.14.1.48/lib/*.so* /usr/lib/aarch64-linux-gnu
# Link required library
RUN cd /usr/aarch64-linux-gnu/lib && ln -sf librt.so.1 librt.so
include/NvInfer.h
+627 -23
View File
@@ -19,7 +19,7 @@
#define NV_INFER_H
#include "NvInferLegacyDims.h"
#include "NvInferRuntime.h"
#include "NvInferRuntime.h" // IWYU pragma: export
//!
//! \mainpage
@@ -106,7 +106,9 @@ enum class LayerType : int32_t
kSQUEEZE = 47, //!< Squeeze Layer.
kUNSQUEEZE = 48, //!< Unsqueeze Layer.
kCUMULATIVE = 49, //!< Cumulative layer.
kDYNAMIC_QUANTIZE = 50, //!< Dynamic Quantize layer.
kDYNAMIC_QUANTIZE = 50, //!< Dynamic Quantize layer.
kATTENTION_INPUT = 51, //!< Attention Input.
kATTENTION_OUTPUT = 52, //!< Attention Output.
};
//!
@@ -117,7 +119,7 @@ enum class LayerType : int32_t
template <>
constexpr inline int32_t EnumMax<LayerType>() noexcept
{
return 51;
return 53;
}
//!
@@ -151,6 +153,7 @@ enum class ActivationType : int32_t
kGELU_TANH = 13 //!< GELU tanh activation: 0.5 * x * (1 + tanh(sqrt(2/pi) * (0.044715F * pow(x, 3) + x)))
};
namespace impl
{
//!
@@ -3645,6 +3648,33 @@ public:
//!
using ILayer::setInput;
//!
//! \brief Set the indices type for the layer.
//!
//! \param type The DataType of the indices tensor.
//!
//! \return true if set successfully, false otherwise.
//!
//! Set the indices (the second output) type of the TopK layer. Valid values are DataType::kINT32 and
//! DataType::kINT64, otherwise an error occurs and the type is not updated.
//!
bool setIndicesType(DataType type) noexcept
{
return mImpl->setIndicesType(type);
}
//!
//! \brief Return the TopK layer indices type.
//!
//! \return indices type set during layer creation or by setIndicesType().
//! The return value is the indices type of the TopK layer.
//! The default value is DataType::kINT32.
//!
DataType getIndicesType() const noexcept
{
return mImpl->getIndicesType();
}
protected:
apiv::VTopKLayer* mImpl;
virtual ~ITopKLayer() noexcept = default;
@@ -3757,7 +3787,7 @@ protected:
//!
//! The input may have type kFLOAT, kHALF, kINT32, or kBOOL.
//!
//! The output is a matrix of type kINT32.
//! The output is a matrix of type kINT32 or kINT64.
//! For an input with dimensions [L1, L2, ..., Lm], the output has dimensions [m,n],
//! where n is the number of non-zero elements. I.e., each column denotes a m-D position.
//!
@@ -3771,6 +3801,34 @@ protected:
//!
class INonZeroLayer : public ILayer
{
public:
//!
//! \brief Set the indices type for the layer.
//!
//! \param type The DataType of the indices tensor.
//!
//! \return true if set successfully, false otherwise.
//!
//! Set the indices (the first output) type of the NonZero layer. Valid values are DataType::kINT32 and
//! DataType::kINT64, otherwise an error occurs and the type is not updated.
//!
bool setIndicesType(DataType type) noexcept
{
return mImpl->setIndicesType(type);
}
//!
//! \brief Return the NonZero layer indices type.
//!
//! \return indices type set during layer creation or by setIndicesType().
//! The return value is the indices type of the NonZero layer.
//! The default value is DataType::kINT32.
//!
DataType getIndicesType() const noexcept
{
return mImpl->getIndicesType();
}
protected:
virtual ~INonZeroLayer() noexcept = default;
apiv::VNonZeroLayer* mImpl;
@@ -6145,10 +6203,12 @@ constexpr inline int32_t EnumMax<BoundingBoxFormat>() noexcept
//! intersection-over-union (IoU) with previously selected boxes is less than or equal to a given threshold.
//! This layer implements NMS per batch item and per class.
//!
//! Per batch item, boxes are initially sorted by their scores without regard to class. Only boxes up to a maximum of the TopK limit are considered for selection (per batch).
//! During selection, only overlapping boxes of the same class are compared, so that overlapping boxes of different classes do not suppress each other.
//! Per batch item, boxes are initially sorted by their scores without regard to class. Only boxes up to a maximum of
//! the TopK limit are considered for selection (per batch). During selection, only overlapping boxes of the same class
//! are compared, so that overlapping boxes of different classes do not suppress each other.
//!
//! For each batch item, the ordering of candidate bounding boxes with the same score is unspecified, but the ordering will be consistent across different runs for the same inputs.
//! For each batch item, the ordering of candidate bounding boxes with the same score is unspecified, but the ordering
//! will be consistent across different runs for the same inputs.
//!
//! The layer has the following inputs, in order of input index:
//!
@@ -6161,19 +6221,21 @@ constexpr inline int32_t EnumMax<BoundingBoxFormat>() noexcept
//! It is a scalar (0D tensor) of type kINT32.
//! * IoUThreshold is the maximum IoU for selected boxes. It is a scalar (0D tensor) of type kFLOAT in the range
//! [0.0f, 1.0f]. It is an optional input with default 0.0f.
//! * ScoreThreshold is the value that a box score must exceed in order to be selected. It is a scalar (0D tensor) of type kFLOAT. It is an optional
//! * ScoreThreshold is the value that a box score must exceed in order to be selected. It is a scalar (0D tensor) of
//! type kFLOAT. It is an optional
//! input with default 0.0f.
//!
//! The layer has the following outputs, in order of output index:
//!
//! * SelectedIndices contains the indices of the selected boxes. It is a linear tensor of type kINT32. It has shape
//! * SelectedIndices contains the indices of the selected boxes. It is a linear tensor of type kINT32 or kINT64. It has
//! shape
//! [NumOutputBoxes, 3]. Each row contains a (batchIndex, classIndex, boxIndex) tuple.
//! The output boxes are sorted in order of increasing batchIndex and then in order of decreasing score within each batchIndex.
//! For each batchIndex, the ordering of output boxes with the same score is unspecified.
//! If MaxOutputBoxesPerClass is a constant input, the maximum number of output boxes is
//! batchSize * numClasses * min(numInputBoundingBoxes, MaxOutputBoxesPerClass).
//! Otherwise, the maximum number of output boxes is batchSize * numClasses * numInputBoundingBoxes.
//! The maximum number of output boxes is used to determine the upper-bound on allocated memory for this output tensor.
//! The output boxes are sorted in order of increasing batchIndex and then in order of decreasing score within each
//! batchIndex. For each batchIndex, the ordering of output boxes with the same score is unspecified. If
//! MaxOutputBoxesPerClass is a constant input, the maximum number of output boxes is batchSize * numClasses *
//! min(numInputBoundingBoxes, MaxOutputBoxesPerClass). Otherwise, the maximum number of output boxes is batchSize *
//! numClasses * numInputBoundingBoxes. The maximum number of output boxes is used to determine the upper-bound on
//! allocated memory for this output tensor.
//! * NumOutputBoxes is the number of output boxes in SelectedIndices. It is a scalar (0D tensor) of type kINT32.
//!
//! \warning There is a hardware-dependent limit K such that only the K highest scoring boxes in each batch item
@@ -6254,6 +6316,33 @@ public:
//!
using ILayer::setInput;
//!
//! \brief Set the indices type for the layer.
//!
//! \param type The DataType of the indices tensor.
//!
//! \return true if set successfully, false otherwise.
//!
//! Set the indices (the first output) type of the NMS layer. Valid values are DataType::kINT32 and
//! DataType::kINT64, otherwise an error occurs and the type is not updated.
//!
bool setIndicesType(DataType type) noexcept
{
return mImpl->setIndicesType(type);
}
//!
//! \brief Return the NMS layer indices type.
//!
//! \return indices type set during layer creation or by setIndicesType().
//! The return value is the indices type of the NMS layer.
//! The default value is DataType::kINT32.
//!
DataType getIndicesType() const noexcept
{
return mImpl->getIndicesType();
}
protected:
apiv::VNMSLayer* mImpl;
virtual ~INMSLayer() noexcept = default;
@@ -6661,6 +6750,409 @@ protected:
virtual ~ICumulativeLayer() noexcept = default;
};
//!
//! \enum AttentionNormalizationOp
//!
//! \brief Enumerates the operations that may be performed by the normalization in the attention subgraph.
//!
enum class AttentionNormalizationOp : int32_t
{
kNONE
= 0, //!< Apply no normalization on the attention scores. Must be used with decomposable=True on pre-Blackwell GPUs
kSOFTMAX = 1, //!< Apply softmax normalization on the attention scores on the `s_kv` dimension.
};
namespace impl
{
//!
//! Maximum number of elements in AttentionNormalizationOp enum.
//!
//! \see AttentionNormalizationOp
//!
template <>
struct EnumMaxImpl<AttentionNormalizationOp>
{
static constexpr int32_t kVALUE = 2;
};
} // namespace impl
//!
//! \class IAttentionBoundaryLayer
//!
//! \brief This is a base class for Attention boundary layers.
//!
//! Boundary layers are used to demarcate the boundaries of IAttention.
//! Typically client code does not deal directly with the boundary layers.
//! However, they are indirectly visible via method `INetworkDefinition::getLayer(int32_t index)`.
//!
class IAttentionBoundaryLayer : public ILayer
{
public:
//!
//! \brief Get a pointer to the IAttention associated with this boundary layer.
//!
IAttention* getAttention() const noexcept
{
return mBoundary->getAttention();
}
protected:
virtual ~IAttentionBoundaryLayer() noexcept = default;
apiv::VAttentionBoundaryLayer* mBoundary;
};
//!
//! \class IAttentionInputLayer
//!
//! \brief This layer represents an input to an attention subgraph.
//!
//! This layer is automatically created when an `IAttention` is created. Clients typically do not
//! deal with the layer directly, but instead specify its input via `addAttention` or `IAttention::setInput`.
//!
//! An IAttentionInputLayer has three to four inputs and one output.
//!
class IAttentionInputLayer : public IAttentionBoundaryLayer
{
public:
//!
//! \brief Append or replace an input of this layer with a specific tensor
//!
//! \param index the index of the input to modify.
//! \param tensor the new input tensor
//!
//! The indices are as follows:
//!
//! Input 0 is the input query tensor.
//! Input 1 is the input key tensor.
//! Input 2 is the input value tensor.
//! Input 3 is the optional mask tensor. setMask should be used instead of setInput
//! Input 4 is the optional normalizationQuantizeScale tensor. setNormalizationQuantizeScale should be used instead
//! of setInput
//!
using ILayer::setInput;
protected:
virtual ~IAttentionInputLayer() noexcept = default;
apiv::VAttentionInputLayer* mImpl;
};
//!
//! \class IAttentionOutputLayer
//!
//! \brief This layer represents an output of an IAttention.
//!
//! This layer is automatically created when an `IAttention` is created. Clients typically do not
//! deal with the layer directly, but instead getting its output via `IAttention::getOutput`.
//!
//! An IAttentionOutputLayer has one input and one output.
//!
class IAttentionOutputLayer : public IAttentionBoundaryLayer
{
public:
protected:
virtual ~IAttentionOutputLayer() noexcept = default;
apiv::VAttentionOutputLayer* mImpl;
};
//!
//! \class IAttention
//!
//! \brief Helper for constructing an attention that consumes query, key and value tensors.
//!
//! An attention subgraph implicitly includes three main components, two MatrixMultiply layers
//! known as BMM1 and BMM2, and one normalization operation which defaults to be a Softmax.
//! By default, IAttention is not decomposable and TensorRT will try to use a single fused kernel, which may be more
//! efficient than if the subgraph is expressed without IAttention. Setting the IAttention to decomposable=True can
//! allow IAttention to be decomposed to use multiple kernels if no fused kernel support found.
//!
//! Query Key Value Mask (optional) NormalizationQuantizeScale (optional)
//! | | | | |
//! | Transpose | | |
//! | | | | |
//! ----BMM1---- | | |
//! | | | |
//! *--------------------------- |
//! | | |
//! Normalization | |
//! | | |
//! *------------------------------------------------
//! | |
//! -------BMM2------
//! |
//! Output
//!
//! The attention has the following inputs, in order of input index:
//!
//! * Query contains the input query. It is a tensor of type kFLOAT, kHALF or kBF16 with
//! shape [batchSize, numHeadsQuery, sequenceLengthQuery, dimHead]
//! * Key contains the input key. It is a tensor of type kFLOAT, kHALF or kBF16 with
//! shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead]
//! * Value contains the input value. It is a tensor of type kFLOAT, kHALF or kBF16 with
//! shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead]
//! * Mask (optional) contains the mask value. It is a tensor of type kBOOL or the same data type of
//! BMM1 output with shape [batchSize, numHeadsQuery, sequenceLengthQuery, sequenceLengthKeyValue]
//! with batchSize and numHeadsQuery broadcastable. For a kBOOL mask, a True value indicates that the corresponding
//! position is allowed to attend. For other data types, the mask values will be added to the BMM1 output, known
//! as an add mask.
//! * NormalizationQuantizeScale (optional) contains the quantization scale for the attention normalization output.
//! It is a tensor of type kFLOAT, kHALF or kBF16 with dimension 0 or 1.
//!
//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.
//!
class IAttention : public INoCopy
{
public:
//!
//! \brief Set the normalization operation for the attention.
//!
//! \see getNormalizationOperation(), AttentionNormalizationOp
//!
//! \return True if the normalization operation is set successfully, false otherwise.
//!
bool setNormalizationOperation(AttentionNormalizationOp op) noexcept
{
return mImpl->setNormalizationOperation(op);
}
//!
//! \brief Get the normalization operation for the attention.
//!
//! \see setNormalizationOperation(), AttentionNormalizationOp
//!
//! \return The normalization operation for the attention. Default is kSOFTMAX.
//!
AttentionNormalizationOp getNormalizationOperation() const noexcept
{
return mImpl->getNormalizationOperation();
}
//!
//! \brief Set whether a mask will be used for the normalization operation.
//!
//! \param mask the mask tensor of type kBOOL or the same data type of
//! BMM1 output with shape [batchSize, sequenceLengthQuery, sequenceLengthKeyValue]. For a kBOOL mask, a True value
//! indicates that the corresponding position is allowed to attend. For other data types, the mask values will
//! be added to the BMM1 output, known as an add mask.
//!
//! \see getMask
//!
//! \return True if the mask is set successfully, false otherwise.
//!
bool setMask(ITensor& mask) noexcept
{
return mImpl->setMask(mask);
}
//!
//! \brief Get the optional mask in attention.
//!
//! \see setMask
//!
//! \return The optional mask in attention, nullptr if no mask is set.
//!
ITensor* getMask() noexcept
{
return mImpl->getMask();
}
//!
//! \brief Set whether the attention will run a causal inference.
//! Cannot be used together with setMask().
//!
//! \see getCausal
//!
//! \return True if the causal inference is set successfully, false otherwise.
//!
bool setCausal(bool isCausal) noexcept
{
return mImpl->setCausal(isCausal);
}
//!
//! \brief Get whether the attention will run a causal inference.
//!
//! \see setCausal
//!
//! \return True if the attention will run a causal inference, false otherwise. Default is false.
//!
bool getCausal() const noexcept
{
return mImpl->getCausal();
}
//!
//! \brief Set whether the attention can be decomposed to use multiple kernels if no fused kernel support found.
//!
//! \see getDecomposable
//!
//! \return True if the decomposable attention is set successfully, false otherwise.
//!
bool setDecomposable(bool decomposable) noexcept
{
return mImpl->setDecomposable(decomposable);
}
//!
//! \brief Get whether the attention can be decomposed to use multiple kernels if no fused kernel support found.
//!
//! \return True if the attention can be decomposed to use multiple kernels by the compiler,
//! false otherwise. Default is false.
//!
//! \see setDecomposable
//!
bool getDecomposable() const noexcept
{
return mImpl->getDecomposable();
}
//!
//! \brief Append or replace an input of this layer with a specific tensor.
//!
//! \param index the index of the input to modify.
//! \param input the new input tensor.
//!
//! The indices are as follows:
//!
//! Input 0 is the input query tensor.
//! Input 1 is the input key tensor.
//! Input 2 is the input value tensor.
//!
//! \return True if the input tensor is set successfully, false otherwise.
//!
bool setInput(int32_t index, ITensor& input) noexcept
{
return mImpl->setInput(index, input);
}
//!
//! \brief Get the number of inputs of IAttention. IAttention has three inputs.
//!
//! \return The number of inputs of IAttention.
int32_t getNbInputs() const noexcept
{
return mImpl->getNbInputs();
}
//!
//! \brief Get the IAttention input corresponding to the given index.
//!
//! \param index The index of the input tensor.
//!
//! \return The input tensor, or nullptr if the index is out of range.
//!
ITensor* getInput(int32_t index) const noexcept
{
return mImpl->getInput(index);
}
//!
//! \brief Get the number of outputs of a layer. IAttention has one output.
//!
int32_t getNbOutputs() const noexcept
{
return mImpl->getNbOutputs();
}
//!
//! \brief Get the IAttention output corresponding to the given index. IAttention has only one output.
//!
//! \param index The index of the output tensor.
//!
//! \return The indexed output tensor, or nullptr if the index is out of range.
//!
ITensor* getOutput(int32_t index) const noexcept
{
return mImpl->getOutput(index);
}
//!
//! \brief Set the name of the attention.
//!
//! The name is used in error diagnostics.
//! This method copies the name string.
//!
//! \warning The string name must be null-terminated, and be at most 4096 bytes including the terminator.
//!
//! \see getName()
//!
//! \return True if the name is set successfully, false otherwise.
//!
bool setName(char const* name) noexcept
{
return mImpl->setName(name);
}
//!
//! \brief Return the name of the attention.
//!
//! \see setName()
//!
//! \return The name of the attention.
//!
char const* getName() const noexcept
{
return mImpl->getName();
}
//!
//! \brief Set the quantization scale for the attention normalization output.
//!
//! \param tensor for quantization scale. Data type must be DataType::kFLOAT, DataType::kHALF or DataType::kBF16.
//! Must be a 0-d or 1-d.
//!
//! \return True if the quantization scale is set successfully, false otherwise.
//!
//! \warning Must be used together with setNormalizationQuantizeToType to set normalization output datatype to
//! DataType::kFP8 or DataType::kINT8.
//!
bool setNormalizationQuantizeScale(ITensor& tensor) noexcept
{
return mImpl->setNormalizationQuantizeScale(tensor);
}
//!
//! \brief Get the quantization scale for the attention normalization output.
//!
//! \return The quantization scale for the attention normalization output or nullptr if no quantization scale is
//! set.
//!
ITensor* getNormalizationQuantizeScale() const noexcept
{
return mImpl->getNormalizationQuantizeScale();
}
//!
//! \brief Set the datatype the attention normalization is quantized to.
//!
//! \param type the datatype the attention normalization is quantized to. Must be one of DataType::kFP8,
//! DataType::kINT8.
//!
//! \return True if the quantization to type is set successfully, false otherwise.
//!
bool setNormalizationQuantizeToType(DataType type) noexcept
{
return mImpl->setNormalizationQuantizeToType(type);
}
//!
//! \brief Get the datatype the attention normalization is quantized to.
//!
//! \return The datatype the attention normalization is quantized to.
//! The default value is DataType::kFLOAT.
//!
//! \warning Must be used after normalization quantization to type is set by setNormalizationQuantizeToType.
DataType getNormalizationQuantizeToType() const noexcept
{
return mImpl->getNormalizationQuantizeToType();
}
protected:
apiv::VAttention* mImpl;
virtual ~IAttention() noexcept = default;
};
//!
//! \class INetworkDefinition
//!
@@ -7112,6 +7604,8 @@ public:
//!
//! Currently only values of K up to 3840 are supported.
//!
//! The default indices tensor (the second output) data type is DataType::kINT32.
//!
//! \param input The input tensor to the layer.
//!
//! \param op Operation to perform.
@@ -7122,19 +7616,53 @@ public:
//! \param reduceAxes The reduction dimensions.
//! The bit in position i of bitmask reduceAxes corresponds to explicit dimension i of the result.
//! E.g., the least significant bit corresponds to the first explicit dimension and the next to least
//! significant bit corresponds to the second explicit dimension.
//!
//! Currently reduceAxes must specify exactly one dimension, and it must be one of the last four dimensions.
//! significant bit corresponds to the second explicit dimension. Currently reduceAxes must specify
//! exactly one dimension, and it must be one of the last four dimensions.
//!
//! \see ITopKLayer
//!
//! \return The new TopK layer, or nullptr if it could not be created.
//!
ITopKLayer* addTopK(ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes) noexcept
//! \deprecated Deprecated in TensorRT 10.14. Superseded by five-argument addTopK.
//!
TRT_DEPRECATED ITopKLayer* addTopK(ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes) noexcept
{
return mImpl->addTopK(input, op, k, reduceAxes);
}
//!
//! \brief Add a TopK layer to the network.
//!
//! The TopK layer has two outputs of the same dimensions. The first contains data values,
//! the second contains index positions for the values. Output values are sorted, largest first
//! for operation kMAX and smallest first for operation kMIN.
//!
//! Currently only values of K up to 3840 are supported.
//!
//! \param input The input tensor to the layer.
//!
//! \param op Operation to perform.
//!
//! \param k The number of elements to keep. For dynamic k, use the setInput() method to pass in k as a tensor
//! instead, which will override the static k value passed here in calculations.
//!
//! \param reduceAxes The reduction dimensions.
//! The bit in position i of bitmask reduceAxes corresponds to explicit dimension i of the result.
//! E.g., the least significant bit corresponds to the first explicit dimension and the next to least
//! significant bit corresponds to the second explicit dimension. Currently reduceAxes must specify
//! exactly one dimension, and it must be one of the last four dimensions.
//!
//! \param indicesType Indices tensor (the second output) data type, must be DataType::kINT32 or DataType::kINT64.
//!
//! \see ITopKLayer
//!
//! \return The new TopK layer, or nullptr if it could not be created.
//!
ITopKLayer* addTopK(ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes, DataType indicesType) noexcept
{
return mImpl->addTopKV2(input, op, k, reduceAxes, indicesType);
}
//!
//! \brief Add gather with mode GatherMode::kDEFAULT and specified axis and nbElementWiseDims=0.
//!
@@ -7211,17 +7739,37 @@ public:
//!
//! \brief Add a nonzero layer to the network.
//!
//! The default indices tensor (the first output) data type is DataType::kINT32.
//!
//! \param input The input tensor to the layer.
//!
//! \see INonZeroLayer
//!
//! \return The new nonzero layer, or nullptr if it could be created.
//! \return The new nonzero layer, or nullptr if it could not be created.
//!
INonZeroLayer* addNonZero(ITensor& input) noexcept
//! \deprecated Deprecated in TensorRT 10.14. Superseded by two-argument addNonZero.
//!
TRT_DEPRECATED INonZeroLayer* addNonZero(ITensor& input) noexcept
{
return mImpl->addNonZero(input);
}
//!
//! \brief Add a nonzero layer to the network.
//!
//! \param input The input tensor to the layer.
//!
//! \param indicesType Indices tensor (the first output) data type, must be DataType::kINT32 or DataType::kINT64.
//!
//! \see INonZeroLayer
//!
//! \return The new nonzero layer, or nullptr if it could not be created.
//!
INonZeroLayer* addNonZero(ITensor& input, DataType indicesType) noexcept
{
return mImpl->addNonZeroV2(input, indicesType);
}
//!
//! \brief Add a constant layer to the network.
//!
@@ -7940,7 +8488,7 @@ public:
//!
//! This layer performs dynamic block quantization of its input tensor and outputs the
//! quantized data and the computed block scale-factors.
//! The block size is currently limited to 16 and the size of the blocked axis must be divisible by 16.
//! The blocked axis dimension size must be divisible by the block size.
//!
//! \param input The input tensor to be quantized. Its data type must be one of DataType::kFLOAT,
//! DataType::kHALF, or DataType::kBF16. Currently only 2D and 3D inputs are supported.
@@ -7999,6 +8547,8 @@ public:
//!
//! \brief Add a non-maximum suppression layer to the network.
//!
//! The default indices tensor (the first output) data type is DataType::kINT32.
//!
//! \param boxes The input boxes tensor to the layer.
//!
//! \param scores The input scores tensor to the layer.
@@ -8009,11 +8559,33 @@ public:
//!
//! \return The new NMS layer, or nullptr if it could not be created.
//!
INMSLayer* addNMS(ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass) noexcept
//! \deprecated Deprecated in TensorRT 10.14. Superseded by four-argument addNMS.
//!
TRT_DEPRECATED INMSLayer* addNMS(ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass) noexcept
{
return mImpl->addNMS(boxes, scores, maxOutputBoxesPerClass);
}
//!
//! \brief Add a non-maximum suppression layer to the network.
//!
//! \param boxes The input boxes tensor to the layer.
//!
//! \param scores The input scores tensor to the layer.
//!
//! \param maxOutputBoxesPerClass The input maxOutputBoxesPerClass tensor to the layer.
//!
//! \param indicesType Indices tensor (the first output) data type, must be DataType::kINT32 or DataType::kINT64.
//!
//! \see INMSLayer
//!
//! \return The new NMS layer, or nullptr if it could not be created.
//!
INMSLayer* addNMS(ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass, DataType indicesType) noexcept
{
return mImpl->addNMSV2(boxes, scores, maxOutputBoxesPerClass, indicesType);
}
//!
//! \brief Add a ReverseSequence layer to the network.
//!
@@ -8079,6 +8651,34 @@ public:
return mImpl->addCumulative(input, axis, operation, exclusive, reverse);
}
//!
//! \brief Add an attention to the network.
//!
//! \param query A 4d input query tensor to the layer.
//! \param key A 4d input key tensor to the layer.
//! \param value A 4d input value tensor to the layer.
//! \param normOp The normalization operation to perform.
//! \param causal Use causual inference or not.
//!
//! query must have shape [batchSize, numHeadsQuery, sequenceLengthQuery, dimHead].
//! key and value must have shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead].
//! pastKey and pastValue must have shape [batchSize, numHeadsKeyValue, sequenceLengthKeyValue, dimHead].
//! normOp defaults to kSOFTMAX isCausal defaults to false.
//!
//! By default, IAttention is not decomposable and TensorRT will try to use a single fused kernel, which may be more
//! efficient than if the subgraph is expressed without IAttention. Setting the IAttention to decomposable=True can
//! allow IAttention to be to use multiple kernels if no fused kernel support found.
//!
//! \see IAttention
//!
//! \return The new attention, or nullptr if it could not be created.
//!
IAttention* addAttention(
ITensor& query, ITensor& key, ITensor& value, AttentionNormalizationOp normOp, bool causal) noexcept
{
return mImpl->addAttention(query, key, value, normOp, causal);
}
//!
//! \brief Return the builder from which this INetworkDefinition was created.
//!
@@ -8812,6 +9412,8 @@ enum class RuntimePlatform : int32_t
//! Designates the target platform for engine execution as Windows AMD64 system. Currently this flag can only be
//! enabled when building engines on Linux AMD64 platforms.
kWINDOWS_AMD64 = 1,
};
namespace impl
@@ -9018,6 +9620,8 @@ enum class BuilderFlag : int32_t
//! For layers that perform einsum:
//! Let n be the leftmost reduction axis. The axes to the left of n are distributive axes.
kDISTRIBUTIVE_INDEPENDENCE = 28,
};
//!
include/NvInferImpl.h
+64 -2
View File
@@ -68,6 +68,7 @@ class IAlgorithmIOInfo;
class IAlgorithmVariant;
#endif // !STRIP_TRT_RTX_INTERNAL_API
class IAssertionLayer;
class IAttention;
class IBuilder;
class IBuilderConfig;
class IConcatenationLayer;
@@ -175,6 +176,7 @@ struct Permutation;
class Weights;
enum class ActivationType : int32_t;
enum class AttentionNormalizationOp : int32_t;
enum class BoundingBoxFormat : int32_t;
enum class BuilderFlag : int32_t;
#if !STRIP_TRT_RTX_INTERNAL_API
@@ -218,6 +220,7 @@ enum class HardwareCompatibilityLevel : int32_t;
enum class ExecutionContextAllocationStrategy : int32_t;
enum class RuntimePlatform : int32_t;
enum class TilingOptimizationLevel : int32_t;
enum class EngineStat : int32_t;
using TacticSources = uint32_t;
@@ -420,6 +423,7 @@ public:
TRT_NODISCARD virtual IExecutionContext* createExecutionContextWithRuntimeConfig(
IRuntimeConfig* runtimeConfig) noexcept = 0;
TRT_NODISCARD virtual IRuntimeConfig* createRuntimeConfig() noexcept = 0;
TRT_NODISCARD virtual int64_t getEngineStat(EngineStat stat) const noexcept = 0;
};
class VExecutionContext : public VRoot
@@ -776,6 +780,8 @@ public:
virtual int32_t getK() const noexcept = 0;
virtual void setReduceAxes(uint32_t reduceAxes) noexcept = 0;
virtual uint32_t getReduceAxes() const noexcept = 0;
virtual bool setIndicesType(DataType type) noexcept = 0;
virtual DataType getIndicesType() const noexcept = 0;
};
class VMatrixMultiplyLayer : public VRoot
@@ -788,6 +794,8 @@ public:
class VNonZeroLayer : public VRoot
{
public:
virtual bool setIndicesType(DataType type) noexcept = 0;
virtual DataType getIndicesType() const noexcept = 0;
};
class VRaggedSoftMaxLayer : public VRoot
@@ -917,6 +925,46 @@ public:
virtual char const* getName() const noexcept = 0;
};
class VAttentionBoundaryLayer : public VRoot
{
public:
virtual IAttention* getAttention() const noexcept = 0;
};
class VAttentionInputLayer : public VRoot
{
public:
};
class VAttentionOutputLayer : public VRoot
{
public:
};
class VAttention : public VRoot
{
public:
TRT_NODISCARD virtual bool setInput(int32_t index, ITensor& input) noexcept = 0;
TRT_NODISCARD virtual int32_t getNbInputs() const noexcept = 0;
TRT_NODISCARD virtual ITensor* getInput(int32_t index) const noexcept = 0;
TRT_NODISCARD virtual int32_t getNbOutputs() const noexcept = 0;
TRT_NODISCARD virtual ITensor* getOutput(int32_t index) const noexcept = 0;
TRT_NODISCARD virtual bool setName(char const* name) noexcept = 0;
TRT_NODISCARD virtual char const* getName() const noexcept = 0;
TRT_NODISCARD virtual bool setNormalizationOperation(AttentionNormalizationOp op) noexcept = 0;
TRT_NODISCARD virtual AttentionNormalizationOp getNormalizationOperation() const noexcept = 0;
TRT_NODISCARD virtual bool setCausal(bool isCausal) noexcept = 0;
TRT_NODISCARD virtual bool getCausal() const noexcept = 0;
TRT_NODISCARD virtual bool setMask(ITensor& mask) noexcept = 0;
TRT_NODISCARD virtual ITensor* getMask() const noexcept = 0;
TRT_NODISCARD virtual bool setDecomposable(bool decomposable) noexcept = 0;
TRT_NODISCARD virtual bool getDecomposable() const noexcept = 0;
TRT_NODISCARD virtual bool setNormalizationQuantizeScale(ITensor& tensor) noexcept = 0;
TRT_NODISCARD virtual ITensor* getNormalizationQuantizeScale() const noexcept = 0;
TRT_NODISCARD virtual bool setNormalizationQuantizeToType(DataType type) noexcept = 0;
TRT_NODISCARD virtual DataType getNormalizationQuantizeToType() const noexcept = 0;
}; // class VAttention
class VSelectLayer : public VRoot
{
};
@@ -1020,6 +1068,8 @@ public:
virtual BoundingBoxFormat getBoundingBoxFormat() const noexcept = 0;
virtual void setTopKBoxLimit(int32_t limit) noexcept = 0;
virtual int32_t getTopKBoxLimit() const noexcept = 0;
virtual bool setIndicesType(DataType type) noexcept = 0;
virtual DataType getIndicesType() const noexcept = 0;
}; // class VNMSLayer
class VReverseSequenceLayer : public VRoot
@@ -1155,11 +1205,23 @@ public:
virtual ISqueezeLayer* addSqueeze(ITensor& input, ITensor& axes) noexcept = 0;
virtual IUnsqueezeLayer* addUnsqueeze(ITensor& input, ITensor& axes) noexcept = 0;
virtual IDynamicQuantizeLayer* addDynamicQuantize(
ITensor& input, int32_t axis, int32_t blockSize, DataType toType, DataType scaleType) noexcept = 0;
ITensor& input, int32_t axis, int32_t blockSize, DataType toType, DataType scaleType) noexcept
= 0;
virtual ICumulativeLayer* addCumulative(
ITensor& input, ITensor& axis, CumulativeOperation operation, bool exclusive, bool reverse) noexcept = 0;
ITensor& input, ITensor& axis, CumulativeOperation operation, bool exclusive, bool reverse) noexcept
= 0;
virtual bool markUnfusedTensorsAsDebugTensors() noexcept = 0;
virtual bool unmarkUnfusedTensorsAsDebugTensors() noexcept = 0;
virtual ITopKLayer* addTopKV2(
ITensor& input, TopKOperation op, int32_t k, uint32_t reduceAxes, DataType indicesType) noexcept
= 0;
virtual INonZeroLayer* addNonZeroV2(ITensor& input, DataType indicesType) noexcept = 0;
virtual INMSLayer* addNMSV2(
ITensor& boxes, ITensor& scores, ITensor& maxOutputBoxesPerClass, DataType indicesType) noexcept
= 0;
virtual IAttention* addAttention(
ITensor& query, ITensor& key, ITensor& value, AttentionNormalizationOp normOp, bool isCausal) noexcept
= 0;
};
#if !STRIP_TRT_RTX_INTERNAL_API
include/NvInferRuntime.h
+67 -4
View File
@@ -24,11 +24,11 @@
//! This is the top-level API file for TensorRT extended runtime library.
//!
#include "NvInferImpl.h"
#include "NvInferImpl.h" // IWYU pragma: export
#define NV_INFER_INTERNAL_INCLUDE 1
#include "NvInferPluginBase.h" // IWYU pragma: exports
#include "NvInferPluginBase.h" // IWYU pragma: export
#undef NV_INFER_INTERNAL_INCLUDE
#include "NvInferRuntimeCommon.h"
#include "NvInferRuntimeCommon.h" // IWYU pragma: export
namespace nvinfer1
{
@@ -1854,6 +1854,7 @@ protected:
//!
using IGpuAllocator = v_1_0::IGpuAllocator;
//!
//! \class IRuntime
//!
@@ -2178,6 +2179,7 @@ public:
return mImpl->getEngineHostCodeAllowed();
}
protected:
apiv::VRuntime* mImpl;
};
@@ -3000,13 +3002,14 @@ enum class SerializationFlag : int32_t
{
kEXCLUDE_WEIGHTS = 0, //!< Exclude the weights that can be refitted.
kEXCLUDE_LEAN_RUNTIME = 1, //!< Exclude the lean runtime.
kINCLUDE_REFIT = 2, //!< Remain refittable if originally so.
};
//! Maximum number of serialization flags in SerializationFlag enum. \see SerializationFlag
template <>
constexpr inline int32_t EnumMax<SerializationFlag>() noexcept
{
return 2;
return 3;
}
//!
@@ -3155,6 +3158,34 @@ protected:
apiv::VRuntimeConfig* mImpl;
}; // class IRuntimeConfig
//!
//! \enum EngineStat
//!
//! \brief The kind of engine statistics that queried from the ICudaEngine.
//!
//! \see ICudaEngine::getEngineStat()
//! \see BuilderFlag::kSTRIP_PLAN
//!
enum class EngineStat : int32_t
{
//! Return the total weight size in bytes.
kTOTAL_WEIGHTS_SIZE = 0,
//! Return the stripped weight size in bytes for engines built with BuilderFlag::kSTRIP_PLAN.
kSTRIPPED_WEIGHTS_SIZE = 1,
};
//!
//! \brief Maximum number of engine statistic kinds in EngineStat enum.
//!
//! \see EngineStat
//!
template <>
constexpr inline int32_t EnumMax<EngineStat>() noexcept
{
return 2;
}
//!
//! \class ICudaEngine
//!
@@ -3827,6 +3858,10 @@ public:
//! Serializing plan file with SerializationFlag::kEXCLUDE_WEIGHTS requires building the engine with kREFIT,
//! kREFIT_IDENTICAL or kREFIT_INDIVIDUAL.
//!
//! The only applicable scenario for SerializationFlag::kINCLUDE_REFIT is when serializing weight-stripping
//! engines without kEXCLUDE_WEIGHTS. By default, the resulting serialized engine is unrefittable. Setting
//! SerializationFlag::kINCLUDE_REFIT ensures that the serialized engine remains refittable.
//!
//! \see IRuntime::deserializeCudaEngine()
//!
IHostMemory* serializeWithConfig(ISerializationConfig& config) const noexcept
@@ -4090,6 +4125,34 @@ public:
return mImpl->getProfileTensorValuesV2(tensorName, profileIndex, select);
}
//!
//! \brief Get engine statistics according to the given enum value.
//!
//! \param stat The kind of statistics to query.
//!
//! If stat is kTOTAL_WEIGHTS_SIZE, the return value is the total weights size in bytes in the engine.
//! If stat is kSTRIPPED_WEIGHTS_SIZE, the return value is the stripped weight size in bytes for engines
//! built with BuilderFlag::kSTRIP_PLAN.
//!
//! When the BuilderFlag::kWEIGHT_STREAMING flag is enabled, engine weights may not be fully copied to the device.
//! The reported total weight size reflects the sum of all weights utilized by the engine,
//! which does not necessarily correspond to the actual GPU memory allocated.
//!
//! \return The kind of statistics specified by EngineStat.
//!
//! \warning if kSTRIPPED_WEIGHTS_SIZE is passed to query a normal engine, this function will
//! return -1 to indicate invalid enum value.
//!
//! \see EngineStat
//! \see BuilderFlag::kWEIGHT_STREAMING
//! \see setWeightStreamingBudget()
//! \see getStreamableWeightsSize()
//!
int64_t getEngineStat(EngineStat stat) const noexcept
{
return mImpl->getEngineStat(stat);
}
protected:
apiv::VCudaEngine* mImpl;
};
include/NvInferRuntimeBase.h
+2
View File
@@ -138,6 +138,8 @@ constexpr int32_t EnumMax() noexcept
//!
//! \enum DataType
//! \brief The type of weights and tensors.
//! The datatypes other than kBOOL, kINT32, and kINT64 are "activation datatypes,"
//! as they often represent values corresponding to inference results.
//!
enum class DataType : int32_t
{
include/NvInferVersion.h
+3 -3
View File
@@ -24,9 +24,9 @@
#define NV_INFER_VERSION_H
#define TRT_MAJOR_ENTERPRISE 10
#define TRT_MINOR_ENTERPRISE 13
#define TRT_PATCH_ENTERPRISE 3
#define TRT_BUILD_ENTERPRISE 9
#define TRT_MINOR_ENTERPRISE 14
#define TRT_PATCH_ENTERPRISE 1
#define TRT_BUILD_ENTERPRISE 48
#define NV_TENSORRT_MAJOR TRT_MAJOR_ENTERPRISE //!< TensorRT major version.
#define NV_TENSORRT_MINOR TRT_MINOR_ENTERPRISE //!< TensorRT minor version.
#define NV_TENSORRT_PATCH TRT_PATCH_ENTERPRISE //!< TensorRT patch version.
include/NvOnnxParser.h
+640
View File
@@ -0,0 +1,640 @@
/*
* 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.
*/
#ifndef NV_ONNX_PARSER_H
#define NV_ONNX_PARSER_H
#include "NvInfer.h"
#include <stddef.h>
#include <string>
#include <vector>
//!
//! \file NvOnnxParser.h
//!
//! This is the API for the ONNX Parser
//!
#define NV_ONNX_PARSER_MAJOR 0
#define NV_ONNX_PARSER_MINOR 1
#define NV_ONNX_PARSER_PATCH 0
static constexpr int32_t NV_ONNX_PARSER_VERSION
= ((NV_ONNX_PARSER_MAJOR * 10000) + (NV_ONNX_PARSER_MINOR * 100) + NV_ONNX_PARSER_PATCH);
//!
//! \typedef SubGraph_t
//!
//! \brief The data structure containing the parsing capability of
//! a set of nodes in an ONNX graph.
//!
typedef std::pair<std::vector<size_t>, bool> SubGraph_t;
//!
//! \typedef SubGraphCollection_t
//!
//! \brief The data structure containing all SubGraph_t partitioned
//! out of an ONNX graph.
//!
typedef std::vector<SubGraph_t> SubGraphCollection_t;
//!
//! \namespace nvonnxparser
//!
//! \brief The TensorRT ONNX parser API namespace
//!
namespace nvonnxparser
{
template <typename T>
constexpr inline int32_t EnumMax() noexcept;
//!
//! \enum ErrorCode
//!
//! \brief The type of error that the parser or refitter may return
//!
enum class ErrorCode : int
{
kSUCCESS = 0,
kINTERNAL_ERROR = 1,
kMEM_ALLOC_FAILED = 2,
kMODEL_DESERIALIZE_FAILED = 3,
kINVALID_VALUE = 4,
kINVALID_GRAPH = 5,
kINVALID_NODE = 6,
kUNSUPPORTED_GRAPH = 7,
kUNSUPPORTED_NODE = 8,
kUNSUPPORTED_NODE_ATTR = 9,
kUNSUPPORTED_NODE_INPUT = 10,
kUNSUPPORTED_NODE_DATATYPE = 11,
kUNSUPPORTED_NODE_DYNAMIC = 12,
kUNSUPPORTED_NODE_SHAPE = 13,
kREFIT_FAILED = 14
};
//!
//! Maximum number of flags in the ErrorCode enum.
//!
//! \see ErrorCode
//!
template <>
constexpr inline int32_t EnumMax<ErrorCode>() noexcept
{
return 14;
}
//!
//! \brief Represents one or more OnnxParserFlag values using binary OR
//! operations, e.g., 1U << OnnxParserFlag::kNATIVE_INSTANCENORM
//!
//! \see IParser::setFlags() and IParser::getFlags()
//!
using OnnxParserFlags = uint32_t;
enum class OnnxParserFlag : int32_t
{
//! Parse the ONNX model into the INetworkDefinition with the intention of using TensorRT's native layer
//! implementation over the plugin implementation for InstanceNormalization nodes.
//! This flag is required when building version-compatible or hardware-compatible engines.
//! This flag is set to be ON by default.
kNATIVE_INSTANCENORM = 0,
//! Enable UINT8 as a quantization data type and asymmetric quantization with non-zero zero-point values
//! in Quantize and Dequantize nodes. This flag is set to be OFF by default.
//! The resulting engine must be built targeting DLA version >= 3.16.
kENABLE_UINT8_AND_ASYMMETRIC_QUANTIZATION_DLA = 1,
};
//!
//! Maximum number of flags in the OnnxParserFlag enum.
//!
//! \see OnnxParserFlag
//!
template <>
constexpr inline int32_t EnumMax<OnnxParserFlag>() noexcept
{
return 2;
}
//!
//! \class IParserError
//!
//! \brief an object containing information about an error
//!
class IParserError
{
public:
//!
//!\brief the error code.
//!
virtual ErrorCode code() const = 0;
//!
//!\brief description of the error.
//!
virtual char const* desc() const = 0;
//!
//!\brief source file in which the error occurred.
//!
virtual char const* file() const = 0;
//!
//!\brief source line at which the error occurred.
//!
virtual int line() const = 0;
//!
//!\brief source function in which the error occurred.
//!
virtual char const* func() const = 0;
//!
//!\brief index of the ONNX model node in which the error occurred.
//!
virtual int node() const = 0;
//!
//!\brief name of the node in which the error occurred.
//!
virtual char const* nodeName() const = 0;
//!
//!\brief name of the node operation in which the error occurred.
//!
virtual char const* nodeOperator() const = 0;
//!
//!\brief A list of the local function names, from the top level down, constituting the current
//! stack trace in which the error occurred. A top-level node that is not inside any
//! local function would return a nullptr.
//!
virtual char const* const* localFunctionStack() const = 0;
//!
//!\brief The size of the stack of local functions at the point where the error occurred.
//! A top-level node that is not inside any local function would correspond to
// a stack size of 0.
//!
virtual int32_t localFunctionStackSize() const = 0;
protected:
virtual ~IParserError() {}
};
//!
//! \class IParser
//!
//! \brief an object for parsing ONNX models into a TensorRT network definition
//!
//! \warning If the ONNX model has a graph output with the same name as a graph input,
//! the output will be renamed by prepending "__".
//!
//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.
//!
class IParser
{
public:
//!
//! \brief Parse a serialized ONNX model into the TensorRT network.
//! This method has very limited diagnostics. If parsing the serialized model
//! fails for any reason (e.g. unsupported IR version, unsupported opset, etc.)
//! it the user responsibility to intercept and report the error.
//! To obtain a better diagnostic, use the parseFromFile method below.
//!
//! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serialized_onnx_model_size Size of the serialized ONNX model
//! in bytes
//! \param model_path Absolute path to the model file for loading external weights if required
//! \return true if the model was parsed successfully
//! \see getNbErrors() getError()
//!
virtual bool parse(
void const* serialized_onnx_model, size_t serialized_onnx_model_size, const char* model_path = nullptr) noexcept
= 0;
//!
//! \brief Parse an onnx model file, which can be a binary protobuf or a text onnx model
//! calls parse method inside.
//!
//! \param onnxModelFile name
//! \param verbosity Level
//!
//! \return true if the model was parsed successfully
//!
//!
virtual bool parseFromFile(const char* onnxModelFile, int verbosity) noexcept = 0;
//!
//! [DEPRECATED] Deprecated in TensorRT 10.1. See supportsModelV2.
//!
//! \brief Check whether TensorRT supports a particular ONNX model.
//! If the function returns True, one can proceed to engine building
//! without having to call \p parse or \p parseFromFile.
//!
//! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serialized_onnx_model_size Size of the serialized ONNX model
//! in bytes
//! \param sub_graph_collection Container to hold supported subgraphs
//! \param model_path Absolute path to the model file for loading external weights if required
//! \return true if the model is supported
//!
TRT_DEPRECATED virtual bool supportsModel(void const* serialized_onnx_model, size_t serialized_onnx_model_size,
SubGraphCollection_t& sub_graph_collection, const char* model_path = nullptr) noexcept = 0;
//!
//! [DEPRECATED] Deprecated in TensorRT 10.13. See loadInitializer().
//!
//!\brief Parse a serialized ONNX model into the TensorRT network
//! with consideration of user provided weights
//!
//! \param serialized_onnx_model Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serialized_onnx_model_size Size of the serialized ONNX model
//! in bytes
//! \return true if the model was parsed successfully
//! \see getNbErrors() getError()
//!
TRT_DEPRECATED virtual bool parseWithWeightDescriptors(
void const* serialized_onnx_model, size_t serialized_onnx_model_size) noexcept = 0;
//!
//!\brief Returns whether the specified operator may be supported by the
//! parser.
//!
//! Note that a result of true does not guarantee that the operator will be
//! supported in all cases (i.e., this function may return false-positives).
//!
//! \param op_name The name of the ONNX operator to check for support
//!
virtual bool supportsOperator(const char* op_name) const noexcept = 0;
//!
//!\brief Get the number of errors that occurred during prior calls to
//! \p parse
//!
//! \see getError() clearErrors() IParserError
//!
virtual int getNbErrors() const noexcept = 0;
//!
//!\brief Get an error that occurred during prior calls to \p parse
//!
//! \see getNbErrors() clearErrors() IParserError
//!
virtual IParserError const* getError(int index) const noexcept = 0;
//!
//!\brief Clear errors from prior calls to \p parse
//!
//! \see getNbErrors() getError() IParserError
//!
virtual void clearErrors() noexcept = 0;
virtual ~IParser() noexcept = default;
//!
//! \brief Query the plugin libraries needed to implement operations used by the parser in a version-compatible
//! engine.
//!
//! This provides a list of plugin libraries on the filesystem needed to implement operations
//! in the parsed network. If you are building a version-compatible engine using this network,
//! provide this list to IBuilderConfig::setPluginsToSerialize to serialize these plugins along
//! with the version-compatible engine, or, if you want to ship these plugin libraries externally
//! to the engine, ensure that IPluginRegistry::loadLibrary is used to load these libraries in the
//! appropriate runtime before deserializing the corresponding engine.
//!
//! \param[out] nbPluginLibs Returns the number of plugin libraries in the array, or -1 if there was an error.
//! \return Array of `nbPluginLibs` C-strings describing plugin library paths on the filesystem if nbPluginLibs > 0,
//! or nullptr otherwise. This array is owned by the IParser, and the pointers in the array are only valid until
//! the next call to parse(), supportsModel(), parseFromFile(), or parseWithWeightDescriptors().
//!
virtual char const* const* getUsedVCPluginLibraries(int64_t& nbPluginLibs) const noexcept = 0;
//!
//! \brief Set the parser flags.
//!
//! The flags are listed in the OnnxParserFlag enum.
//!
//! \param OnnxParserFlags The flags used when parsing an ONNX model.
//!
//! \note This function will override the previous set flags, rather than bitwise ORing the new flag.
//!
//! \see getFlags()
//!
virtual void setFlags(OnnxParserFlags onnxParserFlags) noexcept = 0;
//!
//! \brief Get the parser flags. Defaults to 0.
//!
//! \return The parser flags as a bitmask.
//!
//! \see setFlags()
//!
virtual OnnxParserFlags getFlags() const noexcept = 0;
//!
//! \brief clear a parser flag.
//!
//! clears the parser flag from the enabled flags.
//!
//! \see setFlags()
//!
virtual void clearFlag(OnnxParserFlag onnxParserFlag) noexcept = 0;
//!
//! \brief Set a single parser flag.
//!
//! Add the input parser flag to the already enabled flags.
//!
//! \see setFlags()
//!
virtual void setFlag(OnnxParserFlag onnxParserFlag) noexcept = 0;
//!
//! \brief Returns true if the parser flag is set
//!
//! \see getFlags()
//!
//! \return True if flag is set, false if unset.
//!
virtual bool getFlag(OnnxParserFlag onnxParserFlag) const noexcept = 0;
//!
//!\brief Return the i-th output ITensor object for the ONNX layer "name".
//!
//! Return the i-th output ITensor object for the ONNX layer "name".
//! If "name" is not found or i is out of range, return nullptr.
//! In the case of multiple nodes sharing the same name this function will return
//! the output tensors of the first instance of the node in the ONNX graph.
//!
//! \param name The name of the ONNX layer.
//!
//! \param i The index of the output. i must be in range [0, layer.num_outputs).
//!
virtual nvinfer1::ITensor const* getLayerOutputTensor(char const* name, int64_t i) noexcept = 0;
//!
//! \brief Check whether TensorRT supports a particular ONNX model.
//! If the function returns True, one can proceed to engine building
//! without having to call \p parse or \p parseFromFile.
//! Results can be queried through \p getNbSubgraphs, \p isSubgraphSupported,
//! \p getSubgraphNodes.
//!
//! \param serializedOnnxModel Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serializedOnnxModelSize Size of the serialized ONNX model in bytes
//! \param modelPath Absolute path to the model file for loading external weights if required
//! \return true if the model is supported
//!
virtual bool supportsModelV2(
void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0;
//!
//! \brief Get the number of subgraphs. Calling this function before calling \p supportsModelV2 results in undefined
//! behavior.
//!
//!
//! \return Number of subgraphs.
//!
virtual int64_t getNbSubgraphs() noexcept = 0;
//!
//! \brief Returns whether the subgraph is supported. Calling this function before calling \p supportsModelV2
//! results in undefined behavior.
//!
//!
//! \param index Index of the subgraph.
//! \return Whether the subgraph is supported.
//!
virtual bool isSubgraphSupported(int64_t const index) noexcept = 0;
//!
//! \brief Get the nodes of the specified subgraph. Calling this function before calling \p supportsModelV2 results
//! in undefined behavior.
//!
//!
//! \param index Index of the subgraph.
//! \param subgraphLength Returns the length of the subgraph as reference.
//!
//! \return Pointer to the subgraph nodes array. This pointer is owned by the Parser.
//!
virtual int64_t* getSubgraphNodes(int64_t const index, int64_t& subgraphLength) noexcept = 0;
//!
//! \brief Load a serialized ONNX model into the parser. Unlike the parse(), parseFromFile(), or
//! parseWithWeightDescriptors() functions, this function does not immediately convert the model into a TensorRT
//! INetworkDefinition. Using this function allows users to provide their own initializers for the ONNX model
//! through the loadInitializer() function.
//!
//! Only one model can be loaded at a time. Subsequent calls to loadModelProto() will result in an error.
//!
//! To begin the conversion of the model into a TensorRT INetworkDefinition, use parseModelProto().
//!
//! \param serializedOnnxModel Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serializedOnnxModelSize Size of the serialized ONNX model in bytes.
//! \param modelPath Absolute path to the model file for loading external weights if required.
//! \return true if the model was loaded successfully
//! \see getNbErrors() getError()
//!
virtual bool loadModelProto(
void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0;
//!
//! \brief Prompt the ONNX parser to load an initializer with user-provided binary data.
//! The lifetime of the data must exceed the lifetime of the parser.
//!
//! All user-provided initializers must be provided prior to calling refitModelProto().
//!
//! This function can be called multiple times to specify the names of multiple initializers.
//!
//! Calling this function with an initializer previously specified will overwrite the previous instance.
//!
//!
//! This function will return false if initializer validation fails. Possible validation errors are:
//! * This function was called prior to loadModelProto().
//! * The requested initializer was not found in the model.
//! * The size of the data provided is different from the corresponding initializer in the model.
//!
//! \param name Name of the initializer.
//! \param data Binary data containing the values of the initializer.
//! \param size Size of the initializer in bytes.
//! \return true if the initializer was loaded successfully
//! \see loadModelProto()
//!
virtual bool loadInitializer(char const* name, void const* data, size_t size) noexcept = 0;
//! \brief Begin the parsing and conversion process of the loaded ONNX model into a TensorRT INetworkDefinition.
//!
//! \return true if conversion was successful
//! \see getNbErrors() getError() loadModelProto() loadModelProtoFromFile()
//!
virtual bool parseModelProto() noexcept = 0;
};
//!
//! \class IParserRefitter
//!
//! \brief An interface designed to refit weights from an ONNX model.
//!
//! \warning Do not inherit from this class, as doing so will break forward-compatibility of the API and ABI.
//!
class IParserRefitter
{
public:
//!
//! \brief Load a serialized ONNX model from memory and perform weight refit.
//!
//! \param serializedOnnxModel Pointer to the serialized ONNX model
//! \param serializedOnnxModelSize Size of the serialized ONNX model
//! in bytes
//! \param modelPath Absolute path to the model file for loading external weights if required
//! \return true if all the weights in the engine were refit successfully.
//!
//! The serialized ONNX model must be identical to the one used to generate the engine
//! that will be refit.
//!
virtual bool refitFromBytes(
void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept
= 0;
//!
//! \brief Load and parse a ONNX model from disk and perform weight refit.
//!
//! \param onnxModelFile Path to the ONNX model to load from disk.
//!
//! \return true if the model was loaded successfully, and if all the weights in the engine were refit successfully.
//!
//! The provided ONNX model must be identical to the one used to generate the engine
//! that will be refit.
//!
virtual bool refitFromFile(char const* onnxModelFile) noexcept = 0;
//!
//!\brief Get the number of errors that occurred during prior calls to \p refitFromBytes or \p refitFromFile
//!
//! \see getError() IParserError
//!
virtual int32_t getNbErrors() const noexcept = 0;
//!
//!\brief Get an error that occurred during prior calls to \p refitFromBytes or \p refitFromFile
//!
//! \see getNbErrors() IParserError
//!
virtual IParserError const* getError(int32_t index) const noexcept = 0;
//!
//!\brief Clear errors from prior calls to \p refitFromBytes or \p refitFromFile
//!
//! \see getNbErrors() getError() IParserError
//!
virtual void clearErrors() = 0;
virtual ~IParserRefitter() noexcept = default;
//!
//! \brief Load a serialized ONNX model into the parser. Unlike the refit(), or refitFromFile()
//! functions, this function does not immediately begin the refit process. Using this function
//! allows users to provide their own initializers for the ONNX model through the loadInitializer() function.
//!
//! Only one model can be loaded at a time. Subsequent calls to loadModelProto() will result in an error.
//!
//! To begin the refit process, use refitModelProto().
//!
//! \param serializedOnnxModel Pointer to the serialized ONNX model. Can be freed after this function returns.
//! \param serializedOnnxModelSize Size of the serialized ONNX model in bytes.
//! \param modelPath Absolute path to the model file for loading external weights if required.
//! \return true if the model was loaded successfully
//! \see getNbErrors() getError()
//!
virtual bool loadModelProto(
void const* serializedOnnxModel, size_t serializedOnnxModelSize, char const* modelPath = nullptr) noexcept = 0;
//!
//! \brief Prompt the ONNX refitter to load an initializer with user-provided binary data.
//! The lifetime of the data must exceed the lifetime of the refitter.
//!
//! All user-provided initializers must be provided prior to calling refitModelProto().
//!
//! This function can be called multiple times to specify the names of multiple initializers.
//!
//! Calling this function with an initializer previously specified will overwrite the previous instance.
//!
//! This function will return false if initializer validation fails. Possible validation errors are:
//! * This function was called prior to loadModelProto()
//! * The requested initializer was not found in the model.
//! * The size of the data provided is different from the corresponding initializer in the model.
//!
//! \param name Name of the initializer.
//! \param data Binary data containing the values of the initializer.
//! \param size Size of the initializer in bytes.
//! \return true if the initializer was loaded successfully
//! \see loadModelProto()
//!
virtual bool loadInitializer(char const* name, void const* data, size_t size) noexcept = 0;
//! \brief Begin the refit process from the loaded ONNX model.
//!
//! \return true if refit was successful
//! \see getNbErrors() getError() loadModelProto()
//!
virtual bool refitModelProto() noexcept = 0;
};
} // namespace nvonnxparser
extern "C" TENSORRTAPI void* createNvOnnxParser_INTERNAL(void* network, void* logger, int version) noexcept;
extern "C" TENSORRTAPI void* createNvOnnxParserRefitter_INTERNAL(
void* refitter, void* logger, int32_t version) noexcept;
extern "C" TENSORRTAPI int getNvOnnxParserVersion() noexcept;
namespace nvonnxparser
{
namespace
{
//!
//! \brief Create a new parser object
//!
//! \param network The network definition that the parser will write to
//! \param logger The logger to use
//! \return a new parser object or NULL if an error occurred
//!
//! Any input dimensions that are constant should not be changed after parsing,
//! because correctness of the translation may rely on those constants.
//! Changing a dynamic input dimension, i.e. one that translates to -1 in
//! TensorRT, to a constant is okay if the constant is consistent with the model.
//! Each instance of the parser is designed to only parse one ONNX model once.
//!
//! \see IParser
//!
inline IParser* createParser(nvinfer1::INetworkDefinition& network, nvinfer1::ILogger& logger) noexcept
{
return static_cast<IParser*>(createNvOnnxParser_INTERNAL(&network, &logger, NV_ONNX_PARSER_VERSION));
}
//!
//! \brief Create a new ONNX refitter object
//!
//! \param refitter The Refitter object used to refit the model
//! \param logger The logger to use
//! \return a new ParserRefitter object or NULL if an error occurred
//!
//! \see IParserRefitter
//!
inline IParserRefitter* createParserRefitter(nvinfer1::IRefitter& refitter, nvinfer1::ILogger& logger) noexcept
{
return static_cast<IParserRefitter*>(
createNvOnnxParserRefitter_INTERNAL(&refitter, &logger, NV_ONNX_PARSER_VERSION));
}
} // namespace
} // namespace nvonnxparser
#endif // NV_ONNX_PARSER_H
python/include/impl/NvInferPythonPlugin.h → include/impl/NvInferPythonPlugin.h
-1
View File
@@ -115,7 +115,6 @@ protected:
virtual ~ISymExprs() noexcept = default;
};
//! \enum QuickPluginCreationRequest
//! \brief Communicates preference when a quickly deployable plugin is to be added to the network
enum class QuickPluginCreationRequest : int32_t
parsers/onnx
+1 -1
Submodule parsers/onnx updated: 9a9f7883dd...c72727708d
plugin/CMakeLists.txt
+39 -18
View File
@@ -73,6 +73,7 @@ set(TRT_PLUGIN_NAMES
)
if(${TRT_BUILD_INCLUDE_BERT_QKV_PLUGIN})
list(APPEND TRT_PLUGIN_NAMES
bertQKVToContextPlugin
@@ -117,20 +118,23 @@ endforeach()
target_compile_options(trt_plugins PUBLIC $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
target_compile_options(trt_vc_plugins PUBLIC $<$<COMPILE_LANGUAGE:CUDA>:--expt-relaxed-constexpr>)
target_compile_definitions(trt_vc_plugins PRIVATE
COMPILE_VFC_PLUGIN=1
)
# Create all the library targets, reusing the objects we've compiled in the first step.
add_library(tensorrt_plugins SHARED $<TARGET_OBJECTS:trt_plugins>)
add_library(tensorrt_plugins_internal SHARED $<TARGET_OBJECTS:trt_plugins>)
add_library(tensorrt_plugins_static STATIC $<TARGET_OBJECTS:trt_plugins>)
add_library(tensorrt_vc_plugins SHARED $<TARGET_OBJECTS:trt_vc_plugins>)
add_library(tensorrt_vc_plugins_static STATIC $<TARGET_OBJECTS:trt_vc_plugins>)
add_library(tensorrt_plugins SHARED)
add_library(tensorrt_plugins_internal SHARED)
add_library(tensorrt_plugins_static STATIC)
foreach(lib tensorrt_plugins tensorrt_plugins_internal tensorrt_plugins_static)
target_link_libraries(${lib} PRIVATE trt_plugins)
endforeach()
target_compile_definitions(tensorrt_vc_plugins PRIVATE
COMPILE_VFC_PLUGIN=1
)
target_compile_definitions(tensorrt_vc_plugins_static PRIVATE
COMPILE_VFC_PLUGIN=1
)
add_library(tensorrt_vc_plugins SHARED)
add_library(tensorrt_vc_plugins_static STATIC)
foreach(lib tensorrt_vc_plugins tensorrt_vc_plugins_static)
target_link_libraries(${lib} PRIVATE trt_vc_plugins)
endforeach()
if (NOT MSVC)
set(trt_plugins_link_options
@@ -168,8 +172,10 @@ set_target_properties(
OUTPUT_NAME nvinfer_plugin
VERSION ${TensorRT_VERSION}
SOVERSION ${TRT_MAJOR}
LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports.map)
LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports.map
)
update_windows_output_name(tensorrt_plugins ${TRT_MAJOR} ${TRT_MINOR})
if (NOT MSVC)
set(trt_plugins_internal_link_options
@@ -198,8 +204,10 @@ set_target_properties(
OUTPUT_NAME nvinfer_plugin_internal
VERSION ${TensorRT_VERSION}
SOVERSION ${TRT_MAJOR}
LINK_DEPENDS ${TensorRT_SOURCE_DIR}/Exports-plugin_internal.map)
LINK_DEPENDS ${TensorRT_SOURCE_DIR}/Exports-plugin_internal.map
)
update_windows_output_name(tensorrt_plugins_internal ${TRT_MAJOR} ${TRT_MINOR})
### Static Plugin Setup
set(trt_plugin_static_dependencies
@@ -225,6 +233,8 @@ if(NOT ${TRT_BUILD_ENABLE_STATIC_LIBS})
)
endif()
smoke_test_static_lib(tensorrt_plugins_static)
### VC Plugin Setup
if (NOT MSVC)
set(trt_vc_plugins_link_options
@@ -242,7 +252,7 @@ endif()
# Target properties for tensorrt_vc_plugins
# This library includes a minimal subset of the plugins used for version compatibility.
target_include_directories(tensorrt_vc_plugins PRIVATE ${trt_plugin_include_dirs})
target_link_libraries(tensorrt_vc_plugins PRIVATE ${trt_plugin_dependencies})
target_link_libraries(tensorrt_vc_plugins PRIVATE trt_global_definitions)
target_link_options(tensorrt_vc_plugins PRIVATE ${trt_vc_plugins_link_options})
set_target_properties(
@@ -252,7 +262,10 @@ set_target_properties(
OUTPUT_NAME nvinfer_vc_plugin
VERSION ${TensorRT_VERSION}
SOVERSION ${TRT_MAJOR}
LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports-vfc_plugin.map)
LINK_DEPENDS ${CMAKE_CURRENT_LIST_DIR}/exports-vfc_plugin.map
)
update_windows_output_name(tensorrt_vc_plugins ${TRT_MAJOR} ${TRT_MINOR})
### VC Plugin Static Setup
target_include_directories(tensorrt_vc_plugins_static PRIVATE ${trt_plugin_include_dirs})
@@ -273,16 +286,23 @@ if(NOT ${TRT_BUILD_ENABLE_STATIC_LIBS})
)
endif()
smoke_test_static_lib(tensorrt_vc_plugins_static)
if(${TRT_BUILD_STUB_LIBS} AND NOT MSVC)
create_stub_lib(tensorrt_plugins)
create_stub_lib(tensorrt_vc_plugins)
endif()
installLibraries(
TARGETS tensorrt_plugins tensorrt_plugins_static tensorrt_vc_plugins tensorrt_vc_plugins_static
OPTIONAL
COMPONENT release
COMPONENT external
)
installLibraries(
TARGETS tensorrt_plugins_internal
OPTIONAL
COMPONENT full
COMPONENT internal
)
else() # TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW
@@ -387,6 +407,7 @@ endif()
include_directories(common common/kernels ${CMAKE_SOURCE_DIR}/third_party)
foreach(PLUGIN_ITER ${PLUGIN_LISTS})
include_directories(${PLUGIN_ITER})
add_subdirectory(${PLUGIN_ITER})
plugin/batchedNMSPlugin/gatherNMSOutputs.cu
+1 -1
View File
@@ -174,7 +174,7 @@ struct nmsOutLaunchConfig
, function(function)
{
}
bool operator==(const nmsOutLaunchConfig& other)
bool operator==(nmsOutLaunchConfig const& other) const
{
return t_bbox == other.t_bbox && t_score == other.t_score;
}
plugin/bertQKVToContextPlugin/CMakeLists.txt
+1
View File
@@ -49,3 +49,4 @@ if (NOT DEFINED ${TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW})
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE)
endif()
plugin/bertQKVToContextPlugin/CustomQKVToContextPluginDynamic_PluginConfig.yaml
+2 -2
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 2022-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -69,7 +69,7 @@ versions:
- hidden_size
- num_heads
- has_mask
golden_io_path: "plugin/bertQKVToContextPlugin/CustomQKVToContextPluginDynamic_PluginGoldenIO.json"
golden_io_path: "plugin/CustomQKVToContextPluginDynamic_PluginGoldenIO.json"
abs_tol: 1e-5
rel_tol: 1e-5
fp16_atol: 1e-2
plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPlugin.cpp
+1 -1
View File
@@ -54,7 +54,7 @@ QKVToContextInterleavedPlugin::QKVToContextInterleavedPlugin(std::string const&
, mQkvScale(qkvScale)
, mCtxScale(ctxScale)
{
mSM = getSMVersion();
mSM = getSmVersion();
mUseInt8ScaleMax = static_cast<int32_t>(useInt8ScaleMax);
mUseExplicitInt8 = static_cast<int32_t>(useExplicitInt8);
// variable sequence length is only supported with the fused MHA kernels
plugin/bertQKVToContextPlugin/qkvToContextInt8InterleavedPluginLegacy.cpp
+1 -1
View File
@@ -57,7 +57,7 @@ QKVToContextInterleavedPluginLegacy::QKVToContextInterleavedPluginLegacy(std::st
, mQkvScale(qkvScale)
, mCtxScale(ctxScale)
{
mSM = getSMVersion();
mSM = getSmVersion();
// variable sequence length is only supported with the fused MHA kernels
// we should not override mS!
bool isSMSupported = elem(mSM,
plugin/bertQKVToContextPlugin/qkvToContextPlugin.cpp
+3 -4
View File
@@ -66,7 +66,7 @@ QKVToContextPluginDynamic::QKVToContextPluginDynamic(const std::string name, con
{
mHasImask = static_cast<int32_t>(hasImask);
mSM = getSMVersion();
mSM = getSmVersion();
}
QKVToContextPluginDynamic::QKVToContextPluginDynamic(const std::string name, const DataType type, const int32_t S,
@@ -772,7 +772,7 @@ QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(std::string const name,
, mDqProbs(dqProbs)
, mHdim(HDIM)
{
mSM = getSMVersion();
mSM = getSmVersion();
mUseVarSeqlen = static_cast<int32_t>(varSeqlen);
mUseInt8ScaleMax = static_cast<int32_t>(useInt8ScaleMax);
mHasImask = static_cast<int32_t>(hasImask);
@@ -802,7 +802,7 @@ QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(std::string const name,
, mDqProbs(dqProbs)
, mHdim(HDIM)
{
mSM = getSMVersion();
mSM = getSmVersion();
mUseVarSeqlen = static_cast<int32_t>(varSeqlen);
mUseInt8ScaleMax = static_cast<int32_t>(useInt8ScaleMax);
mHasImask = static_cast<int32_t>(hasImask);
@@ -825,7 +825,6 @@ QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(std::string const name,
mDispatcher->deserialize(runnerStateBuffer, length);
}
IPluginCapability* QKVToContextVarSeqlenPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept
{
try
plugin/bertQKVToContextPlugin/qkvToContextPluginLegacy.cpp
+2 -2
View File
@@ -65,7 +65,7 @@ QKVToContextPluginDynamicLegacy::QKVToContextPluginDynamicLegacy(std::string con
, mDqProbs(dqProbs)
{
mSM = getSMVersion();
mSM = getSmVersion();
}
QKVToContextPluginDynamicLegacy::QKVToContextPluginDynamicLegacy(
@@ -603,7 +603,7 @@ QKVToContextVarSeqlenPluginLegacy::QKVToContextVarSeqlenPluginLegacy(std::string
, mUseVarSeqlen(varSeqlen)
, mUseInt8ScaleMax(useInt8ScaleMax)
{
mSM = getSMVersion();
mSM = getSmVersion();
if (varSeqlen)
{
plugin/clipPlugin/CMakeLists.txt
+1
View File
@@ -21,3 +21,4 @@ add_plugin_source(
clipPlugin.cpp
clipPlugin.h
)
plugin/common/CMakeLists.txt
+6 -1
View File
@@ -14,7 +14,8 @@
# See the License for the specific language governing permissions and
# limitations under the License.
#
add_plugin_source(
# Base source files
set(PLUGIN_COMMON_SOURCES
bboxUtils.h
bertCommon.h
checkMacrosPlugin.cpp
@@ -37,10 +38,14 @@ add_plugin_source(
plugin.cpp
plugin.h
reducedMathPlugin.cpp
scopedCudaStream.h
serialize.hpp
templates.h
)
add_plugin_source(${PLUGIN_COMMON_SOURCES})
add_subdirectory(kernels)
# Promote PLUGIN_SOURCES and PLUGIN_CU_SOURCES added by `add_subdirectory` to this file's parent.
plugin/common/bertCommon.h
+1 -1
View File
@@ -139,7 +139,7 @@ inline int32_t getMHAMaskPackedSize(int32_t smVersion, nvinfer1::DataType dataTy
return packedSize;
}
inline uint32_t getElementSize(nvinfer1::DataType t) noexcept
inline uint32_t getElementSize(nvinfer1::DataType t)
{
switch (t)
{
plugin/common/cudaDriverWrapper.cpp
+84 -12
View File
@@ -50,18 +50,45 @@ CUDADriverWrapper::CUDADriverWrapper()
return ret;
};
*(void**) (&_cuGetErrorName) = load_sym(handle, "cuGetErrorName");
*(void**) (&_cuFuncSetAttribute) = load_sym(handle, "cuFuncSetAttribute");
*(void**) (&_cuLinkComplete) = load_sym(handle, "cuLinkComplete");
*(void**) (&_cuModuleUnload) = load_sym(handle, "cuModuleUnload");
*(void**) (&_cuLinkDestroy) = load_sym(handle, "cuLinkDestroy");
*(void**) (&_cuModuleLoadData) = load_sym(handle, "cuModuleLoadData");
*(void**) (&_cuLinkCreate) = load_sym(handle, "cuLinkCreate_v2");
*(void**) (&_cuModuleGetFunction) = load_sym(handle, "cuModuleGetFunction");
*(void**) (&_cuLinkAddFile) = load_sym(handle, "cuLinkAddFile_v2");
*(void**) (&_cuLinkAddData) = load_sym(handle, "cuLinkAddData_v2");
*(void**) (&_cuLaunchCooperativeKernel) = load_sym(handle, "cuLaunchCooperativeKernel");
*(void**) (&_cuLaunchKernel) = load_sym(handle, "cuLaunchKernel");
_cuGetErrorName = reinterpret_cast<CUresult (*)(CUresult, char const**)>(load_sym(handle, "cuGetErrorName"));
_cuGetErrorString = reinterpret_cast<CUresult (*)(CUresult, char const**)>(load_sym(handle, "cuGetErrorString"));
_cuFuncSetAttribute = reinterpret_cast<CUresult (*)(CUfunction, CUfunction_attribute, int32_t)>(
load_sym(handle, "cuFuncSetAttribute"));
_cuLinkComplete = reinterpret_cast<CUresult (*)(CUlinkState, void**, size_t*)>(load_sym(handle, "cuLinkComplete"));
_cuModuleUnload = reinterpret_cast<CUresult (*)(CUmodule)>(load_sym(handle, "cuModuleUnload"));
_cuLinkDestroy = reinterpret_cast<CUresult (*)(CUlinkState)>(load_sym(handle, "cuLinkDestroy"));
_cuModuleLoadData = reinterpret_cast<CUresult (*)(CUmodule*, void const*)>(load_sym(handle, "cuModuleLoadData"));
_cuLinkCreate = reinterpret_cast<CUresult (*)(uint32_t, CUjit_option*, void**, CUlinkState*)>(
load_sym(handle, "cuLinkCreate_v2"));
_cuModuleGetFunction
= reinterpret_cast<CUresult (*)(CUfunction*, CUmodule, char const*)>(load_sym(handle, "cuModuleGetFunction"));
_cuLinkAddFile
= reinterpret_cast<CUresult (*)(CUlinkState, CUjitInputType, char const*, uint32_t, CUjit_option*, void**)>(
load_sym(handle, "cuLinkAddFile_v2"));
_cuLinkAddData = reinterpret_cast<CUresult (*)(CUlinkState, CUjitInputType, void*, size_t, char const*, uint32_t,
CUjit_option*, void**)>(load_sym(handle, "cuLinkAddData_v2"));
_cuLaunchCooperativeKernel = reinterpret_cast<CUresult (*)(CUfunction, uint32_t, uint32_t, uint32_t, uint32_t,
uint32_t, uint32_t, uint32_t, CUstream, void**)>(load_sym(handle, "cuLaunchCooperativeKernel"));
_cuLaunchKernel = reinterpret_cast<CUresult (*)(CUfunction, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t,
uint32_t, uint32_t, CUstream, void**, void**)>(load_sym(handle, "cuLaunchKernel"));
#if CUDA_VERSION >= 11060
_cuLaunchKernelEx
= reinterpret_cast<CUresult (*)(CUlaunchConfig const* config, CUfunction f, void** kernelParams, void** extra)>(
dllGetSym(handle, "cuLaunchKernelEx"));
#endif
#if CUDA_VERSION >= 12000
_cuTensorMapEncodeTiled
= reinterpret_cast<CUresult (*)(CUtensorMap*, CUtensorMapDataType, cuuint32_t, void const*, cuuint64_t const*,
cuuint64_t const*, cuuint32_t const*, cuuint32_t const*, CUtensorMapInterleave, CUtensorMapSwizzle,
CUtensorMapL2promotion, CUtensorMapFloatOOBfill)>(load_sym(handle, "cuTensorMapEncodeTiled"));
#endif
_cuMemcpyDtoH = reinterpret_cast<CUresult (*)(void*, CUdeviceptr, size_t)>(load_sym(handle, "cuMemcpyDtoH_v2"));
_cuDeviceGetAttribute = reinterpret_cast<CUresult (*)(int32_t*, CUdevice_attribute, CUdevice)>(
load_sym(handle, "cuDeviceGetAttribute"));
#if CUDA_VERSION >= 12000
_cuOccupancyMaxActiveClusters = reinterpret_cast<CUresult (*)(int32_t*, CUfunction, CUlaunchConfig const*)>(
load_sym(handle, "cuOccupancyMaxActiveClusters"));
#endif
}
CUDADriverWrapper::~CUDADriverWrapper()
@@ -74,6 +101,11 @@ CUresult CUDADriverWrapper::cuGetErrorName(CUresult error, char const** pStr) co
return (*_cuGetErrorName)(error, pStr);
}
CUresult CUDADriverWrapper::cuGetErrorString(CUresult error, char const** pStr) const
{
return (*_cuGetErrorString)(error, pStr);
}
CUresult CUDADriverWrapper::cuFuncSetAttribute(CUfunction hfunc, CUfunction_attribute attrib, int32_t value) const
{
return (*_cuFuncSetAttribute)(hfunc, attrib, value);
@@ -137,3 +169,43 @@ CUresult CUDADriverWrapper::cuLaunchKernel(CUfunction f, uint32_t gridDimX, uint
return (*_cuLaunchKernel)(
f, gridDimX, gridDimY, gridDimZ, blockDimX, blockDimY, blockDimZ, sharedMemBytes, hStream, kernelParams, extra);
}
#if CUDA_VERSION >= 11060
CUresult CUDADriverWrapper::cuLaunchKernelEx(
CUlaunchConfig const* config, CUfunction f, void** kernelParams, void** extra) const
{
PLUGIN_ASSERT(_cuLaunchKernelEx != nullptr);
return (*_cuLaunchKernelEx)(config, f, kernelParams, extra);
}
#endif
#if CUDA_VERSION >= 12000
CUresult CUDADriverWrapper::cuTensorMapEncodeTiled(CUtensorMap* tensorMap, CUtensorMapDataType tensorDataType,
cuuint32_t tensorRank, void const* globalAddress, cuuint64_t const* globalDim, cuuint64_t const* globalStrides,
cuuint32_t const* boxDim, cuuint32_t const* elementStrides, CUtensorMapInterleave interleave,
CUtensorMapSwizzle swizzle, CUtensorMapL2promotion l2Promotion, CUtensorMapFloatOOBfill oobFill) const
{
PLUGIN_ASSERT(_cuTensorMapEncodeTiled != nullptr);
return (*_cuTensorMapEncodeTiled)(tensorMap, tensorDataType, tensorRank, globalAddress, globalDim, globalStrides,
boxDim, elementStrides, interleave, swizzle, l2Promotion, oobFill);
}
#endif
CUresult CUDADriverWrapper::cuMemcpyDtoH(void* dstHost, CUdeviceptr srcDevice, size_t ByteCount) const
{
return (*_cuMemcpyDtoH)(dstHost, srcDevice, ByteCount);
}
CUresult CUDADriverWrapper::cuDeviceGetAttribute(int32_t* pi, CUdevice_attribute attrib, CUdevice dev) const
{
return (*_cuDeviceGetAttribute)(pi, attrib, dev);
}
#if CUDA_VERSION >= 12000
CUresult CUDADriverWrapper::cuOccupancyMaxActiveClusters(
int32_t* maxActiveClusters, CUfunction f, CUlaunchConfig const* config) const
{
PLUGIN_ASSERT(_cuOccupancyMaxActiveClusters != nullptr);
return (*_cuOccupancyMaxActiveClusters)(maxActiveClusters, f, config);
}
#endif
plugin/common/cudaDriverWrapper.h
+141 -1
View File
@@ -1,5 +1,5 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* 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");
@@ -21,6 +21,8 @@
#include <cstdint>
#include <cstdio>
#include <cuda.h>
#include <memory>
#include <utility>
#define cuErrCheck(stat, wrap) \
{ \
@@ -42,6 +44,8 @@ public:
CUresult cuGetErrorName(CUresult error, char const** pStr) const;
CUresult cuGetErrorString(CUresult error, char const** pStr) const;
CUresult cuFuncSetAttribute(CUfunction hfunc, CUfunction_attribute attrib, int32_t value) const;
CUresult cuLinkComplete(CUlinkState state, void** cubinOut, size_t* sizeOut) const;
@@ -70,9 +74,29 @@ public:
uint32_t blockDimY, uint32_t blockDimZ, uint32_t sharedMemBytes, CUstream hStream, void** kernelParams,
void** extra) const;
#if CUDA_VERSION >= 11060
CUresult cuLaunchKernelEx(CUlaunchConfig const* config, CUfunction f, void** kernelParams, void** extra) const;
#endif
#if CUDA_VERSION >= 12000
CUresult cuTensorMapEncodeTiled(CUtensorMap* tensorMap, CUtensorMapDataType tensorDataType, cuuint32_t tensorRank,
void const* globalAddress, cuuint64_t const* globalDim, cuuint64_t const* globalStrides,
cuuint32_t const* boxDim, cuuint32_t const* elementStrides, CUtensorMapInterleave interleave,
CUtensorMapSwizzle swizzle, CUtensorMapL2promotion l2Promotion, CUtensorMapFloatOOBfill oobFill) const;
#endif
CUresult cuMemcpyDtoH(void* dstHost, CUdeviceptr srcDevice, size_t ByteCount) const;
CUresult cuDeviceGetAttribute(int32_t* pi, CUdevice_attribute attrib, CUdevice dev) const;
#if CUDA_VERSION >= 12000
CUresult cuOccupancyMaxActiveClusters(int32_t* maxActiveClusters, CUfunction f, CUlaunchConfig const* config) const;
#endif
private:
void* handle;
CUresult (*_cuGetErrorName)(CUresult, char const**);
CUresult (*_cuGetErrorString)(CUresult, char const**);
CUresult (*_cuFuncSetAttribute)(CUfunction, CUfunction_attribute, int32_t);
CUresult (*_cuLinkComplete)(CUlinkState, void**, size_t*);
CUresult (*_cuModuleUnload)(CUmodule);
@@ -88,6 +112,20 @@ private:
CUresult (*_cuLaunchKernel)(CUfunction f, uint32_t gridDimX, uint32_t gridDimY, uint32_t gridDimZ,
uint32_t blockDimX, uint32_t blockDimY, uint32_t blockDimZ, uint32_t sharedMemBytes, CUstream hStream,
void** kernelParams, void** extra);
#if CUDA_VERSION >= 11060
CUresult (*_cuLaunchKernelEx)(CUlaunchConfig const* config, CUfunction f, void** kernelParams, void** extra);
#endif
#if CUDA_VERSION >= 12000
CUresult (*_cuTensorMapEncodeTiled)(CUtensorMap* tensorMap, CUtensorMapDataType tensorDataType,
cuuint32_t tensorRank, void const* globalAddress, cuuint64_t const* globalDim, cuuint64_t const* globalStrides,
cuuint32_t const* boxDim, cuuint32_t const* elementStrides, CUtensorMapInterleave interleave,
CUtensorMapSwizzle swizzle, CUtensorMapL2promotion l2Promotion, CUtensorMapFloatOOBfill oobFill);
#endif
CUresult (*_cuMemcpyDtoH)(void* dstHost, CUdeviceptr srcDevice, size_t ByteCount);
CUresult (*_cuDeviceGetAttribute)(int32_t*, CUdevice_attribute attrib, CUdevice dev);
#if CUDA_VERSION >= 12000
CUresult (*_cuOccupancyMaxActiveClusters)(int32_t*, CUfunction f, CUlaunchConfig const* config);
#endif
};
inline void cuErrCheck_(CUresult stat, CUDADriverWrapper const& wrap, char const* file, int32_t line)
@@ -106,6 +144,108 @@ constexpr int32_t getCudaLibVersionMaj() noexcept
return CUDA_VERSION / 1000U;
}
// RAII wrapper for CUDA module
// Automatically manages CUDA module lifecycle - loading in constructor, unloading in destructor
// Provides safe module management and prevents memory leaks
class CudaModule
{
public:
// Default constructor - creates an uninitialized module
CudaModule() noexcept
: mModule(nullptr)
, mDriverWrapper(nullptr)
, mLoadResult(CUDA_ERROR_NOT_INITIALIZED)
{
}
// Constructor that loads a CUDA module from data
CudaModule(CUDADriverWrapper const& driverWrapper, void const* image)
: mModule(nullptr)
, mDriverWrapper(&driverWrapper)
, mLoadResult(CUDA_ERROR_NOT_INITIALIZED)
{
mLoadResult = mDriverWrapper->cuModuleLoadData(&mModule, image);
}
// Destructor - automatically unloads the module
~CudaModule()
{
if (mModule != nullptr && mDriverWrapper != nullptr)
{
mDriverWrapper->cuModuleUnload(mModule);
}
}
// Delete copy constructor and assignment
CudaModule(CudaModule const&) = delete;
CudaModule& operator=(CudaModule const&) = delete;
// Move constructor
CudaModule(CudaModule&& other) noexcept
: mModule(std::exchange(other.mModule, nullptr))
, mDriverWrapper(std::exchange(other.mDriverWrapper, nullptr))
, mLoadResult(std::exchange(other.mLoadResult, CUDA_ERROR_NOT_INITIALIZED))
{
}
// Move assignment
CudaModule& operator=(CudaModule&& other) noexcept
{
CudaModule tmp{std::move(other)};
std::swap(mModule, tmp.mModule);
std::swap(mDriverWrapper, tmp.mDriverWrapper);
std::swap(mLoadResult, tmp.mLoadResult);
return *this;
}
//! Get the underlying CUDA module handle (`CUmodule`; a raw pointer).
[[nodiscard]] CUmodule get() noexcept
{
return mModule;
}
//! Get the underlying const CUDA module handle (`CUmodule`; a raw pointer).
//! \note Since `CUmodule` is a raw pointer, we remove the pointer from the
//! type, add `const`, and re-add the pointer.
[[nodiscard]] std::remove_pointer_t<CUmodule> const* get() const noexcept
{
return mModule;
}
// Implicit conversion to CUmodule
operator CUmodule() const noexcept
{
return mModule;
}
// Check if the module is valid
bool isValid() const noexcept
{
return mModule != nullptr && mLoadResult == CUDA_SUCCESS;
}
// Get function from module
CUresult getFunction(CUfunction* hfunc, char const* name) const
{
if (mModule != nullptr && mDriverWrapper != nullptr)
{
return mDriverWrapper->cuModuleGetFunction(hfunc, mModule, name);
}
return CUDA_ERROR_INVALID_HANDLE;
}
private:
CUmodule mModule;
CUDADriverWrapper const* mDriverWrapper;
CUresult mLoadResult;
};
// Helper function to create a unique_ptr to CudaModule
inline std::unique_ptr<CudaModule> makeCudaModule(CUDADriverWrapper const& driverWrapper, void const* image)
{
return std::make_unique<CudaModule>(driverWrapper, image);
}
} // namespace nvinfer1
#endif // CUDA_DRIVER_WRAPPER_H
plugin/common/cudnnWrapper.cpp
+2 -2
View File
@@ -91,8 +91,8 @@ void* CudnnWrapper::tryLoadingCudnn(char const* callerPluginName)
static constexpr int32_t kSM_BLACKWELL_100 = 100;
std::string errorMsgCudnnSupport
= "At least one plugin (" + std::string(callerPluginName) + ") that requires cuDNN is being used. TensorRT does not provide cuDNN support for Blackwell (compute capability: 10.0) and later architectures. Detected compute capability: " + std::to_string(nvinfer1::plugin::getSMVersion() / 10) + "." + std::to_string(nvinfer1::plugin::getSMVersion() % 10) + ". Please run on a platform with compute capability < 10.0, or use an alternative to " + std::string(callerPluginName) + ".";
PLUGIN_VALIDATE(nvinfer1::plugin::getSMVersion() < kSM_BLACKWELL_100, errorMsgCudnnSupport.c_str());
= "At least one plugin (" + std::string(callerPluginName) + ") that requires cuDNN is being used. TensorRT does not provide cuDNN support for Blackwell (compute capability: 10.0) and later architectures. Detected compute capability: " + std::to_string(nvinfer1::plugin::getSmVersion() / 10) + "." + std::to_string(nvinfer1::plugin::getSmVersion() % 10) + ". Please run on a platform with compute capability < 10.0, or use an alternative to " + std::string(callerPluginName) + ".";
PLUGIN_VALIDATE(nvinfer1::plugin::getSmVersion() < kSM_BLACKWELL_100, errorMsgCudnnSupport.c_str());
#endif // CUDART_VERSION >= 12070 && CUDNN_MAJOR == 8
void* cudnnLib = dllOpen(kCUDNN_PLUGIN_LIBNAME.c_str());
std::string errorMsg = "Failed to load " + kCUDNN_PLUGIN_LIBNAME + ".";
plugin/common/kernels/allClassNMS.cu
+1 -1
View File
@@ -360,7 +360,7 @@ struct nmsLaunchConfigSSD
, function(function)
{
}
bool operator==(const nmsLaunchConfigSSD& other)
bool operator==(nmsLaunchConfigSSD const& other) const
{
return t_score == other.t_score && t_bbox == other.t_bbox;
}
plugin/common/kernels/bboxDeltas2Proposals.cu
+1 -1
View File
@@ -236,7 +236,7 @@ struct bd2pLaunchConfig
{
}
bool operator==(const bd2pLaunchConfig& other)
bool operator==(bd2pLaunchConfig const& other) const
{
return t_deltas == other.t_deltas && l_deltas == other.l_deltas && t_proposals == other.t_proposals && l_proposals == other.l_proposals && t_scores == other.t_scores && l_scores == other.l_scores;
}
plugin/common/kernels/decodeBBoxes.cu
+1 -1
View File
@@ -355,7 +355,7 @@ struct dbbLaunchConfig
, function(function)
{
}
bool operator==(const dbbLaunchConfig& other)
bool operator==(dbbLaunchConfig const& other) const
{
return t_bbox == other.t_bbox;
}
plugin/common/kernels/gatherTopDetections.cu
+1 -1
View File
@@ -173,7 +173,7 @@ struct gtdLaunchConfig
, function(function)
{
}
bool operator==(const gtdLaunchConfig& other)
bool operator==(gtdLaunchConfig const& other) const
{
return t_bbox == other.t_bbox && t_score == other.t_score;
}
plugin/common/kernels/nmsLayer.cu
+1 -1
View File
@@ -390,7 +390,7 @@ struct nmsLaunchConfig
{
}
bool operator==(const nmsLaunchConfig& other)
bool operator==(nmsLaunchConfig const& other) const
{
return (t_fgScores == other.t_fgScores) && (l_fgScores == other.l_fgScores) && (t_proposals == other.t_proposals) && (l_proposals == other.l_proposals) && (t_rois == other.t_rois);
}
plugin/common/kernels/permuteData.cu
+1 -1
View File
@@ -80,7 +80,7 @@ struct pdLaunchConfig
, function(function)
{
}
bool operator==(const pdLaunchConfig& other)
bool operator==(pdLaunchConfig const& other) const
{
return t_data == other.t_data;
}
plugin/common/kernels/proposalKernel.cu
+1 -1
View File
@@ -420,7 +420,7 @@ struct nmsLaunchConfig
{
}
bool operator==(const nmsLaunchConfig& other)
bool operator==(nmsLaunchConfig const& other) const
{
return (t_fgScores == other.t_fgScores) && (l_fgScores == other.l_fgScores)
&& (t_proposals == other.t_proposals) && (l_proposals == other.l_proposals)
plugin/common/kernels/roiPooling.cu
+1 -1
View File
@@ -285,7 +285,7 @@ struct roiFwdLaunchConfig
{
}
bool operator==(const roiFwdLaunchConfig& other)
bool operator==(roiFwdLaunchConfig const& other) const
{
return (t_rois == other.t_rois)
&& (t_featureMap == other.t_featureMap)
plugin/common/kernels/sortScoresPerClass.cu
+1 -1
View File
@@ -218,7 +218,7 @@ struct sspcLaunchConfig
, function(function)
{
}
bool operator==(const sspcLaunchConfig& other)
bool operator==(sspcLaunchConfig const& other) const
{
return t_score == other.t_score;
}
plugin/common/kernels/sortScoresPerImage.cu
+1 -1
View File
@@ -79,7 +79,7 @@ struct sspiLaunchConfig
, function(function)
{
}
bool operator==(const sspiLaunchConfig& other)
bool operator==(sspiLaunchConfig const& other) const
{
return t_score == other.t_score;
}
plugin/common/plugin.h
+3 -3
View File
@@ -107,7 +107,7 @@ OutType read(BufferType const*& buffer)
return val;
}
inline int32_t getTrtSMVersionDec(int32_t majorVersion, int32_t minorVersion)
inline int32_t getTrtSmVersionDec(int32_t majorVersion, int32_t minorVersion)
{
return majorVersion * 10 + minorVersion;
}
@@ -132,12 +132,12 @@ struct DeviceComputeCapability
}
};
inline int32_t getSMVersion()
inline int32_t getSmVersion()
{
int32_t device{-1};
PLUGIN_CHECK_CUDA(cudaGetDevice(&device));
auto const cc = DeviceComputeCapability::forDevice(device);
return getTrtSMVersionDec(cc.major, cc.minor);
return getTrtSmVersionDec(cc.major, cc.minor);
}
// Check that all required field names are present in the PluginFieldCollection.
plugin/common/scopedCudaStream.h
+120
View File
@@ -0,0 +1,120 @@
/*
* Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef TRT_SCOPED_CUDA_STREAM_H
#define TRT_SCOPED_CUDA_STREAM_H
#include <cstdint>
#include <cuda_runtime_api.h>
#include <memory>
#include <stdexcept>
namespace nvinfer1
{
namespace pluginInternal
{
// RAII wrapper for CUDA stream
// Automatically manages CUDA stream lifecycle - creation in constructor, destruction in destructor
// Provides safe stream management and prevents memory leaks
class ScopedCudaStream
{
public:
// Constructor that creates a new CUDA stream with default flags
ScopedCudaStream()
: mStream(nullptr)
{
cudaError_t result = cudaStreamCreate(&mStream);
if (result != cudaSuccess)
{
throw std::runtime_error("Failed to create CUDA stream: " + std::string(cudaGetErrorString(result)));
}
}
// Constructor that creates a new CUDA stream with custom flags
explicit ScopedCudaStream(uint32_t const flags)
: mStream(nullptr)
{
cudaError_t result = cudaStreamCreateWithFlags(&mStream, flags);
if (result != cudaSuccess)
{
mStream = nullptr;
}
}
// Destructor - automatically destroys the stream
~ScopedCudaStream()
{
if (mStream != nullptr)
{
cudaStreamDestroy(mStream);
}
}
// Delete copy constructor and assignment
ScopedCudaStream(ScopedCudaStream const&) = delete;
ScopedCudaStream& operator=(ScopedCudaStream const&) = delete;
// Move constructor
ScopedCudaStream(ScopedCudaStream&& other) noexcept
: mStream(std::exchange(other.mStream, nullptr))
{
}
// Move assignment
ScopedCudaStream& operator=(ScopedCudaStream&& other) noexcept
{
ScopedCudaStream tmp{std::move(other)};
std::swap(mStream, tmp.mStream);
return *this;
}
// Get the underlying CUDA stream handle
cudaStream_t get() const noexcept
{
return mStream;
}
// Implicit conversion to cudaStream_t
operator cudaStream_t() const noexcept
{
return mStream;
}
// Check if the stream is valid
bool isValid() const noexcept
{
return mStream != nullptr;
}
private:
cudaStream_t mStream;
};
// Helper function to create a unique_ptr to ScopedCudaStream
inline std::unique_ptr<ScopedCudaStream> makeScopedCudaStream(uint32_t const flags = cudaStreamDefault)
{
if (flags == cudaStreamDefault)
{
return std::make_unique<ScopedCudaStream>();
}
return std::make_unique<ScopedCudaStream>(flags);
}
} // namespace pluginInternal
} // namespace nvinfer1
#endif // TRT_SCOPED_CUDA_STREAM_H
plugin/cropAndResizePlugin/CMakeLists.txt
+1
View File
@@ -21,3 +21,4 @@ add_plugin_source(
cropAndResizePluginLegacy.cpp
cropAndResizePluginLegacy.h
)
plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginConfig.yaml
+2 -2
View File
@@ -52,7 +52,7 @@ versions:
attributes_required:
- crop_height
- crop_width
golden_io_path: "plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginGoldenIO.json"
golden_io_path: "plugin/CropAndResizeDynamic_PluginGoldenIO.json"
abs_tol: 1e-6
rel_tol: 1e-6
fp16_atol: 1e-3
@@ -111,7 +111,7 @@ versions:
attributes_required:
- crop_height
- crop_width
golden_io_path: "plugin/cropAndResizePlugin/CropAndResizeDynamic_PluginGoldenIO.json"
golden_io_path: "plugin/CropAndResizeDynamic_PluginGoldenIO.json"
abs_tol: 1e-6
rel_tol: 1e-6
fp16_atol: 1e-3
plugin/detectionLayerPlugin/CMakeLists.txt
+1
View File
@@ -19,3 +19,4 @@ add_plugin_source(
detectionLayerPlugin.cpp
detectionLayerPlugin.h
)
plugin/disentangledAttentionPlugin/CMakeLists.txt
+1
View File
@@ -23,3 +23,4 @@ add_plugin_source(
disentangledAttentionPluginLegacy.h
disentangledKernel.cu
)
plugin/disentangledAttentionPlugin/DisentangledAttentionPlugin_PluginConfig.yaml
+2 -2
View File
@@ -45,7 +45,7 @@ versions:
attributes_required:
- span
- factor
golden_io_path: "plugin/disentangledAttentionPlugin/DisentangledAttention_PluginGoldenIO.json"
golden_io_path: "plugin/DisentangledAttention_PluginGoldenIO.json"
abs_tol: 1e-5
rel_tol: 1e-5
fp16_atol: 1e-2
@@ -120,7 +120,7 @@ versions:
attributes_required:
- span
- factor
golden_io_path: "plugin/disentangledAttentionPlugin/DisentangledAttention_PluginGoldenIO.json"
golden_io_path: "plugin/DisentangledAttention_PluginGoldenIO.json"
abs_tol: 1e-5
rel_tol: 1e-5
fp16_atol: 1e-2
plugin/efficientNMSPlugin/CMakeLists.txt
+1
View File
@@ -31,3 +31,4 @@ if (NOT DEFINED ${TRT_BUILD_ENABLE_NEW_PLUGIN_FLOW})
set(PLUGIN_SOURCES ${PLUGIN_SOURCES} PARENT_SCOPE)
set(PLUGIN_CU_SOURCES ${PLUGIN_CU_SOURCES} PARENT_SCOPE)
endif()
plugin/efficientNMSPlugin/EfficientNMSPlugin_PluginConfig.yaml
+2 -2
View File
@@ -1,5 +1,5 @@
#
# SPDX-FileCopyrightText: Copyright (c) 2022-2023 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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");
@@ -81,7 +81,7 @@ versions:
- background_class
- score_activation
- box_coding
golden_io_path: "plugin/efficientNMSPlugin/EfficientNMSPlugin_PluginGoldenIO.json"
golden_io_path: "plugin/EfficientNMSPlugin_PluginGoldenIO.json"
abs_tol: 1e-5
rel_tol: 1e-5
configs:
plugin/embLayerNormPlugin/CMakeLists.txt
+2 -2
View File
@@ -9,8 +9,7 @@
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
@@ -28,3 +27,4 @@ add_plugin_source(
embLayerNormVarSeqlenPluginLegacy.cpp
embLayerNormVarSeqlenPluginLegacy.h
)

Some files were not shown because too many files have changed in this diff Show More