chore: import upstream snapshot with attribution

plugin/common/CMakeLists.txt

@@ -0,0 +1,47 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Base source files
+set(PLUGIN_COMMON_SOURCES
+    bboxUtils.h
+    bertCommon.h
+    checkMacrosPlugin.cpp
+    checkMacrosPlugin.h
+    common.cuh
+    cub_helper.h
+    cublasLtWrapper.cpp
+    cublasLtWrapper.h
+    cublasWrapper.cpp
+    cublasWrapper.h
+    cudaDriverWrapper.cpp
+    cudaDriverWrapper.h
+    cudnnWrapper.cpp
+    cudnnWrapper.h
+    dimsHelpers.h
+    half.h
+    mrcnn_config.h
+    nmsHelper.cpp
+    plugin.cpp
+    plugin.h
+    reducedMathPlugin.cpp
+    scopedCudaStream.h
+    serialize.hpp
+    templates.h
+)
+
+add_plugin_source(${PLUGIN_COMMON_SOURCES})
+
+add_subdirectory(kernels)

plugin/common/bboxUtils.h

@@ -0,0 +1,78 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef TRT_BBOX_UTILS_H
+#define TRT_BBOX_UTILS_H
+
+#include "common/plugin.h"
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename T>
+struct Bbox
+{
+    T xmin, ymin, xmax, ymax;
+    Bbox(T xmin, T ymin, T xmax, T ymax)
+        : xmin(xmin)
+        , ymin(ymin)
+        , xmax(xmax)
+        , ymax(ymax)
+    {
+    }
+    Bbox() = default;
+};
+
+template <typename T>
+struct BboxInfo
+{
+    T conf_score;
+    int32_t label;
+    int32_t bbox_idx;
+    bool kept;
+    BboxInfo(T conf_score, int32_t label, int32_t bbox_idx, bool kept)
+        : conf_score(conf_score)
+        , label(label)
+        , bbox_idx(bbox_idx)
+        , kept(kept)
+    {
+    }
+    BboxInfo() = default;
+};
+
+template <typename TFloat>
+bool operator<(Bbox<TFloat> const& lhs, Bbox<TFloat> const& rhs)
+{
+    return lhs.x1 < rhs.x1;
+}
+
+template <typename TFloat>
+bool operator==(Bbox<TFloat> const& lhs, Bbox<TFloat> const& rhs)
+{
+    return lhs.x1 == rhs.x1 && lhs.y1 == rhs.y1 && lhs.x2 == rhs.x2 && lhs.y2 == rhs.y2;
+}
+// }}}
+
+int8_t* alignPtr(int8_t* ptr, uintptr_t to);
+
+int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize);
+
+size_t dataTypeSize(nvinfer1::DataType dtype);
+
+void setUniformOffsets(cudaStream_t stream, int32_t num_segments, int32_t offset, int32_t* d_offsets);
+} // namespace plugin
+} // namespace nvinfer1
+#endif

plugin/common/bertCommon.h

@@ -0,0 +1,540 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cuda.h>
+#if CUDA_VERSION >= 10010
+
+#ifndef BERT_COMMON_H
+#define BERT_COMMON_H
+
+#include "NvInfer.h"
+#include "NvInferRuntimeCommon.h"
+#include "common/checkMacrosPlugin.h"
+#include "common/cublasWrapper.h"
+#include "common/plugin.h"
+#include <cuda_fp16.h>
+
+#include <algorithm>
+#include <cassert>
+#include <cuda_runtime_api.h>
+#include <memory>
+#include <numeric>
+#include <stdexcept>
+#include <vector>
+
+#define TRT_UNUSED (void)
+
+#define BERT_PRINT_DEBUG_MSG 0
+
+#if BERT_PRINT_DEBUG_MSG
+#define BERT_DEBUG_MSG(msg) (gLogVerbose << (msg) << std::endl)
+#define BERT_DEBUG_VALUE(key, value) (gLogVerbose << key << value << std::endl)
+#else
+#define BERT_DEBUG_MSG(msg) TRT_UNUSED(msg)
+#define BERT_DEBUG_VALUE(key, value)                                                                                   \
+    TRT_UNUSED(key);                                                                                                   \
+    TRT_UNUSED(value)
+#endif
+
+using half = __half;
+
+constexpr uint32_t BDIM = 1; // batch dimension
+constexpr uint32_t SDIM = 0; // seq len dimension
+constexpr uint32_t HDIM = 2; // hidden dimension
+
+constexpr int32_t kSM_75 = 75;
+constexpr int32_t kSM_80 = 80;
+constexpr int32_t kSM_86 = 86;
+constexpr int32_t kSM_87 = 87;
+constexpr int32_t kSM_89 = 89;
+constexpr int32_t kSM_90 = 90;
+constexpr int32_t kSM_100 = 100;
+constexpr int32_t kSM_120 = 120;
+
+// For full mask mode, we must produce the compressed mask format expected by the fused attention path. Currently, only
+// two sequence lengths are supported. We hard code the sizes here.
+// The number of threads per CTA: warps_m * warps_n * warps_k * 32;
+constexpr size_t threadsPerCta128 = 2 * 2 * 32;
+constexpr size_t threadsPerCta384 = 1 * 8 * 32;
+
+// The number of xmmas in the M dimension. We use one uint32_t per XMMA in the M dimension: (s + 16*warps_m - 1)
+// / (16*warps_m);
+constexpr size_t xmmasM128 = 4;
+constexpr size_t xmmasM384 = 24;
+
+// Packed mask size per batch. Layout is XMMAS_M * THREADS_PER_CTA.
+constexpr size_t unfusedMaskSize = 1;
+constexpr size_t packedMaskSize64 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize96 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize128 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize384 = xmmasM384 * threadsPerCta384;
+
+namespace nvinfer1
+{
+namespace pluginInternal
+{
+template <typename T>
+struct CudaDeleter
+{
+    void operator()(T* buf)
+    {
+        PLUGIN_CUASSERT(cudaFree(buf));
+    }
+};
+
+} // namespace pluginInternal
+namespace plugin
+{
+
+namespace bert
+{
+
+//! \brief Checks if the first argument matches any of the list items.
+//! \return True if v is a member of list.
+template <typename TElem, typename Container = std::initializer_list<TElem>>
+bool elem(TElem const& v, Container const& list)
+{
+    return std::any_of(std::begin(list), std::end(list), [&v](TElem const& t) { return t == v; });
+}
+
+inline int32_t getMHAMaskPackedSize(int32_t smVersion, nvinfer1::DataType dataType, int32_t sequenceLength)
+{
+    // this code must match EmbLayerNormPluginDynamic::getOutputDimensions in embLayerNormPlugin.cpp
+    int32_t packedSize = unfusedMaskSize;
+    bool const isSmOK = elem(smVersion, {kSM_75, kSM_80, kSM_86, kSM_87, kSM_89, kSM_90, kSM_100, kSM_120});
+    bool isPrecisionOK = (dataType == nvinfer1::DataType::kINT8 || dataType == nvinfer1::DataType::kHALF);
+    if (isSmOK && isPrecisionOK)
+    {
+        if (sequenceLength == 64)
+        {
+            packedSize = packedMaskSize64;
+        }
+        else if (sequenceLength == 96)
+        {
+            packedSize = packedMaskSize96;
+        }
+        else if (sequenceLength == 128)
+        {
+            packedSize = packedMaskSize128;
+        }
+        else if (sequenceLength == 384)
+        {
+            packedSize = packedMaskSize384;
+        }
+    }
+    return packedSize;
+}
+
+inline uint32_t getElementSize(nvinfer1::DataType t)
+{
+    switch (t)
+    {
+    case nvinfer1::DataType::kINT64: return 8;
+    case nvinfer1::DataType::kINT32:
+    case nvinfer1::DataType::kFLOAT: return 4;
+    case nvinfer1::DataType::kBF16:
+    case nvinfer1::DataType::kHALF: return 2;
+    case nvinfer1::DataType::kBOOL:
+    case nvinfer1::DataType::kUINT8:
+    case nvinfer1::DataType::kINT8:
+    case nvinfer1::DataType::kFP8: return 1;
+    case nvinfer1::DataType::kINT4:
+    case nvinfer1::DataType::kFP4:
+    case nvinfer1::DataType::kE8M0: PLUGIN_FAIL("Element size is not implemented for sub-byte data-types");
+    }
+    return 0;
+}
+
+inline int64_t getWeightsSize(nvinfer1::Weights const& w, nvinfer1::DataType type)
+{
+    return w.count * getElementSize(type);
+}
+
+inline int64_t volume(nvinfer1::Dims const& d)
+{
+    return std::accumulate(d.d, d.d + d.nbDims, int64_t{1}, std::multiplies<int64_t>{});
+}
+
+//! Check if the hardware supports BERT Multi-Head Attention plugins
+//! The plugin calls precompiled cubins (compiled from fmha_v2/xmma kernels)
+//! that are SM-specific.
+inline bool doesHwSupportBertMHAPlugin() noexcept
+{
+    int32_t device{-1};
+    cudaGetDevice(&device);
+    int32_t smMajor{0};
+    int32_t smMinor{0};
+    cudaDeviceGetAttribute(&smMajor, cudaDevAttrComputeCapabilityMajor, device);
+    cudaDeviceGetAttribute(&smMinor, cudaDevAttrComputeCapabilityMinor, device);
+    int32_t smVersion = (smMajor << 4) | (smMinor);
+    // Turing and above
+    static constexpr int32_t kSM_TURING_HEX{0x75};
+    static constexpr int32_t kSM_BLACKWELL_100_HEX{0xA0};
+    static constexpr int32_t kSM_BLACKWELL_120_HEX{0xC0};
+    static constexpr int32_t kSM_ORIN_HEX{0x87};
+    bool isAuto = smVersion == kSM_ORIN_HEX;
+    bool isSm100OrLower = smVersion >= kSM_TURING_HEX && smVersion <= kSM_BLACKWELL_100_HEX;
+    bool isHardwareSupported = (isSm100OrLower || smVersion == kSM_BLACKWELL_120_HEX) && !isAuto;
+
+    return isHardwareSupported;
+}
+
+template <typename IntType>
+constexpr IntType ceildiv(IntType a, IntType b)
+{
+    return (a + b - 1) / b;
+}
+template <typename IntType>
+constexpr IntType alignTo(IntType a, IntType b)
+{
+    return ceildiv(a, b) * b;
+}
+
+template <typename T>
+inline T* deserToDev(char const*& buffer, size_t nbElem)
+{
+    void* dev{nullptr};
+    const size_t len = sizeof(T) * nbElem;
+    PLUGIN_CUASSERT(cudaMalloc(&dev, len));
+    PLUGIN_CUASSERT(cudaMemcpy(dev, buffer, len, cudaMemcpyHostToDevice));
+
+    buffer += len;
+    return static_cast<T*>(dev);
+}
+
+template <typename T>
+inline void serFromDev(char*& buffer, T const* data, size_t nbElem)
+{
+    const size_t len = sizeof(T) * nbElem;
+    PLUGIN_CUASSERT(cudaMemcpy(buffer, static_cast<void const*>(data), len, cudaMemcpyDeviceToHost));
+    buffer += len;
+}
+
+template <typename T>
+inline T* devToDev(T const* data, size_t nbElem)
+{
+    void* dev{nullptr};
+    const size_t len = sizeof(T) * nbElem;
+    PLUGIN_CUASSERT(cudaMalloc(&dev, len));
+    PLUGIN_CUASSERT(cudaMemcpy(dev, static_cast<void const*>(data), len, cudaMemcpyDeviceToDevice));
+    return static_cast<T*>(dev);
+}
+
+template <typename T>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemm(nvinfer1::pluginInternal::cublasHandle_t handle,
+    nvinfer1::pluginInternal::cublasOperation_t transa, nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m,
+    int32_t n, int32_t k, const T alpha, T const* A, int32_t lda, T const* B, int32_t ldb, const T beta, T* C,
+    int32_t ldc);
+
+template <>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemm(nvinfer1::pluginInternal::cublasHandle_t handle,
+    nvinfer1::pluginInternal::cublasOperation_t transa, nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m,
+    int32_t n, int32_t k, float const alpha, float const* A, int32_t lda, float const* B, int32_t ldb, float const beta,
+    float* C, int32_t ldc)
+{
+    nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
+    return wrapper.cublasSgemm(handle, transa, transb, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
+}
+
+template <>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemm(nvinfer1::pluginInternal::cublasHandle_t handle,
+    nvinfer1::pluginInternal::cublasOperation_t transa, nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m,
+    int32_t n, int32_t k, const half alpha, half const* A, int32_t lda, half const* B, int32_t ldb, const half beta,
+    half* C, int32_t ldc)
+{
+    nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
+    return wrapper.cublasHgemm(handle, transa, transb, m, n, k, &alpha, A, lda, B, ldb, &beta, C, ldc);
+}
+
+template <typename T>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatchedEx(
+    nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
+    nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const T alpha, T const* A,
+    int32_t lda, int64_t strideA, T const* B, int32_t ldb, int64_t strideB, const T beta, T* C, int32_t ldc,
+    int64_t strideC, int32_t batchCount, nvinfer1::pluginInternal::cublasGemmAlgo_t algo);
+
+template <>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatchedEx(
+    nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
+    nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, float const alpha,
+    float const* A, int32_t lda, int64_t strideA, float const* B, int32_t ldb, int64_t strideB, float const beta,
+    float* C, int32_t ldc, int64_t strideC, int32_t batchCount, nvinfer1::pluginInternal::cublasGemmAlgo_t algo)
+{
+    nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
+    return wrapper.cublasGemmStridedBatchedEx(handle, transa, transb, m, n, k, &alpha, A, CUDA_R_32F, lda, strideA, B,
+        CUDA_R_32F, ldb, strideB, &beta, C, CUDA_R_32F, ldc, strideC, batchCount, CUDA_R_32F, algo);
+}
+
+template <>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatchedEx(
+    nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
+    nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const half alpha,
+    half const* A, int32_t lda, int64_t strideA, half const* B, int32_t ldb, int64_t strideB, const half beta, half* C,
+    int32_t ldc, int64_t strideC, int32_t batchCount, nvinfer1::pluginInternal::cublasGemmAlgo_t algo)
+{
+    nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
+    return wrapper.cublasGemmStridedBatchedEx(handle, transa, transb, m, n, k, &alpha, A, CUDA_R_16F, lda, strideA, B,
+        CUDA_R_16F, ldb, strideB, &beta, C, CUDA_R_16F, ldc, strideC, batchCount, CUDA_R_16F, algo);
+}
+
+template <typename T>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatched(
+    nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
+    nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const T alpha, T const* A,
+    int32_t lda, int64_t strideA, T const* B, int32_t ldb, int64_t strideB, const T beta, T* C, int32_t ldc,
+    int64_t strideC, int32_t batchCount);
+
+template <>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatched(
+    nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
+    nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, float const alpha,
+    float const* A, int32_t lda, int64_t strideA, float const* B, int32_t ldb, int64_t strideB, float const beta,
+    float* C, int32_t ldc, int64_t strideC, int32_t batchCount)
+{
+    nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
+    return wrapper.cublasSgemmStridedBatched(
+        handle, transa, transb, m, n, k, &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, strideC, batchCount);
+}
+
+template <>
+nvinfer1::pluginInternal::cublasStatus_t inline cublasGemmStridedBatched(
+    nvinfer1::pluginInternal::cublasHandle_t handle, nvinfer1::pluginInternal::cublasOperation_t transa,
+    nvinfer1::pluginInternal::cublasOperation_t transb, int32_t m, int32_t n, int32_t k, const half alpha,
+    half const* A, int32_t lda, int64_t strideA, half const* B, int32_t ldb, int64_t strideB, const half beta, half* C,
+    int32_t ldc, int64_t strideC, int32_t batchCount)
+{
+    nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
+    return wrapper.cublasHgemmStridedBatched(
+        handle, transa, transb, m, n, k, &alpha, A, lda, strideA, B, ldb, strideB, &beta, C, ldc, strideC, batchCount);
+}
+
+struct CublasConfigHelper
+{
+    nvinfer1::pluginInternal::cublasPointerMode_t pm;
+    nvinfer1::pluginInternal::cublasMath_t mm;
+    nvinfer1::pluginInternal::cublasHandle_t cublas;
+    nvinfer1::pluginInternal::CublasWrapper& wrapper = nvinfer1::pluginInternal::getCublasWrapper();
+    CublasConfigHelper(nvinfer1::pluginInternal::cublasHandle_t cublas_)
+        : cublas(cublas_)
+    {
+        PLUGIN_CUBLASASSERT(wrapper.cublasGetPointerMode(cublas, &pm));
+        PLUGIN_CUBLASASSERT(wrapper.cublasGetMathMode(cublas, &mm));
+        PLUGIN_CUBLASASSERT(wrapper.cublasSetPointerMode(cublas, nvinfer1::pluginInternal::CUBLAS_POINTER_MODE_HOST));
+        PLUGIN_CUBLASASSERT(wrapper.cublasSetMathMode(cublas, nvinfer1::pluginInternal::CUBLAS_TENSOR_OP_MATH));
+    }
+    ~CublasConfigHelper()
+    {
+        wrapper.cublasSetMathMode(cublas, mm);
+        wrapper.cublasSetPointerMode(cublas, pm);
+    }
+};
+
+template <typename T>
+using cuda_unique_ptr = std::unique_ptr<T, pluginInternal::CudaDeleter<T>>;
+
+template <typename T>
+using cuda_shared_ptr = std::shared_ptr<T>;
+
+template <typename T>
+void make_cuda_shared(cuda_shared_ptr<T>& ptr, void* cudaMem)
+{
+    ptr.reset(static_cast<T*>(cudaMem), pluginInternal::CudaDeleter<T>());
+}
+
+struct WeightsWithOwnership : public nvinfer1::Weights
+{
+    WeightsWithOwnership()
+    {
+        values = nullptr;
+        count = 0;
+    }
+    ~WeightsWithOwnership()
+    {
+        operator delete[](const_cast<void*>(values));
+    }
+
+    WeightsWithOwnership(WeightsWithOwnership const&) = delete;
+    WeightsWithOwnership operator=(WeightsWithOwnership const&) = delete;
+    WeightsWithOwnership(WeightsWithOwnership const&&) = delete;
+    WeightsWithOwnership operator=(WeightsWithOwnership const&&) = delete;
+
+    void convertAndCopy(nvinfer1::Weights const& src, nvinfer1::DataType type)
+    {
+        this->type = type;
+        this->count = src.count;
+
+        if (type == nvinfer1::DataType::kFLOAT)
+        {
+            auto destBuf = new float[src.count];
+            this->values = destBuf;
+
+            if (src.type == nvinfer1::DataType::kFLOAT)
+            {
+                BERT_DEBUG_MSG("Float Weights(Host) => Float Array(Host)");
+                std::copy_n(static_cast<float const*>(src.values), src.count, destBuf);
+            }
+            else
+            {
+                PLUGIN_ASSERT(src.type == nvinfer1::DataType::kHALF);
+
+                BERT_DEBUG_MSG("Half Weights(Host) => Float Array(Host)");
+                auto const s = static_cast<half const*>(src.values);
+                auto d = static_cast<float*>(const_cast<void*>(this->values));
+
+                for (auto it = 0; it < src.count; it++)
+                {
+                    d[it] = __half2float(s[it]);
+                }
+            }
+        }
+        else if (type == nvinfer1::DataType::kHALF)
+        {
+            auto destBuf = new half[src.count];
+            this->values = destBuf;
+
+            if (src.type == nvinfer1::DataType::kHALF)
+            {
+                BERT_DEBUG_MSG("Half Weights(Host) => Half Array(Host)");
+                std::copy_n(static_cast<half const*>(src.values), src.count, destBuf);
+            }
+            else
+            {
+                PLUGIN_ASSERT(src.type == nvinfer1::DataType::kFLOAT);
+
+                BERT_DEBUG_MSG("Float Weights(Host) => Half Array(Host)");
+                auto const s = static_cast<float const*>(src.values);
+                auto d = static_cast<half*>(const_cast<void*>(this->values));
+
+                for (auto it = 0; it < src.count; it++)
+                {
+                    d[it] = __float2half(s[it]);
+                }
+            }
+        }
+        else
+        {
+            throw std::runtime_error("Unsupported DataType specified for plugin.");
+        }
+    }
+
+    void convertAndCopy(char const*& srcBuf, size_t count, nvinfer1::DataType type) noexcept
+    {
+        this->type = type;
+        this->count = count;
+        auto const nbBytes = getWeightsSize(*this, type);
+        auto destBuf = new char[nbBytes];
+        this->values = destBuf;
+
+        std::copy_n(srcBuf, nbBytes, destBuf);
+        srcBuf += nbBytes;
+    }
+};
+
+template <typename T>
+inline void copyToDevice(WeightsWithOwnership& hostWeights, size_t nbBytes, cuda_unique_ptr<T>& cudaWeights)
+{
+    if (hostWeights.values)
+    {
+        void* cudaMem{nullptr};
+        PLUGIN_CUASSERT(cudaMalloc(&cudaMem, nbBytes));
+        PLUGIN_CUASSERT(cudaMemcpy(cudaMem, hostWeights.values, nbBytes, cudaMemcpyHostToDevice));
+        cudaWeights.reset(static_cast<T*>(cudaMem));
+    }
+}
+
+inline void convertAndCopyToDevice(nvinfer1::Weights const& src, float* destDev)
+{
+
+    size_t wordSize = sizeof(float);
+    size_t nbBytes = src.count * wordSize;
+    if (src.type == nvinfer1::DataType::kFLOAT)
+    {
+        BERT_DEBUG_MSG("Float Weights(Host) => Float Array(Device)");
+        PLUGIN_CUASSERT(cudaMemcpy(destDev, src.values, nbBytes, cudaMemcpyHostToDevice));
+    }
+    else
+    {
+        BERT_DEBUG_MSG("Half Weights(Host) => Float Array(Device)");
+        std::vector<float> tmp(src.count);
+        half const* values = reinterpret_cast<half const*>(src.values);
+
+        for (size_t it = 0; it < tmp.size(); it++)
+        {
+            tmp[it] = __half2float(values[it]);
+        }
+
+        PLUGIN_CUASSERT(cudaMemcpy(destDev, tmp.data(), nbBytes, cudaMemcpyHostToDevice));
+    }
+}
+
+inline void convertAndCopyToDevice(nvinfer1::Weights const& src, half* destDev)
+{
+    size_t wordSize = sizeof(half);
+    size_t nbBytes = src.count * wordSize;
+    if (src.type == nvinfer1::DataType::kHALF)
+    {
+        BERT_DEBUG_MSG("Half Weights(Host) => Half Array(Device)");
+        PLUGIN_CUASSERT(cudaMemcpy(destDev, src.values, nbBytes, cudaMemcpyHostToDevice));
+    }
+    else
+    {
+        BERT_DEBUG_MSG("Float Weights(Host) => Half Array(Device)");
+        std::vector<half> tmp(src.count);
+        float const* values = reinterpret_cast<float const*>(src.values);
+
+        for (size_t it = 0; it < tmp.size(); it++)
+        {
+            tmp[it] = __float2half(values[it]);
+        }
+        PLUGIN_CUASSERT(cudaMemcpy(destDev, tmp.data(), nbBytes, cudaMemcpyHostToDevice));
+    }
+}
+
+inline nvinfer1::DataType fieldTypeToDataType(const nvinfer1::PluginFieldType ftype)
+{
+    switch (ftype)
+    {
+    case nvinfer1::PluginFieldType::kFLOAT32:
+    {
+        BERT_DEBUG_MSG("PluginFieldType is Float32");
+        return nvinfer1::DataType::kFLOAT;
+    }
+    case nvinfer1::PluginFieldType::kFLOAT16:
+    {
+        BERT_DEBUG_MSG("PluginFieldType is Float16");
+        return nvinfer1::DataType::kHALF;
+    }
+    case nvinfer1::PluginFieldType::kINT32:
+    {
+        BERT_DEBUG_MSG("PluginFieldType is Int32");
+        return nvinfer1::DataType::kINT32;
+    }
+    case nvinfer1::PluginFieldType::kINT8:
+    {
+        BERT_DEBUG_MSG("PluginFieldType is Int8");
+        return nvinfer1::DataType::kINT8;
+    }
+    default: throw std::invalid_argument("No corresponding datatype for plugin field type");
+    }
+}
+
+} // namespace bert
+} // namespace plugin
+} // namespace nvinfer1
+#endif // BERT_COMMON_H
+
+#endif // CUDA_VERSION >= 10010

plugin/common/checkMacrosPlugin.cpp

@@ -0,0 +1,76 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "common/checkMacrosPlugin.h"
+#include "common/cublasWrapper.h"
+#include "vc/vfcCommon.h"
+#include <cstdlib>
+#include <cuda_runtime.h>
+#include <string_view>
+
+using namespace nvinfer1::pluginInternal;
+
+namespace nvinfer1::plugin
+{
+
+namespace
+{
+//! \return Static identifier string for \p status, or an empty view if unrecognized.
+[[nodiscard]] constexpr std::string_view cublasStatusName(cublasStatus_t status)
+{
+    switch (status)
+    {
+    case CUBLAS_STATUS_SUCCESS: return "CUBLAS_STATUS_SUCCESS";
+    case CUBLAS_STATUS_NOT_INITIALIZED: return "CUBLAS_STATUS_NOT_INITIALIZED";
+    case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED";
+    case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE";
+    case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH";
+    case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR";
+    case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED";
+    case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR";
+    case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED";
+    case CUBLAS_STATUS_LICENSE_ERROR: return "CUBLAS_STATUS_LICENSE_ERROR";
+    }
+    return {};
+}
+} // namespace
+
+// break-pointable
+void throwCublasError(
+    char const* file, char const* function, int32_t line, int32_t status, std::optional<std::string> msg)
+{
+    if (!msg)
+    {
+        msg.emplace(cublasStatusName(static_cast<cublasStatus_t>(status)));
+    }
+    CublasError error(file, function, line, status, std::move(msg));
+    error.log(gLogError);
+    // NOLINTNEXTLINE(misc-throw-by-value-catch-by-reference)
+    throw error;
+}
+
+// break-pointable
+void throwCudnnError(
+    char const* file, char const* function, int32_t line, int32_t status, std::optional<std::string> msg)
+{
+    CudnnError error(file, function, line, status, std::move(msg));
+    error.log(gLogError);
+    // NOLINTNEXTLINE(misc-throw-by-value-catch-by-reference)
+    throw error;
+}
+
+} // namespace nvinfer1::plugin

plugin/common/checkMacrosPlugin.h

@@ -0,0 +1,108 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef CHECK_MACROS_PLUGIN_H
+#define CHECK_MACROS_PLUGIN_H
+
+#include "NvInfer.h"
+#include "common/cudnnWrapper.h"
+#include "vc/checkMacrosPlugin.h"
+#include <mutex>
+#include <optional>
+#include <sstream>
+#include <string>
+#include <utility>
+
+#ifdef _MSC_VER
+#define FN_NAME __FUNCTION__
+#else
+#define FN_NAME __func__
+#endif
+
+namespace nvinfer1
+{
+namespace plugin
+{
+
+//! \throw CudnnError carrying \p msg, after logging it via \c gLogError.
+//! Parameter semantics match \c nvinfer1::plugin::throwCudaError.
+[[noreturn]] void throwCudnnError(char const* file, char const* function, int32_t line, int32_t status,
+    std::optional<std::string> msg = std::nullopt);
+
+//! \throw CublasError carrying \p msg (or a default message derived from \p status if \p msg is
+//!     \c std::nullopt), after logging it via \c gLogError. An explicit empty string is preserved.
+//! Parameter semantics match \c nvinfer1::plugin::throwCudaError.
+[[noreturn]] void throwCublasError(char const* file, char const* function, int32_t line, int32_t status,
+    std::optional<std::string> msg = std::nullopt);
+
+//! \c TRTException specialization for cuDNN failures.
+class CudnnError : public TRTException
+{
+public:
+    //! \see TRTException::TRTException. Sets the subsystem name to \c "Cudnn".
+    CudnnError(char const* fl, char const* fn, int32_t ln, int32_t stat, std::optional<std::string> msg = std::nullopt)
+        : TRTException(fl, fn, ln, stat, std::move(msg), "Cudnn")
+    {
+    }
+};
+
+//! \c TRTException specialization for cuBLAS failures.
+class CublasError : public TRTException
+{
+public:
+    //! \see TRTException::TRTException. Sets the subsystem name to \c "cuBLAS".
+    CublasError(char const* fl, char const* fn, int32_t ln, int32_t stat, std::optional<std::string> msg = std::nullopt)
+        : TRTException(fl, fn, ln, stat, std::move(msg), "cuBLAS")
+    {
+    }
+};
+
+} // namespace plugin
+
+} // namespace nvinfer1
+
+#define PLUGIN_CHECK_CUDNN(call)                                                                                       \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        cudnnStatus_t status = call;                                                                                   \
+        if (status != CUDNN_STATUS_SUCCESS)                                                                            \
+        {                                                                                                              \
+            return status;                                                                                             \
+        }                                                                                                              \
+    } while (0)
+
+#define PLUGIN_CUBLASASSERT(status_)                                                                                   \
+    {                                                                                                                  \
+        auto s_ = status_;                                                                                             \
+        if (s_ != nvinfer1::pluginInternal::CUBLAS_STATUS_SUCCESS)                                                     \
+        {                                                                                                              \
+            nvinfer1::plugin::throwCublasError(__FILE__, FN_NAME, __LINE__, s_);                                       \
+        }                                                                                                              \
+    }
+
+#define PLUGIN_CUDNNASSERT(status_)                                                                                    \
+    {                                                                                                                  \
+        auto s_ = status_;                                                                                             \
+        if (s_ != CUDNN_STATUS_SUCCESS)                                                                                \
+        {                                                                                                              \
+            nvinfer1::pluginInternal::CudnnWrapper& wrapper                                                            \
+                = nvinfer1::pluginInternal::getCudnnWrapper(/* plugin caller name */ nullptr);                         \
+            const char* msg = wrapper.cudnnGetErrorString(s_);                                                         \
+            nvinfer1::plugin::throwCudnnError(__FILE__, FN_NAME, __LINE__, s_, msg);                                   \
+        }                                                                                                              \
+    }
+
+#endif /*CHECK_MACROS_PLUGIN_H*/

plugin/common/common.cuh

@@ -0,0 +1,529 @@
+/*
+ * 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 COMMON_CUH
+#define COMMON_CUH
+
+// TODO: Remove WAR once issue resolved in CUB (CUDA 12.6+?)
+
+#ifndef CUDA_VERSION
+#include <cuda.h>
+#endif // CUDA_VERSION
+
+#include "common/cublasWrapper.h"
+#include "common/cubCcclCompat.h"
+#include <cfloat>
+
+#define HDI inline __host__ __device__
+
+using kv_float = cub::KeyValuePair<float, float>;
+using kv_half = cub::KeyValuePair<half, half>;
+using kv_half2 = cub::KeyValuePair<half2, half2>;
+
+template <typename T>
+__device__ inline T tanh(const T& x);
+
+template <typename T>
+__device__ inline T rsqrt(const T& x);
+
+template <typename T>
+__device__ inline T exp(const T& x);
+
+
+// Float32 Operations
+template <>
+__device__ inline float tanh(const float& x)
+{
+    return tanhf(x);
+}
+
+template <>
+__device__ inline float rsqrt(const float& x)
+{
+    return rsqrtf(x);
+}
+
+template <>
+__device__ inline float exp(const float& x)
+{
+    return expf(x);
+}
+
+
+__device__ inline kv_float operator+(const kv_float& a, const kv_float& b)
+{
+    return kv_float(a.key + b.key, a.value + b.value);
+}
+
+
+// Half Operations
+
+__device__ inline half2 __hadd2_with_fallback(const half2 a, const half2 b)
+{
+#if __CUDA_ARCH__ >= 530
+    return __hadd2(a, b);
+#else
+    float2 out{};
+    out.x = __half2float(a.x) + __half2float(b.x);
+    out.y = __half2float(a.y) + __half2float(b.y);
+    return __float22half2_rn(out);
+#endif
+}
+
+#if __CUDA_ARCH__ < 530
+
+template <typename T>
+__device__ inline T operator+(const T& a, const T& b);
+
+template <typename T>
+__device__ inline T operator*(const T& a, const T& b);
+
+
+template <>
+__device__ inline half2 operator+(const half2& a, const half2& b)
+{
+    return __hadd2_with_fallback(a, b);
+}
+
+template <>
+__device__ inline half2 operator*(const half2& a, const half2& b)
+{
+    float2 out{};
+    out.x = __half2float(a.x) * __half2float(b.x);
+    out.y = __half2float(a.y) * __half2float(b.y);
+    return __float22half2_rn(out);
+}
+
+template <typename T>
+__device__ inline T operator+(const T& a, const T& b);
+
+
+template <typename T>
+__device__ inline T operator/(const T& a, const T& b);
+
+
+template <typename T>
+__device__ inline T& operator+=(T& a, const T& b);
+
+template <typename T>
+__device__ inline T operator-(const T& a, const T& b);
+
+template <typename T>
+__device__ inline T operator*(const T& a, const T& b);
+
+template <typename T>
+__device__ inline bool operator>(const T& a, const T& b);
+
+template <typename T>
+__device__ inline bool operator>=(const T& a, const T& b);
+
+template <typename T>
+__device__ inline bool operator<(const T& a, const T& b);
+
+template <typename T>
+__device__ inline bool operator<=(const T& a, const T& b);
+
+template <>
+__device__ inline half operator+(const half& a, const half& b)
+{
+    return __float2half(__half2float(a) + __half2float(b));
+}
+
+template <>
+__device__ inline half& operator+=(half& a, const half& b)
+{
+    a = __float2half(__half2float(a) + __half2float(b));
+    return a;
+}
+
+template <>
+__device__ inline half operator-(const half& a, const half& b)
+{
+    return __float2half(__half2float(a) - __half2float(b));
+}
+
+template <>
+__device__ inline half operator*(const half& a, const half& b)
+{
+    return __float2half(__half2float(a) * __half2float(b));
+}
+
+template <>
+__device__ inline half operator/(const half& a, const half& b)
+{
+    return __float2half(__half2float(a) / __half2float(b));
+}
+
+template <>
+__device__ inline bool operator>(const half& a, const half& b)
+{
+    return __half2float(a) > __half2float(b);
+}
+
+template <>
+__device__ inline bool operator>=(const half& a, const half& b)
+{
+    return __half2float(a) >= __half2float(b);
+}
+
+template <>
+__device__ inline bool operator<(const half& a, const half& b)
+{
+    return __half2float(a) < __half2float(b);
+}
+
+template <>
+__device__ inline bool operator<=(const half& a, const half& b)
+{
+    return __half2float(a) <= __half2float(b);
+}
+
+#endif
+
+template <>
+__device__ inline half tanh(const half& x)
+{
+    const float tmp = tanhf(__half2float(x));
+    return __float2half(tmp);
+}
+
+template <>
+__device__ inline half2 tanh(const half2& x)
+{
+    // at the moment, there is no half2 tanh builtin
+    float2 tmp = (__half22float2(x));
+    tmp.x = tanhf(tmp.x);
+    tmp.y = tanhf(tmp.y);
+    return __float22half2_rn(tmp);
+}
+
+template <>
+__device__ inline half rsqrt(const half& x)
+{
+#if __CUDA_ARCH__ >= 530
+    return hrsqrt(x);
+#else
+    return __float2half(rsqrt(__half2float(x)));
+#endif
+}
+
+template <>
+__device__ inline half exp(const half& x)
+{
+#if __CUDA_ARCH__ >= 530
+    return hexp(x);
+#else
+    return __float2half(exp(__half2float(x)));
+#endif
+}
+
+__device__ inline kv_half operator+(const kv_half& a, const kv_half& b)
+{
+    const half2 a2 = __halves2half2(a.key, a.value);
+    const half2 b2 = __halves2half2(b.key, b.value);
+    const half2 res = __hadd2_with_fallback(a2, b2);
+    return kv_half(res.x, res.y);
+}
+
+__device__ inline kv_half2 operator+(const kv_half2& a, const kv_half2& b)
+{
+    return kv_half2(__hadd2_with_fallback(a.key, b.key), __hadd2_with_fallback(a.value, b.value));
+}
+
+
+// Helper Functions
+
+template <typename T>
+using kvp = cub::KeyValuePair<T, T>;
+
+template <typename T, typename R, typename P, int32_t TPB>
+__device__ inline void layerNorm(
+    const kvp<R>& threadData, const int32_t ld, const int32_t offset, const P* beta, const P* gamma, T* output)
+{
+    // Assuming threadData is already divided by ld
+
+    using BlockReduce = cub::BlockReduce<kvp<R>, TPB>;
+    __shared__ typename BlockReduce::TempStorage temp_storage;
+    __shared__ R mu;     // mean
+    __shared__ R rsigma; // 1 / std.dev.
+
+    const auto sumKV = BlockReduce(temp_storage).Reduce(threadData, [](auto const& lhs, auto const& rhs){return lhs + rhs;});
+
+    if (threadIdx.x == 0)
+    {
+        mu = sumKV.key;
+        rsigma = rsqrt(sumKV.value - mu * mu);
+    }
+    __syncthreads();
+
+    for (int32_t i = threadIdx.x; i < ld; i += TPB)
+    {
+        const int32_t idx = offset + i;
+        const R val = output[idx];
+        const R g(gamma[i]);
+        const R b(beta[i]);
+        output[idx] = g * (val - mu) * rsigma + b;
+    }
+}
+
+template <typename T, typename P, int32_t TPB>
+__device__ inline void layerNormSmall(
+    const T val, const kvp<T>& threadData, const int32_t ld, const int32_t idx, const P* beta, const P* gamma, T* output)
+{
+    // Assuming threadData is already divided by ld
+    // Small settings: the block covers the leading dimension TPB >= ld. The input
+    // value is available in a register
+
+    using BlockReduce = cub::BlockReduce<kvp<T>, TPB>;
+    __shared__ typename BlockReduce::TempStorage temp_storage;
+    __shared__ T mu;     // mean
+    __shared__ T rsigma; // 1 / std.dev.
+
+    const auto sumKV = BlockReduce(temp_storage).Reduce(threadData, [](auto const& lhs, auto const& rhs){return lhs + rhs;});
+
+    if (threadIdx.x == 0)
+    {
+        mu = sumKV.key;
+        rsigma = rsqrt(sumKV.value - mu * mu);
+    }
+    __syncthreads();
+
+    if (threadIdx.x < ld)
+    {
+        const T g(gamma[threadIdx.x]);
+        const T b(beta[threadIdx.x]);
+        output[idx] = g * (val - mu) * rsigma + b;
+    }
+}
+
+template <typename T, unsigned TPB>
+__device__ inline void scaledSoftmaxSmall(
+    const int32_t ld, const int32_t lastValid, const float rsqrtHeadSize, const T* input, T* output)
+{
+
+    using BlockReduce = cub::BlockReduce<float, TPB>;
+
+    __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+    __shared__ float rZ;
+    __shared__ float fMax;
+
+    const int32_t offset = (blockIdx.y * gridDim.x + blockIdx.x) * ld;
+
+    const float w(rsqrtHeadSize);
+    float threadData(-FLT_MAX);
+
+    const int32_t idx = offset + threadIdx.x;
+    if (threadIdx.x < lastValid)
+    {
+        threadData = input[idx];
+    }
+
+    const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, compat::getCudaMaxOp());
+    if (threadIdx.x == 0)
+    {
+        fMax = maxElem;
+    }
+    __syncthreads();
+
+    if (threadIdx.x < lastValid)
+    {
+        threadData = exp((threadData - fMax) * w);
+    }
+    else
+    {
+        threadData = 0;
+    }
+
+    const auto Z = BlockReduce(tmpStorage).Reduce(threadData, [](auto const& lhs, auto const& rhs){return lhs + rhs;});
+
+    if (threadIdx.x == 0)
+    {
+        rZ = (1.f) / Z;
+    }
+    __syncthreads();
+
+    if (threadIdx.x < ld)
+    {
+        float const val = (threadIdx.x < lastValid) ? threadData * rZ : 0.F;
+        output[idx] = static_cast<T>(val);
+    }
+}
+
+template <typename T, unsigned TPB>
+__device__ inline void scaledSoftmax(
+    const int32_t ld, const int32_t lastValid, const float rsqrtHeadSize, const T* input, T* output)
+{
+    using BlockReduce = cub::BlockReduce<float, TPB>;
+    __shared__ typename BlockReduce::TempStorage tmpStorage;
+
+    __shared__ float rZ;
+    __shared__ float fMax;
+
+    const int32_t offset = (blockIdx.y * gridDim.x + blockIdx.x) * ld;
+
+    const float w(rsqrtHeadSize);
+    float threadData(-FLT_MAX);
+
+    if (lastValid >= blockDim.x)
+    {
+        threadData = 0;
+    }
+    for (int32_t i = threadIdx.x; i < lastValid; i += TPB)
+    {
+        const int32_t idx = offset + i;
+        threadData = max(static_cast<float>(input[idx]), threadData);
+    }
+
+    const float maxElem = BlockReduce(tmpStorage).Reduce(threadData, compat::getCudaMaxOp());
+    if (threadIdx.x == 0)
+    {
+        fMax = maxElem;
+    }
+    __syncthreads();
+
+    threadData = 0;
+
+    for (int32_t i = threadIdx.x; i < lastValid; i += TPB)
+    {
+        const int32_t idx = offset + i;
+        threadData += exp((static_cast<float>(input[idx]) - fMax) * w);
+    }
+
+    const auto Z = BlockReduce(tmpStorage).Reduce(threadData, [](auto const& lhs, auto const& rhs){return lhs + rhs;});
+
+    if (threadIdx.x == 0)
+    {
+        rZ = 1.f / Z;
+    }
+    __syncthreads();
+
+    for (int32_t i = threadIdx.x; i < ld; i += TPB)
+    {
+        const int32_t idx = offset + i;
+        const float val = (i < lastValid) ? exp((static_cast<float>(input[idx]) - fMax) * w) * rZ : 0.f;
+        output[idx] = T(val);
+    }
+}
+
+template <typename IntType>
+constexpr HDI IntType ceildiv(IntType a, IntType b)
+{
+    return (a + b - 1) / b;
+}
+template <typename IntType>
+constexpr HDI IntType alignTo(IntType a, IntType b)
+{
+    return ceildiv(a, b) * b;
+}
+
+template <int32_t VPT>
+struct BytesToType;
+
+template <>
+struct BytesToType<2>
+{
+    using type = uint16_t;
+};
+template <>
+struct BytesToType<4>
+{
+    using type = uint32_t;
+};
+template <>
+struct BytesToType<8>
+{
+    using type = uint64_t;
+};
+template <>
+struct BytesToType<16>
+{
+    using type = float4;
+};
+
+template <int32_t Bytes>
+__device__ inline void copy(const void* local, void* data)
+{
+    using T = typename BytesToType<Bytes>::type;
+
+    const T* in = static_cast<const T*>(local);
+    T* out = static_cast<T*>(data);
+    *out = *in;
+}
+
+template <typename T>
+__device__ inline T myExp(const T x);
+
+template <>
+__device__ inline half myExp<half>(const half x)
+{
+    return exp(x);
+}
+template <>
+__device__ inline float myExp<float>(const float x)
+{
+    return __expf(x);
+}
+
+static inline __device__ uint32_t float4_to_char4(float x,
+                                                  float y,
+                                                  float z,
+                                                  float w) {
+  uint32_t dst;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 720
+  uint32_t a; asm volatile("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(a) : "f"(x));
+  uint32_t b; asm volatile("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(b) : "f"(y));
+  uint32_t c; asm volatile("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(c) : "f"(z));
+  uint32_t d; asm volatile("cvt.rni.sat.s32.f32 %0, %1;\n" : "=r"(d) : "f"(w));
+
+  asm volatile("cvt.pack.sat.s8.s32.b32 %0, %1, %2,  0;\n" : "=r"(dst) : "r"(d), "r"(c));
+  asm volatile("cvt.pack.sat.s8.s32.b32 %0, %1, %2, %0;\n" : "+r"(dst) : "r"(b), "r"(a));
+#else
+  char4 tmp;
+  tmp.x = x;
+  tmp.y = y;
+  tmp.z = z;
+  tmp.w = w;
+  dst = reinterpret_cast<const uint32_t&>(tmp);
+#endif
+  return dst;
+}
+
+inline __device__ char quantize(const float x, const float qScale)
+{
+    int32_t tmpq = __float2int_rn(qScale * x);  // scale and round
+    char tmpq8 = min(127, max(-127, tmpq)); // clip and cast
+    return tmpq8;
+}
+
+inline __device__ void ldg(const int8_t* input, uint4& data)
+{
+    data = *reinterpret_cast<const uint4*>(input);
+}
+
+inline __device__ void stg(int8_t* output, uint4& data)
+{
+    *reinterpret_cast<uint4*>(output) = data;
+}
+
+inline __device__ uint32_t pack4(const float (&hdata)[4], const float qScale)
+{
+    return float4_to_char4(hdata[0] * qScale, hdata[1] * qScale , hdata[2] * qScale, hdata[3] * qScale);
+}
+
+
+#endif // #ifndef COMMON_CUH

plugin/common/cubCcclCompat.h

@@ -0,0 +1,48 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_CUB_CCCL_COMPAT_H
+#define TRT_CUB_CCCL_COMPAT_H
+
+#include <cub/cub.cuh>
+#include <cuda.h>
+#include <cuda/std/functional>
+
+namespace compat
+{
+
+// Min operation compatibility wrapper
+__host__ __device__ __forceinline__ auto getCudaMinOp()
+{
+    return cuda::minimum();
+}
+
+// Max operation compatibility wrapper
+__host__ __device__ __forceinline__ auto getCudaMaxOp()
+{
+    return cuda::maximum();
+}
+
+// Sum operation compatibility wrapper
+__host__ __device__ __forceinline__ auto getCudaSumOp()
+{
+    return cuda::std::plus<>();
+}
+
+} // namespace compat
+
+#endif // TRT_CUB_CCCL_COMPAT_H

plugin/common/cub_helper.h

@@ -0,0 +1,36 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+// TODO: Remove WAR once issue resolved in CUB (CUDA 12.6+?)
+
+#ifndef CUDA_VERSION
+#include <cuda.h>
+#endif // CUDA_VERSION
+
+#include "common/kernels/kernel.h"
+#include <cub/cub.cuh>
+template <typename KeyT, typename ValueT>
+size_t cubSortPairsWorkspaceSize(int32_t num_items, int32_t num_segments)
+{
+    size_t temp_storage_bytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairsDescending((void*) NULL, temp_storage_bytes, (KeyT const*) NULL,
+        (KeyT*) NULL, (ValueT const*) NULL, (ValueT*) NULL,
+        num_items,    // # items
+        num_segments, // # segments
+        (int32_t const*) NULL, (int32_t const*) NULL);
+    return temp_storage_bytes;
+}

plugin/common/cublasLtWrapper.cpp

@@ -0,0 +1,230 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "cublasLtWrapper.h"
+#include "common/checkMacrosPlugin.h"
+#include "cudaDriverWrapper.h"
+
+#if defined(_WIN32)
+#if !defined(WIN32_LEAN_AND_MEAN)
+#define WIN32_LEAN_AND_MEAN
+#endif // defined(WIN32_LEAN_AND_MEAN)
+#include <windows.h>
+#define dllOpen(name) (void*) LoadLibraryA(name)
+#define dllClose(handle) FreeLibrary(static_cast<HMODULE>(handle))
+#define dllGetSym(handle, name) GetProcAddress(static_cast<HMODULE>(handle), name)
+auto const kCUBLASLT_PLUGIN_LIBNAME
+    = std::string{"cublasLt64_"} + std::to_string(nvinfer1::getCudaLibVersionMaj()) + ".dll";
+#else // defined(_WIN32)
+#include <dlfcn.h>
+#define dllOpen(name) dlopen(name, RTLD_LAZY)
+#define dllClose(handle) dlclose(handle)
+#define dllGetSym(handle, name) dlsym(handle, name)
+auto const kCUBLASLT_PLUGIN_LIBNAME = std::string{"libcublasLt.so."} + std::to_string(nvinfer1::getCudaLibVersionMaj());
+#endif // defined(_WIN32)
+
+namespace nvinfer1::pluginInternal
+{
+using namespace nvinfer1;
+
+// If tryLoadingCublasLt failed, the CublasLtWrapper object won't be created.
+CublasLtWrapper::CublasLtWrapper(bool initHandle)
+    : mLibrary(tryLoadingCublasLt())
+{
+    PLUGIN_VALIDATE(mLibrary != nullptr);
+    auto load_sym = [](void* handle, char const* name) {
+        void* ret = dllGetSym(handle, name);
+        std::string loadError = "Fail to load symbol " + std::string(name) + " from the cublasLt library.";
+        PLUGIN_VALIDATE(ret != nullptr, loadError.c_str());
+        return ret;
+    };
+    *(void**) (&_cublasLtCreate) = load_sym(mLibrary, "cublasLtCreate");
+    *(void**) (&_cublasLtDestroy) = load_sym(mLibrary, "cublasLtDestroy");
+    *(void**) (&_cublasLtMatmul) = load_sym(mLibrary, "cublasLtMatmul");
+    *(void**) (&_cublasLtMatmulDescCreate) = load_sym(mLibrary, "cublasLtMatmulDescCreate");
+    *(void**) (&_cublasLtMatmulDescDestroy) = load_sym(mLibrary, "cublasLtMatmulDescDestroy");
+    *(void**) (&_cublasLtMatmulPreferenceCreate) = load_sym(mLibrary, "cublasLtMatmulPreferenceCreate");
+    *(void**) (&_cublasLtMatmulPreferenceDestroy) = load_sym(mLibrary, "cublasLtMatmulPreferenceDestroy");
+    *(void**) (&_cublasLtMatmulPreferenceSetAttribute) = load_sym(mLibrary, "cublasLtMatmulPreferenceSetAttribute");
+    *(void**) (&_cublasLtMatmulAlgoInit) = load_sym(mLibrary, "cublasLtMatmulAlgoInit");
+    *(void**) (&_cublasLtMatmulAlgoCheck) = load_sym(mLibrary, "cublasLtMatmulAlgoCheck");
+    *(void**) (&_cublasLtMatmulAlgoGetIds) = load_sym(mLibrary, "cublasLtMatmulAlgoGetIds");
+    *(void**) (&_cublasLtMatrixLayoutCreate) = load_sym(mLibrary, "cublasLtMatrixLayoutCreate");
+    *(void**) (&_cublasLtMatrixLayoutDestroy) = load_sym(mLibrary, "cublasLtMatrixLayoutDestroy");
+    *(void**) (&_cublasLtMatrixLayoutSetAttribute) = load_sym(mLibrary, "cublasLtMatrixLayoutSetAttribute");
+    *(void**) (&_cublasLtMatmulAlgoConfigSetAttribute) = load_sym(mLibrary, "cublasLtMatmulAlgoConfigSetAttribute");
+    *(void**) (&_cublasLtMatmulAlgoConfigGetAttribute) = load_sym(mLibrary, "cublasLtMatmulAlgoConfigGetAttribute");
+    *(void**) (&_cublasLtMatmulAlgoCapGetAttribute) = load_sym(mLibrary, "cublasLtMatmulAlgoCapGetAttribute");
+    *(void**) (&_cublasLtMatmulDescSetAttribute) = load_sym(mLibrary, "cublasLtMatmulDescSetAttribute");
+
+    if (initHandle)
+    {
+        PLUGIN_VALIDATE(cublasLtCreate(&mHandle) == CUBLAS_STATUS_SUCCESS, "Could not create cublasLt handle.");
+        PLUGIN_VALIDATE(mHandle != nullptr);
+    }
+}
+
+CublasLtWrapper::~CublasLtWrapper()
+{
+    if (mHandle != nullptr)
+    {
+        PLUGIN_VALIDATE(cublasLtDestroy(mHandle) == CUBLAS_STATUS_SUCCESS, "Could not destroy cublas handle.");
+        mHandle = nullptr;
+    }
+
+    if (mLibrary != nullptr)
+    {
+        dllClose(mLibrary);
+    }
+}
+
+void* CublasLtWrapper::tryLoadingCublasLt()
+{
+    void* cublasLtLib = dllOpen(kCUBLASLT_PLUGIN_LIBNAME.c_str());
+    std::string errorMsg = "Failed to load " + kCUBLASLT_PLUGIN_LIBNAME + ".";
+    PLUGIN_VALIDATE(cublasLtLib != nullptr, errorMsg.c_str());
+    return cublasLtLib;
+}
+
+cublasLtContext* CublasLtWrapper::getCublasLtHandle()
+{
+    return mHandle;
+}
+
+bool CublasLtWrapper::isValid() const
+{
+    return mHandle != nullptr;
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtCreate(cublasLtHandle_t* handle)
+{
+    return (*_cublasLtCreate)(handle);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtDestroy(cublasLtHandle_t handle)
+{
+    return (*_cublasLtDestroy)(handle);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmul(cublasLtHandle_t lightHandle, cublasLtMatmulDesc_t computeDesc,
+    void const* alpha, void const* A, cublasLtMatrixLayout_t Adesc, void const* B, cublasLtMatrixLayout_t Bdesc,
+    void const* beta, void const* C, cublasLtMatrixLayout_t Cdesc, void* D, cublasLtMatrixLayout_t Ddesc,
+    cublasLtMatmulAlgo_t const* algo, void* workspace, size_t workspaceSizeInBytes, cudaStream_t stream)
+{
+    return (*_cublasLtMatmul)(lightHandle, computeDesc, alpha, A, Adesc, B, Bdesc, beta, C, Cdesc, D, Ddesc, algo,
+        workspace, workspaceSizeInBytes, stream);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulDescCreate(
+    cublasLtMatmulDesc_t* matmulDesc, cublasComputeType_t computeType, cudaDataType_t scaleType)
+{
+    return (*_cublasLtMatmulDescCreate)(matmulDesc, computeType, scaleType);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulDescDestroy(cublasLtMatmulDesc_t matmulDesc)
+{
+    return (*_cublasLtMatmulDescDestroy)(matmulDesc);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulPreferenceCreate(cublasLtMatmulPreference_t* pref)
+{
+    return (*_cublasLtMatmulPreferenceCreate)(pref);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulPreferenceDestroy(cublasLtMatmulPreference_t pref)
+{
+    return (*_cublasLtMatmulPreferenceDestroy)(pref);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulPreferenceSetAttribute(
+    cublasLtMatmulPreference_t pref, cublasLtMatmulPreferenceAttributes_t attr, void const* buf, size_t sizeInBytes)
+{
+    return (*_cublasLtMatmulPreferenceSetAttribute)(pref, attr, buf, sizeInBytes);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoInit(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
+    cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype, cudaDataType_t Dtype,
+    int algoId, cublasLtMatmulAlgo_t* algo)
+{
+    return (*_cublasLtMatmulAlgoInit)(lightHandle, computeType, scaleType, Atype, Btype, Ctype, Dtype, algoId, algo);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoCheck(cublasLtHandle_t lightHandle,
+    cublasLtMatmulDesc_t operationDesc, cublasLtMatrixLayout_t Adesc, cublasLtMatrixLayout_t Bdesc,
+    cublasLtMatrixLayout_t Cdesc, cublasLtMatrixLayout_t Ddesc, cublasLtMatmulAlgo_t const* algo,
+    cublasLtMatmulHeuristicResult_t* result)
+{
+    return (*_cublasLtMatmulAlgoCheck)(lightHandle, operationDesc, Adesc, Bdesc, Cdesc, Ddesc, algo, result);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoGetIds(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
+    cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype, cudaDataType_t Dtype,
+    int requestedAlgoCount, int algoIdsArray[], int* returnAlgoCount)
+{
+    return (*_cublasLtMatmulAlgoGetIds)(lightHandle, computeType, scaleType, Atype, Btype, Ctype, Dtype,
+        requestedAlgoCount, algoIdsArray, returnAlgoCount);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatrixLayoutCreate(
+    cublasLtMatrixLayout_t* matLayout, cudaDataType type, uint64_t rows, uint64_t cols, int64_t ld)
+{
+    return (*_cublasLtMatrixLayoutCreate)(matLayout, type, rows, cols, ld);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatrixLayoutDestroy(cublasLtMatrixLayout_t matLayout)
+{
+    return (*_cublasLtMatrixLayoutDestroy)(matLayout);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatrixLayoutSetAttribute(
+    cublasLtMatrixLayout_t matLayout, cublasLtMatrixLayoutAttribute_t attr, void const* buf, size_t sizeInBytes)
+{
+    return (*_cublasLtMatrixLayoutSetAttribute)(matLayout, attr, buf, sizeInBytes);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoConfigGetAttribute(cublasLtMatmulAlgo_t const* algo,
+    cublasLtMatmulAlgoConfigAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten)
+{
+    return (*_cublasLtMatmulAlgoConfigGetAttribute)(algo, attr, buf, sizeInBytes, sizeWritten);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoConfigSetAttribute(
+    cublasLtMatmulAlgo_t* algo, cublasLtMatmulAlgoConfigAttributes_t attr, void const* buf, size_t sizeInBytes)
+{
+    return (*_cublasLtMatmulAlgoConfigSetAttribute)(algo, attr, buf, sizeInBytes);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulAlgoCapGetAttribute(cublasLtMatmulAlgo_t const* algo,
+    cublasLtMatmulAlgoCapAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten)
+{
+    return (*_cublasLtMatmulAlgoCapGetAttribute)(algo, attr, buf, sizeInBytes, sizeWritten);
+}
+
+cublasStatus_t CublasLtWrapper::cublasLtMatmulDescSetAttribute(
+    cublasLtMatmulDesc_t matmulDesc, cublasLtMatmulDescAttributes_t attr, void const* buf, size_t sizeInBytes)
+{
+    return (*_cublasLtMatmulDescSetAttribute)(matmulDesc, attr, buf, sizeInBytes);
+}
+
+CublasLtWrapper& getCublasLtWrapper()
+{
+    // Initialize a global cublasLtWrapper instance to be used to call cublasLt functions.
+    static CublasLtWrapper sGCublasLtWrapper;
+    return sGCublasLtWrapper;
+}
+
+} // namespace nvinfer1::pluginInternal

plugin/common/cublasLtWrapper.h

@@ -0,0 +1,327 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_PLUGIN_CUBLASLT_WRAPPER_H
+#define TRT_PLUGIN_CUBLASLT_WRAPPER_H
+
+#include "NvInferPlugin.h"
+#include "cublasWrapper.h"
+#include <cuda_fp16.h>
+#include <library_types.h>
+#include <string>
+
+extern "C"
+{
+    struct cublasLtContext;
+}
+
+namespace nvinfer1
+{
+namespace pluginInternal
+{
+
+struct cublasLtMatmulAlgo
+{
+    uint64_t data[8];
+};
+
+using cublasLtMatmulAlgo_t = cublasLtMatmulAlgo;
+
+struct cublasLtMatrixLayoutOpaque
+{
+    uint64_t data[8];
+};
+
+struct cublasLtMatmulPreferenceOpaque
+{
+    uint64_t data[8];
+};
+
+struct cublasLtMatmulDescOpaque
+{
+    uint64_t data[32];
+};
+
+struct cublasLtMatmulHeuristicResult
+{
+    cublasLtMatmulAlgo_t algo;
+    size_t workspaceSize;
+    cublasStatus_t state;
+    float wavesCount;
+    int reserved[4];
+};
+
+/* Copy of CUBLASLT cublasLtReductionScheme_t */
+enum cublasLtReductionScheme
+{
+    CUBLASLT_REDUCTION_SCHEME_NONE = 0,
+    CUBLASLT_REDUCTION_SCHEME_INPLACE = 1,
+    CUBLASLT_REDUCTION_SCHEME_COMPUTE_TYPE = 2,
+    CUBLASLT_REDUCTION_SCHEME_OUTPUT_TYPE = 4,
+    CUBLASLT_REDUCTION_SCHEME_MASK = 0x7,
+};
+
+/* Copy of CUBLASLT cublasLtMatrixLayoutAttribute_t */
+enum cublasLtMatrixLayoutAttribute
+{
+    CUBLASLT_MATRIX_LAYOUT_TYPE = 0,
+    CUBLASLT_MATRIX_LAYOUT_ORDER = 1,
+    CUBLASLT_MATRIX_LAYOUT_ROWS = 2,
+    CUBLASLT_MATRIX_LAYOUT_COLS = 3,
+    CUBLASLT_MATRIX_LAYOUT_LD = 4,
+    CUBLASLT_MATRIX_LAYOUT_BATCH_COUNT = 5,
+    CUBLASLT_MATRIX_LAYOUT_STRIDED_BATCH_OFFSET = 6,
+    CUBLASLT_MATRIX_LAYOUT_PLANE_OFFSET = 7,
+};
+
+enum cublasLtMatmulAlgoConfigAttributes
+{
+    CUBLASLT_ALGO_CONFIG_ID = 0,
+    CUBLASLT_ALGO_CONFIG_TILE_ID = 1,
+    CUBLASLT_ALGO_CONFIG_SPLITK_NUM = 2,
+    CUBLASLT_ALGO_CONFIG_REDUCTION_SCHEME = 3,
+    CUBLASLT_ALGO_CONFIG_CTA_SWIZZLING = 4,
+    CUBLASLT_ALGO_CONFIG_CUSTOM_OPTION = 5,
+    CUBLASLT_ALGO_CONFIG_STAGES_ID = 6,
+    CUBLASLT_ALGO_CONFIG_INNER_SHAPE_ID = 7,
+    CUBLASLT_ALGO_CONFIG_CLUSTER_SHAPE_ID = 8,
+};
+
+enum cublasLtMatmulPreferenceAttributes
+{
+    CUBLASLT_MATMUL_PREF_SEARCH_MODE = 0,
+    CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES = 1,
+    CUBLASLT_MATMUL_PREF_REDUCTION_SCHEME_MASK = 3,
+    CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_A_BYTES = 5,
+    CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_B_BYTES = 6,
+    CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_C_BYTES = 7,
+    CUBLASLT_MATMUL_PREF_MIN_ALIGNMENT_D_BYTES = 8,
+    CUBLASLT_MATMUL_PREF_MAX_WAVES_COUNT = 9,
+    CUBLASLT_MATMUL_PREF_IMPL_MASK = 12,
+};
+
+enum cublasLtMatmulTile
+{
+    CUBLASLT_MATMUL_TILE_UNDEFINED = 0,
+    CUBLASLT_MATMUL_TILE_8x8 = 1,
+    CUBLASLT_MATMUL_TILE_8x16 = 2,
+    CUBLASLT_MATMUL_TILE_16x8 = 3,
+    CUBLASLT_MATMUL_TILE_8x32 = 4,
+    CUBLASLT_MATMUL_TILE_16x16 = 5,
+    CUBLASLT_MATMUL_TILE_32x8 = 6,
+    CUBLASLT_MATMUL_TILE_8x64 = 7,
+    CUBLASLT_MATMUL_TILE_16x32 = 8,
+    CUBLASLT_MATMUL_TILE_32x16 = 9,
+    CUBLASLT_MATMUL_TILE_64x8 = 10,
+    CUBLASLT_MATMUL_TILE_32x32 = 11,
+    CUBLASLT_MATMUL_TILE_32x64 = 12,
+    CUBLASLT_MATMUL_TILE_64x32 = 13,
+    CUBLASLT_MATMUL_TILE_32x128 = 14,
+    CUBLASLT_MATMUL_TILE_64x64 = 15,
+    CUBLASLT_MATMUL_TILE_128x32 = 16,
+    CUBLASLT_MATMUL_TILE_64x128 = 17,
+    CUBLASLT_MATMUL_TILE_128x64 = 18,
+    CUBLASLT_MATMUL_TILE_64x256 = 19,
+    CUBLASLT_MATMUL_TILE_128x128 = 20,
+    CUBLASLT_MATMUL_TILE_256x64 = 21,
+    CUBLASLT_MATMUL_TILE_64x512 = 22,
+    CUBLASLT_MATMUL_TILE_128x256 = 23,
+    CUBLASLT_MATMUL_TILE_256x128 = 24,
+    CUBLASLT_MATMUL_TILE_512x64 = 25,
+    CUBLASLT_MATMUL_TILE_64x96 = 26,
+    CUBLASLT_MATMUL_TILE_96x64 = 27,
+    CUBLASLT_MATMUL_TILE_96x128 = 28,
+    CUBLASLT_MATMUL_TILE_128x160 = 29,
+    CUBLASLT_MATMUL_TILE_160x128 = 30,
+    CUBLASLT_MATMUL_TILE_192x128 = 31,
+    CUBLASLT_MATMUL_TILE_128x192 = 32,
+    CUBLASLT_MATMUL_TILE_128x96 = 33,
+    CUBLASLT_MATMUL_TILE_32x256 = 34,
+    CUBLASLT_MATMUL_TILE_256x32 = 35,
+    CUBLASLT_MATMUL_TILE_END
+};
+
+enum cublasLtMatmulAlgoCapAttributes
+{
+    CUBLASLT_ALGO_CAP_SPLITK_SUPPORT = 0,
+    CUBLASLT_ALGO_CAP_REDUCTION_SCHEME_MASK = 1,
+    CUBLASLT_ALGO_CAP_CTA_SWIZZLING_SUPPORT = 2,
+    CUBLASLT_ALGO_CAP_STRIDED_BATCH_SUPPORT = 3,
+    CUBLASLT_ALGO_CAP_OUT_OF_PLACE_RESULT_SUPPORT = 4,
+    CUBLASLT_ALGO_CAP_UPLO_SUPPORT = 5,
+    CUBLASLT_ALGO_CAP_TILE_IDS = 6,
+    CUBLASLT_ALGO_CAP_CUSTOM_OPTION_MAX = 7,
+    CUBLASLT_ALGO_CAP_CUSTOM_MEMORY_ORDER = 10,
+    CUBLASLT_ALGO_CAP_POINTER_MODE_MASK = 11,
+    CUBLASLT_ALGO_CAP_EPILOGUE_MASK = 12,
+    CUBLASLT_ALGO_CAP_STAGES_IDS = 13,
+    CUBLASLT_ALGO_CAP_LD_NEGATIVE = 14,
+    CUBLASLT_ALGO_CAP_NUMERICAL_IMPL_FLAGS = 15,
+    CUBLASLT_ALGO_CAP_MIN_ALIGNMENT_A_BYTES = 16,
+    CUBLASLT_ALGO_CAP_MIN_ALIGNMENT_B_BYTES = 17,
+    CUBLASLT_ALGO_CAP_MIN_ALIGNMENT_C_BYTES = 18,
+    CUBLASLT_ALGO_CAP_MIN_ALIGNMENT_D_BYTES = 19,
+    CUBLASLT_ALGO_CAP_ATOMIC_SYNC = 20,
+};
+
+enum cublasLtMatmulDescAttributes
+{
+    CUBLASLT_MATMUL_DESC_COMPUTE_TYPE = 0,
+    CUBLASLT_MATMUL_DESC_SCALE_TYPE = 1,
+    CUBLASLT_MATMUL_DESC_POINTER_MODE = 2,
+    CUBLASLT_MATMUL_DESC_TRANSA = 3,
+    CUBLASLT_MATMUL_DESC_TRANSB = 4,
+    CUBLASLT_MATMUL_DESC_TRANSC = 5,
+    CUBLASLT_MATMUL_DESC_FILL_MODE = 6,
+    CUBLASLT_MATMUL_DESC_EPILOGUE = 7,
+    CUBLASLT_MATMUL_DESC_BIAS_POINTER = 8,
+    CUBLASLT_MATMUL_DESC_BIAS_BATCH_STRIDE = 10,
+    CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER = 11,
+    CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD = 12,
+    CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_BATCH_STRIDE = 13,
+    CUBLASLT_MATMUL_DESC_ALPHA_VECTOR_BATCH_STRIDE = 14,
+    CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET = 15,
+    CUBLASLT_MATMUL_DESC_A_SCALE_POINTER = 17,
+    CUBLASLT_MATMUL_DESC_B_SCALE_POINTER = 18,
+    CUBLASLT_MATMUL_DESC_C_SCALE_POINTER = 19,
+    CUBLASLT_MATMUL_DESC_D_SCALE_POINTER = 20,
+    CUBLASLT_MATMUL_DESC_AMAX_D_POINTER = 21,
+    CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_DATA_TYPE = 22,
+    CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_SCALE_POINTER = 23,
+    CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_AMAX_POINTER = 24,
+    CUBLASLT_MATMUL_DESC_FAST_ACCUM = 25,
+    CUBLASLT_MATMUL_DESC_BIAS_DATA_TYPE = 26,
+    CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_NUM_CHUNKS_D_ROWS = 27,
+    CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_NUM_CHUNKS_D_COLS = 28,
+    CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_IN_COUNTERS_POINTER = 29,
+    CUBLASLT_MATMUL_DESC_ATOMIC_SYNC_OUT_COUNTERS_POINTER = 30,
+};
+
+using cublasLtMatrixLayoutOpaque_t = cublasLtMatrixLayoutOpaque;
+using cublasLtMatrixLayout_t = cublasLtMatrixLayoutOpaque_t*;
+using cublasLtMatmulPreferenceOpaque_t = cublasLtMatmulPreferenceOpaque;
+using cublasLtMatmulPreference_t = cublasLtMatmulPreferenceOpaque_t*;
+using cublasLtMatmulDescOpaque_t = cublasLtMatmulDescOpaque;
+using cublasLtMatmulDesc_t = cublasLtMatmulDescOpaque_t*;
+using cublasLtMatmulHeuristicResult_t = cublasLtMatmulHeuristicResult;
+using cublasLtReductionScheme_t = cublasLtReductionScheme;
+using cublasLtHandle_t = struct cublasLtContext*;
+using cublasLtMatrixLayoutAttribute_t = cublasLtMatrixLayoutAttribute;
+using cublasLtMatmulAlgoConfigAttributes_t = cublasLtMatmulAlgoConfigAttributes;
+using cublasLtMatmulPreferenceAttributes_t = cublasLtMatmulPreferenceAttributes;
+using cublasLtMatmulTile_t = cublasLtMatmulTile;
+using cublasLtMatmulAlgoCapAttributes_t = cublasLtMatmulAlgoCapAttributes;
+using cublasLtMatmulDescAttributes_t = cublasLtMatmulDescAttributes;
+
+/* Copy of CUBLASLT_NUMERICAL_IMPL_FLAGS_FMA */
+constexpr auto CUBLASLT_NUMERICAL_IMPL_FLAGS_FMA = 0x01ull << 0;
+/* Copy of CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA */
+constexpr auto CUBLASLT_NUMERICAL_IMPL_FLAGS_HMMA = 0x02ull << 0;
+
+class CublasLtWrapper
+{
+public:
+    CublasLtWrapper(bool initHandle = false);
+    ~CublasLtWrapper();
+
+    cublasLtContext* getCublasLtHandle();
+    bool isValid() const;
+
+    cublasStatus_t cublasLtCreate(cublasLtHandle_t* handle);
+    cublasStatus_t cublasLtDestroy(cublasLtHandle_t handle);
+    cublasStatus_t cublasLtMatmul(cublasLtHandle_t lightHandle, cublasLtMatmulDesc_t computeDesc, void const* alpha,
+        void const* A, cublasLtMatrixLayout_t Adesc, void const* B, cublasLtMatrixLayout_t Bdesc, void const* beta,
+        void const* C, cublasLtMatrixLayout_t Cdesc, void* D, cublasLtMatrixLayout_t Ddesc,
+        cublasLtMatmulAlgo_t const* algo, void* workspace, size_t workspaceSizeInBytes, cudaStream_t stream);
+    cublasStatus_t cublasLtMatmulDescCreate(
+        cublasLtMatmulDesc_t* matmulDesc, cublasComputeType_t computeType, cudaDataType_t scaleType);
+    cublasStatus_t cublasLtMatmulDescDestroy(cublasLtMatmulDesc_t matmulDesc);
+    cublasStatus_t cublasLtMatmulPreferenceCreate(cublasLtMatmulPreference_t* pref);
+    cublasStatus_t cublasLtMatmulPreferenceDestroy(cublasLtMatmulPreference_t pref);
+    cublasStatus_t cublasLtMatmulPreferenceSetAttribute(cublasLtMatmulPreference_t pref,
+        cublasLtMatmulPreferenceAttributes_t attr, void const* buf, size_t sizeInBytes);
+    cublasStatus_t cublasLtMatmulAlgoInit(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
+        cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype,
+        cudaDataType_t Dtype, int algoId, cublasLtMatmulAlgo_t* algo);
+    cublasStatus_t cublasLtMatmulAlgoCheck(cublasLtHandle_t lightHandle, cublasLtMatmulDesc_t operationDesc,
+        cublasLtMatrixLayout_t Adesc, cublasLtMatrixLayout_t Bdesc, cublasLtMatrixLayout_t Cdesc,
+        cublasLtMatrixLayout_t Ddesc, cublasLtMatmulAlgo_t const* algo, cublasLtMatmulHeuristicResult_t* result);
+    cublasStatus_t cublasLtMatmulAlgoGetIds(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
+        cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype,
+        cudaDataType_t Dtype, int requestedAlgoCount, int algoIdsArray[], int* returnAlgoCount);
+    cublasStatus_t cublasLtMatrixLayoutCreate(
+        cublasLtMatrixLayout_t* matLayout, cudaDataType type, uint64_t rows, uint64_t cols, int64_t ld);
+    cublasStatus_t cublasLtMatrixLayoutDestroy(cublasLtMatrixLayout_t matLayout);
+    cublasStatus_t cublasLtMatrixLayoutSetAttribute(
+        cublasLtMatrixLayout_t matLayout, cublasLtMatrixLayoutAttribute_t attr, void const* buf, size_t sizeInBytes);
+    cublasStatus_t cublasLtMatmulAlgoConfigGetAttribute(cublasLtMatmulAlgo_t const* algo,
+        cublasLtMatmulAlgoConfigAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten);
+    cublasStatus_t cublasLtMatmulAlgoConfigSetAttribute(
+        cublasLtMatmulAlgo_t* algo, cublasLtMatmulAlgoConfigAttributes_t attr, void const* buf, size_t sizeInBytes);
+    cublasStatus_t cublasLtMatmulAlgoCapGetAttribute(cublasLtMatmulAlgo_t const* algo,
+        cublasLtMatmulAlgoCapAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten);
+    cublasStatus_t cublasLtMatmulDescSetAttribute(
+        cublasLtMatmulDesc_t matmulDesc, cublasLtMatmulDescAttributes_t attr, void const* buf, size_t sizeInBytes);
+
+private:
+    void* mLibrary{nullptr};
+    cublasLtContext* mHandle{nullptr};
+    void* tryLoadingCublasLt();
+
+    cublasStatus_t (*_cublasLtCreate)(cublasLtHandle_t*);
+    cublasStatus_t (*_cublasLtDestroy)(cublasLtHandle_t);
+    cublasStatus_t (*_cublasLtMatmul)(cublasLtHandle_t lightHandle, cublasLtMatmulDesc_t computeDesc, void const* alpha,
+        void const* A, cublasLtMatrixLayout_t Adesc, void const* B, cublasLtMatrixLayout_t Bdesc, void const* beta,
+        void const* C, cublasLtMatrixLayout_t Cdesc, void* D, cublasLtMatrixLayout_t Ddesc,
+        cublasLtMatmulAlgo_t const* algo, void* workspace, size_t workspaceSizeInBytes, cudaStream_t stream);
+    cublasStatus_t (*_cublasLtMatmulDescCreate)(
+        cublasLtMatmulDesc_t* matmulDesc, cublasComputeType_t computeType, cudaDataType_t scaleType);
+    cublasStatus_t (*_cublasLtMatmulDescDestroy)(cublasLtMatmulDesc_t matmulDesc);
+    cublasStatus_t (*_cublasLtMatmulPreferenceCreate)(cublasLtMatmulPreference_t* pref);
+    cublasStatus_t (*_cublasLtMatmulPreferenceDestroy)(cublasLtMatmulPreference_t pref);
+    cublasStatus_t (*_cublasLtMatmulPreferenceSetAttribute)(cublasLtMatmulPreference_t pref,
+        cublasLtMatmulPreferenceAttributes_t attr, void const* buf, size_t sizeInBytes);
+    cublasStatus_t (*_cublasLtMatmulAlgoInit)(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
+        cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype,
+        cudaDataType_t Dtype, int algoId, cublasLtMatmulAlgo_t* algo);
+    cublasStatus_t (*_cublasLtMatmulAlgoCheck)(cublasLtHandle_t lightHandle, cublasLtMatmulDesc_t operationDesc,
+        cublasLtMatrixLayout_t Adesc, cublasLtMatrixLayout_t Bdesc, cublasLtMatrixLayout_t Cdesc,
+        cublasLtMatrixLayout_t Ddesc, cublasLtMatmulAlgo_t const* algo, cublasLtMatmulHeuristicResult_t* result);
+    cublasStatus_t (*_cublasLtMatmulAlgoGetIds)(cublasLtHandle_t lightHandle, cublasComputeType_t computeType,
+        cudaDataType_t scaleType, cudaDataType_t Atype, cudaDataType_t Btype, cudaDataType_t Ctype,
+        cudaDataType_t Dtype, int requestedAlgoCount, int algoIdsArray[], int* returnAlgoCount);
+    cublasStatus_t (*_cublasLtMatrixLayoutCreate)(
+        cublasLtMatrixLayout_t* matLayout, cudaDataType type, uint64_t rows, uint64_t cols, int64_t ld);
+    cublasStatus_t (*_cublasLtMatrixLayoutDestroy)(cublasLtMatrixLayout_t matLayout);
+    cublasStatus_t (*_cublasLtMatrixLayoutSetAttribute)(
+        cublasLtMatrixLayout_t matLayout, cublasLtMatrixLayoutAttribute_t attr, void const* buf, size_t sizeInBytes);
+    cublasStatus_t (*_cublasLtMatmulAlgoConfigGetAttribute)(cublasLtMatmulAlgo_t const* algo,
+        cublasLtMatmulAlgoConfigAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten);
+    cublasStatus_t (*_cublasLtMatmulAlgoConfigSetAttribute)(
+        cublasLtMatmulAlgo_t* algo, cublasLtMatmulAlgoConfigAttributes_t attr, void const* buf, size_t sizeInBytes);
+    cublasStatus_t (*_cublasLtMatmulAlgoCapGetAttribute)(cublasLtMatmulAlgo_t const* algo,
+        cublasLtMatmulAlgoCapAttributes_t attr, void* buf, size_t sizeInBytes, size_t* sizeWritten);
+    cublasStatus_t (*_cublasLtMatmulDescSetAttribute)(
+        cublasLtMatmulDesc_t matmulDesc, cublasLtMatmulDescAttributes_t attr, void const* buf, size_t sizeInBytes);
+};
+
+CublasLtWrapper& getCublasLtWrapper();
+
+} // namespace pluginInternal
+} // namespace nvinfer1
+
+#endif // TRT_PLUGIN_CUBLASLT_WRAPPER_H

plugin/common/cublasWrapper.cpp

@@ -0,0 +1,225 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "cublasWrapper.h"
+#include "common/checkMacrosPlugin.h"
+#include "cudaDriverWrapper.h"
+
+#if defined(_WIN32)
+#if !defined(WIN32_LEAN_AND_MEAN)
+#define WIN32_LEAN_AND_MEAN
+// Ensure that macros appearing in multiple files are only defined once.
+#endif // defined(WIN32_LEAN_AND_MEAN)
+#include <windows.h>
+#define dllOpen(name) (void*) LoadLibraryA(name)
+#define dllClose(handle) FreeLibrary(static_cast<HMODULE>(handle))
+#define dllGetSym(handle, name) GetProcAddress(static_cast<HMODULE>(handle), name)
+auto const kCUBLAS_PLUGIN_LIBNAME
+    = std::string{"cublas64_"} + std::to_string(nvinfer1::getCudaLibVersionMaj()) + ".dll";
+#else // defined(_WIN32)
+#include <dlfcn.h>
+#define dllOpen(name) dlopen(name, RTLD_LAZY)
+#define dllClose(handle) dlclose(handle)
+#define dllGetSym(handle, name) dlsym(handle, name)
+auto const kCUBLAS_PLUGIN_LIBNAME = std::string{"libcublas.so."} + std::to_string(nvinfer1::getCudaLibVersionMaj());
+#endif // defined(_WIN32)
+
+namespace nvinfer1::pluginInternal
+{
+using namespace nvinfer1;
+
+// If tryLoadingCublas failed, the CublasWrapper object won't be created.
+CublasWrapper::CublasWrapper(bool initHandle)
+    : mLibrary(tryLoadingCublas())
+{
+    PLUGIN_VALIDATE(mLibrary != nullptr);
+    auto load_sym = [](void* handle, char const* name) {
+        void* ret = dllGetSym(handle, name);
+        std::string loadError = "Fail to load symbol " + std::string(name) + " from the cublas library.";
+        PLUGIN_VALIDATE(ret != nullptr, loadError.c_str());
+        return ret;
+    };
+    *(void**) (&_cublasCreate) = load_sym(mLibrary, "cublasCreate_v2");
+    *(void**) (&_cublasDestroy) = load_sym(mLibrary, "cublasDestroy_v2");
+    *(void**) (&_cublasSetStream) = load_sym(mLibrary, "cublasSetStream_v2");
+    *(void**) (&_cublasGetPointerMode) = load_sym(mLibrary, "cublasGetPointerMode_v2");
+    *(void**) (&_cublasSetPointerMode) = load_sym(mLibrary, "cublasSetPointerMode_v2");
+    *(void**) (&_cublasGetMathMode) = load_sym(mLibrary, "cublasGetMathMode");
+    *(void**) (&_cublasSetMathMode) = load_sym(mLibrary, "cublasSetMathMode");
+    *(void**) (&_cublasDscal) = load_sym(mLibrary, "cublasDscal_v2");
+    *(void**) (&_cublasSasum) = load_sym(mLibrary, "cublasSasum_v2");
+    *(void**) (&_cublasScopy) = load_sym(mLibrary, "cublasScopy_v2");
+    *(void**) (&_cublasSscal) = load_sym(mLibrary, "cublasSscal_v2");
+    *(void**) (&_cublasSgemm) = load_sym(mLibrary, "cublasSgemm_v2");
+    *(void**) (&_cublasHgemm) = load_sym(mLibrary, "cublasHgemm");
+    *(void**) (&_cublasHgemmStridedBatched) = load_sym(mLibrary, "cublasHgemmStridedBatched");
+    *(void**) (&_cublasSgemmStridedBatched) = load_sym(mLibrary, "cublasSgemmStridedBatched");
+    *(void**) (&_cublasGemmEx) = load_sym(mLibrary, "cublasGemmEx");
+    *(void**) (&_cublasGemmStridedBatchedEx) = load_sym(mLibrary, "cublasGemmStridedBatchedEx");
+
+    if (initHandle)
+    {
+        PLUGIN_VALIDATE(cublasCreate(&mHandle) == CUBLAS_STATUS_SUCCESS, "Could not create cublas handle.");
+        PLUGIN_VALIDATE(mHandle != nullptr);
+    }
+}
+
+CublasWrapper::~CublasWrapper()
+{
+    if (mHandle != nullptr)
+    {
+        PLUGIN_VALIDATE(cublasDestroy(mHandle) == CUBLAS_STATUS_SUCCESS, "Could not destroy cublas handle.");
+        mHandle = nullptr;
+    }
+
+    if (mLibrary != nullptr)
+    {
+        dllClose(mLibrary);
+    }
+}
+
+void* CublasWrapper::tryLoadingCublas()
+{
+    void* cublasLib = dllOpen(kCUBLAS_PLUGIN_LIBNAME.c_str());
+    std::string errorMsg = "Failed to load " + kCUBLAS_PLUGIN_LIBNAME + ".";
+    PLUGIN_VALIDATE(cublasLib != nullptr, errorMsg.c_str());
+    return cublasLib;
+}
+
+cublasContext* CublasWrapper::getCublasHandle()
+{
+    return mHandle;
+}
+
+bool CublasWrapper::isValid() const
+{
+    return mHandle != nullptr;
+}
+
+cublasStatus_t CublasWrapper::cublasCreate(cublasContext** handle)
+{
+    return (*_cublasCreate)(handle);
+}
+
+cublasStatus_t CublasWrapper::cublasDestroy(cublasContext* handle)
+{
+    return (*_cublasDestroy)(handle);
+}
+
+cublasStatus_t CublasWrapper::cublasSetStream(cublasHandle_t handle, cudaStream_t streamId)
+{
+    return (*_cublasSetStream)(handle, streamId);
+}
+
+cublasStatus_t CublasWrapper::cublasGetPointerMode(cublasHandle_t handle, cublasPointerMode_t* mode)
+{
+    return (*_cublasGetPointerMode)(handle, mode);
+}
+
+cublasStatus_t CublasWrapper::cublasSetPointerMode(cublasHandle_t handle, cublasPointerMode_t mode)
+{
+    return (*_cublasSetPointerMode)(handle, mode);
+}
+
+cublasStatus_t CublasWrapper::cublasGetMathMode(cublasHandle_t handle, cublasMath_t* mode)
+{
+    return (*_cublasGetMathMode)(handle, mode);
+}
+
+cublasStatus_t CublasWrapper::cublasSetMathMode(cublasHandle_t handle, cublasMath_t mode)
+{
+    return (*_cublasSetMathMode)(handle, mode);
+}
+
+cublasStatus_t CublasWrapper::cublasDscal(cublasHandle_t handle, int n, float const* alpha, float* x, int incx)
+{
+    return (*_cublasDscal)(handle, n, alpha, x, incx);
+}
+
+cublasStatus_t CublasWrapper::cublasSasum(cublasHandle_t handle, int n, float const* x, int incx, float* result)
+{
+    return (*_cublasSasum)(handle, n, x, incx, result);
+}
+
+cublasStatus_t CublasWrapper::cublasScopy(cublasHandle_t handle, int n, float const* x, int incx, float* y, int incy)
+{
+    return (*_cublasScopy)(handle, n, x, incx, y, incy);
+}
+
+cublasStatus_t CublasWrapper::cublasSscal(cublasHandle_t handle, int n, float const* alpha, float* x, int incx)
+{
+    return (*_cublasSscal)(handle, n, alpha, x, incx);
+}
+
+cublasStatus_t CublasWrapper::cublasSgemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+    int m, int n, int k, float const* alpha, float const* A, int lda, float const* B, int ldb, float const* beta,
+    float* C, int ldc)
+{
+    return (*_cublasSgemm)(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
+}
+
+cublasStatus_t CublasWrapper::cublasHgemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+    int m, int n, int k, __half const* alpha, __half const* A, int lda, __half const* B, int ldb, __half const* beta,
+    __half* C, int ldc)
+{
+    return (*_cublasHgemm)(handle, transa, transb, m, n, k, alpha, A, lda, B, ldb, beta, C, ldc);
+}
+
+cublasStatus_t CublasWrapper::cublasHgemmStridedBatched(cublasHandle_t handle, cublasOperation_t transa,
+    cublasOperation_t transb, int m, int n, int k, __half const* alpha, __half const* A, int lda, long long int strideA,
+    __half const* B, int ldb, long long int strideB, __half const* beta, __half* C, int ldc, long long int strideC,
+    int batchCount)
+{
+    return (*_cublasHgemmStridedBatched)(
+        handle, transa, transb, m, n, k, alpha, A, lda, strideA, B, ldb, strideB, beta, C, ldc, strideC, batchCount);
+}
+
+cublasStatus_t CublasWrapper::cublasSgemmStridedBatched(cublasHandle_t handle, cublasOperation_t transa,
+    cublasOperation_t transb, int m, int n, int k, float const* alpha, float const* A, int lda, long long int strideA,
+    float const* B, int ldb, long long int strideB, float const* beta, float* C, int ldc, long long int strideC,
+    int batchCount)
+{
+    return (*_cublasSgemmStridedBatched)(
+        handle, transa, transb, m, n, k, alpha, A, lda, strideA, B, ldb, strideB, beta, C, ldc, strideC, batchCount);
+}
+
+cublasStatus_t CublasWrapper::cublasGemmEx(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+    int m, int n, int k, void const* alpha, void const* A, cudaDataType Atype, int lda, void const* B,
+    cudaDataType Btype, int ldb, void const* beta, void* C, cudaDataType Ctype, int ldc, cudaDataType computeType,
+    cublasGemmAlgo_t algo)
+{
+    return (*_cublasGemmEx)(
+        handle, transa, transb, m, n, k, alpha, A, Atype, lda, B, Btype, ldb, beta, C, Ctype, ldc, computeType, algo);
+}
+
+cublasStatus_t CublasWrapper::cublasGemmStridedBatchedEx(cublasHandle_t handle, cublasOperation_t transa,
+    cublasOperation_t transb, int m, int n, int k, void const* alpha, void const* A, cudaDataType Atype, int lda,
+    long long int strideA, void const* B, cudaDataType Btype, int ldb, long long int strideB, void const* beta, void* C,
+    cudaDataType Ctype, int ldc, long long int strideC, int batchCount, cudaDataType computeType, cublasGemmAlgo_t algo)
+{
+    return (*_cublasGemmStridedBatchedEx)(handle, transa, transb, m, n, k, alpha, A, Atype, lda, strideA, B, Btype, ldb,
+        strideB, beta, C, Ctype, ldc, strideC, batchCount, computeType, algo);
+}
+
+CublasWrapper& getCublasWrapper()
+{
+    // Initialize a global cublasWrapper instance to be used to call cublas functions.
+    static CublasWrapper sGCublasWrapper;
+    return sGCublasWrapper;
+}
+
+} // namespace nvinfer1::pluginInternal

plugin/common/cublasWrapper.h

@@ -0,0 +1,234 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_PLUGIN_CUBLAS_WRAPPER_H
+#define TRT_PLUGIN_CUBLAS_WRAPPER_H
+
+#include "NvInferPlugin.h"
+#include <cuda_fp16.h>
+#include <library_types.h>
+#include <string>
+
+namespace nvinfer1
+{
+namespace pluginInternal
+{
+/* Copy of CUBLAS status type returns */
+enum CublasStatus
+{
+    CUBLAS_STATUS_SUCCESS = 0,
+    CUBLAS_STATUS_NOT_INITIALIZED = 1,
+    CUBLAS_STATUS_ALLOC_FAILED = 3,
+    CUBLAS_STATUS_INVALID_VALUE = 7,
+    CUBLAS_STATUS_ARCH_MISMATCH = 8,
+    CUBLAS_STATUS_MAPPING_ERROR = 11,
+    CUBLAS_STATUS_EXECUTION_FAILED = 13,
+    CUBLAS_STATUS_INTERNAL_ERROR = 14,
+    CUBLAS_STATUS_NOT_SUPPORTED = 15,
+    CUBLAS_STATUS_LICENSE_ERROR = 16
+};
+
+/* Copy of CUBLAS math types*/
+enum CublasMath
+{
+    CUBLAS_DEFAULT_MATH = 0,
+    CUBLAS_TENSOR_OP_MATH = 1,
+    CUBLAS_PEDANTIC_MATH = 2,
+    CUBLAS_TF32_TENSOR_OP_MATH = 3,
+    CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION = 16,
+};
+
+/* Copy of CUBLAS operation types*/
+enum cublasOperation
+{
+    CUBLAS_OP_N = 0,
+    CUBLAS_OP_T = 1,
+    CUBLAS_OP_C = 2,
+    CUBLAS_OP_HERMITAN = 2,
+    CUBLAS_OP_CONJG = 3
+};
+
+/* Copy of CUBLAS pointer mode types*/
+enum cublasPointerMode
+{
+    CUBLAS_POINTER_MODE_HOST = 0,
+    CUBLAS_POINTER_MODE_DEVICE = 1
+};
+
+/* Copy of CUBLAS compute types*/
+enum cublasComputeType
+{
+    CUBLAS_COMPUTE_16F = 64,
+    CUBLAS_COMPUTE_16F_PEDANTIC = 65,
+    CUBLAS_COMPUTE_32F = 68,
+    CUBLAS_COMPUTE_32F_PEDANTIC = 69,
+    CUBLAS_COMPUTE_32F_FAST_16F = 74,
+    CUBLAS_COMPUTE_32F_FAST_16BF = 75,
+    CUBLAS_COMPUTE_32F_FAST_TF32 = 77,
+    CUBLAS_COMPUTE_64F = 70,
+    CUBLAS_COMPUTE_64F_PEDANTIC = 71,
+    CUBLAS_COMPUTE_32I = 72,
+    CUBLAS_COMPUTE_32I_PEDANTIC = 73,
+};
+
+/* Copy of CUBLAS GEMM algorithm types*/
+enum cublasGemmAlgo
+{
+    CUBLAS_GEMM_DFALT = -1,
+    CUBLAS_GEMM_DEFAULT = -1,
+    CUBLAS_GEMM_ALGO0 = 0,
+    CUBLAS_GEMM_ALGO1 = 1,
+    CUBLAS_GEMM_ALGO2 = 2,
+    CUBLAS_GEMM_ALGO3 = 3,
+    CUBLAS_GEMM_ALGO4 = 4,
+    CUBLAS_GEMM_ALGO5 = 5,
+    CUBLAS_GEMM_ALGO6 = 6,
+    CUBLAS_GEMM_ALGO7 = 7,
+    CUBLAS_GEMM_ALGO8 = 8,
+    CUBLAS_GEMM_ALGO9 = 9,
+    CUBLAS_GEMM_ALGO10 = 10,
+    CUBLAS_GEMM_ALGO11 = 11,
+    CUBLAS_GEMM_ALGO12 = 12,
+    CUBLAS_GEMM_ALGO13 = 13,
+    CUBLAS_GEMM_ALGO14 = 14,
+    CUBLAS_GEMM_ALGO15 = 15,
+    CUBLAS_GEMM_ALGO16 = 16,
+    CUBLAS_GEMM_ALGO17 = 17,
+    CUBLAS_GEMM_ALGO18 = 18,
+    CUBLAS_GEMM_ALGO19 = 19,
+    CUBLAS_GEMM_ALGO20 = 20,
+    CUBLAS_GEMM_ALGO21 = 21,
+    CUBLAS_GEMM_ALGO22 = 22,
+    CUBLAS_GEMM_ALGO23 = 23,
+    CUBLAS_GEMM_DEFAULT_TENSOR_OP = 99,
+    CUBLAS_GEMM_DFALT_TENSOR_OP = 99,
+    CUBLAS_GEMM_ALGO0_TENSOR_OP = 100,
+    CUBLAS_GEMM_ALGO1_TENSOR_OP = 101,
+    CUBLAS_GEMM_ALGO2_TENSOR_OP = 102,
+    CUBLAS_GEMM_ALGO3_TENSOR_OP = 103,
+    CUBLAS_GEMM_ALGO4_TENSOR_OP = 104,
+    CUBLAS_GEMM_ALGO5_TENSOR_OP = 105,
+    CUBLAS_GEMM_ALGO6_TENSOR_OP = 106,
+    CUBLAS_GEMM_ALGO7_TENSOR_OP = 107,
+    CUBLAS_GEMM_ALGO8_TENSOR_OP = 108,
+    CUBLAS_GEMM_ALGO9_TENSOR_OP = 109,
+    CUBLAS_GEMM_ALGO10_TENSOR_OP = 110,
+    CUBLAS_GEMM_ALGO11_TENSOR_OP = 111,
+    CUBLAS_GEMM_ALGO12_TENSOR_OP = 112,
+    CUBLAS_GEMM_ALGO13_TENSOR_OP = 113,
+    CUBLAS_GEMM_ALGO14_TENSOR_OP = 114,
+    CUBLAS_GEMM_ALGO15_TENSOR_OP = 115
+};
+
+using cublasStatus_t = CublasStatus;
+using cublasMath_t = CublasMath;
+using cublasOperation_t = cublasOperation;
+using cublasPointerMode_t = cublasPointerMode;
+using cublasComputeType_t = cublasComputeType;
+using cublasGemmAlgo_t = cublasGemmAlgo;
+using cublasDataType_t = cudaDataType;
+using cublasHandle_t = struct cublasContext*;
+
+class CublasWrapper
+{
+public:
+    CublasWrapper(bool initHandle = false);
+    ~CublasWrapper();
+
+    cublasContext* getCublasHandle();
+    bool isValid() const;
+
+    cublasStatus_t cublasCreate(cublasContext** handle);
+    cublasStatus_t cublasDestroy(cublasContext* handle);
+    cublasStatus_t cublasSetStream(cublasHandle_t handle, cudaStream_t streamId);
+    cublasStatus_t cublasGetPointerMode(cublasHandle_t handle, cublasPointerMode_t* mode);
+    cublasStatus_t cublasSetPointerMode(cublasHandle_t handle, cublasPointerMode_t mode);
+    cublasStatus_t cublasGetMathMode(cublasHandle_t handle, cublasMath_t* mode);
+    cublasStatus_t cublasSetMathMode(cublasHandle_t handle, cublasMath_t mode);
+    cublasStatus_t cublasDscal(cublasHandle_t handle, int n, float const* alpha, float* x, int incx);
+    cublasStatus_t cublasSasum(cublasHandle_t handle, int n, float const* x, int incx, float* result);
+    cublasStatus_t cublasScopy(cublasHandle_t handle, int n, float const* x, int incx, float* y, int incy);
+    cublasStatus_t cublasSscal(cublasHandle_t handle, int n, float const* alpha, float* x, int incx);
+    cublasStatus_t cublasSgemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n,
+        int k, float const* alpha, float const* A, int lda, float const* B, int ldb, float const* beta, float* C,
+        int ldc);
+    cublasStatus_t cublasHgemm(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n,
+        int k, __half const* alpha, __half const* A, int lda, __half const* B, int ldb, __half const* beta, __half* C,
+        int ldc);
+    cublasStatus_t cublasHgemmStridedBatched(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+        int m, int n, int k, __half const* alpha, __half const* A, int lda, long long int strideA, __half const* B,
+        int ldb, long long int strideB, __half const* beta, __half* C, int ldc, long long int strideC, int batchCount);
+    cublasStatus_t cublasSgemmStridedBatched(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+        int m, int n, int k, float const* alpha, float const* A, int lda, long long int strideA, float const* B,
+        int ldb, long long int strideB, float const* beta, float* C, int ldc, long long int strideC, int batchCount);
+    cublasStatus_t cublasGemmEx(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m, int n,
+        int k, void const* alpha, void const* A, cudaDataType Atype, int lda, void const* B, cudaDataType Btype,
+        int ldb, void const* beta, void* C, cudaDataType Ctype, int ldc, cudaDataType computeType,
+        cublasGemmAlgo_t algo);
+    cublasStatus_t cublasGemmStridedBatchedEx(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb,
+        int m, int n, int k, void const* alpha, void const* A, cudaDataType Atype, int lda, long long int strideA,
+        void const* B, cudaDataType Btype, int ldb, long long int strideB, void const* beta, void* C,
+        cudaDataType Ctype, int ldc, long long int strideC, int batchCount, cudaDataType computeType,
+        cublasGemmAlgo_t algo);
+
+private:
+    void* mLibrary{nullptr};
+    cublasContext* mHandle{nullptr};
+    void* tryLoadingCublas();
+
+    cublasStatus_t (*_cublasCreate)(cublasContext**);
+    cublasStatus_t (*_cublasDestroy)(cublasContext*);
+    cublasStatus_t (*_cublasSetStream)(cublasHandle_t handle, cudaStream_t streamId);
+    cublasStatus_t (*_cublasGetPointerMode)(cublasHandle_t handle, cublasPointerMode_t* mode);
+    cublasStatus_t (*_cublasSetPointerMode)(cublasHandle_t handle, cublasPointerMode_t mode);
+    cublasStatus_t (*_cublasGetMathMode)(cublasHandle_t handle, cublasMath_t* mode);
+    cublasStatus_t (*_cublasSetMathMode)(cublasHandle_t handle, cublasMath_t mode);
+    cublasStatus_t (*_cublasDscal)(cublasHandle_t handle, int n, float const* alpha, float* x, int incx);
+    cublasStatus_t (*_cublasSasum)(cublasHandle_t handle, int n, float const* x, int incx, float* result);
+    cublasStatus_t (*_cublasScopy)(cublasHandle_t handle, int n, float const* x, int incx, float* y, int incy);
+    cublasStatus_t (*_cublasSscal)(cublasHandle_t handle, int n, float const* alpha, float* x, int incx);
+    cublasStatus_t (*_cublasSgemm)(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m,
+        int n, int k, float const* alpha, float const* A, int lda, float const* B, int ldb, float const* beta, float* C,
+        int ldc);
+    cublasStatus_t (*_cublasHgemm)(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m,
+        int n, int k, __half const* alpha, __half const* A, int lda, __half const* B, int ldb, __half const* beta,
+        __half* C, int ldc);
+    cublasStatus_t (*_cublasHgemmStridedBatched)(cublasHandle_t handle, cublasOperation_t transa,
+        cublasOperation_t transb, int m, int n, int k, __half const* alpha, __half const* A, int lda,
+        long long int strideA, __half const* B, int ldb, long long int strideB, __half const* beta, __half* C, int ldc,
+        long long int strideC, int batchCount);
+    cublasStatus_t (*_cublasSgemmStridedBatched)(cublasHandle_t handle, cublasOperation_t transa,
+        cublasOperation_t transb, int m, int n, int k, float const* alpha, float const* A, int lda,
+        long long int strideA, float const* B, int ldb, long long int strideB, float const* beta, float* C, int ldc,
+        long long int strideC, int batchCount);
+    cublasStatus_t (*_cublasGemmEx)(cublasHandle_t handle, cublasOperation_t transa, cublasOperation_t transb, int m,
+        int n, int k, void const* alpha, void const* A, cudaDataType Atype, int lda, void const* B, cudaDataType Btype,
+        int ldb, void const* beta, void* C, cudaDataType Ctype, int ldc, cudaDataType computeType,
+        cublasGemmAlgo_t algo);
+    cublasStatus_t (*_cublasGemmStridedBatchedEx)(cublasHandle_t handle, cublasOperation_t transa,
+        cublasOperation_t transb, int m, int n, int k, void const* alpha, void const* A, cudaDataType Atype, int lda,
+        long long int strideA, void const* B, cudaDataType Btype, int ldb, long long int strideB, void const* beta,
+        void* C, cudaDataType Ctype, int ldc, long long int strideC, int batchCount, cudaDataType computeType,
+        cublasGemmAlgo_t algo);
+};
+
+CublasWrapper& getCublasWrapper();
+
+} // namespace pluginInternal
+} // namespace nvinfer1
+
+#endif // TRT_PLUGIN_CUBLAS_WRAPPER_H

plugin/common/cudaDriverWrapper.cpp

@@ -0,0 +1,211 @@
+/*
+ * 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.
+ */
+#define CUDA_LIB_NAME "cuda"
+
+#if defined(_WIN32)
+#if !defined(WIN32_LEAN_AND_MEAN)
+#define WIN32_LEAN_AND_MEAN
+#endif // defined(WIN32_LEAN_AND_MEAN)
+#include <windows.h>
+#define dllOpen(name) (void*) LoadLibraryA("nv" name ".dll")
+#define dllClose(handle) FreeLibrary(static_cast<HMODULE>(handle))
+#define dllGetSym(handle, name) GetProcAddress(static_cast<HMODULE>(handle), name)
+#else // defined(_WIN32)
+#include <dlfcn.h>
+#define dllOpen(name) dlopen("lib" name ".so.1", RTLD_LAZY)
+#define dllClose(handle) dlclose(handle)
+#define dllGetSym(handle, name) dlsym(handle, name)
+#endif // defined(_WIN32)
+
+#include "common/cudaDriverWrapper.h"
+#include "common/plugin.h"
+#include <cstdint>
+#include <cstdio>
+#include <cuda.h>
+
+using namespace nvinfer1;
+
+CUDADriverWrapper::CUDADriverWrapper()
+{
+    handle = dllOpen(CUDA_LIB_NAME);
+    PLUGIN_ASSERT(handle != nullptr);
+
+    auto load_sym = [](void* handle, char const* name) {
+        void* ret = dllGetSym(handle, name);
+        PLUGIN_ASSERT(ret != nullptr);
+        return ret;
+    };
+
+    _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()
+{
+    dllClose(handle);
+}
+
+CUresult CUDADriverWrapper::cuGetErrorName(CUresult error, char const** pStr) const
+{
+    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);
+}
+
+CUresult CUDADriverWrapper::cuLinkComplete(CUlinkState state, void** cubinOut, size_t* sizeOut) const
+{
+    return (*_cuLinkComplete)(state, cubinOut, sizeOut);
+}
+
+CUresult CUDADriverWrapper::cuModuleUnload(CUmodule hmod) const
+{
+    return (*_cuModuleUnload)(hmod);
+}
+
+CUresult CUDADriverWrapper::cuLinkDestroy(CUlinkState state) const
+{
+    return (*_cuLinkDestroy)(state);
+}
+
+CUresult CUDADriverWrapper::cuModuleLoadData(CUmodule* module, void const* image) const
+{
+    return (*_cuModuleLoadData)(module, image);
+}
+
+CUresult CUDADriverWrapper::cuLinkCreate(
+    uint32_t numOptions, CUjit_option* options, void** optionValues, CUlinkState* stateOut) const
+{
+    return (*_cuLinkCreate)(numOptions, options, optionValues, stateOut);
+}
+
+CUresult CUDADriverWrapper::cuModuleGetFunction(CUfunction* hfunc, CUmodule hmod, char const* name) const
+{
+    return (*_cuModuleGetFunction)(hfunc, hmod, name);
+}
+
+CUresult CUDADriverWrapper::cuLinkAddFile(CUlinkState state, CUjitInputType type, char const* path, uint32_t numOptions,
+    CUjit_option* options, void** optionValues) const
+{
+    return (*_cuLinkAddFile)(state, type, path, numOptions, options, optionValues);
+}
+
+CUresult CUDADriverWrapper::cuLinkAddData(CUlinkState state, CUjitInputType type, void* data, size_t size,
+    char const* name, uint32_t numOptions, CUjit_option* options, void** optionValues) const
+{
+    return (*_cuLinkAddData)(state, type, data, size, name, numOptions, options, optionValues);
+}
+
+CUresult CUDADriverWrapper::cuLaunchCooperativeKernel(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) const
+{
+    return (*_cuLaunchCooperativeKernel)(
+        f, gridDimX, gridDimY, gridDimZ, blockDimX, blockDimY, blockDimZ, sharedMemBytes, hStream, kernelParams);
+}
+
+CUresult CUDADriverWrapper::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) const
+{
+    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

@@ -0,0 +1,251 @@
+/*
+ * 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 CUDA_DRIVER_WRAPPER_H
+#define CUDA_DRIVER_WRAPPER_H
+
+#include <cstdint>
+#include <cstdio>
+#include <cuda.h>
+#include <memory>
+#include <utility>
+
+#define cuErrCheck(stat, wrap)                                                                                         \
+    {                                                                                                                  \
+        nvinfer1::cuErrCheck_((stat), wrap, __FILE__, __LINE__);                                                       \
+    }
+
+namespace nvinfer1
+{
+class CUDADriverWrapper
+{
+public:
+    CUDADriverWrapper();
+
+    ~CUDADriverWrapper();
+
+    // Delete default copy constructor and copy assignment constructor
+    CUDADriverWrapper(CUDADriverWrapper const&) = delete;
+    CUDADriverWrapper& operator=(CUDADriverWrapper const&) = delete;
+
+    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;
+
+    CUresult cuModuleUnload(CUmodule hmod) const;
+
+    CUresult cuLinkDestroy(CUlinkState state) const;
+
+    CUresult cuModuleLoadData(CUmodule* module, void const* image) const;
+
+    CUresult cuLinkCreate(uint32_t numOptions, CUjit_option* options, void** optionValues, CUlinkState* stateOut) const;
+
+    CUresult cuModuleGetFunction(CUfunction* hfunc, CUmodule hmod, char const* name) const;
+
+    CUresult cuLinkAddFile(CUlinkState state, CUjitInputType type, char const* path, uint32_t numOptions,
+        CUjit_option* options, void** optionValues) const;
+
+    CUresult cuLinkAddData(CUlinkState state, CUjitInputType type, void* data, size_t size, char const* name,
+        uint32_t numOptions, CUjit_option* options, void** optionValues) const;
+
+    CUresult cuLaunchCooperativeKernel(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) const;
+
+    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) 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);
+    CUresult (*_cuLinkDestroy)(CUlinkState);
+    CUresult (*_cuLinkCreate)(uint32_t, CUjit_option*, void**, CUlinkState*);
+    CUresult (*_cuModuleLoadData)(CUmodule*, void const*);
+    CUresult (*_cuModuleGetFunction)(CUfunction*, CUmodule, char const*);
+    CUresult (*_cuLinkAddFile)(CUlinkState, CUjitInputType, char const*, uint32_t, CUjit_option*, void**);
+    CUresult (*_cuLinkAddData)(
+        CUlinkState, CUjitInputType, void*, size_t, char const*, uint32_t, CUjit_option*, void**);
+    CUresult (*_cuLaunchCooperativeKernel)(
+        CUfunction, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, uint32_t, CUstream, void**);
+    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)
+{
+    if (stat != CUDA_SUCCESS)
+    {
+        char const* msg = nullptr;
+        wrap.cuGetErrorName(stat, &msg);
+        fprintf(stderr, "CUDA Error: %s %s %d\n", msg, file, line);
+    }
+}
+
+//! Return CUDA major version
+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

@@ -0,0 +1,185 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "cudnnWrapper.h"
+#include "common/checkMacrosPlugin.h"
+#include "common/plugin.h"
+
+#define CUDNN_MAJOR 8
+#if defined(_WIN32)
+#if !defined(WIN32_LEAN_AND_MEAN)
+// Ensure that macros appearing in multiple files are only defined once.
+#define WIN32_LEAN_AND_MEAN
+#endif // defined(WIN32_LEAN_AND_MEAN)
+#include <windows.h>
+#define dllOpen(name) (void*) LoadLibraryA(name)
+#define dllClose(handle) FreeLibrary(static_cast<HMODULE>(handle))
+#define dllGetSym(handle, name) GetProcAddress(static_cast<HMODULE>(handle), name)
+auto const kCUDNN_PLUGIN_LIBNAME = std::string("cudnn64_") + std::to_string(CUDNN_MAJOR) + ".dll";
+#else // defined(_WIN32)
+#include <dlfcn.h>
+#define dllOpen(name) dlopen(name, RTLD_LAZY)
+#define dllClose(handle) dlclose(handle)
+#define dllGetSym(handle, name) dlsym(handle, name)
+auto const kCUDNN_PLUGIN_LIBNAME = std::string("libcudnn.so.") + std::to_string(CUDNN_MAJOR);
+#endif // defined(_WIN32)
+
+namespace nvinfer1::pluginInternal
+{
+// If tryLoadingCudnn failed, the CudnnWrapper object won't be created.
+CudnnWrapper::CudnnWrapper(bool initHandle, char const* callerPluginName)
+    : mLibrary(tryLoadingCudnn(callerPluginName))
+{
+    auto load_sym = [](void* handle, char const* name) {
+        void* ret = dllGetSym(handle, name);
+        std::string loadError = "Fail to load symbol " + std::string(name) + " from the cudnn library.";
+        PLUGIN_VALIDATE(ret != nullptr, loadError.c_str());
+        return ret;
+    };
+
+    *reinterpret_cast<void**>(&_cudnnCreate) = load_sym(mLibrary, "cudnnCreate");
+    *reinterpret_cast<void**>(&_cudnnDestroy) = load_sym(mLibrary, "cudnnDestroy");
+    *reinterpret_cast<void**>(&_cudnnCreateTensorDescriptor) = load_sym(mLibrary, "cudnnCreateTensorDescriptor");
+    *reinterpret_cast<void**>(&_cudnnDestroyTensorDescriptor) = load_sym(mLibrary, "cudnnDestroyTensorDescriptor");
+    *reinterpret_cast<void**>(&_cudnnSetStream) = load_sym(mLibrary, "cudnnSetStream");
+    *reinterpret_cast<void**>(&_cudnnBatchNormalizationForwardTraining)
+        = load_sym(mLibrary, "cudnnBatchNormalizationForwardTraining");
+    *reinterpret_cast<void**>(&_cudnnSetTensor4dDescriptor) = load_sym(mLibrary, "cudnnSetTensor4dDescriptor");
+    *reinterpret_cast<void**>(&_cudnnSetTensorNdDescriptor) = load_sym(mLibrary, "cudnnSetTensorNdDescriptor");
+    *reinterpret_cast<void**>(&_cudnnSetTensorNdDescriptorEx) = load_sym(mLibrary, "cudnnSetTensorNdDescriptorEx");
+    *reinterpret_cast<void**>(&_cudnnDeriveBNTensorDescriptor) = load_sym(mLibrary, "cudnnDeriveBNTensorDescriptor");
+    *reinterpret_cast<void**>(&_cudnnGetErrorString) = load_sym(mLibrary, "cudnnGetErrorString");
+
+    if (initHandle)
+    {
+        PLUGIN_CUDNNASSERT(cudnnCreate(&mHandle));
+        PLUGIN_VALIDATE(mHandle != nullptr);
+    }
+}
+
+CudnnWrapper::~CudnnWrapper()
+{
+    if (mHandle != nullptr)
+    {
+        PLUGIN_CUDNNASSERT(cudnnDestroy(mHandle));
+        mHandle = nullptr;
+    }
+
+    if (mLibrary != nullptr)
+    {
+        dllClose(mLibrary);
+    }
+}
+
+void* CudnnWrapper::tryLoadingCudnn(char const* callerPluginName)
+{
+#if CUDART_VERSION >= 12070 && CUDNN_MAJOR == 8
+    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());
+#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_VALIDATE(cudnnLib != nullptr, errorMsg.c_str());
+    return cudnnLib;
+}
+
+cudnnContext* CudnnWrapper::getCudnnHandle()
+{
+    return mHandle;
+}
+
+bool CudnnWrapper::isValid() const
+{
+    return mHandle != nullptr;
+}
+
+cudnnStatus_t CudnnWrapper::cudnnCreate(cudnnContext** handle)
+{
+    return (*_cudnnCreate)(handle);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnDestroy(cudnnContext* handle)
+{
+    return (*_cudnnDestroy)(handle);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnCreateTensorDescriptor(cudnnTensorDescriptor_t* tensorDesc)
+{
+    return (*_cudnnCreateTensorDescriptor)(tensorDesc);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnDestroyTensorDescriptor(cudnnTensorDescriptor_t tensorDesc)
+{
+    return (*_cudnnDestroyTensorDescriptor)(tensorDesc);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnSetStream(cudnnHandle_t handle, cudaStream_t streamId)
+{
+    return (*_cudnnSetStream)(handle, streamId);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnBatchNormalizationForwardTraining(cudnnHandle_t handle, cudnnBatchNormMode_t mode,
+    void const* alpha, void const* beta, cudnnTensorStruct const* xDesc, void const* x, cudnnTensorStruct const* yDesc,
+    void* y, cudnnTensorStruct const* bnScaleBiasMeanVarDesc, void const* bnScale, void const* bnBias,
+    double exponentialAverageFactor, void* resultRunningMean, void* resultRunningVariance, double epsilon,
+    void* resultSaveMean, void* resultSaveInvVariance)
+{
+    return (*_cudnnBatchNormalizationForwardTraining)(handle, mode, alpha, beta, xDesc, x, yDesc, y,
+        bnScaleBiasMeanVarDesc, bnScale, bnBias, exponentialAverageFactor, resultRunningMean, resultRunningVariance,
+        epsilon, resultSaveMean, resultSaveInvVariance);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnSetTensor4dDescriptor(cudnnTensorDescriptor_t tensorDesc, cudnnTensorFormat_t format,
+    cudnnDataType_t dataType, int n, int c, int h, int w)
+{
+    return (*_cudnnSetTensor4dDescriptor)(tensorDesc, format, dataType, n, c, h, w);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnSetTensorNdDescriptor(
+    cudnnTensorDescriptor_t tensorDesc, cudnnDataType_t dataType, int nbDims, int const dimA[], int const strideA[])
+{
+    return (*_cudnnSetTensorNdDescriptor)(tensorDesc, dataType, nbDims, dimA, strideA);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnSetTensorNdDescriptorEx(cudnnTensorDescriptor_t tensorDesc, cudnnTensorFormat_t format,
+    cudnnDataType_t dataType, int nbDims, int const dimA[])
+{
+    return (*_cudnnSetTensorNdDescriptorEx)(tensorDesc, format, dataType, nbDims, dimA);
+}
+
+cudnnStatus_t CudnnWrapper::cudnnDeriveBNTensorDescriptor(
+    cudnnTensorDescriptor_t derivedBnDesc, cudnnTensorStruct const* xDesc, cudnnBatchNormMode_t mode)
+{
+    return (*_cudnnDeriveBNTensorDescriptor)(derivedBnDesc, xDesc, mode);
+}
+
+char const* CudnnWrapper::cudnnGetErrorString(cudnnStatus_t status)
+{
+    return (*_cudnnGetErrorString)(status);
+}
+
+CudnnWrapper& getCudnnWrapper(char const* callerPluginName)
+{
+    // Initialize a global cublasWrapper instance to be used to call cudnn functions.
+    static CudnnWrapper sGCudnnWrapper{/*initHandle*/ false, callerPluginName};
+    return sGCudnnWrapper;
+}
+
+} // namespace nvinfer1::pluginInternal

plugin/common/cudnnWrapper.h

@@ -0,0 +1,168 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_PLUGIN_CUDNN_WRAPPER_H
+#define TRT_PLUGIN_CUDNN_WRAPPER_H
+
+#include "NvInferPlugin.h"
+#include <functional>
+#include <string>
+
+extern "C"
+{
+    //! Forward declaration of cudnnTensorStruct to use in other interfaces.
+    struct cudnnTensorStruct;
+}
+
+namespace nvinfer1
+{
+namespace pluginInternal
+{
+/*
+ * Copy of the CUDNN return codes
+ */
+enum CudnnStatus
+{
+    CUDNN_STATUS_SUCCESS = 0,
+    CUDNN_STATUS_NOT_INITIALIZED = 1,
+    CUDNN_STATUS_ALLOC_FAILED = 2,
+    CUDNN_STATUS_BAD_PARAM = 3,
+    CUDNN_STATUS_INTERNAL_ERROR = 4,
+    CUDNN_STATUS_INVALID_VALUE = 5,
+    CUDNN_STATUS_ARCH_MISMATCH = 6,
+    CUDNN_STATUS_MAPPING_ERROR = 7,
+    CUDNN_STATUS_EXECUTION_FAILED = 8,
+    CUDNN_STATUS_NOT_SUPPORTED = 9,
+    CUDNN_STATUS_LICENSE_ERROR = 10,
+    CUDNN_STATUS_RUNTIME_PREREQUISITE_MISSING = 11,
+    CUDNN_STATUS_RUNTIME_IN_PROGRESS = 12,
+    CUDNN_STATUS_RUNTIME_FP_OVERFLOW = 13,
+    CUDNN_STATUS_VERSION_MISMATCH = 14,
+};
+
+/*
+ * Copy of the CUDNN cudnnBatchNormMode_t
+ */
+enum cudnnBatchNormMode
+{
+    CUDNN_BATCHNORM_PER_ACTIVATION = 0,
+    CUDNN_BATCHNORM_SPATIAL = 1,
+    CUDNN_BATCHNORM_SPATIAL_PERSISTENT = 2,
+};
+
+/*
+ * Copy of the CUDNN cudnnTensorFormat_t
+ */
+enum cudnnTensorFormat
+{
+    CUDNN_TENSOR_NCHW = 0,
+    CUDNN_TENSOR_NHWC = 1,
+    CUDNN_TENSOR_NCHW_VECT_C = 2,
+};
+
+/*
+ * Copy of CUDNN data type
+ */
+enum cudnnDataType
+{
+    CUDNN_DATA_FLOAT = 0,
+    CUDNN_DATA_DOUBLE = 1,
+    CUDNN_DATA_HALF = 2,
+    CUDNN_DATA_INT8 = 3,
+    CUDNN_DATA_INT32 = 4,
+    CUDNN_DATA_INT8x4 = 5,
+    CUDNN_DATA_UINT8 = 6,
+    CUDNN_DATA_UINT8x4 = 7,
+    CUDNN_DATA_INT8x32 = 8,
+    CUDNN_DATA_BFLOAT16 = 9,
+    CUDNN_DATA_INT64 = 10,
+    CUDNN_DATA_BOOLEAN = 11,
+    CUDNN_DATA_FP8_E4M3 = 12,
+    CUDNN_DATA_FP8_E5M2 = 13,
+    CUDNN_DATA_FAST_FLOAT_FOR_FP8 = 14,
+};
+
+using cudnnStatus_t = CudnnStatus;
+using cudnnBatchNormMode_t = cudnnBatchNormMode;
+using cudnnTensorFormat_t = cudnnTensorFormat;
+using cudnnDataType_t = cudnnDataType;
+
+using cudnnHandle_t = struct cudnnContext*;
+using cudnnTensorDescriptor_t = struct cudnnTensorStruct*;
+
+class CudnnWrapper
+{
+public:
+    explicit CudnnWrapper(bool initHandle = false, char const* callerPluginName = nullptr);
+    ~CudnnWrapper();
+
+    cudnnContext* getCudnnHandle();
+    bool isValid() const;
+
+    /*
+     * Copy of the CUDNN APIs
+     */
+    cudnnStatus_t cudnnCreate(cudnnContext** handle);
+    cudnnStatus_t cudnnDestroy(cudnnContext* handle);
+    cudnnStatus_t cudnnCreateTensorDescriptor(cudnnTensorDescriptor_t* tensorDesc);
+    cudnnStatus_t cudnnDestroyTensorDescriptor(cudnnTensorDescriptor_t tensorDesc);
+    cudnnStatus_t cudnnSetStream(cudnnHandle_t handle, cudaStream_t streamId);
+    cudnnStatus_t cudnnBatchNormalizationForwardTraining(cudnnHandle_t handle, cudnnBatchNormMode_t mode,
+        void const* alpha, void const* beta, cudnnTensorStruct const* xDesc, void const* x,
+        cudnnTensorStruct const* yDesc, void* y, cudnnTensorStruct const* bnScaleBiasMeanVarDesc, void const* bnScale,
+        void const* bnBias, double exponentialAverageFactor, void* resultRunningMean, void* resultRunningVariance,
+        double epsilon, void* resultSaveMean, void* resultSaveInvVariance);
+    cudnnStatus_t cudnnSetTensor4dDescriptor(cudnnTensorDescriptor_t tensorDesc, cudnnTensorFormat_t format,
+        cudnnDataType_t dataType, int n, int c, int h, int w);
+    cudnnStatus_t cudnnSetTensorNdDescriptor(cudnnTensorDescriptor_t tensorDesc, cudnnDataType_t dataType, int nbDims,
+        int const dimA[], int const strideA[]);
+    cudnnStatus_t cudnnSetTensorNdDescriptorEx(cudnnTensorDescriptor_t tensorDesc, cudnnTensorFormat_t format,
+        cudnnDataType_t dataType, int nbDims, int const dimA[]);
+    cudnnStatus_t cudnnDeriveBNTensorDescriptor(
+        cudnnTensorDescriptor_t derivedBnDesc, cudnnTensorStruct const* xDesc, cudnnBatchNormMode_t mode);
+    char const* cudnnGetErrorString(cudnnStatus_t status);
+
+private:
+    void* mLibrary{nullptr};
+    cudnnContext* mHandle{nullptr};
+    void* tryLoadingCudnn(char const*);
+
+    cudnnStatus_t (*_cudnnCreate)(cudnnContext**);
+    cudnnStatus_t (*_cudnnDestroy)(cudnnContext*);
+    cudnnStatus_t (*_cudnnCreateTensorDescriptor)(cudnnTensorDescriptor_t*);
+    cudnnStatus_t (*_cudnnDestroyTensorDescriptor)(cudnnTensorDescriptor_t);
+    cudnnStatus_t (*_cudnnSetStream)(cudnnHandle_t, cudaStream_t);
+    cudnnStatus_t (*_cudnnBatchNormalizationForwardTraining)(cudnnHandle_t, cudnnBatchNormMode_t, void const*,
+        void const*, cudnnTensorStruct const*, void const*, cudnnTensorStruct const*, void*, cudnnTensorStruct const*,
+        void const*, void const*, double, void*, void*, double, void*, void*);
+    cudnnStatus_t (*_cudnnSetTensor4dDescriptor)(
+        cudnnTensorDescriptor_t, cudnnTensorFormat_t, cudnnDataType_t, int, int, int, int);
+    cudnnStatus_t (*_cudnnSetTensorNdDescriptor)(cudnnTensorDescriptor_t tensorDesc, cudnnDataType_t dataType,
+        int nbDims, int const dimA[], int const strideA[]);
+    cudnnStatus_t (*_cudnnSetTensorNdDescriptorEx)(cudnnTensorDescriptor_t tensorDesc, cudnnTensorFormat_t format,
+        cudnnDataType_t dataType, int nbDims, int const dimA[]);
+    cudnnStatus_t (*_cudnnDeriveBNTensorDescriptor)(
+        cudnnTensorDescriptor_t, cudnnTensorStruct const*, cudnnBatchNormMode_t);
+    char const* (*_cudnnGetErrorString)(cudnnStatus_t status);
+};
+
+CudnnWrapper& getCudnnWrapper(char const* callerPluginName);
+
+} // namespace pluginInternal
+} // namespace nvinfer1
+
+#endif // TRT_PLUGIN_CUDNN_WRAPPER_H

plugin/common/dimsHelpers.h

@@ -0,0 +1,64 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_PLUGIN_DIMS_HELPERS_H
+#define TRT_PLUGIN_DIMS_HELPERS_H
+
+#include "common/plugin.h" // purely for assertions
+
+#include <algorithm> // all of
+#include <functional>
+#include <numeric>
+
+namespace nvinfer1
+{
+
+namespace pluginInternal
+{
+
+//! Return number of elements in the given dimensions in the range [start, stop).
+//! Does not include padding added for vectorized formats.
+//!
+//! \param dims dimensions whose partial volume needs to be computed
+//! \param start inclusive start axis
+//! \param stop exclusive stop axis
+//!
+//! Expects 0 <= start <= stop <= dims.nbDims.
+//! For i in the range [start,stop), dims.d[i] must be non-negative.
+//!
+inline int64_t volume(Dims const& dims, int32_t start, int32_t stop)
+{
+    // The signature is int32_t start (and not uint32_t start) because int32_t is used
+    // for indexing everywhere
+    ASSERT_PARAM(start >= 0);
+    ASSERT_PARAM(start <= stop);
+    ASSERT_PARAM(stop <= dims.nbDims);
+    ASSERT_PARAM(std::all_of(dims.d + start, dims.d + stop, [](int32_t x) { return x >= 0; }));
+    return std::accumulate(dims.d + start, dims.d + stop, int64_t{1}, std::multiplies<int64_t>{});
+}
+
+//! Shorthand for volume(dims, 0, dims.nbDims).
+inline int64_t volume(Dims const& dims)
+{
+    return volume(dims, 0, dims.nbDims);
+}
+
+} // namespace pluginInternal
+
+} // namespace nvinfer1
+
+#endif // TRT_PLUGIN_DIMS_HELPERS_H

plugin/common/half.h

@@ -0,0 +1,44 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+//
+// Custom wrapper around external half-precision header
+//
+// Header has some "extra parentheses" warnings when different rounding modes are used.
+
+#if defined(__GNUC__)
+#pragma GCC diagnostic push
+#pragma GCC diagnostic ignored "-Wparentheses"
+#endif
+
+#if defined(__clang__)
+#pragma clang diagnostic push
+#pragma clang diagnostic ignored "-Wmismatched-tags"
+#endif
+
+#ifndef HALF_ROUND_TIES_TO_EVEN
+#define HALF_ROUND_TIES_TO_EVEN 1
+#endif
+
+#include "ieee/half.h"
+
+#if defined(__clang__)
+#pragma clang diagnostic pop
+#endif
+
+#if defined(__GNUC__)
+#pragma GCC diagnostic pop
+#endif

plugin/common/kernels/CMakeLists.txt

@@ -0,0 +1,49 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+add_plugin_source(
+    allClassNMS.cu
+    bboxDeltas2Proposals.cu
+    common.cu
+    cropAndResizeKernel.cu
+    decodeBbox3DKernels.cu
+    extractFgScores.cu
+    gatherTopDetections.cu
+    generateAnchors.cu
+    gridAnchorLayer.cu
+    kernel.h
+    lReLU.cu
+    maskRCNNKernels.cu
+    maskRCNNKernels.h
+    nmsLayer.cu
+    normalizeLayer.cu
+    permuteData.cu
+    pillarScatterKernels.cu
+    priorBoxLayer.cu
+    proposalKernel.cu
+    proposalsForward.cu
+    reducedMathPlugin.h
+    regionForward.cu
+    reorgForward.cu
+    roiPooling.cu
+    rproiInferenceFused.cu
+    saturate.h
+    sortScoresPerClass.cu
+    sortScoresPerImage.cu
+    topkLastDim.cu
+    voxelGeneratorKernels.cu
+)

plugin/common/kernels/allClassNMS.cu

@@ -0,0 +1,394 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/bboxUtils.h"
+#include "common/kernels/kernel.h"
+#include <cuda_fp16.h>
+#include <array>
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename T_BBOX>
+__device__ float bboxSize(const Bbox<T_BBOX>& bbox, const bool normalized)
+{
+    if (float(bbox.xmax) < float(bbox.xmin) || float(bbox.ymax) < float(bbox.ymin))
+    {
+        // If bbox is invalid (e.g. xmax < xmin or ymax < ymin), return 0.
+        return 0;
+    }
+    else
+    {
+        float width = float(bbox.xmax) - float(bbox.xmin);
+        float height = float(bbox.ymax) - float(bbox.ymin);
+        if (normalized)
+        {
+            return width * height;
+        }
+        else
+        {
+            // If bbox is not within range [0, 1].
+            return (width + 1.f) * (height + 1.f);
+        }
+    }
+}
+
+template <typename T_BBOX>
+__device__ void intersectBbox(
+    const Bbox<T_BBOX>& bbox1,
+    const Bbox<T_BBOX>& bbox2,
+    Bbox<T_BBOX>* intersect_bbox)
+{
+    if (bbox2.xmin > bbox1.xmax || bbox2.xmax < bbox1.xmin || bbox2.ymin > bbox1.ymax || bbox2.ymax < bbox1.ymin)
+    {
+        // Return [0, 0, 0, 0] if there is no intersection.
+        intersect_bbox->xmin = T_BBOX(0);
+        intersect_bbox->ymin = T_BBOX(0);
+        intersect_bbox->xmax = T_BBOX(0);
+        intersect_bbox->ymax = T_BBOX(0);
+    }
+    else
+    {
+        intersect_bbox->xmin = max(bbox1.xmin, bbox2.xmin);
+        intersect_bbox->ymin = max(bbox1.ymin, bbox2.ymin);
+        intersect_bbox->xmax = min(bbox1.xmax, bbox2.xmax);
+        intersect_bbox->ymax = min(bbox1.ymax, bbox2.ymax);
+    }
+}
+
+
+template <>
+__device__ void intersectBbox<__half>(
+    const Bbox<__half>& bbox1,
+    const Bbox<__half>& bbox2,
+    Bbox<__half>* intersect_bbox)
+{
+    if (float(bbox2.xmin) > float(bbox1.xmax)
+        || float(bbox2.xmax) < float(bbox1.xmin)
+        || float(bbox2.ymin) > float(bbox1.ymax)
+        || float(bbox2.ymax) < float(bbox1.ymin))
+    {
+        // Return [0, 0, 0, 0] if there is no intersection.
+        intersect_bbox->xmin = __half(0);
+        intersect_bbox->ymin = __half(0);
+        intersect_bbox->xmax = __half(0);
+        intersect_bbox->ymax = __half(0);
+    }
+    else
+    {
+        intersect_bbox->xmin = max(float(bbox1.xmin), float(bbox2.xmin));
+        intersect_bbox->ymin = max(float(bbox1.ymin), float(bbox2.ymin));
+        intersect_bbox->xmax = min(float(bbox1.xmax), float(bbox2.xmax));
+        intersect_bbox->ymax = min(float(bbox1.ymax), float(bbox2.ymax));
+    }
+}
+
+
+template <typename T_BBOX>
+__device__ Bbox<T_BBOX> getDiagonalMinMaxSortedBox(const Bbox<T_BBOX>& bbox1)
+{
+    Bbox<T_BBOX> result;
+    result.xmin = min(bbox1.xmin, bbox1.xmax);
+    result.xmax = max(bbox1.xmin, bbox1.xmax);
+
+    result.ymin = min(bbox1.ymin, bbox1.ymax);
+    result.ymax = max(bbox1.ymin, bbox1.ymax);
+    return result;
+}
+
+template <>
+__device__ Bbox<__half> getDiagonalMinMaxSortedBox(const Bbox<__half>& bbox1)
+{
+    Bbox<__half> result;
+    result.xmin = min(float(bbox1.xmin), float(bbox1.xmax));
+    result.xmax = max(float(bbox1.xmin), float(bbox1.xmax));
+
+    result.ymin = min(float(bbox1.ymin), float(bbox1.ymax));
+    result.ymax = max(float(bbox1.ymin), float(bbox1.ymax));
+    return result;
+}
+
+template <typename T_BBOX>
+__device__ float jaccardOverlap(
+    const Bbox<T_BBOX>& bbox1, const Bbox<T_BBOX>& bbox2, const bool normalized, const bool caffeSemantics)
+{
+    Bbox<T_BBOX> intersect_bbox;
+
+    Bbox<T_BBOX> localbbox1 = getDiagonalMinMaxSortedBox(bbox1);
+    Bbox<T_BBOX> localbbox2 = getDiagonalMinMaxSortedBox(bbox2);
+
+    intersectBbox(localbbox1, localbbox2, &intersect_bbox);
+
+    float intersect_width, intersect_height;
+    // Only when using Caffe semantics, IOU calculation adds "1" to width and height if bbox is not normalized.
+    // https://github.com/weiliu89/caffe/blob/ssd/src/caffe/util/bbox_util.cpp#L92-L97
+    if (normalized || !caffeSemantics)
+    {
+        intersect_width = float(intersect_bbox.xmax) - float(intersect_bbox.xmin);
+        intersect_height = float(intersect_bbox.ymax) - float(intersect_bbox.ymin);
+    }
+    else
+    {
+        intersect_width = float(intersect_bbox.xmax) - float(intersect_bbox.xmin) + float(T_BBOX(1));
+        intersect_height = float(intersect_bbox.ymax) - float(intersect_bbox.ymin) + float(T_BBOX(1));
+    }
+    if (intersect_width > 0 && intersect_height > 0)
+    {
+        float intersect_size = intersect_width * intersect_height;
+        float bbox1_size = bboxSize(localbbox1, normalized);
+        float bbox2_size = bboxSize(localbbox2, normalized);
+        return intersect_size / (bbox1_size + bbox2_size - intersect_size);
+    }
+    else
+    {
+        return 0.;
+    }
+}
+
+template <typename T_BBOX>
+__device__ void emptyBboxInfo(
+    BboxInfo<T_BBOX>* bbox_info)
+{
+    bbox_info->conf_score = T_BBOX(0);
+    bbox_info->label = -2; // -1 is used for all labels when shared_location is ture
+    bbox_info->bbox_idx = -1;
+    bbox_info->kept = false;
+}
+/********** new NMS for only score and index array **********/
+
+template <typename T_SCORE, typename T_BBOX, int TSIZE>
+__global__ void allClassNMS_kernel(const int num, const int num_classes, const int num_preds_per_class, const int top_k,
+    const float nms_threshold, const bool share_location, const bool isNormalized,
+    T_BBOX* bbox_data, // bbox_data should be float to preserve location information
+    T_SCORE* beforeNMS_scores, int* beforeNMS_index_array, T_SCORE* afterNMS_scores, int* afterNMS_index_array,
+    bool flipXY, const float score_shift, bool caffeSemantics)
+{
+    //__shared__ bool kept_bboxinfo_flag[CAFFE_CUDA_NUM_THREADS * TSIZE];
+    extern __shared__ bool kept_bboxinfo_flag[];
+
+    for (int i = 0; i < num; i++)
+    {
+        int32_t const offset = i * num_classes * num_preds_per_class + blockIdx.x * num_preds_per_class;
+        // Should not write data beyond [offset, top_k).
+        int32_t const max_idx = offset + top_k;
+        // Should not read beyond [offset, num_preds_per_class).
+        int32_t const max_read_idx = offset + min(top_k, num_preds_per_class);
+        int32_t const bbox_idx_offset = i * num_preds_per_class * (share_location ? 1 : num_classes);
+
+        // local thread data
+        int loc_bboxIndex[TSIZE];
+        Bbox<T_BBOX> loc_bbox[TSIZE];
+
+        // initialize Bbox, Bboxinfo, kept_bboxinfo_flag
+        // Eliminate shared memory RAW hazard
+        __syncthreads();
+#pragma unroll
+        for (int t = 0; t < TSIZE; t++)
+        {
+            const int cur_idx = threadIdx.x + blockDim.x * t;
+            const int item_idx = offset + cur_idx;
+            // Init all output data
+            if (item_idx < max_idx)
+            {
+                // Do not access data if it exceeds read boundary
+                if (item_idx < max_read_idx)
+                {
+                    loc_bboxIndex[t] = beforeNMS_index_array[item_idx];
+                }
+                else
+                {
+                    loc_bboxIndex[t] = -1;
+                }
+
+                if (loc_bboxIndex[t] != -1)
+                {
+                    const int bbox_data_idx = share_location ? (loc_bboxIndex[t] % num_preds_per_class + bbox_idx_offset) : loc_bboxIndex[t];
+
+                    loc_bbox[t].xmin = flipXY ? bbox_data[bbox_data_idx * 4 + 1] : bbox_data[bbox_data_idx * 4 + 0];
+                    loc_bbox[t].ymin = flipXY ? bbox_data[bbox_data_idx * 4 + 0] : bbox_data[bbox_data_idx * 4 + 1];
+                    loc_bbox[t].xmax = flipXY ? bbox_data[bbox_data_idx * 4 + 3] : bbox_data[bbox_data_idx * 4 + 2];
+                    loc_bbox[t].ymax = flipXY ? bbox_data[bbox_data_idx * 4 + 2] : bbox_data[bbox_data_idx * 4 + 3];
+                    kept_bboxinfo_flag[cur_idx] = true;
+                }
+                else
+                {
+                    kept_bboxinfo_flag[cur_idx] = false;
+                }
+            }
+            else
+            {
+                kept_bboxinfo_flag[cur_idx] = false;
+            }
+        }
+
+        // filter out overlapped boxes with lower scores
+        int ref_item_idx = offset;
+
+        int32_t ref_bbox_idx = -1;
+        if (ref_item_idx < max_read_idx)
+        {
+            ref_bbox_idx = share_location
+                ? (beforeNMS_index_array[ref_item_idx] % num_preds_per_class + bbox_idx_offset)
+                : beforeNMS_index_array[ref_item_idx];
+        }
+        while ((ref_bbox_idx != -1) && ref_item_idx < max_read_idx)
+        {
+            Bbox<T_BBOX> ref_bbox;
+            ref_bbox.xmin = flipXY ? bbox_data[ref_bbox_idx * 4 + 1] : bbox_data[ref_bbox_idx * 4 + 0];
+            ref_bbox.ymin = flipXY ? bbox_data[ref_bbox_idx * 4 + 0] : bbox_data[ref_bbox_idx * 4 + 1];
+            ref_bbox.xmax = flipXY ? bbox_data[ref_bbox_idx * 4 + 3] : bbox_data[ref_bbox_idx * 4 + 2];
+            ref_bbox.ymax = flipXY ? bbox_data[ref_bbox_idx * 4 + 2] : bbox_data[ref_bbox_idx * 4 + 3];
+
+            // Eliminate shared memory RAW hazard
+            __syncthreads();
+
+            for (int t = 0; t < TSIZE; t++)
+            {
+                const int cur_idx = threadIdx.x + blockDim.x * t;
+                const int item_idx = offset + cur_idx;
+
+                if ((kept_bboxinfo_flag[cur_idx]) && (item_idx > ref_item_idx))
+                {
+                    if (jaccardOverlap(ref_bbox, loc_bbox[t], isNormalized, caffeSemantics) > nms_threshold)
+                    {
+                        kept_bboxinfo_flag[cur_idx] = false;
+                    }
+                }
+            }
+            __syncthreads();
+
+            do
+            {
+                ref_item_idx++;
+            } while (ref_item_idx < max_read_idx && !kept_bboxinfo_flag[ref_item_idx - offset]);
+
+            // Move to next valid point
+            if (ref_item_idx < max_read_idx)
+            {
+                ref_bbox_idx = share_location
+                    ? (beforeNMS_index_array[ref_item_idx] % num_preds_per_class + bbox_idx_offset)
+                    : beforeNMS_index_array[ref_item_idx];
+            }
+        }
+
+        // store data
+        for (int t = 0; t < TSIZE; t++)
+        {
+            const int cur_idx = threadIdx.x + blockDim.x * t;
+            const int read_item_idx = offset + cur_idx;
+            const int write_item_idx = (i * num_classes * top_k + blockIdx.x * top_k) + cur_idx;
+            /*
+             * If not not keeping the bbox
+             * Set the score to 0
+             * Set the bounding box index to -1
+             */
+            if (read_item_idx < max_idx)
+            {
+                afterNMS_scores[write_item_idx] = kept_bboxinfo_flag[cur_idx] ? T_SCORE(beforeNMS_scores[read_item_idx]) : T_SCORE(score_shift);
+                afterNMS_index_array[write_item_idx] = kept_bboxinfo_flag[cur_idx] ? loc_bboxIndex[t] : -1;
+            }
+        }
+    }
+}
+
+template <typename T_SCORE, typename T_BBOX>
+pluginStatus_t allClassNMS_gpu(cudaStream_t stream, const int num, const int num_classes, const int num_preds_per_class,
+    const int top_k, const float nms_threshold, const bool share_location, const bool isNormalized, void* bbox_data,
+    void* beforeNMS_scores, void* beforeNMS_index_array, void* afterNMS_scores, void* afterNMS_index_array, bool flipXY,
+    const float score_shift, bool caffeSemantics)
+{
+#define P(tsize) allClassNMS_kernel<T_SCORE, T_BBOX, (tsize)>
+
+    void (*kernel[8])(const int, const int, const int, const int, const float, const bool, const bool, T_BBOX*,
+        T_SCORE*, int*, T_SCORE*, int*, bool, const float, bool)
+        = {
+            P(1),
+            P(2),
+            P(3),
+            P(4),
+            P(5),
+            P(6),
+            P(7),
+            P(8),
+        };
+
+    const int BS = 512;
+    const int GS = num_classes;
+    const int t_size = (top_k + BS - 1) / BS;
+
+    kernel[t_size - 1]<<<GS, BS, BS * t_size * sizeof(bool), stream>>>(num, num_classes, num_preds_per_class, top_k,
+        nms_threshold, share_location, isNormalized, (T_BBOX*) bbox_data, (T_SCORE*) beforeNMS_scores,
+        (int*) beforeNMS_index_array, (T_SCORE*) afterNMS_scores, (int*) afterNMS_index_array, flipXY, score_shift,
+        caffeSemantics);
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// allClassNMS LAUNCH CONFIG
+typedef pluginStatus_t (*nmsFunc)(cudaStream_t, const int, const int, const int, const int, const float, const bool,
+    const bool, void*, void*, void*, void*, void*, bool, const float, bool);
+
+struct nmsLaunchConfigSSD
+{
+    DataType t_score;
+    DataType t_bbox;
+    nmsFunc function;
+
+    nmsLaunchConfigSSD(DataType t_score, DataType t_bbox)
+        : t_score(t_score)
+        , t_bbox(t_bbox)
+        , function(nullptr)
+    {
+    }
+    nmsLaunchConfigSSD(DataType t_score, DataType t_bbox, nmsFunc function)
+        : t_score(t_score)
+        , t_bbox(t_bbox)
+        , function(function)
+    {
+    }
+    bool operator==(nmsLaunchConfigSSD const& other) const
+    {
+        return t_score == other.t_score && t_bbox == other.t_bbox;
+    }
+};
+
+static std::array<nmsLaunchConfigSSD, 2> nmsSsdLCOptions = {
+    nmsLaunchConfigSSD(DataType::kFLOAT, DataType::kFLOAT, allClassNMS_gpu<float, float>),
+    nmsLaunchConfigSSD(DataType::kHALF, DataType::kHALF, allClassNMS_gpu<__half, __half>)
+};
+
+pluginStatus_t allClassNMS(cudaStream_t stream, const int num, const int num_classes, const int num_preds_per_class,
+    const int top_k, const float nms_threshold, const bool share_location, const bool isNormalized,
+    const DataType DT_SCORE, const DataType DT_BBOX, void* bbox_data, void* beforeNMS_scores,
+    void* beforeNMS_index_array, void* afterNMS_scores, void* afterNMS_index_array, bool flipXY,
+    const float score_shift, bool caffeSemantics)
+{
+    nmsLaunchConfigSSD lc = nmsLaunchConfigSSD(DT_SCORE, DT_BBOX);
+    for (unsigned i = 0; i < nmsSsdLCOptions.size(); ++i)
+    {
+        if (lc == nmsSsdLCOptions[i])
+        {
+            DEBUG_PRINTF("all class nms kernel %d\n", i);
+            return nmsSsdLCOptions[i].function(stream, num, num_classes, num_preds_per_class, top_k, nms_threshold,
+                share_location, isNormalized, bbox_data, beforeNMS_scores, beforeNMS_index_array, afterNMS_scores,
+                afterNMS_index_array, flipXY, score_shift, caffeSemantics);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/bboxDeltas2Proposals.cu

@@ -0,0 +1,290 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+#include <algorithm>
+#include <array>
+#include <math.h>
+#include <stdio.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+using std::max;
+using std::min;
+
+// BBD2P KERNEL
+template <typename T_DELTAS,
+          DLayout_t L_DELTAS,
+          typename TV_PROPOSALS,
+          DLayout_t L_PROPOSALS,
+          typename T_FGSCORES,
+          DLayout_t L_FGSCORES>
+__global__ void bboxDeltas2Proposals_kernel(
+    int N,
+    int A,
+    int H,
+    int W,
+    const float* __restrict__ anchors,
+    const float* __restrict__ imInfo,
+    int featureStride,
+    float minSize,
+    const T_DELTAS* __restrict__ deltas,
+    TV_PROPOSALS* __restrict__ proposals,
+    T_FGSCORES* __restrict__ scores)
+{
+    int tid = threadIdx.x + blockIdx.x * blockDim.x;
+
+    if (tid < N * A * H * W)
+    { // TODO this can be a loop.
+        // Find out the index of bounding box for the output
+        int cnt = tid;
+        // width index
+        int w = cnt % W;
+        cnt = cnt / W;
+        // height index
+        int h = cnt % H;
+        cnt = cnt / H;
+        // anchor box index
+        int a = cnt % A;
+        cnt = cnt / A;
+        // batch index
+        int n = cnt;
+        int hw = h * W + w;
+
+        // Get the height and width of the original input image
+        float imHeight = imInfo[3 * n];
+        float imWidth = imInfo[3 * n + 1];
+
+        // Point to the right anchor box
+        float4 anchor = ((float4*) anchors)[a];
+        // Get anchor box coordinates
+        float a_ctr_x = anchor.x;
+        float a_ctr_y = anchor.y;
+        float a_w = anchor.z;
+        float a_h = anchor.w;
+
+        // NCHW format
+        // Find out the starting position of the bounding box in the input (predicted bounding box offsets)
+        int id = ((tid - hw) * 4) + hw;
+        T_DELTAS dx;
+        T_DELTAS dy;
+        T_DELTAS dw;
+        T_DELTAS dh;
+        if (L_DELTAS == NCHW)
+        {
+            // The offsets between adjacent coordinates on linear memory is H * W
+            dx = deltas[id];
+            dy = deltas[id + 1 * H * W];
+            dw = deltas[id + 2 * H * W];
+            dh = deltas[id + 3 * H * W];
+        }
+        // NC4HW format
+        else if (L_DELTAS == NC4HW)
+        {
+            dx = deltas[tid * 4 + 0];
+            dy = deltas[tid * 4 + 1];
+            dw = deltas[tid * 4 + 2];
+            dh = deltas[tid * 4 + 3];
+        }
+        /*
+         * Calculate the coordinates of decoded bounding box on the original input image scale
+         * Only works if param.minBoxSize == param.featureStride
+         */
+        float ctr_x = a_ctr_x + w * featureStride;
+        float ctr_y = a_ctr_y + h * featureStride;
+        // float ctr_x = (w + 0.5) * featureStride;
+        // float ctr_y = (h + 0.5) * featureStride;
+
+        /*
+         * Decode the predicted bounding box
+         * The decoded bounding boxes has coordinates of [x_topleft, y_topleft, x_bottomright, y_bottomright]
+         * The units are in pixels
+         */
+        ctr_x = ctr_x + dx * a_w;
+        ctr_y = ctr_y + dy * a_h;
+        float b_w = __expf(dw) * a_w;
+        float b_h = __expf(dh) * a_h;
+        float bx = ctr_x - (b_w / 2);
+        float by = ctr_y - (b_h / 2);
+        float bz = ctr_x + (b_w / 2);
+        float bw = ctr_y + (b_h / 2);
+
+        TV_PROPOSALS bbox;
+        // Make sure that the decoded bouding box go outside of the original input image
+        bbox.x = fminf(fmaxf(bx, 0.0f), imWidth - 1.0f);
+        bbox.y = fminf(fmaxf(by, 0.0f), imHeight - 1.0f);
+        bbox.z = fminf(fmaxf(bz, 0.0f), imWidth - 1.0f);
+        bbox.w = fminf(fmaxf(bw, 0.0f), imHeight - 1.0f);
+
+        // Put the decoded bounding box information to the outputs
+        if (L_PROPOSALS == NC4HW)
+        {
+            proposals[tid] = bbox;
+        }
+
+        int ininf = 0xff800000;
+        float ninf = *(float*) &ininf;
+        // minBoxSize at the original input image scale
+        float scaledMinSize = minSize * imInfo[3 * n + 2];
+        // Set the objectness score to -inf if the predicted bounding box has edgth length less than the minimal box size expected.
+        if (bbox.z - bbox.x + 1 < scaledMinSize || bbox.w - bbox.y + 1 < scaledMinSize)
+        {
+            if (L_FGSCORES == NCHW)
+                scores[tid] = ninf;
+        }
+    }
+}
+
+// BBD2P KERNEL LAUNCHER
+template <typename T_DELTAS,
+          DLayout_t L_DELTAS,
+          typename TV_PROPOSALS,
+          DLayout_t L_PROPOSALS,
+          typename T_FGSCORES,
+          DLayout_t L_FGSCORES>
+pluginStatus_t bboxDeltas2Proposals_gpu(cudaStream_t stream,
+                                       int N,
+                                       int A,
+                                       int H,
+                                       int W,
+                                       const float* imInfo,
+                                       int featureStride,
+                                       float minBoxSize,
+                                       const float* anchors,
+                                       const void* deltas,
+                                       void* propos,
+                                       void* scores)
+{
+    const int BS = 32;
+    const int GS = ((N * A * H * W) + BS - 1) / BS;
+
+    bboxDeltas2Proposals_kernel<T_DELTAS, L_DELTAS, TV_PROPOSALS, L_PROPOSALS, T_FGSCORES, L_FGSCORES><<<GS, BS, 0, stream>>>(N, A, H, W,
+                                                                                                                              anchors,
+                                                                                                                              imInfo,
+                                                                                                                              featureStride,
+                                                                                                                              minBoxSize,
+                                                                                                                              (T_DELTAS*) deltas,
+                                                                                                                              (TV_PROPOSALS*) propos,
+                                                                                                                              (T_FGSCORES*) scores);
+
+    DEBUG_PRINTF("&&&& [bboxD2P] POST LAUNCH\n");
+    DEBUG_PRINTF("&&&& [bboxD2P] PROPOS %u\n", hash(propos, N * A * H * W * 4 * sizeof(float)));
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// BBD2P LAUNCH CONFIG {{{
+typedef pluginStatus_t (*bd2pFun)(cudaStream_t,
+                                 int,
+                                 int,
+                                 int,
+                                 int,
+                                 const float*,
+                                 int,
+                                 float,
+                                 const float*,
+                                 const void*,
+                                 void*,
+                                 void*);
+
+struct bd2pLaunchConfig
+{
+    DataType t_deltas;
+    DLayout_t l_deltas;
+    DataType t_proposals;
+    DLayout_t l_proposals;
+    DataType t_scores;
+    DLayout_t l_scores;
+    bd2pFun function;
+
+    bd2pLaunchConfig(DataType t_deltas, DLayout_t l_deltas, DataType t_proposals, DLayout_t l_proposals,
+        DataType t_scores, DLayout_t l_scores)
+        : t_deltas(t_deltas)
+        , l_deltas(l_deltas)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_scores(t_scores)
+        , l_scores(l_scores)
+        , function(nullptr)
+    {
+    }
+
+    bd2pLaunchConfig(DataType t_deltas, DLayout_t l_deltas, DataType t_proposals, DLayout_t l_proposals, DataType t_scores, DLayout_t l_scores, bd2pFun function)
+        : t_deltas(t_deltas)
+        , l_deltas(l_deltas)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_scores(t_scores)
+        , l_scores(l_scores)
+        , function(function)
+    {
+    }
+
+    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;
+    }
+};
+
+#define FLOAT32 nvinfer1::DataType::kFLOAT
+static std::array<bd2pLaunchConfig, 2> bd2pLCOptions = {
+    bd2pLaunchConfig(FLOAT32, NC4HW, FLOAT32, NC4HW, FLOAT32, NCHW, bboxDeltas2Proposals_gpu<float, NC4HW, float4, NC4HW, float, NCHW>),
+    bd2pLaunchConfig(FLOAT32, NCHW, FLOAT32, NC4HW, FLOAT32, NCHW, bboxDeltas2Proposals_gpu<float, NCHW, float4, NC4HW, float, NCHW>)};
+
+// BBD2P
+pluginStatus_t bboxDeltas2Proposals(cudaStream_t stream,
+                                   const int N,
+                                   const int A,
+                                   const int H,
+                                   const int W,
+                                   const int featureStride,
+                                   const float minBoxSize,
+                                   const float* imInfo,
+                                   const float* anchors,
+                                   const DataType t_deltas,
+                                   const DLayout_t l_deltas,
+                                   const void* deltas,
+                                   const DataType t_proposals,
+                                   const DLayout_t l_proposals,
+                                   void* proposals,
+                                   const DataType t_scores,
+                                   const DLayout_t l_scores,
+                                   void* scores)
+{
+    bd2pLaunchConfig lc = bd2pLaunchConfig(t_deltas, l_deltas, t_proposals, l_proposals, t_scores, l_scores);
+    for (unsigned i = 0; i < bd2pLCOptions.size(); i++)
+    {
+        if (lc == bd2pLCOptions[i])
+        {
+            DEBUG_PRINTF("BBD2P kernel %d\n", i);
+            return bd2pLCOptions[i].function(stream,
+                                          N, A, H, W,
+                                          imInfo,
+                                          featureStride,
+                                          minBoxSize,
+                                          anchors,
+                                          deltas,
+                                          proposals,
+                                          scores);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/common.cu

@@ -0,0 +1,175 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "common/bboxUtils.h"
+#include "common/kernels/kernel.h"
+#include "cuda.h"
+#include <cub/cub.cuh>
+#include <stdint.h>
+
+using namespace nvinfer1::pluginInternal;
+
+#define CUDA_MEM_ALIGN 256
+namespace nvinfer1
+{
+namespace plugin
+{
+// HASH
+unsigned int hash(const void* array_, size_t size)
+{
+    // Apply hashing only when debugging RPN codes.
+    if (DEBUG_ENABLE)
+    {
+        const char* array_const;
+        char* array;
+        PLUGIN_CHECK_CUDA(cudaMallocHost((void**) &array, size));
+        PLUGIN_CHECK_CUDA(cudaMemcpy(array, array_, size, cudaMemcpyDeviceToHost));
+        array_const = array;
+        unsigned int hash = 45599;
+        for (size_t i = 0; i < size; i++)
+        {
+            unsigned int value = array_const[i];
+            hash = hash * 1487 + value;
+            hash = hash * 317;
+            hash = hash % 105359;
+        }
+        return hash;
+    }
+    else
+    {
+        return 0;
+    }
+}
+
+// ALIGNPTR
+int8_t* alignPtr(int8_t* ptr, uintptr_t to)
+{
+    uintptr_t addr = (uintptr_t) ptr;
+    if (addr % to)
+    {
+        addr += to - addr % to;
+    }
+    return (int8_t*) addr;
+}
+
+// NEXTWORKSPACEPTR
+int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize)
+{
+    uintptr_t addr = (uintptr_t) ptr;
+    addr += previousWorkspaceSize;
+    return alignPtr((int8_t*) addr, CUDA_MEM_ALIGN);
+}
+
+// CALCULATE TOTAL WORKSPACE SIZE
+size_t calculateTotalWorkspaceSize(size_t* workspaces, int count)
+{
+    size_t total = 0;
+    for (int i = 0; i < count; i++)
+    {
+        total += workspaces[i];
+        if (workspaces[i] % CUDA_MEM_ALIGN)
+        {
+            total += CUDA_MEM_ALIGN - (workspaces[i] % CUDA_MEM_ALIGN);
+        }
+    }
+    return total;
+}
+
+using nvinfer1::DataType;
+
+// DATA TYPE SIZE
+size_t dataTypeSize(const DataType dtype)
+{
+    switch (dtype)
+    {
+    case DataType::kINT8: return sizeof(char);
+    case DataType::kHALF: return sizeof(short);
+    case DataType::kFLOAT: return sizeof(float);
+    case DataType::kBF16:
+    case DataType::kINT64: PLUGIN_FAIL("Unsupported data type");
+    default: return 0;
+    }
+}
+
+// CUB
+/*
+size_t cubSortFloatIntPairsWorkspaceSize(int num_items, int num_segments)
+{
+    size_t temp_storage_bytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+    (int *)NULL, temp_storage_bytes,
+    (const float *)NULL, (float *)NULL,
+    (const int *)NULL, (int *)NULL,
+    num_items,     // # items
+    num_segments,  // # segments
+    (const int *)NULL, (const int *)NULL);
+    return temp_storage_bytes;
+}
+
+size_t cubSortFloatBboxInfoPairsWorkspaceSize(int num_items, int num_segments)
+{
+    size_t temp_storage_bytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+    (int *)NULL, temp_storage_bytes,
+    (const float *)NULL, (float *)NULL,
+    (const BboxInfo<float> *)NULL, (BboxInfo<float> *)NULL,
+    num_items,     // # items
+    num_segments,  // # segments
+    (const int *)NULL, (const int *)NULL);
+    return temp_storage_bytes;
+}
+*/
+
+template <unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta) __global__
+    void setUniformOffsets_kernel(const int num_segments, const int offset, int* d_offsets)
+{
+    const int idx = blockIdx.x * nthds_per_cta + threadIdx.x;
+    if (idx <= num_segments)
+        d_offsets[idx] = idx * offset;
+}
+
+void setUniformOffsets(cudaStream_t stream, const int num_segments, const int offset, int* d_offsets)
+{
+    const int BS = 32;
+    const int GS = (num_segments + 1 + BS - 1) / BS;
+    setUniformOffsets_kernel<BS><<<GS, BS, 0, stream>>>(num_segments, offset, d_offsets);
+}
+
+const char* cublasGetErrorString(cublasStatus_t error)
+{
+    switch (error)
+    {
+    case CUBLAS_STATUS_SUCCESS: return "CUBLAS_STATUS_SUCCESS";
+    case CUBLAS_STATUS_NOT_INITIALIZED: return "CUBLAS_STATUS_NOT_INITIALIZED";
+    case CUBLAS_STATUS_ALLOC_FAILED: return "CUBLAS_STATUS_ALLOC_FAILED";
+    case CUBLAS_STATUS_INVALID_VALUE: return "CUBLAS_STATUS_INVALID_VALUE";
+    case CUBLAS_STATUS_ARCH_MISMATCH: return "CUBLAS_STATUS_ARCH_MISMATCH";
+    case CUBLAS_STATUS_MAPPING_ERROR: return "CUBLAS_STATUS_MAPPING_ERROR";
+    case CUBLAS_STATUS_EXECUTION_FAILED: return "CUBLAS_STATUS_EXECUTION_FAILED";
+    case CUBLAS_STATUS_INTERNAL_ERROR: return "CUBLAS_STATUS_INTERNAL_ERROR";
+#if CUDA_VERSION >= 6000
+    case CUBLAS_STATUS_NOT_SUPPORTED: return "CUBLAS_STATUS_NOT_SUPPORTED";
+#endif
+#if CUDA_VERSION >= 6050
+    case CUBLAS_STATUS_LICENSE_ERROR: return "CUBLAS_STATUS_LICENSE_ERROR";
+#endif
+    }
+    return "Unknown cublas status";
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/cropAndResizeKernel.cu

@@ -0,0 +1,146 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename T>
+__global__ void cropAndResizeKernel(const int nthreads, const T* image_ptr, const float* boxes_ptr,
+                                    int num_boxes, int batch, int image_height, int image_width,
+                                    int crop_height, int crop_width, int depth,
+                                    float extrapolation_value, float* crops_ptr)
+{
+    for (int out_idx = threadIdx.x + blockIdx.x * blockDim.x ; out_idx < nthreads;
+            out_idx += blockDim.x * gridDim.x)
+    {
+        int idx =  out_idx;
+        const int x = idx % crop_width;
+        idx /= crop_width;
+        const int y = idx % crop_height;
+        idx /= crop_height;
+        const int d = idx % depth;
+        const int b = idx / depth;
+        const float y1 = boxes_ptr[b * 4];
+        const float x1 = boxes_ptr[b * 4 + 1];
+        const float y2 = boxes_ptr[b * 4 + 2];
+        const float x2 = boxes_ptr[b * 4 + 3];
+        //each image has num_boxes of boxes, so we simply divide to get the box index.
+        const int b_in = b / num_boxes;
+
+        if (b_in < 0 || b_in >= batch)
+        {
+            continue;
+        }
+
+        const float height_scale =
+            (crop_height > 1) ? (y2 - y1) * (image_height - 1) / (crop_height - 1)
+            : 0;
+        const float width_scale =
+            (crop_width > 1) ? (x2 - x1) * (image_width - 1) / (crop_width - 1) : 0;
+        const float in_y = (crop_height > 1)
+                            ? y1 * (image_height - 1) + y * height_scale
+                            : 0.5 * (y1 + y2) * (image_height - 1);
+
+        if (in_y < 0 || in_y > image_height - 1)
+        {
+            crops_ptr[out_idx] = extrapolation_value;
+            continue;
+        }
+
+        const float in_x = (crop_width > 1)
+                           ? x1 * (image_width - 1) + x * width_scale
+                           : 0.5 * (x1 + x2) * (image_width - 1);
+
+        if (in_x < 0 || in_x > image_width - 1)
+        {
+            crops_ptr[out_idx] = extrapolation_value;
+            continue;
+        }
+
+        const int top_y_index = floorf(in_y);
+        const int bottom_y_index = ceilf(in_y);
+        const float y_lerp = in_y - top_y_index;
+        const int left_x_index = floorf(in_x);
+        const int right_x_index = ceilf(in_x);
+        const float x_lerp = in_x - left_x_index;
+        const float top_left(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                               top_y_index) *
+                                               image_width +
+                                               left_x_index]));
+        const float top_right(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                                top_y_index) *
+                                                image_width +
+                                                right_x_index]));
+        const float bottom_left(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                               bottom_y_index) *
+                                               image_width +
+                                               left_x_index]));
+        const float bottom_right(static_cast<float>(
+                                     image_ptr[((b_in * depth + d) * image_height +
+                                               bottom_y_index) *
+                                               image_width +
+                                               right_x_index]));
+        const float top = top_left + (top_right - top_left) * x_lerp;
+        const float bottom = bottom_left + (bottom_right - bottom_left) * x_lerp;
+        crops_ptr[out_idx] = top + (bottom - top) * y_lerp;
+    }
+}
+
+
+
+
+
+int cropAndResizeInference(
+    cudaStream_t stream,
+    int n,
+    const void* image,
+    const void* rois,
+    int batch_size,
+    int input_height,
+    int input_width,
+    int num_boxes,
+    int crop_height,
+    int crop_width,
+    int depth,
+    void* output)
+{
+    int output_volume = batch_size * num_boxes * crop_height * crop_width * depth;
+    int block_size = 1024;
+    int grid_size = (output_volume + block_size - 1 ) / block_size;
+    cropAndResizeKernel<float> <<< grid_size, block_size, 0, stream>>>(output_volume,
+            static_cast<const float*>(image),
+            static_cast<const float*>(rois),
+            num_boxes,
+            batch_size,
+            input_height,
+            input_width,
+            crop_height,
+            crop_width,
+            depth,
+            0.0f,
+            static_cast<float*>(output));
+    return 0;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/decodeBbox3DKernels.cu

@@ -0,0 +1,188 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <iostream>
+#include <cuda_runtime_api.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+#define checkCudaErrors(status_)                                  \
+{                                                                 \
+  auto const status = status_;                                    \
+  if (status != 0)                                                \
+  {                                                               \
+    std::cout << "Cuda failure: " << cudaGetErrorString(status)   \
+              << " at line " << __LINE__                          \
+              << " in file " << __FILE__                          \
+              << " error status: " << status                      \
+              << std::endl;                                       \
+              exit(EXIT_FAILURE);                                           \
+    }                                                             \
+}
+
+#ifndef M_PI
+    #define M_PI 3.14159265358979323846
+#endif
+
+__device__ float sigmoid(const float x) { return 1.0f / (1.0f + expf(-x)); }
+
+__global__ void postprocess_kernal(const float *cls_input,
+                                        float const* box_input,
+                                        const float *dir_cls_input,
+                                        float *anchors,
+                                        float *anchors_bottom_height,
+                                        float *bndbox_output,
+                                        int *object_counter,
+                                        const float min_x_range,
+                                        const float max_x_range,
+                                        const float min_y_range,
+                                        const float max_y_range,
+                                        const int feature_x_size,
+                                        const int feature_y_size,
+                                        const int num_anchors,
+                                        const int num_classes,
+                                        const int num_box_values,
+                                        const float score_thresh,
+                                        const float dir_offset,
+                                        const float dir_limit_offset,
+                                        const int num_dir_bins)
+{
+  int max_box_num = feature_x_size * feature_y_size * num_anchors;
+  int loc_index =blockIdx.x;
+  int batch_idx = blockIdx.x / (feature_x_size * feature_y_size);
+  int loc_index_in_frame = blockIdx.x % (feature_x_size * feature_y_size);
+  int ith_anchor = threadIdx.x;
+  if (ith_anchor >= num_anchors)
+  {
+      return;
+  }
+  int col = loc_index_in_frame % feature_x_size;
+  int row = loc_index_in_frame / feature_x_size;
+  float x_offset = min_x_range + col * (max_x_range - min_x_range) / (feature_x_size - 1);
+  float y_offset = min_y_range + row * (max_y_range - min_y_range) / (feature_y_size - 1);
+  int cls_offset = loc_index * num_classes * num_anchors + ith_anchor * num_classes;
+  float dev_cls[2] = {-1, 0};
+  const float *scores = cls_input + cls_offset;
+  float max_score = sigmoid(scores[0]);
+  int cls_id = 0;
+  for (int i = 1; i < num_classes; i++) {
+    float cls_score = sigmoid(scores[i]);
+    if (cls_score > max_score) {
+      max_score = cls_score;
+      cls_id = i;
+    }
+  }
+  dev_cls[0] = static_cast<float>(cls_id);
+  dev_cls[1] = max_score;
+  if (dev_cls[1] >= score_thresh)
+  {
+    int box_offset = loc_index * num_anchors * num_box_values + ith_anchor * num_box_values;
+    int dir_cls_offset = loc_index * num_anchors * 2 + ith_anchor * 2;
+    float *anchor_ptr = anchors + ith_anchor * 4;
+    float z_offset = anchor_ptr[2] / 2 + anchors_bottom_height[ith_anchor / 2];
+    float anchor[7] = {x_offset, y_offset, z_offset, anchor_ptr[0], anchor_ptr[1], anchor_ptr[2], anchor_ptr[3]};
+    float const* box_encodings = box_input + box_offset;
+    float xa = anchor[0];
+    float ya = anchor[1];
+    float za = anchor[2];
+    float dxa = anchor[3];
+    float dya = anchor[4];
+    float dza = anchor[5];
+    float ra = anchor[6];
+    float diagonal = sqrtf(dxa * dxa + dya * dya);
+    float be0 = box_encodings[0] * diagonal + xa;
+    float be1 = box_encodings[1] * diagonal + ya;
+    float be2 = box_encodings[2] * dza + za;
+    float be3 = expf(box_encodings[3]) * dxa;
+    float be4 = expf(box_encodings[4]) * dya;
+    float be5 = expf(box_encodings[5]) * dza;
+    float be6 = box_encodings[6] + ra;
+    float yaw;
+    int dir_label = dir_cls_input[dir_cls_offset] > dir_cls_input[dir_cls_offset + 1] ? 0 : 1;
+    float period = 2.0f * float(M_PI) / num_dir_bins;
+    float val = be6 - dir_offset;
+    float dir_rot = val - floor(val / period + dir_limit_offset) * period;
+    yaw = dir_rot + dir_offset + period * dir_label;
+    int resCount = atomicAdd(object_counter + batch_idx, 1);
+    float *data = bndbox_output + (batch_idx * max_box_num + resCount) * 9;
+    data[0] = be0;
+    data[1] = be1;
+    data[2] = be2;
+    data[3] = be3;
+    data[4] = be4;
+    data[5] = be5;
+    data[6] = yaw;
+    data[7] = dev_cls[0];
+    data[8] = dev_cls[1];
+  }
+}
+
+
+void  decodeBbox3DLaunch(
+  const int batch_size,
+  const float *cls_input,
+  const float *box_input,
+  const float *dir_cls_input,
+  float *anchors,
+  float *anchors_bottom_height,
+  float *bndbox_output,
+  int *object_counter,
+  const float min_x_range,
+  const float max_x_range,
+  const float min_y_range,
+  const float max_y_range,
+  const int feature_x_size,
+  const int feature_y_size,
+  const int num_anchors,
+  const int num_classes,
+  const int num_box_values,
+  const float score_thresh,
+  const float dir_offset,
+  const float dir_limit_offset,
+  const int num_dir_bins,
+  cudaStream_t stream)
+{
+  int bev_size = batch_size * feature_x_size * feature_y_size;
+  dim3 threads (num_anchors);
+  dim3 blocks (bev_size);
+  postprocess_kernal<<<blocks, threads, 0, stream>>>
+                (cls_input,
+                 box_input,
+                 dir_cls_input,
+                 anchors,
+                 anchors_bottom_height,
+                 bndbox_output,
+                 object_counter,
+                 min_x_range,
+                 max_x_range,
+                 min_y_range,
+                 max_y_range,
+                 feature_x_size,
+                 feature_y_size,
+                 num_anchors,
+                 num_classes,
+                 num_box_values,
+                 score_thresh,
+                 dir_offset,
+                 dir_limit_offset,
+                 num_dir_bins);
+  checkCudaErrors(cudaGetLastError());
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/extractFgScores.cu

@@ -0,0 +1,81 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename T>
+pluginStatus_t extractFgScores_gpu(cudaStream_t stream, int N, int A, int H, int W, const void* scores, void* fgScores)
+{
+    // Copy all the objectness scores for one batch
+    size_t size = A * H * W * sizeof(T);
+    for (int n = 0; n < N; n++)
+    {
+        // Find out the starting pointer of the objectness scores in the input
+        size_t offset_ld = (n * 2 + 1) * A * H * W;
+        // Find out the starting pointer of the objectness scores in the output
+        size_t offset_st = n * A * H * W;
+        CSC(cudaMemcpyAsync(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size, cudaMemcpyDeviceToDevice, stream), STATUS_FAILURE);
+    }
+
+    return STATUS_SUCCESS;
+}
+
+template <typename T>
+pluginStatus_t extractFgScores_cpu(int N,
+                                  int A,
+                                  int H,
+                                  int W,
+                                  const void* scores,
+                                  void* fgScores)
+{
+    size_t size = A * H * W * sizeof(T);
+    for (int n = 0; n < N; n++)
+    {
+        size_t offset_ld = (n * 2 + 1) * A * H * W;
+        size_t offset_st = n * A * H * W;
+        memcpy(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size);
+    }
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t extractFgScores(cudaStream_t stream,
+                              const int N,
+                              const int A,
+                              const int H,
+                              const int W,
+                              const DataType t_scores,
+                              const DLayout_t l_scores,
+                              const void* scores,
+                              const DataType t_fgScores,
+                              const DLayout_t l_fgScores,
+                              void* fgScores)
+{
+    if (l_fgScores != NCHW || l_scores != NCHW)
+        return STATUS_BAD_PARAM;
+
+    if (t_fgScores != DataType::kFLOAT)
+        return STATUS_BAD_PARAM;
+
+    if (t_scores != DataType::kFLOAT)
+        return STATUS_BAD_PARAM;
+
+    return extractFgScores_gpu<float>(stream, N, A, H, W, scores, fgScores);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/gatherTopDetections.cu

@@ -0,0 +1,230 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+#include "common/plugin.h"
+#include <cuda_fp16.h>
+#include <array>
+
+
+namespace nvinfer1
+{
+namespace plugin
+{
+inline __device__ __half minus_fb(const __half& a, const __half& b)
+{
+#if __CUDA_ARCH__ >= 530
+    return a - b;
+#else
+    return __float2half(__half2float(a) - __half2float(b));
+#endif
+}
+
+inline __device__ float minus_fb(const float & a, const float & b) {
+    return a - b;
+}
+
+template <typename T_BBOX, typename T_SCORE, unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta)
+    __global__ void gatherTopDetections_kernel(
+        const bool shareLocation,
+        const int numImages,
+        const int numPredsPerClass,
+        const int numClasses,
+        const int topK,
+        const int keepTopK,
+        const int* indices,
+        const T_SCORE* scores,
+        const T_BBOX* bboxData,
+        int* keepCount,
+        T_BBOX* topDetections,
+        const T_SCORE score_shift)
+{
+    if (keepTopK > topK)
+        return;
+    for (int i = blockIdx.x * nthds_per_cta + threadIdx.x;
+         i < numImages * keepTopK;
+         i += gridDim.x * nthds_per_cta)
+    {
+        const int imgId = i / keepTopK;
+        const int detId = i % keepTopK;
+        const int offset = imgId * numClasses * topK;
+        const int index = indices[offset + detId];
+        const T_SCORE score = scores[offset + detId];
+        /*
+         * It is also likely that there is "bad bounding boxes" in the keepTopK bounding boxes.
+         * We set the bounding boxes parameters as the parameters shown below.
+         * These data will only show up at the end of keepTopK bounding boxes since the bounding boxes were sorted previously.
+         * It is also not going to affect the count of valid bounding boxes (keepCount).
+         * These data will probably never be used (because we have keepCount).
+         */
+        if (index == -1)
+        {
+            topDetections[i * 7] = imgId;  // image id
+            topDetections[i * 7 + 1] = -1; // label
+            topDetections[i * 7 + 2] = 0;  // confidence score
+            // score==0 will not pass the VisualizeBBox check
+            topDetections[i * 7 + 3] = 0;   // bbox xmin
+            topDetections[i * 7 + 4] = 0;   // bbox ymin
+            topDetections[i * 7 + 5] = 0;   // bbox xmax
+            topDetections[i * 7 + 6] = 0;   // bbox ymax
+        }
+        else
+        {
+            const int bboxOffset = imgId * (shareLocation ? numPredsPerClass : (numClasses * numPredsPerClass));
+            const int bboxId = ((shareLocation ? (index % numPredsPerClass)
+                        : index % (numClasses * numPredsPerClass)) + bboxOffset) * 4;
+            topDetections[i * 7] = imgId;                                                            // image id
+            topDetections[i * 7 + 1] = (index % (numClasses * numPredsPerClass)) / numPredsPerClass; // label
+            topDetections[i * 7 + 2] = score;                                                        // confidence score
+            // subtract 1.0 score shift we added in sortScorePerClass
+            topDetections[i * 7 + 2] = minus_fb(topDetections[i * 7 + 2], score_shift);
+            const T_BBOX xMin = bboxData[bboxId];
+            const T_BBOX yMin = bboxData[bboxId + 1];
+            const T_BBOX xMax = bboxData[bboxId + 2];
+            const T_BBOX yMax = bboxData[bboxId + 3];
+            // clipped bbox xmin
+            topDetections[i * 7 + 3] = __saturatef(xMin);
+            // clipped bbox ymin
+            topDetections[i * 7 + 4] = __saturatef(yMin);
+            // clipped bbox xmax
+            topDetections[i * 7 + 5] = __saturatef(xMax);
+            // clipped bbox ymax
+            topDetections[i * 7 + 6] = __saturatef(yMax);
+            // Atomic add to increase the count of valid keepTopK bounding boxes
+            // Without having to do manual sync.
+            atomicAdd(&keepCount[i / keepTopK], 1);
+        }
+    }
+}
+
+template <typename T_BBOX, typename T_SCORE>
+pluginStatus_t gatherTopDetections_gpu(
+    cudaStream_t stream,
+    const bool shareLocation,
+    const int numImages,
+    const int numPredsPerClass,
+    const int numClasses,
+    const int topK,
+    const int keepTopK,
+    const void* indices,
+    const void* scores,
+    const void* bboxData,
+    void* keepCount,
+    void* topDetections,
+    const float score_shift
+)
+{
+    CSC(cudaMemsetAsync(keepCount, 0, numImages * sizeof(int), stream), STATUS_FAILURE);
+    const int BS = 32;
+    const int GS = 32;
+    gatherTopDetections_kernel<T_BBOX, T_SCORE, BS><<<GS, BS, 0, stream>>>(shareLocation, numImages, numPredsPerClass,
+                                                                           numClasses, topK, keepTopK,
+                                                                           (int*) indices, (T_SCORE*) scores, (T_BBOX*) bboxData,
+                                                                           (int*) keepCount, (T_BBOX*) topDetections,
+                                                                           T_SCORE(score_shift));
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// gatherTopDetections LAUNCH CONFIG
+typedef pluginStatus_t (*gtdFunc)(cudaStream_t,
+                               const bool,
+                               const int,
+                               const int,
+                               const int,
+                               const int,
+                               const int,
+                               const void*,
+                               const void*,
+                               const void*,
+                               void*,
+                               void*,
+                               const float);
+struct gtdLaunchConfig
+{
+    DataType t_bbox;
+    DataType t_score;
+    gtdFunc function;
+
+    gtdLaunchConfig(DataType t_bbox, DataType t_score)
+        : t_bbox(t_bbox)
+        , t_score(t_score)
+        , function(nullptr)
+    {
+    }
+    gtdLaunchConfig(DataType t_bbox, DataType t_score, gtdFunc function)
+        : t_bbox(t_bbox)
+        , t_score(t_score)
+        , function(function)
+    {
+    }
+    bool operator==(gtdLaunchConfig const& other) const
+    {
+        return t_bbox == other.t_bbox && t_score == other.t_score;
+    }
+};
+
+using nvinfer1::DataType;
+
+static std::array<gtdLaunchConfig, 2> gtdLCOptions = {
+    gtdLaunchConfig(DataType::kFLOAT, DataType::kFLOAT, gatherTopDetections_gpu<float, float>),
+    gtdLaunchConfig(DataType::kHALF, DataType::kHALF, gatherTopDetections_gpu<__half, __half>)
+};
+
+pluginStatus_t gatherTopDetections(
+    cudaStream_t stream,
+    const bool shareLocation,
+    const int numImages,
+    const int numPredsPerClass,
+    const int numClasses,
+    const int topK,
+    const int keepTopK,
+    const DataType DT_BBOX,
+    const DataType DT_SCORE,
+    const void* indices,
+    const void* scores,
+    const void* bboxData,
+    void* keepCount,
+    void* topDetections,
+    const float score_shift)
+{
+    gtdLaunchConfig lc = gtdLaunchConfig(DT_BBOX, DT_SCORE);
+    for (unsigned i = 0; i < gtdLCOptions.size(); ++i)
+    {
+        if (lc == gtdLCOptions[i])
+        {
+            DEBUG_PRINTF("gatherTopDetections kernel %d\n", i);
+            return gtdLCOptions[i].function(stream,
+                                          shareLocation,
+                                          numImages,
+                                          numPredsPerClass,
+                                          numClasses,
+                                          topK,
+                                          keepTopK,
+                                          indices,
+                                          scores,
+                                          bboxData,
+                                          keepCount,
+                                          topDetections,
+                                          score_shift);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/generateAnchors.cu

@@ -0,0 +1,102 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+#include <cstdio>
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t generateAnchors_cpu(
+    int numRatios, float* ratios, int numScales, float* scales, int baseSize, float* anchors)
+{
+#ifdef DEBUG
+    DEBUG_PRINTF("Generating Anchors with:\n");
+    DEBUG_PRINTF("Scales:");
+    for (int s = 0; s < numScales; ++s)
+    {
+        DEBUG_PRINTF("%f\t", scales[s]);
+    }
+    DEBUG_PRINTF("\n");
+    DEBUG_PRINTF("Ratios:");
+    for (int r = 0; r < numRatios; ++r)
+    {
+        DEBUG_PRINTF("%f\t", ratios[r]);
+    }
+    DEBUG_PRINTF("\n");
+#endif
+
+    if ((numScales <= 0) || (numRatios <= 0) || (baseSize <= 0))
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    // Generate parameters for numRatios * numScales general anchor boxes
+    for (int r = 0; r < numRatios; ++r)
+    {
+        for (int s = 0; s < numScales; ++s)
+        {
+            int id = r * numScales + s;
+            float scale = scales[s];
+            float ratio = ratios[r];
+            float bs = baseSize;
+            float ws = round(sqrt((float) (bs * bs) / ratio));
+            float hs = round(ws * ratio);
+            // Width: bs / sqrt(ratio) * scale
+            // Height: bs * sqrt(ratio) * scale
+            ws *= scale;
+            hs *= scale;
+
+            // x_anchor_ctr
+            /*
+             * This value should not useful in this implementation of generating numRatios * numScales general anchor boxes.
+             * Because the center of anchor box in the original input raw image scale will not be dependent on this.
+             */
+            anchors[id * 4] = (bs - 1) / 2;
+            // y_anchor_ctr
+            /*
+             * This value should not useful in this implementation of generating numRatios * numScales general anchor boxes.
+             * Because the center of anchor box in the original input raw image scale will not be dependent on this.
+             */
+            anchors[id * 4 + 1] = (bs - 1) / 2;
+            // w_anchor
+            anchors[id * 4 + 2] = ws;
+            // h_anchor
+            anchors[id * 4 + 3] = hs;
+        }
+    }
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t generateAnchors(cudaStream_t stream,
+                              int numRatios,
+                              float* ratios,
+                              int numScales,
+                              float* scales,
+                              int baseSize,
+                              float* anchors)
+{
+    // Each anchor box has 4 parameters
+    int ac = numRatios * numScales * 4;
+    float* anchors_cpu;
+    CSC(cudaMallocHost((void**) &anchors_cpu, sizeof(float) * ac), STATUS_FAILURE);
+    pluginStatus_t status = generateAnchors_cpu(numRatios, ratios, numScales, scales, baseSize, anchors_cpu);
+    CSC(cudaMemcpyAsync(anchors, anchors_cpu, sizeof(float) * ac, cudaMemcpyHostToDevice, stream), STATUS_FAILURE);
+    CSC(cudaFreeHost(anchors_cpu), STATUS_FAILURE);
+    return status;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/gridAnchorLayer.cu

@@ -0,0 +1,132 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+#include "reducedMathPlugin.h"
+#include <iostream>
+
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+using nvinfer1::plugin::ReducedDivisor;
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__ void gridAnchorKernel(const GridAnchorParameters param,
+    const int numAspectRatios, ReducedDivisor divObj, const float* widths, const float* heights, float* outputData)
+{
+    // output dims: (H, W, param.numMinSize, (1+haveMaxSize+numAR-1), 4)
+    const int dim = param.H * param.W * numAspectRatios;
+    /*
+     * Parameters used to calculate the bounding box coordinates back to input image scale
+     * Normally we calculate the anchorStride = image_input_size (in pixel) / feature_map_size
+     * Here we do not use image_input_size for the moment
+     * Instead we use 1.0
+     * The coordinates calculated are scaled by the input image size.
+     * Most of the coordinates will be in a range of [0, 1], except for the bounding box coordinates going outside of
+     * the image Every coordinate will go back to the pixel coordinates in the input image if being multiplied by
+     * image_input_size.
+     */
+    float anchorStrideH = (1.0F / param.H);
+    float anchorStrideW = (1.0F / param.W);
+    float anchorOffsetH = 0.5F * anchorStrideH;
+    float anchorOffsetW = 0.5F * anchorStrideW;
+
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= dim)
+        return;
+    int arId, currIndex;
+    divObj.divmod(tid, currIndex, arId);
+
+    const int w = currIndex % param.W;
+    const int h = currIndex / param.W;
+
+    // Center coordinates
+    float yC = h * anchorStrideH + anchorOffsetH;
+    float xC = w * anchorStrideW + anchorOffsetW;
+
+    // x_min, y_min
+    float xMin = xC - 0.5 * widths[arId];
+    float yMin = yC - 0.5 * heights[arId];
+
+    // x_max, y_max
+    float xMax = xC + 0.5 * widths[arId];
+    float yMax = yC + 0.5 * heights[arId];
+
+    outputData[tid * 4] = xMin;
+    outputData[tid * 4 + 1] = yMin;
+    outputData[tid * 4 + 2] = xMax;
+    outputData[tid * 4 + 3] = yMax;
+
+    // Remember to move the output cursor
+    float* output = outputData + dim * 4;
+
+    // Simply copying the variance
+    output[tid * 4] = param.variance[0];
+    output[tid * 4 + 1] = param.variance[1];
+    output[tid * 4 + 2] = param.variance[2];
+    output[tid * 4 + 3] = param.variance[3];
+}
+
+pluginStatus_t anchorGridInference(cudaStream_t stream, const GridAnchorParameters param, const int numAspectRatios,
+    const void* widths, const void* heights, void* outputData)
+{
+    const int dim = param.H * param.W * numAspectRatios;
+    ReducedDivisor divObj(numAspectRatios);
+    if (dim > 5120)
+    {
+        const int BS = 128;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    else
+    {
+        const int BS = 32;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t anchorGridInference(cudaStream_t stream, const GridAnchorParameters param, const int numAspectRatios,
+    const void* widths, const void* heights, void* outputData)
+{
+    const int dim = param.H * param.W * numAspectRatios;
+    ReducedDivisor divObj(numAspectRatios);
+    if (dim > 5120)
+    {
+        const int BS = 128;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    else
+    {
+        const int BS = 32;
+        const int GS = (dim + BS - 1) / BS;
+        gridAnchorKernel<BS><<<GS, BS, 0, stream>>>(
+            param, numAspectRatios, divObj, (const float*) widths, (const float*) heights, (float*) outputData);
+    }
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/kernel.h

@@ -0,0 +1,261 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_KERNEL_H
+#define TRT_KERNEL_H
+
+#include "common/plugin.h"
+#include <algorithm>
+#include <cassert>
+#include <cstdio>
+#include <cuda_fp16.h>
+
+#define DEBUG_ENABLE 0
+namespace nvinfer1
+{
+namespace plugin
+{
+typedef enum
+{
+    NCHW = 0,
+    NC4HW = 1,
+    NC32HW = 2
+} DLayout_t;
+#ifndef TRT_RPNLAYER_H
+
+pluginStatus_t allClassNMS(cudaStream_t stream, int32_t num, int32_t num_classes, int32_t num_preds_per_class,
+    int32_t top_k, float nms_threshold, bool share_location, bool isNormalized, nvinfer1::DataType DT_SCORE,
+    nvinfer1::DataType DT_BBOX, void* bbox_data, void* beforeNMS_scores, void* beforeNMS_index_array,
+    void* afterNMS_scores, void* afterNMS_index_array, bool flipXY, float const score_shift, bool caffeSemantics);
+
+pluginStatus_t nmsInference(cudaStream_t stream, int32_t N, int32_t boxesSize, int32_t scoresSize, bool shareLocation,
+    int32_t backgroundLabelId, int32_t numPredsPerClass, int32_t numClasses, int32_t topK, int32_t keepTopK,
+    float scoreThreshold, float iouThreshold, nvinfer1::DataType DT_BBOX, void const* locData,
+    nvinfer1::DataType DT_SCORE, void const* confData, void* keepCount, void* nmsedBoxes, void* nmsedScores,
+    void* nmsedClasses, void* workspace, bool isNormalized = true, bool confSigmoid = false, bool clipBoxes = true,
+    int32_t scoreBits = 16, bool caffeSemantics = true);
+
+pluginStatus_t gatherTopDetections(cudaStream_t stream, bool shareLocation, int32_t numImages, int32_t numPredsPerClass,
+    int32_t numClasses, int32_t topK, int32_t keepTopK, nvinfer1::DataType DT_BBOX, nvinfer1::DataType DT_SCORE,
+    void const* indices, void const* scores, void const* bboxData, void* keepCount, void* topDetections,
+    float const scoreShift);
+
+size_t detectionForwardBBoxDataSize(int32_t N, int32_t C1, nvinfer1::DataType DT_BBOX);
+
+size_t detectionForwardBBoxPermuteSize(bool shareLocation, int32_t N, int32_t C1, nvinfer1::DataType DT_BBOX);
+
+size_t sortScoresPerClassWorkspaceSize(
+    int32_t num, int32_t num_classes, int32_t num_preds_per_class, nvinfer1::DataType DT_CONF);
+
+size_t sortScoresPerImageWorkspaceSize(int32_t num_images, int32_t num_items_per_image, nvinfer1::DataType DT_SCORE);
+
+pluginStatus_t sortScoresPerImage(cudaStream_t stream, int32_t num_images, int32_t num_items_per_image,
+    nvinfer1::DataType DT_SCORE, void* unsorted_scores, void* unsorted_bbox_indices, void* sorted_scores,
+    void* sorted_bbox_indices, void* workspace, int32_t score_bits);
+
+pluginStatus_t sortScoresPerClass(cudaStream_t stream, int32_t num, int32_t num_classes, int32_t num_preds_per_class,
+    int32_t background_label_id, float confidence_threshold, nvinfer1::DataType DT_SCORE, void* conf_scores_gpu,
+    void* index_array_gpu, void* workspace, int32_t const score_bits, float const score_shift);
+
+size_t calculateTotalWorkspaceSize(size_t* workspaces, int32_t count);
+
+char const* cublasGetErrorString(nvinfer1::pluginInternal::cublasStatus_t error);
+
+pluginStatus_t permuteData(cudaStream_t stream, int32_t nthreads, int32_t num_classes, int32_t num_data,
+    int32_t num_dim, nvinfer1::DataType DT_DATA, bool confSigmoid, void const* data, void* new_data);
+
+size_t detectionForwardPreNMSSize(int32_t N, int32_t C2);
+
+size_t detectionForwardPostNMSSize(int32_t N, int32_t numClasses, int32_t topK);
+
+size_t normalizePluginWorkspaceSize(bool acrossSpatial, int32_t C, int32_t H, int32_t W);
+
+pluginStatus_t normalizeInference(cudaStream_t stream, nvinfer1::pluginInternal::cublasHandle_t handle,
+    bool acrossSpatial, bool channelShared, int32_t N, int32_t C, int32_t H, int32_t W, float eps, void const* scale,
+    void const* inputData, void* outputData, void* workspace);
+
+pluginStatus_t scatterNDInference(cudaStream_t stream, int32_t* outputDims, int32_t nOutputDims, int32_t sliceRank,
+    int32_t nRows, int32_t rowSize, int32_t CopySize, int32_t sizeOfElementInBytes, void const* index,
+    void const* updates, void const* data, void* output, void* workspace);
+
+pluginStatus_t priorBoxInference(cudaStream_t stream, nvinfer1::plugin::PriorBoxParameters param, int32_t H, int32_t W,
+    int32_t numPriors, int32_t numAspectRatios, void const* minSize, void const* maxSize, void const* aspectRatios,
+    void* outputData);
+
+pluginStatus_t lReLUInference(cudaStream_t stream, int32_t n, float negativeSlope, void const* input, void* output);
+
+pluginStatus_t reorgInference(cudaStream_t stream, int32_t batch, int32_t C, int32_t H, int32_t W, int32_t stride,
+    void const* input, void* output);
+
+pluginStatus_t anchorGridInference(cudaStream_t stream, nvinfer1::plugin::GridAnchorParameters param,
+    int32_t numAspectRatios, void const* aspectRatios, void const* scales, void* outputData);
+
+pluginStatus_t regionInference(cudaStream_t stream, int32_t batch, int32_t C, int32_t H, int32_t W, int32_t num,
+    int32_t coords, int32_t classes, bool hasSoftmaxTree, nvinfer1::plugin::softmaxTree const* smTree,
+    void const* input, void* output);
+
+// GENERATE ANCHORS
+// For now it takes host pointers - ratios and scales but
+// in GPU MODE anchors should be device pointer
+pluginStatus_t generateAnchors(cudaStream_t stream,
+    int32_t numRatios, // number of ratios
+    float* ratios,     // ratio array
+    int32_t numScales, // number of scales
+    float* scales,     // scale array
+    int32_t baseSize,  // size of the base anchor (baseSize x baseSize)
+    float* anchors);   // output anchors (numRatios x numScales)
+
+// BBD2P
+pluginStatus_t bboxDeltas2Proposals(cudaStream_t stream,
+    int32_t N,                     // batch size
+    int32_t A,                     // number of anchors
+    int32_t H,                     // last feature map H
+    int32_t W,                     // last feature map W
+    int32_t featureStride,         // feature stride
+    float minBoxSize,              // minimum allowed box size before scaling
+    float const* imInfo,           // image info (nrows, ncols, image scale)
+    float const* anchors,          // input anchors
+    nvinfer1::DataType tDeltas,    // type of input deltas
+    DLayout_t lDeltas,             // layout of input deltas
+    void const* deltas,            // input deltas
+    nvinfer1::DataType tProposals, // type of output proposals
+    DLayout_t lProposals,          // layout of output proposals
+    void* proposals,               // output proposals
+    nvinfer1::DataType tScores,    // type of output scores
+    DLayout_t lScores,             // layout of output scores
+    void* scores);                 // output scores (the score associated with too small box will be set to -inf)
+
+// NMS
+pluginStatus_t nms(cudaStream_t stream,
+    int32_t N,                     // batch size
+    int32_t R,                     // number of ROIs (region of interest) per image
+    int32_t preNmsTop,             // number of proposals before applying NMS
+    int32_t nmsMaxOut,             // number of remaining proposals after applying NMS
+    float iouThreshold,            // IoU threshold
+    nvinfer1::DataType tFgScores,  // type of foreground scores
+    DLayout_t lFgScores,           // layout of foreground scores
+    void* fgScores,                // foreground scores
+    nvinfer1::DataType tProposals, // type of proposals
+    DLayout_t lProposals,          // layout of proposals
+    void const* proposals,         // proposals
+    void* workspace,               // workspace
+    nvinfer1::DataType tRois,      // type of ROIs
+    void* rois);                   // ROIs
+
+// WORKSPACE SIZES
+size_t proposalsForwardNMSWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+size_t proposalsForwardBboxWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W);
+
+size_t proposalForwardFgScoresWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W);
+
+size_t proposalsInferenceWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+size_t RPROIInferenceFusedWorkspaceSize(int32_t N, int32_t A, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+// PROPOSALS INFERENCE
+pluginStatus_t proposalsInference(cudaStream_t stream, int32_t N, int32_t A, int32_t H, int32_t W,
+    int32_t featureStride, int32_t preNmsTop, int32_t nmsMaxOut, float iouThreshold, float minBoxSize,
+    float const* imInfo, float const* anchors, nvinfer1::DataType tScores, DLayout_t lScores, void const* scores,
+    nvinfer1::DataType tDeltas, DLayout_t lDeltas, void const* deltas, void* workspace, nvinfer1::DataType tRois,
+    void* rois);
+
+// EXTRACT FG SCORES
+pluginStatus_t extractFgScores(cudaStream_t stream, int32_t N, int32_t A, int32_t H, int32_t W,
+    nvinfer1::DataType tScores, DLayout_t lScores, void const* scores, nvinfer1::DataType tFgScores,
+    DLayout_t lFgScores, void* fgScores);
+
+// ROI INFERENCE
+pluginStatus_t roiInference(cudaStream_t stream,
+    int32_t const R,        // TOTAL number of rois -> ~nmsMaxOut * N
+    int32_t const N,        // Batch size
+    int32_t const C,        // Channels
+    int32_t const H,        // Input feature map H
+    int32_t const W,        // Input feature map W
+    int32_t const poolingH, // Output feature map H
+    int32_t const poolingW, // Output feature map W
+    float const spatialScale, nvinfer1::DataType const tRois, void const* rois, nvinfer1::DataType const tFeatureMap,
+    DLayout_t const lFeatureMap, void const* featureMap, nvinfer1::DataType const tTop, DLayout_t const lTop, void* top,
+    size_t deviceSmemSize);
+
+// ROI FORWARD
+pluginStatus_t roiForward(cudaStream_t stream,
+    int32_t R,        // TOTAL number of rois -> ~nmsMaxOut * N
+    int32_t N,        // Batch size
+    int32_t C,        // Channels
+    int32_t H,        // Input feature map H
+    int32_t W,        // Input feature map W
+    int32_t poolingH, // Output feature map H
+    int32_t poolingW, // Output feature map W
+    float spatialScale, nvinfer1::DataType tRois, void const* rois, nvinfer1::DataType tFeatureMap,
+    DLayout_t lFeatureMap, void const* featureMap, nvinfer1::DataType tTop, DLayout_t lTop, void* top, int32_t* maxIds);
+
+// RP ROI Fused INFERENCE
+pluginStatus_t RPROIInferenceFused(cudaStream_t stream, int32_t N, int32_t A, int32_t C, int32_t H, int32_t W,
+    int32_t poolingH, int32_t poolingW, int32_t featureStride, int32_t preNmsTop, int32_t nmsMaxOut, float iouThreshold,
+    float minBoxSize, float spatialScale, float const* imInfo, float const* anchors, nvinfer1::DataType tScores,
+    DLayout_t lScores, void const* scores, nvinfer1::DataType tDeltas, DLayout_t lDeltas, void const* deltas,
+    nvinfer1::DataType tFeatureMap, DLayout_t lFeatureMap, void const* featureMap, void* workspace,
+    nvinfer1::DataType tRois, void* rois, nvinfer1::DataType tTop, DLayout_t lTop, void* top, size_t deviceSmemSize);
+
+// GENERATE ANCHORS CPU
+pluginStatus_t generateAnchors_cpu(
+    int32_t numRatios, float* ratios, int32_t numScales, float* scales, int32_t baseSize, float* anchors);
+
+int32_t cropAndResizeInference(cudaStream_t stream, int32_t n, void const* image, void const* rois, int32_t batch_size,
+    int32_t input_height, int32_t input_width, int32_t num_boxes, int32_t crop_height, int32_t crop_width,
+    int32_t depth, void* output);
+
+int32_t proposalInference_gpu(cudaStream_t stream, void const* rpn_prob, void const* rpn_regr, int32_t batch_size,
+    int32_t input_height, int32_t input_width, int32_t rpn_height, int32_t rpn_width, int32_t MAX_BOX_NUM,
+    int32_t RPN_PRE_NMS_TOP_N, float* ANCHOR_SIZES, int32_t anc_size_num, float* ANCHOR_RATIOS, int32_t anc_ratio_num,
+    float rpn_std_scaling, int32_t rpn_stride, float bbox_min_size, float nms_iou_threshold, void* workspace,
+    void* output);
+
+size_t _get_workspace_size(
+    int32_t N, int32_t anc_size_num, int32_t anc_ratio_num, int32_t H, int32_t W, int32_t nmsMaxOut);
+
+void decodeBbox3DLaunch(int32_t const batch_size, float const* cls_input, float const* box_input,
+    float const* dir_cls_input, float* anchors, float* anchors_bottom_height, float* bndbox_output,
+    int32_t* object_counter, float const min_x_range, float const max_x_range, float const min_y_range,
+    float const max_y_range, int32_t const feature_x_size, int32_t const feature_y_size, int32_t const num_anchors,
+    int32_t const num_classes, int32_t const num_box_values, float const score_thresh, float const dir_offset,
+    float const dir_limit_offset, int32_t const num_dir_bins, cudaStream_t stream = 0);
+
+template <typename Element>
+int32_t pillarScatterKernelLaunch(int32_t batch_size, int32_t max_pillar_num, int32_t num_features,
+    Element const* pillar_features_data, uint32_t const* coords_data, uint32_t const* params_data, uint32_t featureX,
+    uint32_t featureY, Element* spatial_feature_data, cudaStream_t stream);
+
+void generateVoxels_launch(int32_t batch_size, int32_t max_num_points, float* points, uint32_t* points_size,
+    float min_x_range, float max_x_range, float min_y_range, float max_y_range, float min_z_range, float max_z_range,
+    float pillar_x_size, float pillar_y_size, float pillar_z_size, int32_t grid_y_size, int32_t grid_x_size,
+    int32_t num_point_values, int32_t max_points_per_voxel, uint32_t* mask, float* voxels, cudaStream_t stream);
+
+void generateBaseFeatures_launch(int32_t batch_size, uint32_t* mask, float* voxels, int32_t grid_y_size,
+    int32_t grid_x_size, uint32_t* pillar_num, int32_t max_pillar_num, int32_t max_points_per_voxel,
+    int32_t num_point_values, float* voxel_features, uint32_t* voxel_num_points, uint32_t* coords, cudaStream_t stream);
+
+int32_t generateFeatures_launch(int32_t batch_size, int32_t dense_pillar_num, float* voxel_features,
+    uint32_t* voxel_num_points, uint32_t* coords, uint32_t* params, float voxel_x, float voxel_y, float voxel_z,
+    float range_min_x, float range_min_y, float range_min_z, uint32_t voxel_features_size, uint32_t max_points,
+    uint32_t max_voxels, uint32_t num_point_values, float* features, cudaStream_t stream);
+
+#endif // TRT_RPNLAYER_H
+} // namespace plugin
+} // namespace nvinfer1
+#endif

plugin/common/kernels/lReLU.cu

@@ -0,0 +1,49 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "common/kernels/kernel.h"
+namespace nvinfer1
+{
+namespace plugin
+{
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__
+    void pReLUKernel(const int n, const float negativeSlope, const float* input, float* output)
+{
+    for (int i = blockIdx.x * nthdsPerCTA + threadIdx.x; i < n; i += gridDim.x * nthdsPerCTA)
+    {
+        output[i] = input[i] > 0 ? input[i] : input[i] * negativeSlope;
+    }
+}
+
+pluginStatus_t lReLUGPU(cudaStream_t stream, const int n, const float negativeSlope, const void* input, void* output)
+{
+    const int BS = 512;
+    const int GS = (n + BS - 1) / BS;
+    pReLUKernel<BS><<<GS, BS, 0, stream>>>(n, negativeSlope,
+                                           (const float*) input,
+                                           (float*) output);
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t lReLUInference(
+    cudaStream_t stream, const int n, const float negativeSlope, const void* input, void* output)
+{
+    return lReLUGPU(stream, n, negativeSlope, (const float*) input, (float*) output);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/maskRCNNKernels.h

@@ -0,0 +1,276 @@
+/*
+ * 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 TRT_MASKRCNN_UTILS_H
+#define TRT_MASKRCNN_UTILS_H
+
+#include "NvInfer.h"
+#include "common/plugin.h"
+
+inline size_t nAlignUp(size_t x, size_t align)
+{
+    size_t mask = align - 1;
+    PLUGIN_ASSERT((align & mask) == 0); // power of 2
+    return (x + mask) & (~mask);
+}
+
+inline size_t nAlignDown(size_t x, size_t align)
+{
+    size_t mask = align - 1;
+    PLUGIN_ASSERT((align & mask) == 0); // power of 2
+    return (x) & (~mask);
+}
+
+inline size_t dimVolume(const nvinfer1::Dims& dims)
+{
+    size_t volume = 1;
+    for (int32_t i = 0; i < dims.nbDims; ++i)
+        volume *= dims.d[i];
+
+    return volume;
+}
+
+inline size_t typeSize(const nvinfer1::DataType type)
+{
+    switch (type)
+    {
+    case nvinfer1::DataType::kFLOAT: return sizeof(float);
+    case nvinfer1::DataType::kBF16: return sizeof(uint16_t);
+    case nvinfer1::DataType::kHALF: return sizeof(uint16_t);
+    case nvinfer1::DataType::kINT8: return sizeof(uint8_t);
+    case nvinfer1::DataType::kINT32: return sizeof(int32_t);
+    case nvinfer1::DataType::kINT64: return sizeof(int64_t);
+    case nvinfer1::DataType::kBOOL: return sizeof(bool);
+    case nvinfer1::DataType::kUINT8: return sizeof(uint8_t);
+    case nvinfer1::DataType::kFP8:
+    case nvinfer1::DataType::kINT4:
+    case nvinfer1::DataType::kFP4:
+    case nvinfer1::DataType::kE8M0: PLUGIN_FAIL("Unsupported data type");
+    }
+    return 0;
+}
+
+#define AlignMem(x) nAlignUp(x, 256)
+
+struct RefineNMSParameters
+{
+    int32_t backgroundLabelId, numClasses, keepTopK;
+    float scoreThreshold, iouThreshold;
+};
+
+struct RefineDetectionWorkSpace
+{
+    RefineDetectionWorkSpace(const int32_t batchSize, const int32_t sampleCount, const RefineNMSParameters& param,
+        const nvinfer1::DataType type);
+
+    RefineDetectionWorkSpace() = default;
+
+    nvinfer1::DimsHW argMaxScoreDims;
+    nvinfer1::DimsHW argMaxBboxDims;
+    nvinfer1::DimsHW argMaxLabelDims;
+    nvinfer1::DimsHW sortClassScoreDims;
+    nvinfer1::DimsHW sortClassLabelDims;
+    nvinfer1::DimsHW sortClassSampleIdxDims;
+    nvinfer1::Dims sortClassValidCountDims = {1, {1, 0}};
+    nvinfer1::DimsHW sortClassPosDims;
+    nvinfer1::DimsHW sortNMSMarkDims;
+
+    size_t argMaxScoreOffset = 0;
+    size_t argMaxBboxOffset = 0;
+    size_t argMaxLabelOffset = 0;
+    size_t sortClassScoreOffset = 0;
+    size_t sortClassLabelOffset = 0;
+    size_t sortClassSampleIdxOffset = 0;
+    size_t sortClassValidCountOffset = 0;
+    size_t sortClassPosOffset = 0;
+    size_t sortNMSMarkOffset = 0;
+    size_t totalSize = 0;
+};
+
+struct ProposalWorkSpace
+{
+    ProposalWorkSpace(const int32_t batchSize, const int32_t inputCnt, const int32_t sampleCount,
+        const RefineNMSParameters& param, const nvinfer1::DataType type);
+
+    ProposalWorkSpace() = default;
+
+    nvinfer1::DimsHW preRefineScoreDims;
+    nvinfer1::DimsHW preRefineSortedScoreDims;
+    nvinfer1::DimsHW preRefineBboxDims;
+    nvinfer1::DimsHW argMaxScoreDims;
+    nvinfer1::DimsHW argMaxBboxDims;
+    nvinfer1::DimsHW argMaxLabelDims;
+    nvinfer1::DimsHW sortClassScoreDims;
+    nvinfer1::DimsHW sortClassLabelDims;
+    nvinfer1::DimsHW sortClassSampleIdxDims;
+    nvinfer1::Dims sortClassValidCountDims = {1, {1, 0}};
+    nvinfer1::DimsHW sortClassPosDims;
+    nvinfer1::DimsHW sortNMSMarkDims;
+
+    size_t tempStorageOffset = 0;
+    size_t preRefineScoreOffset = 0;
+    size_t preRefineSortedScoreOffset = 0;
+    size_t preRefineBboxOffset = 0;
+    size_t argMaxScoreOffset = 0;
+    size_t argMaxBboxOffset = 0;
+    size_t argMaxLabelOffset = 0;
+    size_t sortClassScoreOffset = 0;
+    size_t sortClassLabelOffset = 0;
+    size_t sortClassSampleIdxOffset = 0;
+    size_t sortClassValidCountOffset = 0;
+    size_t sortClassPosOffset = 0;
+    size_t sortNMSMarkOffset = 0;
+    size_t totalSize = 0;
+};
+
+struct MultilevelProposeROIWorkSpace
+{
+    MultilevelProposeROIWorkSpace(const int32_t batchSize, const int32_t inputCnt, const int32_t sampleCount,
+        const RefineNMSParameters& param, const nvinfer1::DataType type);
+
+    MultilevelProposeROIWorkSpace() = default;
+
+    nvinfer1::DimsHW preRefineSortedScoreDims;
+    nvinfer1::DimsHW preRefineBboxDims;
+    nvinfer1::DimsHW argMaxScoreDims;
+    nvinfer1::DimsHW argMaxBboxDims;
+    nvinfer1::DimsHW argMaxLabelDims;
+    nvinfer1::DimsHW sortClassScoreDims;
+    nvinfer1::DimsHW sortClassLabelDims;
+    nvinfer1::DimsHW sortClassSampleIdxDims;
+    nvinfer1::Dims sortClassValidCountDims = {1, {1, 0}};
+    nvinfer1::DimsHW sortClassPosDims;
+    nvinfer1::DimsHW sortNMSMarkDims;
+
+    size_t tempStorageOffset = 0;
+    size_t preRefineSortedScoreOffset = 0;
+    size_t preRefineBboxOffset = 0;
+    size_t argMaxScoreOffset = 0;
+    size_t argMaxBboxOffset = 0;
+    size_t argMaxLabelOffset = 0;
+    size_t sortClassScoreOffset = 0;
+    size_t sortClassLabelOffset = 0;
+    size_t sortClassSampleIdxOffset = 0;
+    size_t sortClassValidCountOffset = 0;
+    size_t sortClassPosOffset = 0;
+    size_t sortNMSMarkOffset = 0;
+    size_t totalSize = 0;
+};
+
+struct ConcatTopKWorkSpace
+{
+    ConcatTopKWorkSpace(
+        const int32_t batchSize, const int32_t concatCnt, const int32_t topK, const nvinfer1::DataType inType);
+
+    ConcatTopKWorkSpace() = default;
+
+    nvinfer1::DimsHW concatedScoreDims;
+    nvinfer1::DimsHW concatedBBoxDims;
+    nvinfer1::DimsHW sortedScoreDims;
+    nvinfer1::DimsHW sortedBBoxDims;
+
+    size_t tempStorageOffset = 0;
+    size_t concatedScoreOffset = 0;
+    size_t concatedBBoxOffset = 0;
+    size_t sortedScoreOffset = 0;
+    size_t sortedBBoxOffset = 0;
+    size_t totalSize = 0;
+};
+
+cudaError_t RefineBatchClassNMS(cudaStream_t stream, int32_t N, int32_t samples, nvinfer1::DataType dtype,
+    const RefineNMSParameters& param, const RefineDetectionWorkSpace& refineOffset, void* workspace,
+    const void* inScores, const void* inDelta, const void* inCountValid, const void* inROI, void* outDetections);
+
+cudaError_t DetectionPostProcess(cudaStream_t stream, int32_t N, int32_t samples, const float* regWeight,
+    const float inputHeight, const float inputWidth, nvinfer1::DataType dtype, const RefineNMSParameters& param,
+    const RefineDetectionWorkSpace& refineOffset, void* workspace, const void* inScores, const void* inDelta,
+    const void* inCountValid, const void* inROI, void* outDetections);
+
+cudaError_t proposalRefineBatchClassNMS(cudaStream_t stream, int32_t N,
+    int32_t inputCnt, // candidate anchors
+    int32_t samples,  // preNMS_topK
+    nvinfer1::DataType dtype, const RefineNMSParameters& param, const ProposalWorkSpace& proposalOffset,
+    void* workspace, const void* inScores, const void* inDelta, const void* inCountValid, const void* inAnchors,
+    void* outProposals);
+
+// inScores: [N, anchorsCnt, 1]
+// inDelta: [N, anchorsCnt, 4]
+// outScores: [N, topK, 1]
+// outBbox: [N, topK, 4]
+cudaError_t MultilevelPropose(cudaStream_t stream, int32_t N,
+    int32_t inputCnt, // candidate anchors number among feature map
+    int32_t samples,  // pre nms cnt
+    const float* regWeight, const float inputHeight, const float inputWidth, nvinfer1::DataType dtype,
+    const RefineNMSParameters& param, const MultilevelProposeROIWorkSpace& proposalOffset, void* workspace,
+    const void* inScore, const void* inDelta, void* inCountValid, const void* inAnchors, void* outScores,
+    void* outBbox);
+
+// inScores: [N, topK, 1] * featureCnt
+// inBboxes: [N, topK, 4] * featureCnt
+// outProposals: [N, topK, 4]
+cudaError_t ConcatTopK(cudaStream_t stream, int32_t N, int32_t featureCnt, int32_t topK, nvinfer1::DataType dtype,
+    void* workspace, const ConcatTopKWorkSpace& spaceOffset, void** inScores, void** inBBox, void* outProposals);
+
+cudaError_t DecodeBBoxes(cudaStream_t stream, int32_t N,
+    int32_t samples, // number of anchors per image
+    const float* regWeight, const float inputHeight, const float inputWidth,
+    const void* anchors, // [N, anchors, (y1, x1, y2, x2)]
+    const void* delta,   //[N, anchors, (dy, dx, log(dh), log(dw)]
+    void* outputBbox, nvinfer1::DataType dtype);
+
+cudaError_t ApplyDelta2Bboxes(cudaStream_t stream, int32_t N,
+    int32_t samples,     // number of anchors per image
+    const void* anchors, // [N, anchors, (y1, x1, y2, x2)]
+    const void* delta,   //[N, anchors, (dy, dx, log(dh), log(dw)]
+    void* outputBbox);
+
+struct xy_t
+{
+    int32_t y;
+    int32_t x;
+
+    xy_t()
+        : y(0)
+        , x(0)
+    {
+    }
+    xy_t(int32_t y_, int32_t x_)
+        : y(y_)
+        , x(x_)
+    {
+    }
+};
+// PYRAMID ROIALIGN
+cudaError_t roiAlign(cudaStream_t const stream, int32_t const batchSize, xy_t const imageSize,
+    int32_t const featureCount, int32_t const roiCount, float const firstThreshold, int32_t const transformCoords,
+    bool const absCoords, bool const swapCoords, bool const plusOneCoords, int32_t const samplingRatio,
+    void const* rois, void const* const layers[], xy_t const* layerDims, void* pooled, xy_t const poolDims);
+
+cudaError_t roiAlignHalfCenter(cudaStream_t stream, int32_t batchSize, int32_t featureCount, int32_t roiCount,
+    float firstThreshold,
+
+    int32_t inputHeight, int32_t inputWidth, const void* rois, const void* const layers[], const xy_t* layerDims,
+
+    void* pooled, const xy_t poolDims, const nvinfer1::DataType dtype);
+
+// RESIZE NEAREST
+void resizeNearest(dim3 grid, dim3 block, cudaStream_t stream, int32_t nbatch, float scale, int2 osize,
+    float const* idata, int32_t istride, int32_t ibatchstride, float* odata, int32_t ostride, int32_t obatchstride);
+// SPECIAL SLICE
+void specialSlice(cudaStream_t stream, int32_t batch_size, int32_t boxes_cnt, const void* idata, void* odata);
+
+#endif // TRT_MASKRCNN_UTILS_H

plugin/common/kernels/nmsLayer.cu

@@ -0,0 +1,441 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/bboxUtils.h"
+#include "common/kernels/kernel.h"
+#include "cuda_runtime_api.h"
+#include <algorithm>
+#include <array>
+#include <cub/cub.cuh>
+#include <functional>
+#include <stdint.h>
+#include <stdio.h>
+
+using namespace nvinfer1;
+namespace nvinfer1
+{
+namespace plugin
+{
+// CUB's bug workaround:
+// To work properly for large batch size CUB segmented sort needs ridiculous
+// workspace alignment.
+const uintptr_t ALIGNMENT = 1 << 20;
+
+// IOU
+template <typename TFloat>
+__device__ __host__ inline float IoU(const Bbox<TFloat>& a, const Bbox<TFloat>& b)
+{
+    TFloat left = max(a.xmin, b.xmin), right = min(a.xmax, b.xmax);
+    TFloat top = max(a.ymin, b.ymin), bottom = min(a.ymax, b.ymax);
+    TFloat width = max((TFloat)(right - left + (TFloat) 1.0), (TFloat) 0.0);
+    TFloat height = max((TFloat)(bottom - top + (TFloat) 1.0), (TFloat) 0.0);
+    TFloat interS = width * height;
+    TFloat Sa = (a.xmax - a.xmin + (TFloat) 1) * (a.ymax - a.ymin + (TFloat) 1);
+    TFloat Sb = (b.xmax - b.xmin + (TFloat) 1) * (b.ymax - b.ymin + (TFloat) 1);
+    return (float) interS / (float) (Sa + Sb - interS);
+}
+
+// NMS KERNEL FOR SMALL BATCH SIZE
+template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
+__global__ __launch_bounds__(DIM) void nmsKernel1(const int propSize,
+                                                  Bbox<T_PROPOSALS> const* __restrict__ preNmsProposals,
+                                                  T_ROIS* __restrict__ afterNmsProposals,
+                                                  const int preNmsTopN,
+                                                  const float nmsThres,
+                                                  const int afterNmsTopN)
+{
+    __shared__ bool kept_boxes[TSIZE * DIM];
+    int kept = 0;
+    int batch_offset = blockIdx.x * propSize;
+    int max_box_idx = batch_offset + preNmsTopN;
+    int batch_offset_out = blockIdx.x * afterNmsTopN;
+
+    int flag_idx[TSIZE];
+    int boxes_idx[TSIZE];
+    Bbox<T_PROPOSALS> cur_boxes[TSIZE];
+
+// initialize kept_boxes
+#pragma unroll
+    for (int i = 0; i < TSIZE; i++)
+    {
+        boxes_idx[i] = threadIdx.x + batch_offset + DIM * i;
+        flag_idx[i] = threadIdx.x + DIM * i;
+
+        if (boxes_idx[i] < max_box_idx)
+        {
+            cur_boxes[i] = preNmsProposals[boxes_idx[i]];
+            kept_boxes[flag_idx[i]] = true;
+        }
+        else
+        {
+            kept_boxes[flag_idx[i]] = false;
+            boxes_idx[i] = -1.0f;
+            flag_idx[i] = -1.0f;
+        }
+    }
+
+    int ref_box_idx = 0 + batch_offset;
+
+    // remove the overlapped boxes
+    while ((kept < afterNmsTopN) && (ref_box_idx < max_box_idx))
+    {
+        Bbox<T_PROPOSALS> ref_box;
+        ref_box = preNmsProposals[ref_box_idx];
+
+#pragma unroll
+        for (int i = 0; i < TSIZE; i++)
+        {
+            if (boxes_idx[i] > ref_box_idx)
+            {
+                if (IoU(ref_box, cur_boxes[i]) > nmsThres)
+                {
+                    kept_boxes[flag_idx[i]] = false;
+                }
+            }
+            else if (boxes_idx[i] == ref_box_idx)
+            {
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 0] = ref_box.xmin;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 1] = ref_box.ymin;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 2] = ref_box.xmax;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 3] = ref_box.ymax;
+            }
+        }
+        __syncthreads();
+
+        do
+        {
+            ref_box_idx++;
+        } while (!kept_boxes[ref_box_idx - batch_offset] && ref_box_idx < max_box_idx);
+
+        kept++;
+    }
+}
+
+// NMS KERNEL FOR LARGE BATCH SIZE
+template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
+__global__ __launch_bounds__(DIM) void nmsKernel2(const int propSize,
+                                                  Bbox<T_PROPOSALS> const* __restrict__ proposals,
+                                                  T_ROIS* __restrict__ filtered,
+                                                  const int preNmsTopN,
+                                                  const float nmsThres,
+                                                  const int afterNmsTopN)
+{
+    Bbox<T_PROPOSALS> const* cProposals = proposals + blockIdx.x * propSize;
+
+    Bbox<T_PROPOSALS> t[TSIZE];
+    uint64_t del = 0;
+
+    for (int i = 0; i < TSIZE; i++)
+    {
+        if (i < TSIZE - 1 || i * DIM + threadIdx.x < preNmsTopN)
+        {
+            t[i] = cProposals[i * DIM + threadIdx.x];
+        }
+    }
+
+    __shared__ Bbox<T_PROPOSALS> last;
+    __shared__ bool kept;
+    __shared__ int foundBatch;
+    if (threadIdx.x == 0)
+        foundBatch = 0;
+
+    for (int i = 0; i < TSIZE; i++)
+    {
+        for (int j = 0; j < DIM; j++)
+        {
+            int offset = i * DIM;
+            int index = offset + j;
+            if (index >= preNmsTopN)
+                break;
+
+            __syncthreads();
+
+            if (threadIdx.x == j)
+            {
+                kept = 0 == (del & ((uint64_t) 1 << i));
+                last = t[i];
+
+                if (kept)
+                {
+                    int cnt = blockIdx.x * afterNmsTopN + foundBatch;
+                    filtered[cnt * 4 + 0] = t[i].xmin;
+                    filtered[cnt * 4 + 1] = t[i].ymin;
+                    filtered[cnt * 4 + 2] = t[i].xmax;
+                    filtered[cnt * 4 + 3] = t[i].ymax;
+                    foundBatch++;
+                }
+            }
+
+            __syncthreads();
+
+            if (foundBatch == afterNmsTopN)
+            {
+                return;
+            }
+
+            if (kept)
+            {
+                Bbox<T_PROPOSALS> test = last;
+
+                for (int k = 0; k < TSIZE; k++)
+                {
+                    if (index < k * DIM + threadIdx.x
+                        && IoU<T_PROPOSALS>(test, t[k]) > nmsThres)
+                    {
+                        del |= (uint64_t) 1 << k;
+                    }
+                }
+            }
+        }
+    }
+}
+
+// NMS LAUNCH
+template <typename T_PROPOSALS, DLayout_t L_PROPOSALS, typename T_ROIS>
+pluginStatus_t nmsLaunch(cudaStream_t stream,
+                        const int batch,
+                        const int propSize,
+                        void* proposals,
+                        void* filtered,
+                        const int preNmsTopN,
+                        const float nmsThres,
+                        const int afterNmsTopN)
+{
+    const int blockSize = 1024;
+
+#define P1(tsize) nmsKernel1<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
+#define P2(tsize) nmsKernel2<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
+
+    void (*kernel[64])(int, Bbox<T_PROPOSALS> const*, T_ROIS*, int, float, int) = {
+        P1(1), P1(2), P1(3), P1(4), P1(5), P1(6), P1(7), P1(8), P1(9), P1(10), P1(11), P1(12), P2(13), P2(14), P2(15), P2(16),
+        P2(17), P2(18), P2(19), P2(20), P2(21), P2(22), P2(23), P2(24), P2(25), P2(26), P2(27), P2(28), P2(29), P2(30), P2(31), P2(32),
+        P2(33), P2(34), P2(35), P2(36), P2(37), P2(38), P2(39), P2(40), P2(41), P2(42), P2(43), P2(44), P2(45), P2(46), P2(47), P2(48),
+        P2(49), P2(50), P2(51), P2(52), P2(53), P2(54), P2(55), P2(56), P2(57), P2(58), P2(59), P2(60), P2(61), P2(62), P2(63), P2(64)};
+
+    ASSERT_PARAM(preNmsTopN < 64 * blockSize);
+
+    CSC(cudaMemsetAsync(filtered, 0, batch * afterNmsTopN * 4 * sizeof(T_ROIS), stream), STATUS_FAILURE);
+
+    kernel[(preNmsTopN + blockSize - 1) / blockSize - 1]<<<batch, blockSize, 0, stream>>>(propSize,
+                                                                                          (Bbox<T_PROPOSALS>*) proposals,
+                                                                                          (T_ROIS*) filtered,
+                                                                                          preNmsTopN,
+                                                                                          nmsThres,
+                                                                                          afterNmsTopN);
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+
+    return STATUS_SUCCESS;
+}
+
+// SET OFFSET
+// Works for up to 2Gi elements (cub's limitation)!
+__global__ void setOffset(int stride, int size, int* output)
+{
+    // One block, because batch size shouldn't be too large.
+    for (int i = threadIdx.x; i < size; i += blockDim.x)
+    {
+        output[i] = i * stride;
+    }
+}
+
+// NMS GPU
+template <typename T_SCORES, typename T_ROIS>
+pluginStatus_t nmsGpu(cudaStream_t stream,
+                     const int N,
+                     const int R,
+                     const int preNmsTop,
+                     const int nmsMaxOut,
+                     const float iouThreshold,
+                     //const float       minBoxSize,
+                     //const float *     imInfo,
+                     void* fgScores,
+                     const void* proposals,
+                     void* workspace,
+                     void* rois)
+{
+    int8_t* vworkspace = alignPtr((int8_t*) workspace, ALIGNMENT);
+
+    DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
+
+    pluginStatus_t error;
+
+    DEBUG_PRINTF("&&&& [NMS] DISCARD\n");
+    DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
+
+    // Generate offsets
+    int* offsets = (int*) vworkspace;
+    setOffset<<<1, 1024, 0, stream>>>(R, N + 1, offsets);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+
+    vworkspace = vworkspace + N + 1;
+    vworkspace = alignPtr(vworkspace, ALIGNMENT);
+
+    // Sort (batched)
+    std::size_t tempStorageBytes = 0;
+
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+        NULL, tempStorageBytes,
+        (T_SCORES*) fgScores, (T_SCORES*) fgScores,
+        (Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposals,
+        N * R, N,
+        offsets, offsets + 1, 0, 8 * sizeof(T_SCORES), stream);
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+
+    T_SCORES* scoresOut = (T_SCORES*) vworkspace;
+    vworkspace = (int8_t*) (scoresOut + N * R);
+    vworkspace = alignPtr(vworkspace, ALIGNMENT);
+    Bbox<T_ROIS>* proposalsOut = (Bbox<T_ROIS>*) vworkspace;
+    vworkspace = (int8_t*) (proposalsOut + N * R);
+    vworkspace = alignPtr(vworkspace, ALIGNMENT);
+
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+        vworkspace, tempStorageBytes,
+        (T_SCORES*) fgScores, (T_SCORES*) scoresOut,
+        (Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposalsOut,
+        N * R, N,
+        offsets, offsets + 1,
+        0, 8 * sizeof(T_SCORES), stream);
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+
+    DEBUG_PRINTF("&&&& [NMS] POST CUB\n");
+    DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposalsOut, N * R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(scoresOut, N * R * sizeof(float)));
+
+    error = nmsLaunch<T_ROIS, NC4HW, T_ROIS>(stream,
+                                             N,
+                                             R,
+                                             proposalsOut,
+                                             rois,
+                                             preNmsTop,
+                                             iouThreshold,
+                                             nmsMaxOut);
+
+    DEBUG_PRINTF("&&&& [NMS] POST LAUNCH\n");
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(rois, N * nmsMaxOut * 4 * sizeof(float)));
+
+    if (error != STATUS_SUCCESS)
+    {
+        return error;
+    }
+
+    return STATUS_SUCCESS;
+}
+
+// NMS LAUNCH CONFIG
+typedef pluginStatus_t (*nmsFun)(cudaStream_t,
+                                const int,   // N
+                                const int,   // R
+                                const int,   // preNmsTop
+                                const int,   // nmsMaxOut
+                                const float, // iouThreshold
+                                //const float,       // minBoxSize
+                                //const float *,     // imInfo
+                                void*,       // fgScores
+                                const void*, // proposals,
+                                void*,       // workspace,
+                                void*);      // rois
+
+struct nmsLaunchConfig
+{
+    DataType t_fgScores;
+    DLayout_t l_fgScores;
+    DataType t_proposals;
+    DLayout_t l_proposals;
+    DataType t_rois;
+    nmsFun function;
+
+    nmsLaunchConfig(DataType t_fgScores,
+                    DLayout_t l_fgScores,
+                    DataType t_proposals,
+                    DLayout_t l_proposals,
+                    DataType t_rois,
+                    nmsFun function)
+        : t_fgScores(t_fgScores)
+        , l_fgScores(l_fgScores)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_rois(t_rois)
+        , function(function)
+    {
+    }
+
+    nmsLaunchConfig(DataType t_fgScores,
+                    DLayout_t l_fgScores,
+                    DataType t_proposals,
+                    DLayout_t l_proposals,
+                    DataType t_rois)
+        : t_fgScores(t_fgScores)
+        , l_fgScores(l_fgScores)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_rois(t_rois)
+    {
+    }
+
+    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);
+    }
+};
+
+#define FLOAT32 nvinfer1::DataType::kFLOAT
+static std::array<nmsLaunchConfig, 1> nmsLCOptions = {
+    nmsLaunchConfig(FLOAT32, NCHW, FLOAT32, NC4HW, FLOAT32, nmsGpu<float, float>)};
+
+// NMS
+pluginStatus_t nms(cudaStream_t stream,
+                  const int N,
+                  const int R,
+                  const int preNmsTop,
+                  const int nmsMaxOut,
+                  const float iouThreshold,
+                  const DataType t_fgScores,
+                  const DLayout_t l_fgScores,
+                  void* fgScores,
+                  const DataType t_proposals,
+                  const DLayout_t l_proposals,
+                  const void* proposals,
+                  void* workspace,
+                  const DataType t_rois,
+                  void* rois)
+{
+
+    nmsLaunchConfig lc(t_fgScores, l_fgScores, t_proposals, l_proposals, t_rois);
+    for (unsigned i = 0; i < nmsLCOptions.size(); i++)
+    {
+        if (nmsLCOptions[i] == lc)
+        {
+            DEBUG_PRINTF("NMS KERNEL %d\n", i);
+            return nmsLCOptions[i].function(stream,
+                                        N, R,
+                                        preNmsTop,
+                                        nmsMaxOut,
+                                        iouThreshold,
+                                        fgScores,
+                                        proposals,
+                                        workspace,
+                                        rois);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/normalizeLayer.cu

@@ -0,0 +1,294 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/bboxUtils.h"
+#include "common/kernels/kernel.h"
+#include "common/cublasWrapper.h"
+
+using namespace nvinfer1::pluginInternal;
+
+namespace nvinfer1
+{
+namespace plugin
+{
+#define CUBLAS_CHECK(condition)                                                                                        \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        cublasStatus_t status = condition;                                                                             \
+        if (status != CUBLAS_STATUS_SUCCESS)                                                                           \
+        {                                                                                                              \
+            printf("%s %d CUBLAS FAIL %s\n", __FILE__, __LINE__, cublasGetErrorString(status));                        \
+        }                                                                                                              \
+    } while (0)
+
+size_t normalizePluginWorkspaceSize(bool acrossSpatial, int C, int H, int W)
+{
+    if (acrossSpatial)
+        return sizeof(float) * C * H * W;
+    else
+        return (size_t) 0;
+}
+
+template <unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta)
+    __global__ void normalizeNotAcrossSpatialKernel(
+        const bool channelShared,
+        const int N,
+        const int C,
+        const int H,
+        const int W,
+        const float eps,
+        const float* scale,
+        float* inputData,
+        float* outputData)
+{
+    const int dim = C * H * W;
+    const int spatialDim = H * W;
+    const int tile = 32;
+    const int numTile = (spatialDim + tile - 1) / tile;
+    for (int n = blockIdx.x; n < N * numTile; n += gridDim.x)
+    {
+        float* input = inputData + (n / numTile) * dim;
+        float* output = outputData + (n / numTile) * dim;
+        __shared__ float sum[tile];
+        float localsum = 0.0F;
+        for (int i = threadIdx.x; i < tile; i += nthds_per_cta)
+        {
+            sum[i] = 0.0F;
+        }
+        __syncthreads();
+        for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+        {
+            int row = i / tile;
+            int col = (n % numTile) * tile + i % tile;
+            float data = 0.0F;
+            if (col < spatialDim)
+                data = input[row * spatialDim + col];
+            localsum += data * data;
+        }
+        atomicAdd(&sum[threadIdx.x & 31], localsum);
+        __syncthreads();
+        for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+        {
+            int row = i / tile;
+            int col = (n % numTile) * tile + i % tile;
+            if (col < spatialDim)
+            {
+                int offset = row * spatialDim + col;
+                output[offset] = input[offset] / sqrt(sum[threadIdx.x & 31] + eps);
+            }
+        }
+        if (channelShared)
+        {
+            for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+            {
+                int row = i / tile;
+                int col = (n % numTile) * tile + i % tile;
+                if (col < spatialDim)
+                    output[row * spatialDim + col] *= scale[0];
+            }
+        }
+        else
+        {
+            for (int i = threadIdx.x; i < C * tile; i += nthds_per_cta)
+            {
+                int row = i / tile;
+                int col = (n % numTile) * tile + i % tile;
+                if (col < spatialDim)
+                    output[row * spatialDim + col] *= scale[row];
+            }
+        }
+    }
+}
+
+pluginStatus_t normalizeNotAcrossSpatialGpu(
+    cudaStream_t stream,
+    const bool channelShared,
+    const int N,
+    const int C,
+    const int H,
+    const int W,
+    const float eps,
+    const void* scale,
+    const void* inputData,
+    void* outputData)
+{
+    const int BS = 128;
+    const int GS = 256;
+    // assumes warp size == 32
+    PLUGIN_ASSERT(BS % 32 == 0);
+    normalizeNotAcrossSpatialKernel<BS><<<GS, BS, 0, stream>>>(
+        channelShared, N, C, H, W, eps, (const float*) scale, (float*) inputData, (float*) outputData);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+__global__ void squareKernel(
+    const int n,
+    const float* x,
+    float* y)
+{
+    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
+         i < n; i += gridDim.x * blockDim.x)
+    {
+        y[i] = x[i] * x[i];
+    }
+}
+
+__global__ void scalChannelKernel(
+    const int n,
+    const int spatialDim,
+    const float* inputData,
+    const float* scale,
+    float* outputData)
+{
+    for (int i = blockIdx.x * blockDim.x + threadIdx.x;
+         i < n; i += gridDim.x * blockDim.x)
+    {
+        // scale factors are indepedent across different channels
+        // scale[i / spatialDim]: find the right scale factor for specific channels
+        outputData[i] = inputData[i] / scale[i / spatialDim];
+    }
+}
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t normalizeInference(
+    cudaStream_t stream,
+    cublasHandle_t handle,
+    const bool acrossSpatial,
+    const bool channelShared,
+    const int N,
+    const int C,
+    const int H,
+    const int W,
+    const float eps,
+    const void* scale,
+    const void* inputData,
+    void* outputData,
+    void* workspace)
+{
+    CublasWrapper& mCublasWrapper = getCublasWrapper();
+    const int dim = C * H * W;
+    // Normalization is conducted for each sample from the batch indepdently
+    if (acrossSpatial)
+    {
+        float* input = (float*) const_cast<void*>(inputData);
+        float* output = (float*) outputData;
+        float* buffer = (float*) workspace;
+        for (int n = 0; n < N; ++n)
+        {
+            // Take the square of each element in the input
+            squareKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, input, buffer);
+            float normsqr = 0.0F;
+            // Sum up all the squared elements
+            CUBLAS_CHECK(mCublasWrapper.cublasSasum(handle, dim, buffer, 1, &normsqr));
+            // Make a copy of the input to the output
+            CUBLAS_CHECK(mCublasWrapper.cublasScopy(handle, dim, input, 1, output, 1));
+            // Calculate the inverse of the square root of the sum
+            // Use eps to prevent being divided by zero
+            normsqr = 1 / sqrt(normsqr + eps);
+            // Scale all the outputs by normsqr
+            CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, &normsqr, output, 1));
+            // If channel shared is true, scale all the outputs
+            if (channelShared)
+            {
+                CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, (float*) scale, output, 1));
+            }
+            // Use different scale factors for different channels
+            else
+            {
+                // scale the output according to channels
+                scalChannelKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, H * W, output, (float*) scale, output);
+            }
+            // Move cursors
+            input += dim;
+            output += dim;
+        }
+        return STATUS_SUCCESS;
+    }
+    // Normalization ignoring the batch
+    else
+    {
+        return normalizeNotAcrossSpatialGpu(stream, channelShared, N, C, H, W, eps, scale, inputData, outputData);
+    }
+}
+} // namespace plugin
+} // namespace nvinfer1
+
+pluginStatus_t normalizeInference(
+    cudaStream_t stream,
+    cublasHandle_t handle,
+    const bool acrossSpatial,
+    const bool channelShared,
+    const int N,
+    const int C,
+    const int H,
+    const int W,
+    const float eps,
+    const void* scale,
+    const void* inputData,
+    void* outputData,
+    void* workspace)
+{
+    const int dim = C * H * W;
+    // Normalization is conducted for each sample from the batch indepdently
+    if (acrossSpatial)
+    {
+        float* input = (float*) const_cast<void*>(inputData);
+        float* output = (float*) outputData;
+        float* buffer = (float*) workspace;
+        CublasWrapper& mCublasWrapper = getCublasWrapper();
+        for (int n = 0; n < N; ++n)
+        {
+            // Take the square of each element in the input
+            squareKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, input, buffer);
+            float normsqr = 0.0F;
+            // Sum up all the squared elements
+            CUBLAS_CHECK(mCublasWrapper.cublasSasum(handle, dim, buffer, 1, &normsqr));
+            // Make a copy of the input to the output
+            CUBLAS_CHECK(mCublasWrapper.cublasScopy(handle, dim, input, 1, output, 1));
+            // Calculate the inverse of the square root of the sum
+            // Use eps to prevent being divided by zero
+            normsqr = 1 / sqrt(normsqr + eps);
+            // Scale all the outputs by normsqr
+            CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, &normsqr, output, 1));
+            // If channel shared is true, scale all the outputs
+            if (channelShared)
+            {
+                CUBLAS_CHECK(mCublasWrapper.cublasSscal(handle, dim, (float*) scale, output, 1));
+            }
+            // Use different scale factors for different channels
+            else
+            {
+                // scale the output according to channels
+                scalChannelKernel<<<(dim + 511) / 512, 512, 0, stream>>>(dim, H * W, output, (float*) scale, output);
+            }
+            // Move cursors
+            input += dim;
+            output += dim;
+        }
+        return STATUS_SUCCESS;
+    }
+    // Normalization ignoring the batch
+    else
+    {
+        return normalizeNotAcrossSpatialGpu(stream, channelShared, N, C, H, W, eps, scale, inputData, outputData);
+    }
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/permuteData.cu

@@ -0,0 +1,114 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+#include <array>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+
+template <typename Dtype, unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta) __global__ void permuteData_kernel(const int nthreads, const int num_classes,
+    const int num_data, const int num_dim, bool confSigmoid, const Dtype* data, Dtype* new_data)
+{
+    // data format: [batch_size, num_data, num_classes, num_dim]
+    for (int index = blockIdx.x * nthds_per_cta + threadIdx.x; index < nthreads; index += nthds_per_cta * gridDim.x)
+    {
+        const int i = index % num_dim;
+        const int c = (index / num_dim) % num_classes;
+        const int d = (index / num_dim / num_classes) % num_data;
+        const int n = index / num_dim / num_classes / num_data;
+        const int new_index = ((n * num_classes + c) * num_data + d) * num_dim + i;
+        float result = data[index];
+        if (confSigmoid)
+            result = exp(result) / (1 + exp(result));
+
+        new_data[new_index] = result;
+    }
+    // new data format: [batch_size, num_classes, num_data, num_dim]
+}
+
+template <typename Dtype>
+pluginStatus_t permuteData_gpu(
+    cudaStream_t stream,
+    const int nthreads,
+    const int num_classes,
+    const int num_data,
+    const int num_dim,
+    bool confSigmoid,
+    const void* data,
+    void* new_data)
+{
+    const int BS = 512;
+    const int GS = (nthreads + BS - 1) / BS;
+    permuteData_kernel<Dtype, BS><<<GS, BS, 0, stream>>>(nthreads, num_classes, num_data, num_dim, confSigmoid,
+                                                         (const Dtype*) data, (Dtype*) new_data);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// permuteData LAUNCH CONFIG
+typedef pluginStatus_t (*pdFunc)(cudaStream_t, const int, const int, const int, const int, bool, const void*, void*);
+
+struct pdLaunchConfig
+{
+    DataType t_data;
+    pdFunc function;
+
+    pdLaunchConfig(DataType t_data)
+        : t_data(t_data)
+        , function(nullptr)
+    {
+    }
+    pdLaunchConfig(DataType t_data, pdFunc function)
+        : t_data(t_data)
+        , function(function)
+    {
+    }
+    bool operator==(pdLaunchConfig const& other) const
+    {
+        return t_data == other.t_data;
+    }
+};
+
+static std::array<pdLaunchConfig, 2> pdLCOptions = {
+    pdLaunchConfig(DataType::kFLOAT, permuteData_gpu<float>), pdLaunchConfig(DataType::kHALF, permuteData_gpu<__half>)};
+
+pluginStatus_t permuteData(cudaStream_t stream, const int nthreads, const int num_classes, const int num_data,
+    const int num_dim, const DataType DT_DATA, bool confSigmoid, const void* data, void* new_data)
+{
+    pdLaunchConfig lc = pdLaunchConfig(DT_DATA);
+    for (unsigned i = 0; i < pdLCOptions.size(); ++i)
+    {
+        if (lc == pdLCOptions[i])
+        {
+            DEBUG_PRINTF("permuteData kernel %d\n", i);
+            return pdLCOptions[i].function(stream,
+                                         nthreads,
+                                         num_classes,
+                                         num_data,
+                                         num_dim,
+                                         confSigmoid,
+                                         data,
+                                         new_data);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/pillarScatterKernels.cu

@@ -0,0 +1,114 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <iostream>
+#include <cuda_runtime_api.h>
+#include <cuda_fp16.h>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+const int PILLARS_PER_BLOCK = 64;
+const int PILLAR_FEATURE_SIZE = 64;
+
+template <typename Element>
+__global__ void scatterBEV_kernel(const Element *pillar_features_data,
+          const unsigned int *coords_data, const unsigned int *params_data,
+          unsigned int featureX, unsigned int featureY,
+          Element *spatial_feature_data)
+{
+    int pillar_idx = blockIdx.x * PILLARS_PER_BLOCK + threadIdx.x;
+    int valid_pillars_inBlock = PILLARS_PER_BLOCK;
+    const int num_pillars = params_data[0];
+    int valid_blocks = (num_pillars+PILLARS_PER_BLOCK-1)/PILLARS_PER_BLOCK;
+    if(blockIdx.x >= valid_blocks) return;
+    if(blockIdx.x == (valid_blocks-1)) {
+      valid_pillars_inBlock = num_pillars % PILLARS_PER_BLOCK;
+    }
+    valid_pillars_inBlock = (valid_pillars_inBlock==0) ? PILLARS_PER_BLOCK : valid_pillars_inBlock;
+    __shared__ Element pillarSM[PILLARS_PER_BLOCK][PILLAR_FEATURE_SIZE]; //pillar*64
+    for (int i = 0; i < valid_pillars_inBlock; i++)
+    {
+      pillarSM[i][threadIdx.x] = pillar_features_data[ (blockIdx.x * PILLARS_PER_BLOCK +i)*PILLAR_FEATURE_SIZE + threadIdx.x];
+    }
+    __syncthreads();
+    if(pillar_idx >= num_pillars) return;
+    int4 coord = ((const int4 *)coords_data)[pillar_idx];
+    int x = coord.w;
+    int y = coord.z;
+    for (int i = 0; i < PILLAR_FEATURE_SIZE; i++)
+    {
+      spatial_feature_data[i*featureY*featureX + y*featureX + x] = pillarSM[threadIdx.x][i];
+    }
+}
+
+template <typename Element>
+int pillarScatterKernelLaunch(
+  int batch_size,
+  int max_pillar_num,
+  int num_features,
+  const Element *pillar_features_data,
+  const unsigned int *coords_data,
+  const unsigned int *params_data,
+  unsigned int featureX, unsigned int featureY,
+  Element *spatial_feature_data,
+  cudaStream_t stream)
+{
+    dim3 blocks( (featureX*featureY+PILLARS_PER_BLOCK-1)/PILLARS_PER_BLOCK);
+    dim3 threads(PILLARS_PER_BLOCK);
+    for (int b = 0; b < batch_size; b++) {
+      scatterBEV_kernel<Element><<<blocks, threads, 0, stream>>>
+        (pillar_features_data + b*max_pillar_num*num_features,
+         coords_data + b*max_pillar_num*4,
+         params_data + b,
+         featureX,
+         featureY,
+         spatial_feature_data + b*num_features*featureX*featureY
+        );
+      auto err = cudaGetLastError();
+      if (cudaSuccess != err) {
+          fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
+          return -1;
+      }
+    }
+    return 0;
+}
+
+template int pillarScatterKernelLaunch<half>(
+  int batch_size,
+  int max_pillar_num,
+  int num_features,
+  const half *pillar_features_data,
+  const unsigned int *coords_data,
+  const unsigned int *params_data,
+  unsigned int featureX, unsigned int featureY,
+  half *spatial_feature_data,
+  cudaStream_t stream);
+
+template int pillarScatterKernelLaunch<float>(
+  int batch_size,
+  int max_pillar_num,
+  int num_features,
+  const float *pillar_features_data,
+  const unsigned int *coords_data,
+  const unsigned int *params_data,
+  unsigned int featureX, unsigned int featureY,
+  float *spatial_feature_data,
+  cudaStream_t stream);
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/priorBoxLayer.cu

@@ -0,0 +1,202 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "NvInferPluginUtils.h"
+#include "common/kernels/kernel.h"
+#include "reducedMathPlugin.h"
+#include <iostream>
+
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+using nvinfer1::plugin::ReducedDivisor;
+
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__ void priorBoxKernel(PriorBoxParameters param, const int H, const int W,
+    const int numPriors, const int numAspectRatios, const float* minSize, const float* maxSize,
+    const float* aspectRatios, float* outputData)
+{
+    // output dims: (H, W, param.numMinSize, (1+haveMaxSize+numAR-1), 4)
+    const int dim = H * W * numPriors;
+    const bool haveMaxSize = param.numMaxSize > 0;
+    const int dimAR = (haveMaxSize ? 1 : 0) + numAspectRatios;
+    for (int i = blockIdx.x * nthdsPerCTA + threadIdx.x;
+         i < dim; i += gridDim.x * nthdsPerCTA)
+    {
+        const int w = (i / numPriors) % W;
+        const int h = (i / numPriors) / W;
+        // Usually param.offset == 0.5
+        // Calucate the center of prior box at the input image scale
+        const float centerX = (w + param.offset) * param.stepW;
+        const float centerY = (h + param.offset) * param.stepH;
+        // Minimum size index
+        const int minSizeId = (i / dimAR) % param.numMinSize;
+        // Aspect ratio index
+        const int arId = i % dimAR;
+        // Generate square pior box of aspect ratio of 1.0, edge length of minSize[minSizeId]
+        if (arId == 0)
+        {
+            const float boxW = minSize[minSizeId];
+            const float boxH = boxW;
+            float x, y, z, w;
+            // Calculate [x_topleft, y_topleft, x_bottomright, y_bottomright]
+            // Coordinates were scaled to [0, 1] against the width or height of the original input image
+            x = (centerX - boxW / 2.0f) / param.imgW;
+            y = (centerY - boxH / 2.0f) / param.imgH;
+            z = (centerX + boxW / 2.0f) / param.imgW;
+            w = (centerY + boxH / 2.0f) / param.imgH;
+            // If we decided to clip the prior box make sure all the bounding box are inside the original input image
+            if (param.clip)
+            {
+                x = min(max(x, 0.0f), 1.0f);
+                y = min(max(y, 0.0f), 1.0f);
+                z = min(max(z, 0.0f), 1.0f);
+                w = min(max(w, 0.0f), 1.0f);
+            }
+            // Copy the bounding box coordinates to output
+            outputData[i * 4] = x;
+            outputData[i * 4 + 1] = y;
+            outputData[i * 4 + 2] = z;
+            outputData[i * 4 + 3] = w;
+        }
+        // If have maxSize
+        // Generate square pior box for aspect ratio of 1.0, edge length of sqrt(minSize[minSizeId] * maxSize[minSizeId])
+        // Described in SSD paper page 6
+        else if (haveMaxSize && arId == 1)
+        {
+            const float boxW = sqrt(minSize[minSizeId] * maxSize[minSizeId]);
+            const float boxH = boxW;
+            float x, y, z, w;
+            x = (centerX - boxW / 2.0f) / param.imgW;
+            y = (centerY - boxH / 2.0f) / param.imgH;
+            z = (centerX + boxW / 2.0f) / param.imgW;
+            w = (centerY + boxH / 2.0f) / param.imgH;
+            if (param.clip)
+            {
+                x = min(max(x, 0.0f), 1.0f);
+                y = min(max(y, 0.0f), 1.0f);
+                z = min(max(z, 0.0f), 1.0f);
+                w = min(max(w, 0.0f), 1.0f);
+            }
+            outputData[i * 4] = x;
+            outputData[i * 4 + 1] = y;
+            outputData[i * 4 + 2] = z;
+            outputData[i * 4 + 3] = w;
+        }
+        // Generate other bouding boxes with aspect ratios of not one.
+        else
+        {
+            const int arOffset = haveMaxSize ? arId - 1 : arId; // skip aspectRatios[0] which is 1
+            const float boxW = minSize[minSizeId] * sqrt(aspectRatios[arOffset]);
+            const float boxH = minSize[minSizeId] / sqrt(aspectRatios[arOffset]);
+            float x, y, z, w;
+            x = (centerX - boxW / 2.0f) / param.imgW;
+            y = (centerY - boxH / 2.0f) / param.imgH;
+            z = (centerX + boxW / 2.0f) / param.imgW;
+            w = (centerY + boxH / 2.0f) / param.imgH;
+            if (param.clip)
+            {
+                x = min(max(x, 0.0f), 1.0f);
+                y = min(max(y, 0.0f), 1.0f);
+                z = min(max(z, 0.0f), 1.0f);
+                w = min(max(w, 0.0f), 1.0f);
+            }
+            outputData[i * 4] = x;
+            outputData[i * 4 + 1] = y;
+            outputData[i * 4 + 2] = z;
+            outputData[i * 4 + 3] = w;
+        }
+    }
+    // Simply copy variance to from the parameter to output
+    float* output = outputData + dim * 4;
+    for (int i = blockIdx.x * nthdsPerCTA + threadIdx.x;
+         i < dim; i += gridDim.x * nthdsPerCTA)
+    {
+        float x, y, z, w;
+        x = param.variance[0];
+        y = param.variance[1];
+        z = param.variance[2];
+        w = param.variance[3];
+        output[i * 4] = x;
+        output[i * 4 + 1] = y;
+        output[i * 4 + 2] = z;
+        output[i * 4 + 3] = w;
+    }
+}
+
+pluginStatus_t priorBoxGpu(
+    cudaStream_t stream,
+    const PriorBoxParameters param,
+    const int H,
+    const int W,
+    const int numPriors,
+    const int numAspectRatios,
+    const void* minSize,
+    const void* maxSize,
+    const void* aspectRatios,
+    void* outputData)
+{
+    const int dim = H * W * numPriors;
+    if (dim > 5120)
+    {
+        const int BS = 128;
+        const int GS = (dim + BS - 1) / BS;
+        priorBoxKernel<BS><<<GS, BS, 0, stream>>>(param, H, W, numPriors, numAspectRatios,
+                                                  (const float*) minSize, (const float*) maxSize,
+                                                  (const float*) aspectRatios, (float*) outputData);
+        CSC(cudaGetLastError(), STATUS_FAILURE);
+        return STATUS_SUCCESS;
+    }
+    else
+    {
+        const int BS = 32;
+        const int GS = (dim + BS - 1) / BS;
+        priorBoxKernel<BS><<<GS, BS, 0, stream>>>(param, H, W, numPriors, numAspectRatios,
+                                                  (const float*) minSize, (const float*) maxSize,
+                                                  (const float*) aspectRatios, (float*) outputData);
+        CSC(cudaGetLastError(), STATUS_FAILURE);
+        return STATUS_SUCCESS;
+    }
+}
+
+pluginStatus_t priorBoxInference(cudaStream_t stream, const PriorBoxParameters param, const int H, const int W,
+    const int numPriors, const int numAspectRatios, const void* minSize, const void* maxSize, const void* aspectRatios,
+    void* outputData)
+{
+    PLUGIN_ASSERT(param.numMaxSize >= 0);
+    if (param.numMaxSize)
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios, minSize, maxSize, aspectRatios, outputData);
+    else
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios,
+                           minSize, nullptr, aspectRatios, outputData);
+}
+
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t priorBoxInference(cudaStream_t stream, const PriorBoxParameters param, const int H, const int W,
+    const int numPriors, const int numAspectRatios, const void* minSize, const void* maxSize, const void* aspectRatios,
+    void* outputData)
+{
+    PLUGIN_ASSERT(param.numMaxSize >= 0);
+    if (param.numMaxSize)
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios, minSize, maxSize, aspectRatios, outputData);
+    else
+        return priorBoxGpu(stream, param, H, W, numPriors, numAspectRatios,
+                           minSize, nullptr, aspectRatios, outputData);
+}
+} // namespace nvinfer1
+} // namespace plugin

plugin/common/kernels/proposalKernel.cu

@@ -0,0 +1,693 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "NvInfer.h"
+#include "common/plugin.h"
+#include <algorithm>
+#include <assert.h>
+#include <cub/cub.cuh>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <functional>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <vector>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+#define PLUGIN_CHECK_CUDA(call)                                                                                        \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        cudaError_t status = call;                                                                                     \
+        if (status != cudaSuccess)                                                                                     \
+        {                                                                                                              \
+            return status;                                                                                             \
+        }                                                                                                              \
+    } while (0)
+template <typename TFloat>
+struct Bbox
+{
+    TFloat x1, y1, x2, y2;
+};
+
+typedef nvinfer1::DataType DType_t;
+
+typedef enum
+{
+    NCHW = 0,
+    NC4HW = 1
+} DLayout_t;
+
+typedef pluginStatus_t frcnnStatus_t;
+
+#define DEBUG_RPN_ENABLE 0
+
+#define FRCNN_ASSERT_PARAM(exp)                                                         \
+    do                                                                                  \
+    {                                                                                   \
+        if (!(exp))                                                                     \
+        {                                                                               \
+            DEBUG_FPRINTF(stderr, "Bad param - " #exp ", %s:%d\n", __FILE__, __LINE__); \
+            return STATUS_BAD_PARAM;                                                    \
+        }                                                                               \
+    } while (0)
+
+#define DEBUG_FPRINTF(...)        \
+    do                            \
+    {                             \
+        if (DEBUG_RPN_ENABLE)     \
+        {                         \
+            fprintf(__VA_ARGS__); \
+        }                         \
+    } while (0)
+
+#define CUDA_MEM_ALIGN 256
+
+unsigned int hash(const void* array_, size_t size);
+int8_t* alignPtr(int8_t* ptr, uintptr_t to);
+__global__ void setOffset(int stride, int size, int* output);
+frcnnStatus_t nms(cudaStream_t stream,
+    const int N,
+    const int R,
+    const int preNmsTop,
+    const int nmsMaxOut,
+    const float iouThreshold,
+    const DType_t t_fgScores,
+    const DLayout_t l_fgScores,
+    void* fgScores,
+    const DType_t t_proposals,
+    const DLayout_t l_proposals,
+    const void* proposals,
+    void* workspace,
+    const DType_t t_rois,
+    void* rois);
+int8_t* nextWorkspacePtr(int8_t* ptr, uintptr_t previousWorkspaceSize);
+
+
+template <typename TFloat>
+__device__ __host__ inline float IoU(const Bbox<TFloat>& a, const Bbox<TFloat>& b)
+{
+    TFloat left = max(a.x1, b.x1), right = min(a.x2, b.x2);
+    TFloat top = max(a.y1, b.y1), bottom = min(a.y2, b.y2);
+    TFloat width = max((TFloat)(right - left + (TFloat) 1.0), (TFloat) 0.0);
+    TFloat height = max((TFloat)(bottom - top + (TFloat) 1.0), (TFloat) 0.0);
+    TFloat interS = width * height;
+    TFloat Sa = (a.x2 - a.x1 + (TFloat) 1) * (a.y2 - a.y1 + (TFloat) 1);
+    TFloat Sb = (b.x2 - b.x1 + (TFloat) 1) * (b.y2 - b.y1 + (TFloat) 1);
+    return (float) interS / (float) (Sa + Sb - interS);
+}
+
+
+// NMS KERNEL FOR SMALL BATCH SIZE {{{
+template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
+__global__ __launch_bounds__(DIM) void nmsKernel1(const int propSize,
+        Bbox<T_PROPOSALS> const* __restrict__ preNmsProposals,
+        T_ROIS* __restrict__ afterNmsProposals,
+        const int preNmsTopN,
+        const float nmsThres,
+        const int afterNmsTopN)
+{
+    __shared__ bool kept_boxes[TSIZE * DIM];
+    int kept = 0;
+    int batch_offset = blockIdx.x * propSize;
+    int max_box_idx = batch_offset + preNmsTopN;
+    int batch_offset_out = blockIdx.x * afterNmsTopN;
+    int flag_idx[TSIZE];
+    int boxes_idx[TSIZE];
+    Bbox<T_PROPOSALS> cur_boxes[TSIZE];
+// initialize kept_boxes
+#pragma unroll
+
+    for (int i = 0; i < TSIZE; i++)
+    {
+        boxes_idx[i] = threadIdx.x + batch_offset + DIM * i;
+        flag_idx[i] = threadIdx.x + DIM * i;
+
+        if (boxes_idx[i] < max_box_idx)
+        {
+            cur_boxes[i] = preNmsProposals[boxes_idx[i]];
+            kept_boxes[flag_idx[i]] = true;
+        }
+        else
+        {
+            kept_boxes[flag_idx[i]] = false;
+            boxes_idx[i] = -1.0f;
+            flag_idx[i] = -1.0f;
+        }
+    }
+
+    int ref_box_idx = 0 + batch_offset;
+
+    // remove the overlapped boxes
+    while ((kept < afterNmsTopN) && (ref_box_idx < max_box_idx))
+    {
+        Bbox<T_PROPOSALS> ref_box;
+        ref_box = preNmsProposals[ref_box_idx];
+#pragma unroll
+
+        for (int i = 0; i < TSIZE; i++)
+        {
+            if (boxes_idx[i] > ref_box_idx)
+            {
+                if (IoU(ref_box, cur_boxes[i]) > nmsThres)
+                {
+                    kept_boxes[flag_idx[i]] = false;
+                }
+            }
+            else if (boxes_idx[i] == ref_box_idx)
+            {
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 0] = ref_box.x1;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 1] = ref_box.y1;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 2] = ref_box.x2;
+                afterNmsProposals[(batch_offset_out + kept) * 4 + 3] = ref_box.y2;
+            }
+        }
+
+        __syncthreads();
+
+        do
+        {
+            ref_box_idx++;
+        }
+        while (!kept_boxes[ref_box_idx - batch_offset] && ref_box_idx < max_box_idx);
+
+        kept++;
+    }
+}
+// }}}
+
+// NMS KERNEL FOR LARGE BATCH SIZE {{{
+template <typename T_PROPOSALS, typename T_ROIS, int DIM, int TSIZE>
+__global__ __launch_bounds__(DIM) void nmsKernel2(const int propSize,
+        Bbox<T_PROPOSALS> const* __restrict__ proposals,
+        T_ROIS* __restrict__ filtered,
+        const int preNmsTopN,
+        const float nmsThres,
+        const int afterNmsTopN)
+{
+    Bbox<T_PROPOSALS> const* cProposals = proposals + blockIdx.x * propSize;
+    Bbox<T_PROPOSALS> t[TSIZE];
+    uint64_t del = 0;
+
+    for (int i = 0; i < TSIZE; i++)
+    {
+        if (i < TSIZE - 1 || i * DIM + threadIdx.x < preNmsTopN)
+        {
+            t[i] = cProposals[i * DIM + threadIdx.x];
+        }
+    }
+
+    __shared__ Bbox<T_PROPOSALS> last;
+    __shared__ bool kept;
+    __shared__ int foundBatch;
+
+    if (threadIdx.x == 0)
+    {
+        foundBatch = 0;
+    }
+
+    for (int i = 0; i < TSIZE; i++)
+    {
+        for (int j = 0; j < DIM; j++)
+        {
+            int offset = i * DIM;
+            int index = offset + j;
+
+            if (index >= preNmsTopN)
+            {
+                break;
+            }
+            __syncthreads();
+
+            if (threadIdx.x == j)
+            {
+                kept = 0 == (del & ((uint64_t) 1 << i));
+                last = t[i];
+
+                if (kept)
+                {
+                    int cnt = blockIdx.x * afterNmsTopN + foundBatch;
+                    filtered[cnt * 4 + 0] = t[i].x1;
+                    filtered[cnt * 4 + 1] = t[i].y1;
+                    filtered[cnt * 4 + 2] = t[i].x2;
+                    filtered[cnt * 4 + 3] = t[i].y2;
+                    foundBatch++;
+                }
+            }
+
+            __syncthreads();
+
+            if (foundBatch == afterNmsTopN)
+            {
+                return;
+            }
+
+            if (kept)
+            {
+                Bbox<T_PROPOSALS> test = last;
+
+                for (int k = 0; k < TSIZE; k++)
+                {
+                    if (index < k * DIM + threadIdx.x
+                            && IoU<T_PROPOSALS>(test, t[k]) > nmsThres)
+                    {
+                        del |= (uint64_t) 1 << k;
+                    }
+                }
+            }
+        }
+    }
+}
+// }}}
+
+// NMS LAUNCH {{{
+template <typename T_PROPOSALS, DLayout_t L_PROPOSALS, typename T_ROIS>
+frcnnStatus_t nmsLaunch(cudaStream_t stream,
+                        const int batch,
+                        const int propSize,
+                        void* proposals,
+                        void* filtered,
+                        const int preNmsTopN,
+                        const float nmsThres,
+                        const int afterNmsTopN)
+{
+    const int blockSize = 1024;
+#define P1(tsize) nmsKernel1<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
+#define P2(tsize) nmsKernel2<T_PROPOSALS, T_ROIS, blockSize, (tsize)>
+    void (*kernel[64])(int, Bbox<T_PROPOSALS> const*, T_ROIS*, int, float, int) =
+    {
+        P1(1), P1(2), P1(3), P1(4), P1(5), P1(6), P1(7), P1(8), P1(9), P1(10), P1(11), P1(12), P2(13), P2(14), P2(15), P2(16),
+        P2(17), P2(18), P2(19), P2(20), P2(21), P2(22), P2(23), P2(24), P2(25), P2(26), P2(27), P2(28), P2(29), P2(30), P2(31), P2(32),
+        P2(33), P2(34), P2(35), P2(36), P2(37), P2(38), P2(39), P2(40), P2(41), P2(42), P2(43), P2(44), P2(45), P2(46), P2(47), P2(48),
+        P2(49), P2(50), P2(51), P2(52), P2(53), P2(54), P2(55), P2(56), P2(57), P2(58), P2(59), P2(60), P2(61), P2(62), P2(63), P2(64)
+    };
+    FRCNN_ASSERT_PARAM(preNmsTopN < 64 * blockSize);
+    CSC(cudaMemsetAsync(filtered, 0, batch * afterNmsTopN * 4 * sizeof(T_ROIS), stream),
+        STATUS_FAILURE);
+    kernel[(preNmsTopN + blockSize - 1) / blockSize - 1] <<< batch, blockSize, 0, stream>>>(propSize,
+            (Bbox<T_PROPOSALS>*) proposals,
+            (T_ROIS*) filtered,
+            preNmsTopN,
+            nmsThres,
+            afterNmsTopN);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+// }}}
+
+// NMS GPU {{{
+template <typename T_SCORES, typename T_ROIS>
+frcnnStatus_t nmsGpu(cudaStream_t stream, const int N, const int R, const int preNmsTop, const int nmsMaxOut,
+    const float iouThreshold, void* fgScores, const void* proposals, void* workspace, void* rois)
+{
+    // CUB's bug workaround:
+    // To work properly for large batch size CUB segmented sort needs ridiculous
+    // workspace alignment.
+    constexpr uintptr_t kALIGNMENT = 1 << 20;
+
+    int8_t* vworkspace = alignPtr((int8_t*) workspace, kALIGNMENT);
+    DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
+    frcnnStatus_t error;
+    DEBUG_PRINTF("&&&& [NMS] DISCARD\n");
+    DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposals, N * R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(fgScores, N * R * sizeof(float)));
+    // Generate offsets
+    int* offsets = (int*) vworkspace;
+    setOffset<<<1, 1024, 0, stream>>>(R, N + 1, offsets);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    vworkspace = vworkspace + N + 1;
+    vworkspace = alignPtr(vworkspace, kALIGNMENT);
+    // Sort (batched)
+    std::size_t tempStorageBytes = 0;
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(NULL, tempStorageBytes, (T_SCORES*) fgScores,
+        (T_SCORES*) fgScores, (Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposals, N * R, N, offsets, offsets + 1, 0,
+        8 * sizeof(T_SCORES), stream);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    T_SCORES* scoresOut = (T_SCORES*) vworkspace;
+    vworkspace = (int8_t*) (scoresOut + N * R);
+    vworkspace = alignPtr(vworkspace, kALIGNMENT);
+    Bbox<T_ROIS>* proposalsOut = (Bbox<T_ROIS>*) vworkspace;
+    vworkspace = (int8_t*) (proposalsOut + N * R);
+    vworkspace = alignPtr(vworkspace, kALIGNMENT);
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(vworkspace, tempStorageBytes, (T_SCORES*) fgScores,
+        (T_SCORES*) scoresOut, (Bbox<T_ROIS>*) proposals, (Bbox<T_ROIS>*) proposalsOut, N * R, N, offsets, offsets + 1,
+        0, 8 * sizeof(T_SCORES), stream);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    DEBUG_PRINTF("&&&& [NMS] POST CUB\n");
+    DEBUG_PRINTF("&&&& [NMS] PROPOSALS %u\n", hash(proposalsOut, N * R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(scoresOut, N * R * sizeof(float)));
+    error = nmsLaunch<T_ROIS, NC4HW, T_ROIS>(stream,
+            N,
+            R,
+            proposalsOut,
+            rois,
+            preNmsTop,
+            iouThreshold,
+            nmsMaxOut);
+    DEBUG_PRINTF("&&&& [NMS] POST LAUNCH\n");
+    DEBUG_PRINTF("&&&& [NMS] SCORES %u\n", hash(rois, N * nmsMaxOut * 4 * sizeof(float)));
+
+    if (error != STATUS_SUCCESS)
+    {
+        return error;
+    }
+
+    return STATUS_SUCCESS;
+}
+// }}}
+
+typedef frcnnStatus_t (*nmsFun)(cudaStream_t,
+                                const int,   // N
+                                const int,   // R
+                                const int,   // preNmsTop
+                                const int,   // nmsMaxOut
+                                const float, // iouThreshold
+                                void*,       // fgScores
+                                const void*, // proposals,
+                                void*,       // workspace,
+                                void*);      // rois
+
+struct nmsLaunchConfig
+{
+    DType_t t_fgScores;
+    DLayout_t l_fgScores;
+    DType_t t_proposals;
+    DLayout_t l_proposals;
+    DType_t t_rois;
+    nmsFun function;
+
+    nmsLaunchConfig(DType_t t_fgScores,
+                    DLayout_t l_fgScores,
+                    DType_t t_proposals,
+                    DLayout_t l_proposals,
+                    DType_t t_rois,
+                    nmsFun function)
+        : t_fgScores(t_fgScores)
+        , l_fgScores(l_fgScores)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_rois(t_rois)
+        , function(function)
+    {
+    }
+
+    nmsLaunchConfig(
+        DType_t t_fgScores, DLayout_t l_fgScores, DType_t t_proposals, DLayout_t l_proposals, DType_t t_rois)
+        : t_fgScores(t_fgScores)
+        , l_fgScores(l_fgScores)
+        , t_proposals(t_proposals)
+        , l_proposals(l_proposals)
+        , t_rois(t_rois)
+        , function(nullptr)
+    {
+    }
+
+    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);
+    }
+};
+
+static std::vector<nmsLaunchConfig> nmsLCVec;
+#define FLOAT32 nvinfer1::DataType::kFLOAT
+
+__global__ void _inverse_transform_gpu(const float* RPN_prob, const float* RPN_regr, int N,
+                                       int INPUT_H, int INPUT_W, int RPN_H, int RPN_W, float RPN_STD_SCALING, int RPN_STRIDE,
+                                       float* ANCHOR_SIZES, int anc_size_num, float* ANCHOR_RATIOS, int anc_ratio_num, float bbox_min_size,
+                                       float* fg_scores, float* proposal_out)
+{
+    int nthreads = N * RPN_H * RPN_W * anc_size_num * anc_ratio_num;
+    int num_ancs = anc_size_num * anc_ratio_num;
+
+    for (int out_idx = threadIdx.x + blockDim.x * blockIdx.x; out_idx < nthreads;
+            out_idx += blockDim.x * gridDim.x)
+    {
+        //input RPN_regr: (N, A4, H, W), thread: (N, A, H, W)
+        int idx = out_idx;
+        int w = idx % RPN_W;
+        idx /= RPN_W;
+        int h = idx % RPN_H;
+        idx /= RPN_H;
+        int a = idx % num_ancs;
+        int n = idx / num_ancs;
+        // normalize by RPN_STD_SCALING
+        int ptr_1 = ((((n * num_ancs) + a) * 4) * RPN_H + h) * RPN_W + w;
+        int ptr_2 = ((((n * num_ancs) + a) * 4 + 1) * RPN_H + h) * RPN_W + w;
+        int ptr_3 = ((((n * num_ancs) + a) * 4 + 2) * RPN_H + h) * RPN_W + w;
+        int ptr_4 = ((((n * num_ancs) + a) * 4 + 3) * RPN_H + h) * RPN_W + w;
+        float tx = RPN_regr[ptr_1] / RPN_STD_SCALING;
+        float ty = RPN_regr[ptr_2] / RPN_STD_SCALING;
+        float tw = RPN_regr[ptr_3] / RPN_STD_SCALING;
+        float th = RPN_regr[ptr_4] / RPN_STD_SCALING;
+        // do inverse transform
+        int ar = a % anc_ratio_num;
+        int as = a / anc_ratio_num;
+        float anchor_w = ANCHOR_SIZES[as] * ANCHOR_RATIOS[ar];
+        float anchor_h = ANCHOR_SIZES[as] / ANCHOR_RATIOS[ar];
+        float anchor_cx = (w + 0.5f) * RPN_STRIDE;
+        float anchor_cy = (h + 0.5f) * RPN_STRIDE;
+        float cx1 = anchor_cx + anchor_w * tx;
+        float cy1 = anchor_cy + anchor_h * ty;
+        float w1 = __expf(tw) * anchor_w;
+        float h1 = __expf(th) * anchor_h;
+        tx = cx1 - w1 / 2.0f;
+        ty = cy1 - h1 / 2.0f;
+        tw = w1;
+        th = h1;
+        tw += tx;
+        th += ty;
+        // clip to min
+        tx = (tx >= 0.0f) ? tx : 0.0f;
+        ty = (ty >= 0.0f) ? ty : 0.0f;
+        tw = (tw >= 0.0f) ? tw : 0.0f;
+        th = (th >= 0.0f) ? th : 0.0f;
+        //clip to max
+        tx = (tx <= INPUT_W - 1.0f) ? tx : (INPUT_W - 1.0f);
+        ty = (ty <= INPUT_H - 1.0f) ? ty : (INPUT_H - 1.0f);
+        tw = (tw <= INPUT_W - 1.0f) ? tw : (INPUT_W - 1.0f);
+        th = (th <= INPUT_H - 1.0f) ? th : (INPUT_H - 1.0f);
+        // filter out small boxes by setting the confidence to -inf
+        int ininf = 0xff800000;
+        float ninf = *(float*) &ininf;
+
+        if (tw - tx <= bbox_min_size || th - ty <= bbox_min_size)
+        {
+            fg_scores[out_idx] = ninf;
+        }
+
+        // copy to proposal_out, output shape: (N, A, H, W, 4)
+        proposal_out[out_idx * 4] = tx;
+        proposal_out[out_idx * 4 + 1] = ty;
+        proposal_out[out_idx * 4 + 2] = tw;
+        proposal_out[out_idx * 4 + 3] = th;
+    }
+}
+
+
+
+cudaError_t _inverse_transform_wrapper(const float* RPN_prob, const float* RPN_regr, int N, int INPUT_H,
+                                int INPUT_W, int RPN_H, int RPN_W, float RPN_STD_SCALING, int RPN_STRIDE,  float* ANCHOR_SIZES,
+                                int anc_size_num, float* ANCHOR_RATIOS, int anc_ratio_num, float bbox_min_size, float* fg_scores,
+                                float* proposal_out, cudaStream_t stream)
+{
+    const int block_size = 1024;
+    const int grid_size = (N * anc_size_num * anc_ratio_num * RPN_H * RPN_W + block_size - 1) /
+                          (block_size);
+    _inverse_transform_gpu <<< grid_size, block_size, 0, stream>>> (RPN_prob, RPN_regr, N, INPUT_H,
+            INPUT_W, RPN_H, RPN_W, RPN_STD_SCALING, RPN_STRIDE, ANCHOR_SIZES, anc_size_num, ANCHOR_RATIOS,
+            anc_ratio_num, bbox_min_size, fg_scores, proposal_out);
+
+    return cudaGetLastError();
+}
+
+size_t _proposalsForwardNMSWorkspaceSize(int N,
+                                        int A,
+                                        int H,
+                                        int W,
+                                        int nmsMaxOut)
+{
+    return N * A * H * W * 5 * 5 * sizeof(float) + (1 << 22);
+}
+
+size_t _proposalsForwardBboxWorkspaceSize(int N, int A, int H, int W)
+{
+    return N * A * H * W * 4 * sizeof(float);
+}
+
+
+size_t _proposalForwardFgScoresWorkspaceSize(int N, int A, int H, int W)
+{
+    return N * A * H * W * sizeof(float);
+}
+
+
+size_t anchors_buf_size(int anc_size_num, int anc_ratio_num)
+{
+    return (anc_size_num + anc_ratio_num) * sizeof(float);
+}
+
+size_t calculateTotalWorkspaceSize(size_t* workspaces, int count);
+
+size_t _get_workspace_size(int N,
+                           int anc_size_num,
+                           int anc_ratio_num,
+                           int H,
+                           int W,
+                           int nmsMaxOut)
+{
+    size_t wss[4];
+    int A = anc_size_num * anc_ratio_num;
+    wss[0] = _proposalsForwardNMSWorkspaceSize(N, A, H, W, nmsMaxOut);
+    wss[1] = _proposalsForwardBboxWorkspaceSize(N, A, H, W);
+    wss[2] = _proposalForwardFgScoresWorkspaceSize(N, A, H, W);
+    wss[3] = anchors_buf_size(anc_size_num, anc_ratio_num);
+    return calculateTotalWorkspaceSize(wss, 4);
+}
+
+
+
+template <typename T>
+frcnnStatus_t extractFgScores_gpu(cudaStream_t stream,
+                                  int N,
+                                  int A,
+                                  int H,
+                                  int W,
+                                  const void* scores,
+                                  void* fgScores)
+{
+    //TODO custom kernel for this
+    size_t size = A * H * W * sizeof(T);
+
+    for (int n = 0; n < N; n++)
+    {
+        size_t offset_ld = n * A * H * W;
+        size_t offset_st = n * A * H * W;
+        CSC(cudaMemcpyAsync(((T*) fgScores) + offset_st, ((T*) scores) + offset_ld, size,
+                            cudaMemcpyDeviceToDevice, stream), STATUS_FAILURE);
+    }
+
+    return STATUS_SUCCESS;
+}
+
+cudaError_t _copy_anchors_to_gpu(cudaStream_t stream, float* ANCHOR_SIZES, int anc_size_num, float* ANCHOR_RATIOS,
+    int anc_ratio_num, void* anchor_size_buf)
+{
+    PLUGIN_CHECK_CUDA(cudaMemcpyAsync(anchor_size_buf, static_cast<void*>(ANCHOR_SIZES), sizeof(float) * anc_size_num,
+        cudaMemcpyHostToDevice, stream));
+    PLUGIN_CHECK_CUDA(cudaMemcpyAsync(static_cast<void*>(static_cast<float*>(anchor_size_buf) + anc_size_num),
+        static_cast<void*>(ANCHOR_RATIOS), sizeof(float) * anc_ratio_num, cudaMemcpyHostToDevice, stream));
+
+    return cudaSuccess;
+}
+
+
+__global__ void _normalize_rois_kernel(float* roi_after_nms, int nthreads, int width, int height)
+{
+    for(int i = threadIdx.x + blockDim.x * blockIdx.x; i < nthreads; i += blockDim.x * gridDim.x)
+    {
+        float x1 = roi_after_nms[i * 4];
+        float y1 = roi_after_nms[i * 4 + 1];
+        float x2 = roi_after_nms[i * 4 + 2];
+        float y2 = roi_after_nms[i * 4 + 3];
+        roi_after_nms[i * 4] = y1 / (height - 1.0f);
+        roi_after_nms[i * 4 + 1] = x1 / (width - 1.0f);
+        roi_after_nms[i * 4 + 2] = y2 / (height - 1.0f);
+        roi_after_nms[i * 4 + 3] = x2 / (width - 1.0f);
+    }
+}
+
+
+
+cudaError_t _normalize_rois(float* roi_after_nms, int n, int max_box_num, int input_width,
+                     int input_height, cudaStream_t stream)
+{
+    const int block_size = 1024;
+    const int grid_size = (n * max_box_num + block_size - 1) / block_size;
+    _normalize_rois_kernel <<< grid_size, block_size, 0, stream>>>(roi_after_nms, n * max_box_num,
+            input_width, input_height);
+
+    return cudaGetLastError();
+}
+
+
+int proposalInference_gpu(
+    cudaStream_t stream,
+    const void* rpn_prob,
+    const void* rpn_regr,
+    int batch_size,
+    int input_height,
+    int input_width,
+    int rpn_height,
+    int rpn_width,
+    int MAX_BOX_NUM,
+    int RPN_PRE_NMS_TOP_N,
+    float* ANCHOR_SIZES,
+    int anc_size_num,
+    float* ANCHOR_RATIOS,
+    int anc_ratio_num,
+    float rpn_std_scaling,
+    int rpn_stride,
+    float bbox_min_size,
+    float nms_iou_threshold,
+    void * workspace,
+    void* output)
+{
+    size_t nmsWorkspaceSize = _proposalsForwardNMSWorkspaceSize(batch_size, anc_size_num * anc_ratio_num,
+                              rpn_height, rpn_width, MAX_BOX_NUM);
+    void* nmsWorkspace = workspace;
+    size_t proposalsSize = _proposalsForwardBboxWorkspaceSize(batch_size, anc_size_num * anc_ratio_num,
+                           rpn_height, rpn_width);
+    const DType_t t_proposals = nvinfer1::DataType::kFLOAT;
+    const DLayout_t l_proposals = NC4HW;
+    void* proposals = nextWorkspacePtr((int8_t*) nmsWorkspace, nmsWorkspaceSize);
+    void* fg_scores = nextWorkspacePtr((int8_t*) proposals, proposalsSize);
+    size_t fg_scores_size = _proposalForwardFgScoresWorkspaceSize(batch_size,
+                            anc_size_num * anc_ratio_num, rpn_height, rpn_width);
+    void* anchor_size_buf = nextWorkspacePtr((int8_t*) fg_scores, fg_scores_size);
+    void* anchor_ratio_buf = static_cast<void*>(static_cast<float*>(anchor_size_buf) + anc_size_num);
+    frcnnStatus_t status;
+    PLUGIN_CHECK_CUDA(
+        _copy_anchors_to_gpu(stream, ANCHOR_SIZES, anc_size_num, ANCHOR_RATIOS, anc_ratio_num, anchor_size_buf));
+
+    status = extractFgScores_gpu<float>(
+        stream, batch_size, anc_size_num * anc_ratio_num, rpn_height, rpn_width, rpn_prob, fg_scores);
+    PLUGIN_ASSERT(status == STATUS_SUCCESS);
+    PLUGIN_CHECK_CUDA(
+        _inverse_transform_wrapper(static_cast<const float*>(rpn_prob), static_cast<const float*>(rpn_regr), batch_size,
+            input_height, input_width, rpn_height, rpn_width, rpn_std_scaling, rpn_stride,
+            static_cast<float*>(anchor_size_buf), anc_size_num, static_cast<float*>(anchor_ratio_buf), anc_ratio_num,
+            bbox_min_size, static_cast<float*>(fg_scores), static_cast<float*>(proposals), stream));
+
+    status = nms(stream, batch_size, anc_size_num * anc_ratio_num * rpn_height * rpn_width, RPN_PRE_NMS_TOP_N,
+        MAX_BOX_NUM, nms_iou_threshold, nvinfer1::DataType::kFLOAT, NCHW, fg_scores, t_proposals, l_proposals,
+        proposals, workspace, nvinfer1::DataType::kFLOAT, output);
+
+    PLUGIN_ASSERT(status == STATUS_SUCCESS);
+
+    PLUGIN_CHECK_CUDA(
+        _normalize_rois(static_cast<float*>(output), batch_size, MAX_BOX_NUM, input_width, input_height, stream));
+
+    return STATUS_SUCCESS;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/proposalsForward.cu

@@ -0,0 +1,161 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/bboxUtils.h"
+#include "common/kernels/kernel.h"
+
+using namespace nvinfer1;
+namespace nvinfer1
+{
+namespace plugin
+{
+// PROPOSALS INFERENCE
+pluginStatus_t proposalsInference(cudaStream_t stream, const int N, const int A, const int H, const int W,
+    const int featureStride, const int preNmsTop, const int nmsMaxOut, const float iouThreshold, const float minBoxSize,
+    const float* imInfo, const float* anchors, const DataType t_scores, const DLayout_t l_scores, const void* scores,
+    const DataType t_deltas, const DLayout_t l_deltas, const void* deltas, void* workspace, const DataType t_rois,
+    void* rois)
+{
+    /*
+     * N: batch size
+     * A: number of anchor boxes per grid cell on feature map
+     * H: height of feature map
+     * W: width of feature map
+     */
+    if (imInfo == NULL || anchors == NULL || scores == NULL || deltas == NULL || workspace == NULL || rois == NULL)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    DEBUG_PRINTF("&&&& IM INFO %u\n", hash(imInfo, N * 3 * sizeof(float)));
+    // anchors: anchor boxes
+    /*
+     * The following line of code looks somewhat incorrect because it sounds like we always have 9 fixed anchor boxes.
+     * The "corrected" implementation should be
+     * DEBUG_PRINTF("&&&& ANCHORS %u\n", A * 4 * sizeof(float)));
+     */
+    DEBUG_PRINTF("&&&& ANCHORS %u\n", hash(anchors, 9 * 4 * sizeof(float)));
+    // scores: objectness of each predicted bounding boxes
+    // 2: softmax, instead of sigmoid, was used for binary objectness classifcation in Faster R-CNN
+    DEBUG_PRINTF("&&&& SCORES  %u\n", hash(scores, N * A * 2 * H * W * sizeof(float)));
+    // deltas: predicted bounding box offsets
+    DEBUG_PRINTF("&&&& DELTAS  %u\n", hash(deltas, N * A * 4 * H * W * sizeof(float)));
+
+    size_t nmsWorkspaceSize = proposalsForwardNMSWorkspaceSize(N, A, H, W, nmsMaxOut);
+
+    void* nmsWorkspace = workspace;
+
+    size_t proposalsSize = proposalsForwardBboxWorkspaceSize(N, A, H, W);
+    const DataType t_proposals = nvinfer1::DataType::kFLOAT;
+    const DLayout_t l_proposals = NC4HW;
+    void* proposals = nextWorkspacePtr((int8_t*) nmsWorkspace, nmsWorkspaceSize);
+
+    const DataType t_fgScores = t_scores;
+    const DLayout_t l_fgScores = NCHW;
+    void* fgScores = nextWorkspacePtr((int8_t*) proposals, proposalsSize);
+
+    pluginStatus_t status;
+
+    /*
+     * Only the second probability value of the objectness (probability of being a object) from the scores will be extracted.
+     * Because the first probability (probability of not being a object) value is redundant.
+     */
+    status = extractFgScores(stream,
+                             N, A, H, W,
+                             t_scores, l_scores, scores,
+                             t_fgScores, l_fgScores, fgScores);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    DEBUG_PRINTF("&&&& FG SCORES %u\n", hash((void*) fgScores, N * A * H * W * sizeof(float)));
+    DEBUG_PRINTF("&&&& DELTAS %u\n", hash((void*) proposals, N * A * H * W * 4 * sizeof(float)));
+
+    /*
+     * Decode predicted bounding boxes.
+     * Decoded predicted bounding boxes were at the raw input image scale.
+     */
+    status = bboxDeltas2Proposals(stream,
+                                  N, A, H, W,
+                                  featureStride,
+                                  minBoxSize,
+                                  imInfo,
+                                  anchors,
+                                  t_deltas, l_deltas, deltas,
+                                  t_proposals, l_proposals, proposals,
+                                  t_fgScores, l_fgScores, fgScores);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    DEBUG_PRINTF("&&&& PROPOSALS %u\n", hash((void*) proposals, N * A * H * W * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& FG SCORES %u\n", hash((void*) fgScores, N * A * H * W * sizeof(float)));
+
+    /*
+     * Non maximum suppression using objectness scores to get the most representative bounding boxes (ROIs).
+     * The rois were at the feature map scale.
+     */
+    status = nms(stream,
+                 N,
+                 A * H * W,
+                 preNmsTop,
+                 nmsMaxOut,
+                 iouThreshold,
+                 t_fgScores, l_fgScores, fgScores,
+                 t_proposals, l_proposals, proposals,
+                 nmsWorkspace,
+                 t_rois, rois);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    DEBUG_PRINTF("&&&& ROIS %u\n", hash((void*) rois, N * nmsMaxOut * 4 * sizeof(float)));
+    return STATUS_SUCCESS;
+}
+
+// WORKSPACE SIZES
+size_t proposalsForwardNMSWorkspaceSize(int N,
+                                        int A,
+                                        int H,
+                                        int W,
+                                        int nmsMaxOut)
+{
+    return N * A * H * W * 5 * 5 * sizeof(float) + (1 << 22);
+}
+
+size_t proposalsForwardBboxWorkspaceSize(int N,
+                                         int A,
+                                         int H,
+                                         int W)
+{
+    return N * A * H * W * 4 * sizeof(float);
+}
+size_t proposalForwardFgScoresWorkspaceSize(int N,
+                                            int A,
+                                            int H,
+                                            int W)
+{
+    return N * A * H * W * sizeof(float);
+}
+
+size_t proposalsInferenceWorkspaceSize(int N,
+                                       int A,
+                                       int H,
+                                       int W,
+                                       int nmsMaxOut)
+{
+    size_t wss[3];
+    wss[0] = proposalsForwardNMSWorkspaceSize(N, A, H, W, nmsMaxOut);
+    wss[1] = proposalsForwardBboxWorkspaceSize(N, A, H, W);
+    wss[2] = proposalForwardFgScoresWorkspaceSize(N, A, H, W);
+    return calculateTotalWorkspaceSize(wss, 3);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/reducedMathPlugin.h

@@ -0,0 +1,106 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef _REDUCED_MATH_PLUGIN_H
+#define _REDUCED_MATH_PLUGIN_H
+#include <cstdint>
+// Dynamically strength-reduced div and mod
+//
+// Ideas taken from Sean Baxter's MGPU library.
+// These classes provide for reduced complexity division and modulus
+// on integers, for the case where the same divisor or modulus will
+// be used repeatedly.
+
+namespace nvinfer1
+{
+namespace plugin
+{
+namespace detail
+{
+
+void findDivisor(int32_t denom, uint32_t& mul_coeff, uint32_t& shift_coeff);
+
+__host__ __device__ __forceinline__ uint32_t umulhi(uint32_t x, uint32_t y)
+{
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 100
+    return __umulhi(x, y);
+#else
+    uint64_t z = (uint64_t) x * (uint64_t) y;
+    return (uint32_t) (z >> 32);
+#endif
+}
+
+// This is a weird implementation that returns div_up(0,1)=0 but
+// div_up(0,2)=1 (wrong) -- just do not use it with a=0.
+__host__ __device__ inline int32_t div_up(int32_t a, int32_t b)
+{
+    return (a - 1) / b + 1;
+}
+
+} // end namespace detail
+
+class ReducedDivisor
+{
+public:
+    ReducedDivisor() {}
+    __host__ __forceinline__ ReducedDivisor(int32_t _y)
+        : y(_y)
+    {
+        detail::findDivisor(y, mul_coeff, shift_coeff);
+    }
+    __host__ __device__ __forceinline__ ReducedDivisor(uint32_t _mul_coeff, uint32_t _shift_coeff, int32_t _y)
+        : mul_coeff(_mul_coeff)
+        , shift_coeff(_shift_coeff)
+        , y(_y)
+    {
+    }
+    __host__ __device__ __forceinline__ int32_t div(int32_t x) const
+    {
+        // if dividing by 1, then findDivisor wouldn't have worked because
+        // mul_coeff would have had to be 2^32, which can't be represented,
+        // so we have to special case that one.
+        return (y != 1) ? detail::umulhi((uint32_t) x, mul_coeff) >> shift_coeff : x;
+    }
+    __host__ __device__ __forceinline__ int32_t mod(int32_t x) const
+    {
+        return x - (div(x) * y);
+    }
+    __host__ __device__ __forceinline__ void divmod(int32_t x, int32_t& q, int32_t& mod) const
+    {
+        q = div(x);
+        mod = x - (q * y);
+    }
+    __host__ __device__ __forceinline__ int32_t get() const
+    {
+        return y;
+    }
+    inline __host__ void get_mul_shift(uint32_t& mul, uint32_t& shift)
+    {
+        mul = mul_coeff;
+        shift = shift_coeff;
+    }
+
+protected:
+    uint32_t mul_coeff{};
+    uint32_t shift_coeff{};
+    int32_t y{};
+};
+
+} // namespace plugin
+
+} // namespace nvinfer1
+#endif /*_REDUCED_MATH_PLUGIN_H*/

plugin/common/kernels/regionForward.cu

@@ -0,0 +1,140 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+namespace nvinfer1
+{
+namespace plugin
+{
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA) __global__ void softmaxKernel(const float* input, const int n, const int batch,
+    const int batchOffset, const int groups, const int groupOffset, const int stride, const float temp, float* output)
+{
+    int id = blockIdx.x * nthdsPerCTA + threadIdx.x;
+    if (id < batch * groups)
+    {
+        int b = id / groups;
+        int g = id % groups;
+        float sum = 0.;
+        // Initialze largest to be the smallest float number.
+        float largest = -3.402823466e+38;
+        int offset = b * batchOffset + g * groupOffset;
+        // Find the largest digits before softmax
+        for (int i = 0; i < n; ++i)
+        {
+            float val = input[i * stride + offset];
+            largest = (val > largest) ? val : largest;
+        }
+        // Softmax for a group of candidate classes
+        // Calculate exponentials
+        for (int i = 0; i < n; ++i)
+        {
+            /*
+             * Here we used a trick to prevent numeric overflow
+             * xm = max{x_1, x_2, ..., x_n}
+             * e^{x_1} / (e^{x_1} + e^{x_2} + e^{x_n}) = e^{x_1 - xm} / (e^{x_1 - xm} + e^{x_2 - xm} + e^{x_n - xm})
+             */
+            float e = exp(input[i * stride + offset] / temp - largest / temp);
+            sum += e;
+            output[i * stride + offset] = e;
+        }
+        // Normalize
+        for (int i = 0; i < n; ++i)
+            output[i * stride + offset] /= sum;
+    }
+}
+
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA)
+    __global__ void activateKernel(float* data,
+                                   const int range)
+{
+    int i = blockIdx.x * nthdsPerCTA + threadIdx.x;
+    // Sigmoid function
+    if (i < range)
+        data[i] = 1. / (1. + exp(-data[i]));
+}
+
+pluginStatus_t regionGPU(
+    cudaStream_t stream,
+    const int batch,
+    const int C,
+    const int H,
+    const int W,
+    const int num,
+    const int coords,
+    const int classes,
+    const bool hasSoftmaxTree,
+    const nvinfer1::plugin::softmaxTree* smTree,
+    const float* input,
+    float* output)
+{
+    const int BS = 512;
+    const int GS1 = (2 * H * W + BS - 1) / BS;
+    const int GS2 = (H * W + BS - 1) / BS;
+    // Applying sigmoid activations
+    for (int b = 0; b < batch; ++b)
+    {
+        for (int n = 0; n < num; ++n)
+        {
+            // Apply sigmoid activation for the encoded center coordinates t_x, and t_y
+            int index = b * C * H * W + n * H * W * (coords + classes + 1);
+            activateKernel<BS><<<GS1, BS, 0, stream>>>(output + index, 2 * H * W);
+            /*
+             * Apply sigmoid for the encoded objectness t_o
+             * + 4 * H * W because we want to skip the first four coordinates t_x, t_y, t_w, t_h
+             * Chaning 4 * H * W to + coords * H * W will not make this function more general
+             * Since we already assumed the information layout in the channels of the input tensor
+             */
+            index = b * C * H * W + n * H * W * (coords + classes + 1) + 4 * H * W;
+            activateKernel<BS><<<GS2, BS, 0, stream>>>(output + index, H * W);
+        }
+    }
+    const int GS3 = (batch * num * H * W + BS - 1) / BS;
+    // Applying softmax activations
+    if (hasSoftmaxTree)
+    {
+        // Softmax for hierarchical classification
+        // The first 5 elements are t_x, t_y, t_w, t_h, t_o which we don't need to apply softmax activation
+        int count = 5;
+        // Only groups and groupSize information is useful for this plugin
+        // Applying softmax activation sequentially for each group of candidate classes
+        for (int i = 0; i < smTree->groups; ++i)
+        {
+            int groupSize = smTree->groupSize[i];
+            softmaxKernel<BS><<<GS3, BS, 0, stream>>>(input + count * H * W, groupSize, batch * num, (C * H * W / num), H * W, 1, H * W, 1., output + count * H * W);
+            count += groupSize;
+        }
+    }
+    else
+    {
+        // Softmax for non-hierarchical classificiation
+        softmaxKernel<BS><<<GS3, BS, 0, stream>>>(input + 5 * H * W, classes, batch * num, (C * H * W / num), H * W, 1, H * W, 1., output + 5 * H * W);
+    }
+
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t regionInference(cudaStream_t stream, const int batch, const int C, const int H, const int W,
+    const int num, const int coords, const int classes, const bool hasSoftmaxTree,
+    const nvinfer1::plugin::softmaxTree* smTree, const void* input, void* output)
+{
+    PLUGIN_CHECK(cudaMemcpyAsync(output, input, batch * C * H * W * sizeof(float), cudaMemcpyDeviceToDevice, stream));
+    return regionGPU(
+        stream, batch, C, H, W, num, coords, classes, hasSoftmaxTree, smTree, (const float*) input, (float*) output);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/reorgForward.cu

@@ -0,0 +1,104 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+#include "reducedMathPlugin.h"
+
+using namespace nvinfer1::plugin; // for ReducedDivisor
+namespace nvinfer1
+{
+namespace plugin
+{
+template <unsigned nthdsPerCTA>
+__launch_bounds__(nthdsPerCTA)
+    __global__ void reorgKernel(
+        const float* input, // input tensor of shape (batch, C, H, W)
+        const int volume,   // note that volumes of input and output tensors are the same
+        ReducedDivisor batch,
+        ReducedDivisor C,
+        ReducedDivisor H,
+        ReducedDivisor W,
+        ReducedDivisor C_out,
+        ReducedDivisor stride,
+        float* output) // output tensor of shape (batch, C * stride * stride, H / stride, W / stride)
+{
+    /*
+     * Reference
+     * https://github.com/pjreddie/darknet/blob/f6d861736038da22c9eb0739dca84003c5a5e275/src/blas_kernels.cu#L370
+     * https://github.com/pjreddie/darknet/blob/f6d861736038da22c9eb0739dca84003c5a5e275/src/blas.c#L9
+     */
+
+    // outIndex is row-major position of input coordinates
+    for (int outIndex = blockIdx.x * nthdsPerCTA + threadIdx.x; outIndex < volume; outIndex += nthdsPerCTA)
+    {
+        int i = outIndex;
+
+        // calculate output coordinates from outIndex
+        int outW, outH, outC;
+        W.divmod(i, i, outW);
+        H.divmod(i, i, outH);
+        C.divmod(i, i, outC);
+        int outN = i;
+
+        // calculate input coordinates based on output coordinates
+        // offset is [0, 1, ..., stride * stride - 1] = posH * stride + posW
+        int offset, inC, posH, posW;
+        C_out.divmod(outC, offset, inC);
+        stride.divmod(offset, posH, posW);
+        int inH = outH * stride.get() + posH;
+        int inW = outW * stride.get() + posW;
+        int inN = outN;
+
+        // inIndex is row-major position of input coordinates
+        int inIndex = inW + W.get() * stride.get() * (inH + H.get() * stride.get() * (inC + C_out.get() * inN));
+
+        output[outIndex] = input[inIndex];
+    }
+}
+
+pluginStatus_t reorgGPU(
+    cudaStream_t stream,
+    const int batch,
+    const int C,
+    const int H,
+    const int W,
+    const int stride,
+    const float* input,
+    float* output)
+{
+    const int BS = 512;                    // number of threads in one block
+    const int volume = batch * C * H * W;  // size of input tensor
+    const int GS = (volume + BS - 1) / BS; // number of blocks to launch, calculated so global number of threads is >= volume
+
+    ReducedDivisor C_out(C / (stride * stride));
+    reorgKernel<BS><<<GS, BS, 0, stream>>>(input, volume, ReducedDivisor(batch), ReducedDivisor(C), ReducedDivisor(H), ReducedDivisor(W), C_out, ReducedDivisor(stride), output);
+    return STATUS_SUCCESS;
+}
+
+pluginStatus_t reorgInference(
+    cudaStream_t stream,
+    const int batch,
+    const int C,
+    const int H,
+    const int W,
+    const int stride,
+    const void* input,
+    void* output)
+{
+    return reorgGPU(stream, batch, C, H, W, stride, (const float*) input, (float*) output);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/roiPooling.cu

@@ -0,0 +1,353 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+#include <algorithm>
+#include <array>
+#include <assert.h>
+#include <cfloat>
+#include <cstdio>
+#include <math.h>
+#include <stdio.h>
+namespace nvinfer1
+{
+namespace plugin
+{
+// This macro is to control shared memory usage. If set to 1, kernel loads the whole feature map
+// into shared memory for reuse; If set to 0, kernel loads data from global memory directly.
+// Roi pooling performance is data dependent. You can test which value is better to your data.
+// If all bboxes are very small, 0 is recommended, otherwise, shared memory will load many unused
+// data; If bboxes have many overlaps, 1 is recommended to avoid duplicate loads.
+// 1 requires larger shared memory size. It may fail if it is larger than CUDA allowed per-block
+// shared memory size upper bound. Then you have to use 0.
+#define ROIPOOLING_FEATURE_MAP_USE_SHMEM 1
+
+template <typename T>
+__device__ T getMax();
+
+template <>
+__device__ __forceinline__ int8_t getMax<int8_t>()
+{
+    return INT8_MAX;
+}
+
+template <>
+__device__ __forceinline__ float getMax<float>()
+{
+    return FLT_MAX;
+}
+
+// ROI POOLING FORWARD KERNEL
+template <typename DATA_T, typename ROI_T, bool INFER_ONLY, bool FM_IN_SMEM>
+__global__ void ROIPoolingForwardKernelAligned(int32_t ROICount, const ROI_T* rois,
+    int32_t N, // feature map size
+    int32_t C, // feature map size
+    int32_t H, // feature map size
+    int32_t W, // feature map size
+    const DATA_T* featureMap, const int32_t poolingH, const int32_t poolingW, const float spatialScale, DATA_T* top,
+    int32_t* maxIds, int32_t fmapStep)
+{
+    extern __shared__ float smem[];
+    DATA_T* feature_shr = (DATA_T*) &smem[0];
+    int* rois_shr = nullptr;
+    if (FM_IN_SMEM)
+    {
+        rois_shr = (int*) &feature_shr[H * W];
+    }
+    else
+    {
+        rois_shr = (int*) &feature_shr[0];
+        feature_shr = nullptr;
+    }
+
+    const int batch = blockIdx.x / C;
+    const int channel = blockIdx.x % C;
+
+    // load ROIs to shared memory
+    for (int j = threadIdx.x; j < ROICount; j += blockDim.x)
+    {
+        int offset = j << 2;
+        float4 roi = reinterpret_cast<float4*>(const_cast<float*>(rois))[batch * ROICount + j];
+        // spatialScale = 1.0 / featureStride
+        // Convert the coordinates to feature map scale
+        rois_shr[offset] = round(roi.x * spatialScale);     //roi_start_w
+        rois_shr[offset + 1] = round(roi.y * spatialScale); //roi_start_h
+        rois_shr[offset + 2] = round(roi.z * spatialScale) - round(roi.x * spatialScale); //roi_length_w
+        rois_shr[offset + 3] = round(roi.w * spatialScale) - round(roi.y * spatialScale); // roi_length_h
+    }
+
+    // NC/xHW
+    int fmapOffset = blockIdx.x / fmapStep * H * W * fmapStep + blockIdx.x % fmapStep;
+
+    // Assumes #CTAs is just enough to cover all channels of all blocks
+    const DATA_T* bottom_data_offset = featureMap + fmapOffset;
+    if (FM_IN_SMEM)
+    {
+        // load the current channel to the shared memory
+        for (int j = threadIdx.x; j < H * W; j += blockDim.x)
+        {
+            feature_shr[j] = bottom_data_offset[j * fmapStep];
+        }
+    }
+    __syncthreads();
+
+    for (int j = threadIdx.x; j < ROICount; j += blockDim.x)
+    {
+        const int offset = j << 2;
+        // Force malformed ROIs to be 1x1
+        int roi_start_w = rois_shr[offset];
+        int roi_start_h = rois_shr[offset + 1];
+        int roi_width = max(rois_shr[offset + 2] + 1, 1);
+        int roi_height = max(rois_shr[offset + 3] + 1, 1);
+        float bin_size_h = static_cast<float>(roi_height) / static_cast<float>(poolingH);
+        float bin_size_w = static_cast<float>(roi_width) / static_cast<float>(poolingW);
+
+        for (int ph = 0; ph < poolingH; ++ph)
+        {
+            for (int pw = 0; pw < poolingW; ++pw)
+            {
+                int hstart = static_cast<int>(floor(static_cast<float>(ph) * bin_size_h));
+                int wstart = static_cast<int>(floor(static_cast<float>(pw) * bin_size_w));
+                int hend = static_cast<int>(ceil(static_cast<float>(ph + 1) * bin_size_h));
+                int wend = static_cast<int>(ceil(static_cast<float>(pw + 1) * bin_size_w));
+
+                // Add roi offsets and clip to input boundaries
+                // In fact, clipping should be done in the RPN, but just in case...
+                hstart = min(max(hstart + roi_start_h, 0), H);
+                hend = min(max(hend + roi_start_h, 0), H);
+                wstart = min(max(wstart + roi_start_w, 0), W);
+                wend = min(max(wend + roi_start_w, 0), W);
+                bool is_empty = (hend <= hstart) || (wend <= wstart);
+
+                // Define an empty pooling region to be zero
+                DATA_T maxval = is_empty ? 0 : -getMax<DATA_T>();
+                int maxId = -1;
+                DATA_T data = 0;
+                for (int h = hstart; h < hend; ++h)
+                {
+                    for (int w = wstart; w < wend; ++w)
+                    {
+                        int index = h * W + w;
+                        if (FM_IN_SMEM)
+                        {
+                            data = feature_shr[index];
+                        }
+                        else
+                        {
+                            data = bottom_data_offset[index * fmapStep];
+                        }
+                        if (data > maxval)
+                        {
+                            maxval = data;
+                            maxId = index;
+                        }
+                    }
+                }
+                top[(((batch * ROICount + j) * C + channel) * poolingH + ph) * poolingW + pw] = maxval;
+                if (!INFER_ONLY)
+                {
+                    maxIds[(((batch * ROICount + j) * C + channel) * poolingH + ph) * poolingW + pw] = maxId;
+                }
+            } //for:pw
+        }     //for:ph
+    }         // for:j
+}
+
+template <typename DATA_T, DLayout_t DATA_L, typename ROI_T, bool INFER_ONLY>
+pluginStatus_t ROIPoolingForwardKernelAlignedLauncher(cudaStream_t stream,
+    const int R,        // TOTAL number of rois -> ~nmsMaxOut * N
+    const int N,        // Batch size
+    const int C,        // Channels
+    const int H,        // Input feature map H
+    const int W,        // Input feature map W
+    const int poolingH, // Output feature map H
+    const int poolingW, // Output feature map W
+    const float spatialScale, const void* rois, const void* featureMap, void* top, int* maxIds, size_t deviceSmemSize)
+{
+    size_t roiShmemSize = (R / N) * 4 * sizeof(ROI_T);
+
+#if ROIPOOLING_FEATURE_MAP_USE_SHMEM
+    size_t shmemSize = H * W * sizeof(DATA_T) + roiShmemSize;
+    const bool fmap_in_shmem = true;
+#else
+    size_t shmemSize = roiShmemSize;
+    const bool fmap_in_shmem = false;
+#endif
+
+    if (shmemSize > deviceSmemSize)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    // in the aligned version of ROI Pooling R should always be a multiple of N
+    PLUGIN_ASSERT(R % N == 0);
+
+    // NC/xHW
+    int32_t fmapStep = 1;
+    switch(DATA_L)
+    {
+    case NCHW: fmapStep = 1; break;
+    case NC4HW:
+        fmapStep = 4;
+        PLUGIN_ASSERT((N * C) % 4 == 0);
+        break;
+    case NC32HW:
+        fmapStep = 32;
+        PLUGIN_ASSERT((N * C) % 32 == 0);
+        break;
+    default: PLUGIN_ASSERT(false);
+    }
+
+    if (shmemSize > 48 * 1024)
+    {
+        PLUGIN_CHECK(cudaFuncSetAttribute(&ROIPoolingForwardKernelAligned<DATA_T, ROI_T, INFER_ONLY, true>,
+            cudaFuncAttributeMaxDynamicSharedMemorySize, shmemSize));
+    }
+    ROIPoolingForwardKernelAligned<DATA_T, ROI_T, INFER_ONLY, fmap_in_shmem><<<N * C, 256, shmemSize, stream>>>(R / N,
+        (const ROI_T*) rois,
+        N, // feature map size
+        C, // feature map size
+        H, // feature map size
+        W, // feature map size
+        (const DATA_T*) featureMap, poolingH, poolingW, spatialScale, (DATA_T*) top, maxIds, fmapStep);
+
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// ROI POOLING LAUNCH CONFIG
+
+typedef pluginStatus_t (*roiFwd)(cudaStream_t,
+                                const int,   //R, // TOTAL number of rois -> ~nmsMaxOut * N
+                                const int,   //N, // Batch size
+                                const int,   //C, // Channels
+                                const int,   //H, // Input feature map H
+                                const int,   //W, // Input feature map W
+                                const int,   //poolingH, // Output feature map H
+                                const int,   //poolingW, // Output feature map W
+                                const float, //spatialScale,
+                                const void*, //rois,
+                                const void*, //featureMap,
+                                void*,       //top
+                                int*,        //maxIds
+                                size_t);     //device shmem size
+
+// struct
+struct roiFwdLaunchConfig
+{
+    DataType t_rois;
+    DataType t_featureMap;
+    DLayout_t l_featureMap;
+    DataType t_top;
+    DLayout_t l_top;
+    bool inferOnly;
+    roiFwd function;
+
+    roiFwdLaunchConfig(
+        DataType t_rois, DataType t_featureMap, DLayout_t l_featureMap, DataType t_top, DLayout_t l_top, bool inferOnly)
+        : t_rois(t_rois)
+        , t_featureMap(t_featureMap)
+        , l_featureMap(l_featureMap)
+        , t_top(t_top)
+        , l_top(l_top)
+        , inferOnly(inferOnly)
+        , function(nullptr)
+    {
+    }
+
+    roiFwdLaunchConfig(DataType t_rois,
+                       DataType t_featureMap,
+                       DLayout_t l_featureMap,
+                       DataType t_top,
+                       DLayout_t l_top,
+                       bool inferOnly,
+                       roiFwd function)
+        : t_rois(t_rois)
+        , t_featureMap(t_featureMap)
+        , l_featureMap(l_featureMap)
+        , t_top(t_top)
+        , l_top(l_top)
+        , inferOnly(inferOnly)
+        , function(function)
+    {
+    }
+
+    bool operator==(roiFwdLaunchConfig const& other) const
+    {
+        return (t_rois == other.t_rois)
+            && (t_featureMap == other.t_featureMap)
+            && (l_featureMap == other.l_featureMap)
+            && (t_top == other.t_top)
+            && (l_top == other.l_top)
+            && (inferOnly == other.inferOnly);
+    }
+};
+
+#define FLOAT32 nvinfer1::DataType::kFLOAT
+#define INT8 nvinfer1::DataType::kINT8
+static std::array<roiFwdLaunchConfig, 6> roiFwdLCOptions = {
+    roiFwdLaunchConfig(FLOAT32, FLOAT32, NCHW, FLOAT32, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<float, NCHW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, FLOAT32, NCHW, FLOAT32, NCHW, false, ROIPoolingForwardKernelAlignedLauncher<float, NCHW, float, false>),
+    roiFwdLaunchConfig(FLOAT32, INT8, NCHW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NCHW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, INT8, NC4HW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NC4HW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, INT8, NC32HW, INT8, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<int8_t, NC32HW, float, true>),
+    roiFwdLaunchConfig(FLOAT32, FLOAT32, NC4HW, FLOAT32, NCHW, true, ROIPoolingForwardKernelAlignedLauncher<float, NC4HW, float, true>)};
+
+// ROI INFERENCE
+pluginStatus_t roiInference(cudaStream_t stream,
+                           const int R,        // TOTAL number of rois -> ~nmsMaxOut * N
+                           const int N,        // Batch size
+                           const int C,        // Channels
+                           const int H,        // Input feature map H
+                           const int W,        // Input feature map W
+                           const int poolingH, // Output feature map H
+                           const int poolingW, // Output feature map W
+                           const float spatialScale,
+                           const nvinfer1::DataType t_rois,
+                           const void* rois,
+                           const nvinfer1::DataType t_featureMap,
+                           const DLayout_t l_featureMap,
+                           const void* featureMap,
+                           const nvinfer1::DataType t_top,
+                           const DLayout_t l_top,
+                           void* top,
+                           size_t deviceSmemSize)
+{
+    if (featureMap == NULL || rois == NULL || top == NULL)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    DEBUG_PRINTF("&&&& ROIS %u\n", hash(rois, R * 4 * sizeof(float)));
+    DEBUG_PRINTF("&&&& FMAP %u\n", hash(featureMap, N * C * H * W * sizeof(float)));
+
+    roiFwdLaunchConfig rflc = roiFwdLaunchConfig(t_rois, t_featureMap, l_featureMap, t_top, l_top, true);
+    ASSERT_PARAM(R > 0);
+
+    for (unsigned i = 0; i < roiFwdLCOptions.size(); i++)
+    {
+        if (rflc == roiFwdLCOptions[i])
+        {
+            DEBUG_PRINTF("$$$$ ROI KERNEL %d\n", i);
+            return roiFwdLCOptions[i].function(stream,
+                                                R, N, C, H, W, poolingH, poolingW,
+                                                spatialScale, rois, featureMap, top, NULL, deviceSmemSize);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/rproiInferenceFused.cu

@@ -0,0 +1,97 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/kernels/kernel.h"
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+namespace nvinfer1
+{
+namespace plugin
+{
+pluginStatus_t RPROIInferenceFused(cudaStream_t stream, const int N, const int A, const int C, const int H, const int W,
+    const int poolingH, const int poolingW, const int featureStride, const int preNmsTop, const int nmsMaxOut,
+    const float iouThreshold, const float minBoxSize, const float spatialScale, const float* imInfo,
+    const float* anchors, const DataType t_scores, const DLayout_t l_scores, const void* scores,
+    const DataType t_deltas, const DLayout_t l_deltas, const void* deltas, const DataType t_featureMap,
+    const DLayout_t l_featureMap, const void* featureMap, void* workspaces, const DataType t_rois, void* rois,
+    const DataType t_top, const DLayout_t l_top, void* top, size_t deviceSmemSize)
+{
+    if (imInfo == NULL || anchors == NULL || scores == NULL || deltas == NULL || featureMap == NULL
+        || workspaces == NULL || rois == NULL || top == NULL)
+    {
+        return STATUS_BAD_PARAM;
+    }
+
+    pluginStatus_t status;
+    // Region proposal inference
+    // Getting the region of interests (ROIs) bounding box coordinates from region proposals using non maximum suppression (NMS)
+    status = proposalsInference(stream,
+                                N,
+                                A,
+                                H,
+                                W,
+                                featureStride,
+                                preNmsTop,
+                                nmsMaxOut,
+                                iouThreshold,
+                                minBoxSize,
+                                imInfo,
+                                anchors,
+                                t_scores,
+                                l_scores,
+                                scores,
+                                t_deltas,
+                                l_deltas,
+                                deltas,
+                                workspaces,
+                                t_rois,
+                                rois);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+    // ROI inference
+    // ROI pooling for ROIs
+    status = roiInference(stream,
+                          N * nmsMaxOut, // TOTAL number of rois -> ~nmsMaxOut * N
+                          N,             // Batch size
+                          C,             // Channels
+                          H,             // Input feature map H
+                          W,             // Input feature map W
+                          poolingH,      // Output feature map H
+                          poolingW,      // Output feature map W
+                          spatialScale,
+                          t_rois,
+                          rois,
+                          t_featureMap,
+                          l_featureMap,
+                          featureMap,
+                          t_top,
+                          l_top,
+                          top,
+                          deviceSmemSize);
+    ASSERT_FAILURE(status == STATUS_SUCCESS);
+
+    return STATUS_SUCCESS;
+}
+
+size_t RPROIInferenceFusedWorkspaceSize(int N,
+                                        int A,
+                                        int H,
+                                        int W,
+                                        int nmsMaxOut)
+{
+    return proposalsInferenceWorkspaceSize(N, A, H, W, nmsMaxOut);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/saturate.h

@@ -0,0 +1,41 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef TRT_SATURATE_H
+#define TRT_SATURATE_H
+
+#include <array>
+#include <cuda_fp16.h>
+
+template <typename T_BBOX>
+__device__ T_BBOX saturate(T_BBOX v)
+{
+    return max(min(v, T_BBOX(1)), T_BBOX(0));
+}
+
+template <>
+inline __device__ __half saturate(__half v)
+{
+#if __CUDA_ARCH__ >= 800
+    return __hmax(__hmin(v, __half(1)), __half(0));
+#elif __CUDA_ARCH__ >= 530
+    return __hge(v, __half(1)) ? __half(1) : (__hle(v, __half(0)) ? __half(0) : v);
+#else
+    return max(min(v, float(1)), float(0));
+#endif
+}
+
+#endif // TRT_SATURATE_H

plugin/common/kernels/sortScoresPerClass.cu

@@ -0,0 +1,297 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/bboxUtils.h"
+#include "common/cub_helper.h"
+#include "common/kernels/kernel.h"
+#include "cub/cub.cuh"
+#include <array>
+namespace nvinfer1
+{
+namespace plugin
+{
+
+inline __device__ __half add_fb(const __half& a, const __half& b)
+{
+#if __CUDA_ARCH__ >= 530
+    return a + b;
+#else
+    return __float2half(__half2float(a) + __half2float(b));
+#endif
+}
+
+// overload for float
+inline __device__ float add_fb(const float & a, const float & b) {
+    return a + b;
+}
+
+inline __device__ bool ge_fb(const __half & a, const __half & b) {
+#if __CUDA_ARCH__ >= 530
+    return a >= b;
+#else
+    return __half2float(a) >= __half2float(b);
+#endif
+}
+
+template <typename T_SCORE, unsigned nthds_per_cta>
+__launch_bounds__(nthds_per_cta)
+    __global__ void prepareSortData(
+        const int num,
+        const int num_classes,
+        const int num_preds_per_class,
+        const int background_label_id,
+        const float confidence_threshold,
+        T_SCORE* conf_scores_gpu,
+        T_SCORE* temp_scores,
+        T_SCORE score_shift,
+        int* temp_idx,
+        int* d_offsets)
+{
+    // Prepare scores data for sort
+    const int cur_idx = blockIdx.x * nthds_per_cta + threadIdx.x;
+    const int numPredsPerBatch = num_classes * num_preds_per_class;
+    T_SCORE clip_val = T_SCORE(float(score_shift) + 1.f - 1.f / 1024.f);
+    if (cur_idx < numPredsPerBatch)
+    {
+        const int class_idx = cur_idx / num_preds_per_class;
+        for (int i = 0; i < num; i++)
+        {
+            const int targetIdx = i * numPredsPerBatch + cur_idx;
+            const T_SCORE score = conf_scores_gpu[targetIdx];
+
+            // "Clear" background labeled score and index
+            // Because we do not care about background
+            if (class_idx == background_label_id)
+            {
+                // Set scores to 0
+                // Set label = -1
+                // add shift of 1.0 to normalize the score values
+                // to the range [1, 2). 
+                // add a constant shift to scores will not change the sort
+                // result, but will help reduce the computation because
+                // we only need to sort the mantissa part of the floating-point
+                // numbers
+                temp_scores[targetIdx] = score_shift;
+                temp_idx[targetIdx] = -1;
+                conf_scores_gpu[targetIdx] = score_shift;
+            }
+            // "Clear" scores lower than threshold
+            else
+            {
+                if (float(score) > confidence_threshold)
+                {
+                    // add shift of 1.0 to normalize the score values
+                    // to the range [1, 2). 
+                    // add a constant shift to scores will not change the sort
+                    // result, but will help reduce the computation because
+                    // we only need to sort the mantissa part of the floating-point
+                    // numbers
+                    temp_scores[targetIdx] = add_fb(score, score_shift);
+                    if (float(score_shift) > 0.f && (ge_fb(temp_scores[targetIdx], clip_val)))
+                        temp_scores[targetIdx] = clip_val;
+                    temp_idx[targetIdx] = cur_idx + i * numPredsPerBatch;
+                }
+                else
+                {
+                    // Set scores to 0
+                    // Set label = -1
+                    // add shift of 1.0 to normalize the score values
+                    // to the range [1, 2). 
+                    // add a constant shift to scores will not change the sort
+                    // result, but will help reduce the computation because
+                    // we only need to sort the mantissa part of the floating-point
+                    // numbers
+                    temp_scores[targetIdx] = score_shift;
+                    temp_idx[targetIdx] = -1;
+                    conf_scores_gpu[targetIdx] = score_shift;
+                    // TODO: HERE writing memory too many times
+                }
+            }
+
+            if ((cur_idx % num_preds_per_class) == 0)
+            {
+                const int offset_ct = i * num_classes + cur_idx / num_preds_per_class;
+                d_offsets[offset_ct] = offset_ct * num_preds_per_class;
+                // set the last element in d_offset
+                if (blockIdx.x == 0 && threadIdx.x == 0)
+                    d_offsets[num * num_classes] = num * numPredsPerBatch;
+            }
+        }
+    }
+}
+
+template <typename T_SCORE>
+pluginStatus_t sortScoresPerClass_gpu(
+    cudaStream_t stream,
+    const int num,
+    const int num_classes,
+    const int num_preds_per_class,
+    const int background_label_id,
+    const float confidence_threshold,
+    void* conf_scores_gpu,
+    void* index_array_gpu,
+    void* workspace,
+    const int score_bits,
+    const float score_shift
+)
+{
+    const int num_segments = num * num_classes;
+    void* temp_scores = workspace;
+    const int arrayLen = num * num_classes * num_preds_per_class;
+    void* temp_idx = nextWorkspacePtr((int8_t*) temp_scores, arrayLen * sizeof(T_SCORE));
+    void* d_offsets = nextWorkspacePtr((int8_t*) temp_idx, arrayLen * sizeof(int));
+    size_t cubOffsetSize = (num_segments + 1) * sizeof(int);
+    void* cubWorkspace = nextWorkspacePtr((int8_t*) d_offsets, cubOffsetSize);
+
+    const int BS = 512;
+    const int GS = (num_classes * num_preds_per_class + BS - 1) / BS;
+    // prepare the score, index, and offsets for CUB radix sort
+    // also normalize the scores to the range [1, 2)
+    // so we only need to sort the mantissa of floating-point numbers
+    // since their sign bit and exponential bits are identical
+    // we will subtract the 1.0 shift in gatherTopDetections()
+    prepareSortData<T_SCORE, BS><<<GS, BS, 0, stream>>>(num, num_classes, num_preds_per_class,
+                                                        background_label_id, confidence_threshold,
+                                                        (T_SCORE*) conf_scores_gpu,
+                                                        (T_SCORE*) temp_scores,
+                                                        T_SCORE(score_shift),
+                                                        (int*) temp_idx,
+                                                        (int*) d_offsets);
+
+    size_t temp_storage_bytes = cubSortPairsWorkspaceSize<T_SCORE, int>(arrayLen, num_segments);
+    size_t begin_bit = 0;
+    size_t end_bit = sizeof(T_SCORE) * 8;
+    if (sizeof(T_SCORE) == 2 && score_bits > 0 && score_bits <= 10)
+    {
+        // only sort score_bits in 10 mantissa bits.
+        end_bit = 10;
+        begin_bit = end_bit - score_bits;
+    }
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+        cubWorkspace, temp_storage_bytes,
+        (const T_SCORE*) (temp_scores), (T_SCORE*) (conf_scores_gpu),
+        (const int*) (temp_idx), (int*) (index_array_gpu),
+        arrayLen, num_segments,
+        (const int*) d_offsets, (const int*) d_offsets + 1,
+        begin_bit, end_bit,
+        stream);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// sortScoresPerClass LAUNCH CONFIG
+typedef pluginStatus_t (*sspcFunc)(cudaStream_t,
+                                const int,
+                                const int,
+                                const int,
+                                const int,
+                                const float,
+                                void*,
+                                void*,
+                                void*,
+                                const int,
+                                const float);
+struct sspcLaunchConfig
+{
+    DataType t_score;
+    sspcFunc function;
+
+    sspcLaunchConfig(DataType t_score)
+        : t_score(t_score)
+    {
+    }
+    sspcLaunchConfig(DataType t_score, sspcFunc function)
+        : t_score(t_score)
+        , function(function)
+    {
+    }
+    bool operator==(sspcLaunchConfig const& other) const
+    {
+        return t_score == other.t_score;
+    }
+};
+
+static std::array<sspcLaunchConfig, 2> sspcLCOptions = {
+    sspcLaunchConfig(DataType::kFLOAT, sortScoresPerClass_gpu<float>),
+    sspcLaunchConfig(DataType::kHALF, sortScoresPerClass_gpu<__half>)
+};
+
+pluginStatus_t sortScoresPerClass(
+    cudaStream_t stream,
+    const int num,
+    const int num_classes,
+    const int num_preds_per_class,
+    const int background_label_id,
+    const float confidence_threshold,
+    const DataType DT_SCORE,
+    void* conf_scores_gpu,
+    void* index_array_gpu,
+    void* workspace,
+    const int score_bits,
+    const float score_shift
+)
+{
+    sspcLaunchConfig lc = sspcLaunchConfig(DT_SCORE);
+    for (unsigned i = 0; i < sspcLCOptions.size(); ++i)
+    {
+        if (lc == sspcLCOptions[i])
+        {
+            DEBUG_PRINTF("sortScoresPerClass kernel %d\n", i);
+            return sspcLCOptions[i].function(stream,
+                                           num,
+                                           num_classes,
+                                           num_preds_per_class,
+                                           background_label_id,
+                                           confidence_threshold,
+                                           conf_scores_gpu,
+                                           index_array_gpu,
+                                           workspace,
+                                           score_bits,
+                                           score_shift);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+
+size_t sortScoresPerClassWorkspaceSize(
+    const int num,
+    const int num_classes,
+    const int num_preds_per_class,
+    const DataType DT_CONF)
+{
+    size_t wss[4];
+    const int arrayLen = num * num_classes * num_preds_per_class;
+    wss[0] = arrayLen * dataTypeSize(DT_CONF);      // temp scores
+    wss[1] = arrayLen * sizeof(int);                // temp indices
+    wss[2] = (num * num_classes + 1) * sizeof(int); // offsets
+    if (DT_CONF == DataType::kFLOAT)
+    {
+        wss[3] = cubSortPairsWorkspaceSize<float, int>(arrayLen, num * num_classes); // cub workspace
+    }
+    else if (DT_CONF == DataType::kHALF)
+    {
+        wss[3] = cubSortPairsWorkspaceSize<__half, int>(arrayLen, num * num_classes); // cub workspace
+    }
+    else
+    {
+        printf("SCORE type not supported\n");
+        return (size_t) -1;
+    }
+
+    return calculateTotalWorkspaceSize(wss, 4);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/sortScoresPerImage.cu

@@ -0,0 +1,151 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include "common/bboxUtils.h"
+#include "common/cub_helper.h"
+#include "common/kernels/kernel.h"
+#include "cub/cub.cuh"
+#include <array>
+namespace nvinfer1
+{
+namespace plugin
+{
+
+template <typename T_SCORE>
+pluginStatus_t sortScoresPerImage_gpu(cudaStream_t stream, const int num_images, const int num_items_per_image,
+    void* unsorted_scores, void* unsorted_bbox_indices, void* sorted_scores, void* sorted_bbox_indices, void* workspace,
+    int score_bits)
+{
+    void* d_offsets = workspace;
+    void* cubWorkspace = nextWorkspacePtr((int8_t*) d_offsets, (num_images + 1) * sizeof(int));
+
+    setUniformOffsets(stream, num_images, num_items_per_image, (int*) d_offsets);
+
+    const int arrayLen = num_images * num_items_per_image;
+    size_t temp_storage_bytes = cubSortPairsWorkspaceSize<T_SCORE, int>(arrayLen, num_images);
+    size_t begin_bit = 0;
+    size_t end_bit = sizeof(T_SCORE) * 8;
+    if (sizeof(T_SCORE) == 2 && score_bits > 0 && score_bits <= 10)
+    {
+        end_bit = 10;
+        begin_bit = end_bit - score_bits;
+    }
+    cub::DeviceSegmentedRadixSort::SortPairsDescending(
+        cubWorkspace, temp_storage_bytes,
+        (const T_SCORE*) (unsorted_scores), (T_SCORE*) (sorted_scores),
+        (const int*) (unsorted_bbox_indices), (int*) (sorted_bbox_indices),
+        arrayLen, num_images,
+        (const int*) d_offsets, (const int*) d_offsets + 1,
+        begin_bit, end_bit,
+        stream);
+    CSC(cudaGetLastError(), STATUS_FAILURE);
+    return STATUS_SUCCESS;
+}
+
+// sortScoresPerImage LAUNCH CONFIG
+typedef pluginStatus_t (*sspiFunc)(cudaStream_t,
+                                const int,
+                                const int,
+                                void*,
+                                void*,
+                                void*,
+                                void*,
+                                void*,
+                                int);
+struct sspiLaunchConfig
+{
+    DataType t_score;
+    sspiFunc function;
+
+    sspiLaunchConfig(DataType t_score)
+        : t_score(t_score)
+    {
+    }
+    sspiLaunchConfig(DataType t_score, sspiFunc function)
+        : t_score(t_score)
+        , function(function)
+    {
+    }
+    bool operator==(sspiLaunchConfig const& other) const
+    {
+        return t_score == other.t_score;
+    }
+};
+
+static std::array<sspiLaunchConfig, 2> sspiLCOptions = {
+    sspiLaunchConfig(DataType::kFLOAT, sortScoresPerImage_gpu<float>),
+    sspiLaunchConfig(DataType::kHALF, sortScoresPerImage_gpu<__half>),
+};
+
+pluginStatus_t sortScoresPerImage(
+    cudaStream_t stream,
+    const int num_images,
+    const int num_items_per_image,
+    const DataType DT_SCORE,
+    void* unsorted_scores,
+    void* unsorted_bbox_indices,
+    void* sorted_scores,
+    void* sorted_bbox_indices,
+    void* workspace,
+    int score_bits
+)
+{
+    sspiLaunchConfig lc = sspiLaunchConfig(DT_SCORE);
+    for (unsigned i = 0; i < sspiLCOptions.size(); ++i)
+    {
+        if (lc == sspiLCOptions[i])
+        {
+            DEBUG_PRINTF("sortScoresPerImage kernel %d\n", i);
+            return sspiLCOptions[i].function(stream,
+                                           num_images,
+                                           num_items_per_image,
+                                           unsorted_scores,
+                                           unsorted_bbox_indices,
+                                           sorted_scores,
+                                           sorted_bbox_indices,
+                                           workspace,
+                                           score_bits);
+        }
+    }
+    return STATUS_BAD_PARAM;
+}
+
+size_t sortScoresPerImageWorkspaceSize(
+    const int num_images,
+    const int num_items_per_image,
+    const DataType DT_SCORE)
+{
+    const int arrayLen = num_images * num_items_per_image;
+    size_t wss[2];
+    wss[0] = (num_images + 1) * sizeof(int); // offsets
+    if (DT_SCORE == DataType::kFLOAT)
+    {
+        wss[1] = cubSortPairsWorkspaceSize<float, int>(arrayLen, num_images); // cub workspace
+    }
+    else if (DT_SCORE == DataType::kHALF)
+    {
+        wss[1] = cubSortPairsWorkspaceSize<__half, int>(arrayLen, num_images); // cub workspace
+    }
+    else
+    {
+        printf("SCORE type not supported.\n");
+        return (size_t) -1;
+    }
+
+    return calculateTotalWorkspaceSize(wss, 2);
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/kernels/voxelGeneratorKernels.cu

@@ -0,0 +1,412 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <iostream>
+#include <cuda_runtime_api.h>
+namespace nvinfer1
+{
+namespace plugin
+{
+__global__ void generateVoxels_kernel(
+        int max_num_points,
+        float *points, unsigned int* points_size,
+        float min_x_range, float max_x_range,
+        float min_y_range, float max_y_range,
+        float min_z_range, float max_z_range,
+        float pillar_x_size, float pillar_y_size, float pillar_z_size,
+        int grid_y_size, int grid_x_size, int num_point_values,
+        int max_points_per_voxel,
+        unsigned int *mask, float *voxels)
+{
+  int point_idx = blockIdx.x * blockDim.x + threadIdx.x;
+  int batch_idx = point_idx / max_num_points;
+  int point_idx_in_frame = point_idx % max_num_points;
+  if(point_idx_in_frame >= points_size[batch_idx]) return;
+  float px = points[num_point_values * point_idx];
+  float py = points[num_point_values * point_idx + 1];
+  float pz = points[num_point_values * point_idx + 2];
+  float pw = points[num_point_values * point_idx + 3];
+  float pt;
+  if (num_point_values == 5) {
+    pt = points[num_point_values * point_idx + 4];
+  }
+  if(px<min_x_range||px>=max_x_range
+    || py<min_y_range||py>=max_y_range
+    || pz<min_z_range||pz>=max_z_range) return;
+  int voxel_idx = floorf((px - min_x_range)/pillar_x_size);
+  int voxel_idy = floorf((py - min_y_range)/pillar_y_size);
+  unsigned int voxel_index = (batch_idx * grid_y_size + voxel_idy) * grid_x_size + voxel_idx;
+  unsigned int point_id = atomicAdd(&(mask[voxel_index]), 1);
+  if(point_id >= max_points_per_voxel) return;
+  float *address = voxels + (voxel_index*max_points_per_voxel + point_id)*num_point_values;
+  atomicExch(address+0, px);
+  atomicExch(address+1, py);
+  atomicExch(address+2, pz);
+  atomicExch(address+3, pw);
+  if (num_point_values == 5) {
+    atomicExch(address+4, pt);
+  }
+}
+
+__global__ void generateBaseFeatures_kernel(
+        int batch_size,
+        unsigned int *mask, float *voxels,
+        int grid_y_size, int grid_x_size,
+        unsigned int *pillar_num,
+        int max_pillar_num,
+        int max_points_per_voxel,
+        int num_point_values,
+        float *voxel_features,
+        unsigned int *voxel_num_points,
+        unsigned int *coords)
+{
+  int voxel_id = blockIdx.x * blockDim.x + threadIdx.x;
+  int voxel_idx = voxel_id % grid_x_size;
+  int voxel_idy = (voxel_id / grid_x_size) % grid_y_size;
+  int batch_id = voxel_id / (grid_y_size * grid_x_size);
+  if (batch_id >= batch_size) return;
+  unsigned int count = mask[voxel_id];
+  if( !(count>0) ) return;
+  count = count<max_points_per_voxel?count:max_points_per_voxel;
+  int current_pillarId = 0;
+  current_pillarId = atomicAdd(pillar_num + batch_id, 1);
+  voxel_num_points[batch_id * grid_y_size * grid_x_size + current_pillarId] = count;
+  int4 coord = {0, 0, voxel_idy, voxel_idx};
+  ((int4*)coords)[batch_id * max_pillar_num + current_pillarId] = coord;
+  for (int i=0; i<count; i++){
+    int inIndex = voxel_id*max_points_per_voxel + i;
+    int outIndex = (batch_id * grid_x_size * grid_y_size + current_pillarId)*max_points_per_voxel + i;
+    if (num_point_values == 4) {
+      ((float4*)voxel_features)[outIndex] = ((float4*)voxels)[inIndex];
+    }
+    else if (num_point_values == 5) {
+      for(int k=0; k<5;k++)
+          voxel_features[5 * outIndex + k] = voxels[5 * inIndex + k];
+    }
+  }
+}
+
+void generateVoxels_launch(
+        int batch_size, int max_num_points,
+        float *points, unsigned int* points_size,
+        float min_x_range, float max_x_range,
+        float min_y_range, float max_y_range,
+        float min_z_range, float max_z_range,
+        float pillar_x_size, float pillar_y_size, float pillar_z_size,
+        int grid_y_size, int grid_x_size, int num_point_values,
+        int max_points_per_voxel,
+        unsigned int *mask, float *voxels,
+        cudaStream_t stream)
+{
+  int threadNum = 256;
+  dim3 blocks((batch_size * max_num_points + threadNum - 1) / threadNum);
+  dim3 threads(threadNum);
+  generateVoxels_kernel<<<blocks, threads, 0, stream>>>
+      (max_num_points,
+        points, points_size,
+        min_x_range, max_x_range,
+        min_y_range, max_y_range,
+        min_z_range, max_z_range,
+        pillar_x_size, pillar_y_size, pillar_z_size,
+        grid_y_size, grid_x_size, num_point_values,
+        max_points_per_voxel,
+        mask, voxels);
+}
+
+void generateBaseFeatures_launch(
+        int batch_size,
+        unsigned int *mask, float *voxels,
+        int grid_y_size, int grid_x_size,
+        unsigned int *pillar_num,
+        int max_pillar_num,
+        int max_points_per_voxel,
+        int num_point_values,
+        float *voxel_features,
+        unsigned int *voxel_num_points,
+        unsigned int *coords,
+        cudaStream_t stream)
+{
+  int blockSize = 1024;
+  dim3 threads(blockSize);
+  dim3 blocks((batch_size * grid_y_size * grid_x_size + blockSize - 1) / blockSize);
+  generateBaseFeatures_kernel<<<blocks, threads, 0, stream>>>
+      (
+        batch_size,
+        mask, voxels, grid_y_size, grid_x_size,
+        pillar_num,
+        max_pillar_num,
+        max_points_per_voxel,
+        num_point_values,
+        voxel_features,
+        voxel_num_points,
+        coords
+      );
+}
+
+__global__ void generateFeatures_kernel(
+    int batch_size,
+    int dense_pillar_num,
+    float* voxel_features,
+    unsigned int* voxel_num_points,
+    unsigned int* coords, unsigned int *params,
+    float voxel_x, float voxel_y, float voxel_z,
+    float range_min_x, float range_min_y, float range_min_z,
+    unsigned int voxel_features_size, unsigned int max_points,
+    unsigned int max_voxels,
+    float* features)
+{
+    int warp_size = max_points;
+    int pillar_idx = blockIdx.x * 4 + threadIdx.x/warp_size;
+    int point_idx = threadIdx.x % warp_size;
+    // In case the actual number of points is less than warp_size
+    // E.g., warp_size=32, max_points=20
+    if (point_idx >= max_points) return;
+    int batch_idx = pillar_idx / max_voxels;
+    if (batch_idx >= batch_size) return;
+    int pillar_idx_in_frame = pillar_idx % max_voxels;
+    int dense_pillar_idx = pillar_idx_in_frame + dense_pillar_num * batch_idx;
+    int pillar_idx_inBlock = threadIdx.x/warp_size;
+    // Limit number of voxels to max_voxels
+    unsigned int num_pillars = params[batch_idx] > max_voxels ? max_voxels : params[batch_idx];
+    // Update max_voxel to actual number
+    if (pillar_idx_in_frame == 0 && point_idx == 0) {
+      params[batch_idx] = num_pillars;
+    }
+    if (pillar_idx_in_frame >= num_pillars) return;
+
+    //load src
+    __shared__ float pillarSM[4][64][5]; // up to 64 points per pillar
+    __shared__ float4 pillarSumSM[4]; //4*4
+    __shared__ int4 cordsSM[4]; //4*4
+    __shared__ int pointsNumSM[4]; //4
+    __shared__ float pillarOutSM[4][64][11]; // up to 11 features per point
+
+    if (point_idx == 0) {
+      pointsNumSM[pillar_idx_inBlock] = voxel_num_points[dense_pillar_idx];
+      cordsSM[pillar_idx_inBlock] = ((int4*)coords)[dense_pillar_idx];
+      pillarSumSM[pillar_idx_inBlock] = {0,0,0,0};
+    }
+    for(int k=0; k<5; k++) {
+      pillarSM[pillar_idx_inBlock][point_idx][k] = voxel_features[5 * (dense_pillar_idx*max_points + point_idx) + k];
+    }
+    __syncthreads();
+    //calculate sm
+    if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+      atomicAdd(&(pillarSumSM[pillar_idx_inBlock].x),  pillarSM[pillar_idx_inBlock][point_idx][0]);
+      atomicAdd(&(pillarSumSM[pillar_idx_inBlock].y),  pillarSM[pillar_idx_inBlock][point_idx][1]);
+      atomicAdd(&(pillarSumSM[pillar_idx_inBlock].z),  pillarSM[pillar_idx_inBlock][point_idx][2]);
+    }
+    __syncthreads();
+    //feature-mean
+    float4 mean;
+    float validPoints = pointsNumSM[pillar_idx_inBlock];
+    mean.x = pillarSumSM[pillar_idx_inBlock].x / validPoints;
+    mean.y = pillarSumSM[pillar_idx_inBlock].y / validPoints;
+    mean.z = pillarSumSM[pillar_idx_inBlock].z / validPoints;
+    mean.x  = pillarSM[pillar_idx_inBlock][point_idx][0] - mean.x;
+    mean.y  = pillarSM[pillar_idx_inBlock][point_idx][1] - mean.y;
+    mean.z  = pillarSM[pillar_idx_inBlock][point_idx][2] - mean.z;
+    //calculate offset
+    float x_offset = voxel_x / 2.0f + cordsSM[pillar_idx_inBlock].w * voxel_x + range_min_x;
+    float y_offset = voxel_y / 2.0f + cordsSM[pillar_idx_inBlock].z * voxel_y + range_min_y;
+    float z_offset = voxel_z / 2.0f + cordsSM[pillar_idx_inBlock].y * voxel_z + range_min_z;
+    //feature-offset
+    float4 center;
+    center.x  = pillarSM[pillar_idx_inBlock][point_idx][0] - x_offset;
+    center.y  = pillarSM[pillar_idx_inBlock][point_idx][1] - y_offset;
+    center.z  = pillarSM[pillar_idx_inBlock][point_idx][2] - z_offset;
+    //store output
+    if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+      for(int k=0; k<5; k++)
+        pillarOutSM[pillar_idx_inBlock][point_idx][k] = pillarSM[pillar_idx_inBlock][point_idx][k];
+      pillarOutSM[pillar_idx_inBlock][point_idx][5] = mean.x;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 1] = mean.y;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 2] = mean.z;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 3] = center.x;
+      pillarOutSM[pillar_idx_inBlock][point_idx][5 + 4] = center.y;
+      if (5 + 5 < voxel_features_size)
+        pillarOutSM[pillar_idx_inBlock][point_idx][warp_size + 5] = center.z;
+    } else {
+      for (int k = 0; k < voxel_features_size; k++)
+        pillarOutSM[pillar_idx_inBlock][point_idx][k] = 0;
+    }
+    __syncthreads();
+    for(int i = 0; i < voxel_features_size; i ++) {
+      int outputSMId = pillar_idx_inBlock*64*11 + point_idx * 11 + i;
+      int outputId = pillar_idx*max_points*voxel_features_size + point_idx * voxel_features_size + i;
+      features[outputId] = ((float*)pillarOutSM)[outputSMId] ;
+    }
+
+}
+
+__global__ void generateFeatures_kernel_4x(
+  int batch_size,
+  int dense_pillar_num,
+  float* voxel_features,
+  unsigned int* voxel_num_points, unsigned int* coords,
+  unsigned int *params,
+  float voxel_x, float voxel_y, float voxel_z,
+  float range_min_x, float range_min_y, float range_min_z,
+  unsigned int voxel_features_size, unsigned int max_points,
+  unsigned int max_voxels,
+  float* features)
+{
+  int warp_size = max_points;
+  int pillar_idx = blockIdx.x * 4 + threadIdx.x / warp_size;
+  int point_idx = threadIdx.x % warp_size;
+  // In case the actual number of points is less than warp_size
+  // E.g., warp_size=32, max_points=20
+  if (point_idx >= max_points) return;
+  int batch_idx = pillar_idx / max_voxels;
+  if (batch_idx >= batch_size) return;
+  int pillar_idx_in_frame = pillar_idx % max_voxels;
+  int dense_pillar_idx = pillar_idx_in_frame + dense_pillar_num * batch_idx;
+  int pillar_idx_inBlock = threadIdx.x / warp_size;
+  // Limit number of voxels to max_voxels
+  unsigned int num_pillars = params[batch_idx] > max_voxels ? max_voxels : params[batch_idx];
+  // Update max_voxel to actual number
+  if (pillar_idx_in_frame == 0 && point_idx == 0) {
+    params[batch_idx] = num_pillars;
+  }
+  if (pillar_idx_in_frame >= num_pillars) return;
+  //load src
+  __shared__ float4 pillarSM[4][64]; // up to 64 points per pillar
+  __shared__ float4 pillarSumSM[4]; //4*4
+  __shared__ int4 cordsSM[4]; //4*4
+  __shared__ int pointsNumSM[4]; //4
+  __shared__ float pillarOutSM[4][64][11]; // up to 11 output features per point
+
+  if (point_idx == 0) {
+    pointsNumSM[pillar_idx_inBlock] = voxel_num_points[dense_pillar_idx];
+    cordsSM[pillar_idx_inBlock] = ((int4*)coords)[pillar_idx];
+    pillarSumSM[pillar_idx_inBlock] = {0,0,0,0};
+  }
+  pillarSM[pillar_idx_inBlock][point_idx] = ((float4*)voxel_features)[dense_pillar_idx*max_points + point_idx];
+  __syncthreads();
+  //calculate sm
+  if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+    atomicAdd(&(pillarSumSM[pillar_idx_inBlock].x),  pillarSM[pillar_idx_inBlock][point_idx].x);
+    atomicAdd(&(pillarSumSM[pillar_idx_inBlock].y),  pillarSM[pillar_idx_inBlock][point_idx].y);
+    atomicAdd(&(pillarSumSM[pillar_idx_inBlock].z),  pillarSM[pillar_idx_inBlock][point_idx].z);
+  }
+  __syncthreads();
+  //feature-mean
+  float4 mean;
+  float validPoints = pointsNumSM[pillar_idx_inBlock];
+  mean.x = pillarSumSM[pillar_idx_inBlock].x / validPoints;
+  mean.y = pillarSumSM[pillar_idx_inBlock].y / validPoints;
+  mean.z = pillarSumSM[pillar_idx_inBlock].z / validPoints;
+  mean.x  = pillarSM[pillar_idx_inBlock][point_idx].x - mean.x;
+  mean.y  = pillarSM[pillar_idx_inBlock][point_idx].y - mean.y;
+  mean.z  = pillarSM[pillar_idx_inBlock][point_idx].z - mean.z;
+  //calculate offset
+  float x_offset = voxel_x / 2.0f + cordsSM[pillar_idx_inBlock].w * voxel_x + range_min_x;
+  float y_offset = voxel_y / 2.0f + cordsSM[pillar_idx_inBlock].z * voxel_y + range_min_y;
+  float z_offset = voxel_z / 2.0f + cordsSM[pillar_idx_inBlock].y * voxel_z + range_min_z;
+  //feature-offset
+  float4 center;
+  center.x  = pillarSM[pillar_idx_inBlock][point_idx].x - x_offset;
+  center.y  = pillarSM[pillar_idx_inBlock][point_idx].y - y_offset;
+  center.z  = pillarSM[pillar_idx_inBlock][point_idx].z - z_offset;
+  //store output
+  if (point_idx < pointsNumSM[pillar_idx_inBlock]) {
+    pillarOutSM[pillar_idx_inBlock][point_idx][0] = pillarSM[pillar_idx_inBlock][point_idx].x;
+    pillarOutSM[pillar_idx_inBlock][point_idx][1] = pillarSM[pillar_idx_inBlock][point_idx].y;
+    pillarOutSM[pillar_idx_inBlock][point_idx][2] = pillarSM[pillar_idx_inBlock][point_idx].z;
+    pillarOutSM[pillar_idx_inBlock][point_idx][3] = pillarSM[pillar_idx_inBlock][point_idx].w;
+    pillarOutSM[pillar_idx_inBlock][point_idx][4] = mean.x;
+    pillarOutSM[pillar_idx_inBlock][point_idx][5] = mean.y;
+    pillarOutSM[pillar_idx_inBlock][point_idx][6] = mean.z;
+    pillarOutSM[pillar_idx_inBlock][point_idx][7] = center.x;
+    pillarOutSM[pillar_idx_inBlock][point_idx][8] = center.y;
+    pillarOutSM[pillar_idx_inBlock][point_idx][9] = center.z;
+
+  } else {
+    pillarOutSM[pillar_idx_inBlock][point_idx][0] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][1] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][2] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][3] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][4] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][5] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][6] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][7] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][8] = 0;
+    pillarOutSM[pillar_idx_inBlock][point_idx][9] = 0;
+  }
+  __syncthreads();
+  for(int i = 0; i < voxel_features_size; i ++) {
+    int outputSMId = pillar_idx_inBlock*64*11 + point_idx * 11 + i;
+    int outputId = pillar_idx*max_points*voxel_features_size + point_idx * voxel_features_size + i;
+    features[outputId] = ((float*)pillarOutSM)[outputSMId] ;
+  }
+
+}
+
+int generateFeatures_launch(
+    int batch_size,
+    int dense_pillar_num,
+    float* voxel_features,
+    unsigned int* voxel_num_points,
+    unsigned int* coords,
+    unsigned int *params,
+    float voxel_x, float voxel_y, float voxel_z,
+    float range_min_x, float range_min_y, float range_min_z,
+    unsigned int voxel_features_size, unsigned int max_points,
+    unsigned int max_voxels, unsigned int num_point_values,
+    float* features,
+    cudaStream_t stream)
+{
+    unsigned int warp_size = max_points;
+    dim3 blocks((batch_size * max_voxels + 3) / 4);
+    dim3 threads(4*warp_size);
+    if (num_point_values == 4) {
+      generateFeatures_kernel_4x<<<blocks, threads, 0, stream>>>
+      (batch_size,
+      dense_pillar_num,
+      voxel_features,
+      voxel_num_points,
+      coords,
+      params,
+      voxel_x, voxel_y, voxel_z,
+      range_min_x, range_min_y, range_min_z,
+      voxel_features_size, max_points,
+      max_voxels,
+      features);
+    }
+    else {
+      generateFeatures_kernel<<<blocks, threads, 0, stream>>>
+      (batch_size,
+      dense_pillar_num,
+      voxel_features,
+      voxel_num_points,
+      coords,
+      params,
+      voxel_x, voxel_y, voxel_z,
+      range_min_x, range_min_y, range_min_z,
+      voxel_features_size, max_points,
+      max_voxels,
+      features);
+    }
+    auto err = cudaGetLastError();
+    if (cudaSuccess != err) {
+        fprintf(stderr, "CUDA kernel failed : %s\n", cudaGetErrorString(err));
+        exit(-1);
+    }
+    return err;
+}
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/mrcnn_config.h

@@ -0,0 +1,177 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_PLUGIN_MASKRCNN_CONFIG_H
+#define TRT_PLUGIN_MASKRCNN_CONFIG_H
+#include "NvInfer.h"
+#include <string>
+#include <vector>
+
+namespace MaskRCNNConfig
+{
+static const nvinfer1::Dims3 IMAGE_SHAPE{3, 1024, 1024};
+
+// Pooled ROIs
+static int32_t const POOL_SIZE = 7;
+static int32_t const MASK_POOL_SIZE = 14;
+
+// Threshold to determine the mask area out of final convolution output
+static float const MASK_THRESHOLD = 0.5F;
+
+// Bounding box refinement standard deviation for RPN and final detections.
+static float const RPN_BBOX_STD_DEV[] = {0.1F, 0.1F, 0.2F, 0.2F};
+static float const BBOX_STD_DEV[] = {0.1F, 0.1F, 0.2F, 0.2F};
+
+// Max number of final detections
+static int32_t const DETECTION_MAX_INSTANCES = 100;
+
+// Minimum probability value to accept a detected instance
+// ROIs below this threshold are skipped
+static float const DETECTION_MIN_CONFIDENCE = 0.7F;
+
+// Non-maximum suppression threshold for detection
+static float const DETECTION_NMS_THRESHOLD = 0.3F;
+
+// The strides of each layer of the FPN Pyramid. These values
+// are based on a Resnet101 backbone.
+static const std::vector<float> BACKBONE_STRIDES = {4.F, 8.F, 16.F, 32.F, 64.F};
+
+// Size of the fully-connected layers in the classification graph
+static int32_t const FPN_CLASSIF_FC_LAYERS_SIZE = 1024;
+
+// Size of the top-down layers used to build the feature pyramid
+static int32_t const TOP_DOWN_PYRAMID_SIZE = 256;
+
+// Number of classification classes (including background)
+static int32_t const NUM_CLASSES = 1 + 80; // COCO has 80 classes
+
+// Length of square anchor side in pixels
+static const std::vector<float> RPN_ANCHOR_SCALES = {32.F, 64.F, 128.F, 256.F, 512.F};
+
+// Ratios of anchors at each cell (width/height)
+// A value of 1 represents a square anchor, and 0.5 is a wide anchor
+static float const RPN_ANCHOR_RATIOS[] = {0.5F, 1.F, 2.F};
+
+// Anchor stride
+// If 1 then anchors are created for each cell in the backbone feature map.
+// If 2, then anchors are created for every other cell, and so on.
+static int32_t const RPN_ANCHOR_STRIDE = 1;
+
+// Although Python impementation uses 6000,
+//  TRT fails if this number larger than kMAX_TOPK_K defined in engine/checkMacros.h
+static int32_t const MAX_PRE_NMS_RESULTS = 1024; // 3840;
+
+// Non-max suppression threshold to filter RPN proposals.
+// You can increase this during training to generate more propsals.
+static float const RPN_NMS_THRESHOLD = 0.7F;
+
+// ROIs kept after non-maximum suppression (training and inference)
+static int32_t const POST_NMS_ROIS_INFERENCE = 1000;
+
+// COCO Class names
+static const std::vector<std::string> CLASS_NAMES = {
+    "BG",
+    "person",
+    "bicycle",
+    "car",
+    "motorcycle",
+    "airplane",
+    "bus",
+    "train",
+    "truck",
+    "boat",
+    "traffic light",
+    "fire hydrant",
+    "stop sign",
+    "parking meter",
+    "bench",
+    "bird",
+    "cat",
+    "dog",
+    "horse",
+    "sheep",
+    "cow",
+    "elephant",
+    "bear",
+    "zebra",
+    "giraffe",
+    "backpack",
+    "umbrella",
+    "handbag",
+    "tie",
+    "suitcase",
+    "frisbee",
+    "skis",
+    "snowboard",
+    "sports ball",
+    "kite",
+    "baseball bat",
+    "baseball glove",
+    "skateboard",
+    "surfboard",
+    "tennis racket",
+    "bottle",
+    "wine glass",
+    "cup",
+    "fork",
+    "knife",
+    "spoon",
+    "bowl",
+    "banana",
+    "apple",
+    "sandwich",
+    "orange",
+    "broccoli",
+    "carrot",
+    "hot dog",
+    "pizza",
+    "donut",
+    "cake",
+    "chair",
+    "couch",
+    "potted plant",
+    "bed",
+    "dining table",
+    "toilet",
+    "tv",
+    "laptop",
+    "mouse",
+    "remote",
+    "keyboard",
+    "cell phone",
+    "microwave",
+    "oven",
+    "toaster",
+    "sink",
+    "refrigerator",
+    "book",
+    "clock",
+    "vase",
+    "scissors",
+    "teddy bear",
+    "hair drier",
+    "toothbrush",
+};
+
+static const std::string MODEL_NAME = "mrcnn_nchw.uff";
+static const std::string MODEL_INPUT = "input_image";
+static const nvinfer1::Dims3 MODEL_INPUT_SHAPE = IMAGE_SHAPE;
+static const std::vector<std::string> MODEL_OUTPUTS = {"mrcnn_detection", "mrcnn_mask/Sigmoid"};
+static const nvinfer1::Dims2 MODEL_DETECTION_SHAPE{DETECTION_MAX_INSTANCES, 6};
+static const nvinfer1::Dims4 MODEL_MASK_SHAPE{DETECTION_MAX_INSTANCES, NUM_CLASSES, 28, 28};
+} // namespace MaskRCNNConfig
+#endif // TRT_PLUGIN_MASKRCNN_CONFIG_H

plugin/common/nmsHelper.cpp

@@ -0,0 +1,66 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "common/plugin.h"
+#include <algorithm>
+#include <cuda_fp16.h>
+
+namespace nvinfer1::plugin
+{
+
+size_t detectionForwardBBoxDataSize(int32_t N, int32_t C1, DataType DT_BBOX)
+{
+    if (DT_BBOX == DataType::kFLOAT)
+    {
+        return N * C1 * sizeof(float);
+    }
+    if (DT_BBOX == DataType::kHALF)
+    {
+        return N * C1 * sizeof(__half);
+    }
+
+    printf("Only FP32/FP16 type bounding boxes are supported.\n");
+    return (size_t) -1;
+}
+
+size_t detectionForwardBBoxPermuteSize(bool shareLocation, int32_t N, int32_t C1, DataType DT_BBOX)
+{
+    if (DT_BBOX == DataType::kFLOAT)
+    {
+        return shareLocation ? 0 : N * C1 * sizeof(float);
+    }
+    if (DT_BBOX == DataType::kHALF)
+    {
+        return shareLocation ? 0 : N * C1 * sizeof(__half);
+    }
+
+    printf("Only FP32/FP16 type bounding boxes are supported.\n");
+    return (size_t) -1;
+}
+
+size_t detectionForwardPreNMSSize(int32_t N, int32_t C2)
+{
+    PLUGIN_ASSERT(sizeof(float) == sizeof(int32_t));
+    return N * C2 * sizeof(float);
+}
+
+size_t detectionForwardPostNMSSize(int32_t N, int32_t numClasses, int32_t topK)
+{
+    PLUGIN_ASSERT(sizeof(float) == sizeof(int32_t));
+    return N * numClasses * topK * sizeof(float);
+}
+} // namespace nvinfer1::plugin

plugin/common/plugin.cpp

@@ -0,0 +1,185 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "common/plugin.h"
+#include <type_traits>
+
+namespace nvinfer1
+{
+namespace pluginInternal
+{
+
+// Helper to create per-context cudnn/cublas singleton managed by std::shared_ptr.
+// Unlike conventional singletons, singleton created with this will be released
+// when not needed, instead of on process exit.
+// Objects of this class shall always be declared static / global, and shall never own cudnn/cublas handle
+// resources.
+template <typename T>
+class PerContextPluginHandleSingletonCreator
+{
+public:
+    // creator returning std::unique_ptr is by design.
+    // It forces separation of memory for T and memory for control blocks.
+    // So when T is released, but we still have observer weak_ptr in mObservers, the T mem block can be released.
+    // creator itself must not own cudnn/cublas handle resources. Only the object it creates can.
+    template <typename CreatorFunction,
+        typename = std::enable_if_t<std::is_invocable_r_v<std::unique_ptr<T>, CreatorFunction>>>
+    explicit PerContextPluginHandleSingletonCreator(CreatorFunction&& cublasCreator)
+        : PerContextPluginHandleSingletonCreator([creator = std::forward<CreatorFunction>(cublasCreator)](char const*) {
+            return creator();
+        }) //< Adapt `cublasCreator` by ignoring the `char const*` argument.
+    {
+    }
+
+    explicit PerContextPluginHandleSingletonCreator(std::function<std::unique_ptr<T>(char const*)> cudnnCreator)
+        : mCreator(std::move(cudnnCreator))
+    {
+    }
+
+    // \param executionContextIdentifier Unique pointer to identify contexts having overlapping lifetime.
+    // \param callerPluginName Optional name of the plugin that invokes this creator
+    [[nodiscard]] std::shared_ptr<T> operator()(
+        void* executionContextIdentifier, char const* callerPluginName = nullptr)
+    {
+        std::lock_guard<std::mutex> lk{mMutex};
+        auto& observer = mObservers[executionContextIdentifier];
+        std::shared_ptr<T> result = observer.lock();
+        if (result == nullptr)
+        {
+            auto deleter = [this, executionContextIdentifier](T* obj) {
+                delete obj;
+                // Clears observer to avoid growth of mObservers, in case users create/destroy
+                // plugin handle contexts frequently.
+                std::shared_ptr<T> observedObjHolder;
+                // The destructor of observedObjHolder may attempt to acquire a lock on mMutex.
+                // To avoid deadlock, it's critical to release the lock here held by lk first,
+                // before destroying observedObjHolder. Hence observedObjHolder must be declared
+                // before lk.
+                std::lock_guard<std::mutex> lk{mMutex};
+                // Must check observer again because another thread may create new instance for
+                // this ctx just before we lock mMutex. We can't infer that the observer is
+                // stale from the fact that obj is destroyed, because shared_ptr ref-count
+                // checking and observer removing are not in one atomic operation, and the
+                // observer may be changed to observe another instance.
+                if (auto it = mObservers.find(executionContextIdentifier); it != mObservers.end())
+                {
+                    if (observedObjHolder = it->second.lock(); observedObjHolder == nullptr)
+                    {
+                        mObservers.erase(it);
+                    }
+                }
+            };
+            // Note, if `std::shared_ptr<T>{...}` throws, the deleter is still called. See
+            // https://en.cppreference.com/w/cpp/memory/shared_ptr/shared_ptr.html
+            result = std::shared_ptr<T>{mCreator(callerPluginName).release(), std::move(deleter)};
+
+            // Update the per-context observer with the new resource
+            observer = result;
+        }
+
+        return result;
+    }
+
+private:
+    std::function<std::unique_ptr<T>(char const*)> mCreator;
+    mutable std::mutex mMutex;
+    // cudnn/cublas handle resources are per-context.
+    std::unordered_map</*contextIdentifier*/ void*, std::weak_ptr<T>> mObservers;
+}; // class PerContextPluginHandleSingletonCreator
+
+std::unique_ptr<CudnnWrapper> createPluginCudnnWrapperImpl(char const* callerPluginName)
+{
+    // callerPluginName is used to enrich downstream cudnn error message with caller plugin info
+    return std::make_unique<CudnnWrapper>(/*initHandle*/ true, callerPluginName);
+}
+
+std::unique_ptr<CublasWrapper> createPluginCublasWrapperImpl()
+{
+    return std::make_unique<CublasWrapper>(/*initHandle*/ true);
+}
+
+static PerContextPluginHandleSingletonCreator<CudnnWrapper> gCreatePluginCudnnHandleWrapper(
+    createPluginCudnnWrapperImpl);
+static PerContextPluginHandleSingletonCreator<CublasWrapper> gCreatePluginCublasHandleWrapper(
+    createPluginCublasWrapperImpl);
+
+std::shared_ptr<CudnnWrapper> createPluginCudnnWrapper(void* executionContextIdentifier, char const* callerPluginName)
+{
+    return gCreatePluginCudnnHandleWrapper(executionContextIdentifier, callerPluginName);
+}
+
+std::shared_ptr<CublasWrapper> createPluginCublasWrapper(void* executionContextIdentifier)
+{
+    return gCreatePluginCublasHandleWrapper(executionContextIdentifier);
+}
+
+} // namespace pluginInternal
+
+namespace plugin
+{
+
+void validateRequiredAttributesExist(std::set<std::string> requiredFieldNames, PluginFieldCollection const* fc)
+{
+    for (int32_t i = 0; i < fc->nbFields; i++)
+    {
+        requiredFieldNames.erase(fc->fields[i].name);
+    }
+    if (!requiredFieldNames.empty())
+    {
+        std::stringstream msg{};
+        msg << "PluginFieldCollection missing required fields: {";
+        char const* separator = "";
+        for (auto const& field : requiredFieldNames)
+        {
+            msg << separator << field;
+            separator = ", ";
+        }
+        msg << "}";
+        std::string msg_str = msg.str();
+        PLUGIN_ERROR(msg_str.c_str());
+    }
+}
+
+int32_t dimToInt32(int64_t d)
+{
+    if (d < std::numeric_limits<int32_t>::min() || d > std::numeric_limits<int32_t>::max())
+    {
+        std::stringstream msg{};
+        msg << "Plugin cannot handle dimension outside of int32_t range: " << d;
+        std::string msg_str = msg.str();
+        PLUGIN_ERROR(msg_str.c_str());
+    }
+    return static_cast<int32_t>(d);
+}
+
+bool supportsMemPoolsHelper()
+{
+    int32_t device;
+    PLUGIN_CUASSERT(cudaGetDevice(&device));
+    int32_t value;
+    PLUGIN_CUASSERT(cudaDeviceGetAttribute(&value, cudaDevAttrMemoryPoolsSupported, device));
+    return value != 0;
+}
+
+bool supportsMemPools()
+{
+    static bool sResult = supportsMemPoolsHelper();
+    return sResult;
+}
+
+} // namespace plugin
+} // namespace nvinfer1

plugin/common/plugin.h

@@ -0,0 +1,269 @@
+/*
+ * 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 TRT_PLUGIN_H
+#define TRT_PLUGIN_H
+#include "NvInferPlugin.h"
+#include "common/checkMacrosPlugin.h"
+#include "cublasWrapper.h"
+#include "cudnnWrapper.h"
+#include <cstring>
+#include <cuda_runtime.h>
+#include <functional>
+#include <iostream>
+#include <memory>
+#include <set>
+#include <sstream>
+#include <string>
+#include <unordered_map>
+
+// Enumerator for status
+typedef enum
+{
+    STATUS_SUCCESS = 0,
+    STATUS_FAILURE = 1,
+    STATUS_BAD_PARAM = 2,
+    STATUS_NOT_SUPPORTED = 3,
+    STATUS_NOT_INITIALIZED = 4
+} pluginStatus_t;
+
+namespace nvinfer1
+{
+
+namespace pluginInternal
+{
+
+class BasePlugin : public IPluginV2
+{
+protected:
+    void setPluginNamespace(char const* libNamespace) noexcept override
+    {
+        mNamespace = libNamespace;
+    }
+
+    char const* getPluginNamespace() const noexcept override
+    {
+        return mNamespace.c_str();
+    }
+
+    std::string mNamespace;
+};
+
+class BaseCreator : public IPluginCreator
+{
+public:
+    void setPluginNamespace(char const* libNamespace) noexcept override
+    {
+        mNamespace = libNamespace;
+    }
+
+    char const* getPluginNamespace() const noexcept override
+    {
+        return mNamespace.c_str();
+    }
+
+protected:
+    std::string mNamespace;
+};
+
+std::shared_ptr<nvinfer1::pluginInternal::CudnnWrapper> createPluginCudnnWrapper(
+    void* executionContextIdentifier, char const* callerPluginName);
+std::shared_ptr<nvinfer1::pluginInternal::CublasWrapper> createPluginCublasWrapper(void* executionContextIdentifier);
+} // namespace pluginInternal
+
+namespace plugin
+{
+
+// Write values into buffer
+template <typename Type, typename BufferType>
+void write(BufferType*& buffer, Type const& val)
+{
+    static_assert(sizeof(BufferType) == 1, "BufferType must be a 1 byte type.");
+    std::memcpy(buffer, &val, sizeof(Type));
+    buffer += sizeof(Type);
+}
+
+// Read values from buffer
+template <typename OutType, typename BufferType>
+OutType read(BufferType const*& buffer)
+{
+    static_assert(sizeof(BufferType) == 1, "BufferType must be a 1 byte type.");
+    OutType val{};
+    std::memcpy(&val, static_cast<void const*>(buffer), sizeof(OutType));
+    buffer += sizeof(OutType);
+    return val;
+}
+
+inline int32_t getTrtSmVersionDec(int32_t majorVersion, int32_t minorVersion)
+{
+    return majorVersion * 10 + minorVersion;
+}
+
+//! Represents the compute capability of a device.
+//! This pertains to virtual architectures represented by the intermediate PTX format.
+//! This is distinct from the SM version.
+//! See https://forums.developer.nvidia.com/t/how-should-i-use-correctly-the-sm-xx-and-compute-xx/219160
+struct DeviceComputeCapability
+{
+    int32_t major{};
+    int32_t minor{};
+
+    //! \return the compute capability of the CUDA device with the given \p deviceIndex.
+    [[nodiscard]] static DeviceComputeCapability forDevice(int32_t deviceIndex)
+    {
+        int32_t major{0};
+        int32_t minor{0};
+        PLUGIN_CUASSERT(cudaDeviceGetAttribute(&major, cudaDevAttrComputeCapabilityMajor, deviceIndex));
+        PLUGIN_CUASSERT(cudaDeviceGetAttribute(&minor, cudaDevAttrComputeCapabilityMinor, deviceIndex));
+        return {major, minor};
+    }
+};
+
+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);
+}
+
+// Check that all required field names are present in the PluginFieldCollection.
+// If not, throw a PluginError with a message stating which fields are missing.
+void validateRequiredAttributesExist(std::set<std::string> requiredFieldNames, PluginFieldCollection const* fc);
+
+template <typename Dtype>
+struct CudaBind
+{
+    size_t mSize;
+    void* mPtr;
+
+    CudaBind(size_t size)
+    {
+        mSize = size;
+        PLUGIN_CUASSERT(cudaMalloc(&mPtr, sizeof(Dtype) * mSize));
+    }
+
+    ~CudaBind()
+    {
+        if (mPtr != nullptr)
+        {
+            PLUGIN_CUASSERT(cudaFree(mPtr));
+            mPtr = nullptr;
+        }
+    }
+};
+
+// Convert a 64-bit dimension to a 32-bit dimension.
+// Throw exception if it doesn't fit.
+int32_t dimToInt32(int64_t);
+
+// Helper function to determine whether memory pool support is available on the device.
+bool supportsMemPoolsHelper();
+
+// Wrapper function around the helper to keep the result in a static variable to avoid mulitple calls to CUDA APIs.
+bool supportsMemPools();
+} // namespace plugin
+} // namespace nvinfer1
+
+#ifndef DEBUG
+
+#define PLUGIN_CHECK(status)                                                                                           \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        if (status != 0)                                                                                               \
+            exit(EXIT_FAILURE);                                                                                                   \
+    } while (0)
+
+#define ASSERT_PARAM(exp)                                                                                              \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        if (!(exp))                                                                                                    \
+            return STATUS_BAD_PARAM;                                                                                   \
+    } while (0)
+
+#define ASSERT_FAILURE(exp)                                                                                            \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        if (!(exp))                                                                                                    \
+            return STATUS_FAILURE;                                                                                     \
+    } while (0)
+
+#define CSC(call, err)                                                                                                 \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        cudaError_t cudaStatus = call;                                                                                 \
+        if (cudaStatus != cudaSuccess)                                                                                 \
+        {                                                                                                              \
+            return err;                                                                                                \
+        }                                                                                                              \
+    } while (0)
+
+#define DEBUG_PRINTF(...)                                                                                              \
+    do                                                                                                                 \
+    {                                                                                                                  \
+    } while (0)
+
+#else
+
+#define ASSERT_PARAM(exp)                                                                                              \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        if (!(exp))                                                                                                    \
+        {                                                                                                              \
+            fprintf(stderr, "Bad param - " #exp ", %s:%d\n", __FILE__, __LINE__);                                      \
+            return STATUS_BAD_PARAM;                                                                                   \
+        }                                                                                                              \
+    } while (0)
+
+#define ASSERT_FAILURE(exp)                                                                                            \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        if (!(exp))                                                                                                    \
+        {                                                                                                              \
+            fprintf(stderr, "Failure - " #exp ", %s:%d\n", __FILE__, __LINE__);                                        \
+            return STATUS_FAILURE;                                                                                     \
+        }                                                                                                              \
+    } while (0)
+
+#define CSC(call, err)                                                                                                 \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        cudaError_t cudaStatus = call;                                                                                 \
+        if (cudaStatus != cudaSuccess)                                                                                 \
+        {                                                                                                              \
+            printf("%s %d CUDA FAIL %s\n", __FILE__, __LINE__, cudaGetErrorString(cudaStatus));                        \
+            return err;                                                                                                \
+        }                                                                                                              \
+    } while (0)
+
+#define PLUGIN_CHECK(status)                                                                                           \
+    {                                                                                                                  \
+        if (status != 0)                                                                                               \
+        {                                                                                                              \
+            DEBUG_PRINTF("%s %d CUDA FAIL %s\n", __FILE__, __LINE__, cudaGetErrorString(status));                      \
+            exit(EXIT_FAILURE);                                                                                                 \
+        }                                                                                                              \
+    }
+
+#define DEBUG_PRINTF(...)                                                                                              \
+    do                                                                                                                 \
+    {                                                                                                                  \
+        printf(__VA_ARGS__);                                                                                           \
+    } while (0)
+
+#endif // DEBUG
+
+#endif // TRT_PLUGIN_H

plugin/common/reducedMathPlugin.cpp

@@ -0,0 +1,81 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#include <cstdint>
+
+namespace nvinfer1::plugin::detail
+{
+
+// Count leading zeros - start from most significant bit.
+int32_t clz(int32_t x)
+{
+    for (int32_t i = 31; i >= 0; --i)
+    {
+        if ((1U << i) & x)
+        {
+            return 31 - i;
+        }
+    }
+    return 32;
+}
+
+#define CUDNN_IS_POW_2(x) (0 == ((x) & ((x) -1)))
+
+int32_t find_log_2(int32_t x, bool round_up = false)
+{
+    int32_t a = 31 - clz(x);
+    if (round_up)
+    {
+        a += !CUDNN_IS_POW_2(x);
+    }
+    return a;
+}
+
+void findDivisor(int32_t denom, uint32_t& mul_coeff, uint32_t& shift_coeff)
+{
+    if (denom == 0)
+    {
+        return;
+    }
+    if (denom == 1)
+    {
+        // if dividing by 1, reduced math doesn't work because mul_coeff would
+        // need to be 2^32, which doesn't fit into uint32_t.  the div()
+        // routine handles this special case separately.
+        mul_coeff = 0;
+        shift_coeff = 0;
+        return;
+    }
+    // To express the division N/D in terms of a multiplication, what we first
+    // imagine is simply N*(1/D).  However, 1/D will always evaluate to 0 (for D>1),
+    // so we need another way.  There's nothing that says we have to use exactly
+    // the fraction 1/D; instead it could be any X/Y that reduces to 1/D (i.e.,
+    // Y=X*D), or at least to "close enough" to it.  If we pick Y that is a power
+    // of two, then the N*(X/Y) can be N*X followed by a right-shift by some amount.
+    // The power of two we should pick should be at least 2^32, because in the
+    // div() routine we'll use umulhi(), which returns only the upper 32 bits --
+    // this being equivalent to a right-shift by 32.  But we might want a higher
+    // power of two for better accuracy depending on the magnitude of the denominator.
+    // Once we've picked Y, then X [our mul_coeff value] is simply Y/D, rounding up,
+    // and we save shift_coeff as whatever further shift we have to do beyond
+    // what the umulhi() implies.
+    uint32_t p = 31 + find_log_2(denom, true);
+    uint32_t m = ((1ULL << p) + (uint32_t) denom - 1) / (uint32_t) denom;
+    mul_coeff = m;
+    shift_coeff = p - 32;
+}
+
+} // namespace nvinfer1::plugin::detail

plugin/common/scopedCudaStream.h

@@ -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/common/serialize.hpp

@@ -0,0 +1,171 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#pragma once
+
+#include <cstring>
+#include <stdexcept>
+#include <type_traits>
+#include <vector>
+
+#include <iostream>
+using std::cerr;
+using std::cout;
+using std::endl;
+
+template <typename T>
+inline void serialize_value(void** buffer, T const& value);
+
+template <typename T>
+inline void deserialize_value(void const** buffer, size_t* buffer_size, T* value);
+
+namespace
+{
+
+template <typename T, class Enable = void>
+struct Serializer
+{
+};
+
+template <typename T>
+struct Serializer<T, std::enable_if_t<std::is_arithmetic_v<T> || std::is_enum_v<T> || std::is_pod_v<T>>>
+{
+    static size_t serialized_size(T const&)
+    {
+        return sizeof(T);
+    }
+    static void serialize(void** buffer, T const& value)
+    {
+        ::memcpy(*buffer, &value, sizeof(T));
+        reinterpret_cast<char*&>(*buffer) += sizeof(T);
+    }
+    static void deserialize(void const** buffer, size_t* buffer_size, T* value)
+    {
+        if (*buffer_size < sizeof(T))
+        {
+            throw std::runtime_error("Deserialization error: buffer too small for scalar value");
+        }
+        ::memcpy(value, *buffer, sizeof(T));
+        reinterpret_cast<char const*&>(*buffer) += sizeof(T);
+        *buffer_size -= sizeof(T);
+    }
+};
+
+template <>
+struct Serializer<const char*>
+{
+    static size_t serialized_size(const char* value)
+    {
+        return strlen(value) + 1;
+    }
+    static void serialize(void** buffer, const char* value)
+    {
+        ::strcpy(static_cast<char*>(*buffer), value);
+        reinterpret_cast<char*&>(*buffer) += strlen(value) + 1;
+    }
+    static void deserialize(void const** buffer, size_t* buffer_size, char const** value)
+    {
+        *value = static_cast<char const*>(*buffer);
+        size_t const data_size = strnlen(*value, *buffer_size) + 1;
+        if (*buffer_size < data_size)
+        {
+            throw std::runtime_error("Deserialization error: buffer too small for C string");
+        }
+        reinterpret_cast<char const*&>(*buffer) += data_size;
+        *buffer_size -= data_size;
+    }
+};
+
+template <typename T>
+struct Serializer<std::vector<T>, std::enable_if_t<std::is_arithmetic_v<T> || std::is_enum_v<T> || std::is_pod_v<T>>>
+{
+    static size_t serialized_size(std::vector<T> const& value)
+    {
+        return sizeof(value.size()) + value.size() * sizeof(T);
+    }
+    static void serialize(void** buffer, std::vector<T> const& value)
+    {
+        serialize_value(buffer, value.size());
+        size_t nbyte = value.size() * sizeof(T);
+        ::memcpy(*buffer, value.data(), nbyte);
+        reinterpret_cast<char*&>(*buffer) += nbyte;
+    }
+    static void deserialize(void const** buffer, size_t* buffer_size, std::vector<T>* value)
+    {
+        size_t size;
+        deserialize_value(buffer, buffer_size, &size);
+        // Single division-based check covers both integer overflow (size*sizeof(T) wraps)
+        // and out-of-bounds memcpy, and must happen before resize() to prevent DoS.
+        if (size > *buffer_size / sizeof(T))
+        {
+            throw std::runtime_error("Deserialization error: vector size exceeds available buffer");
+        }
+        size_t const nbyte = size * sizeof(T);
+        value->resize(size);
+        ::memcpy(value->data(), *buffer, nbyte);
+        reinterpret_cast<char const*&>(*buffer) += nbyte;
+        *buffer_size -= nbyte;
+    }
+};
+
+template <>
+struct Serializer<std::string>
+{
+    static size_t serialized_size(std::string const& value)
+    {
+        return sizeof(value.size()) + value.size();
+    }
+    static void serialize(void** buffer, std::string const& value)
+    {
+        size_t nbyte = value.size();
+        serialize_value(buffer, nbyte);
+        ::memcpy(*buffer, value.data(), nbyte);
+        reinterpret_cast<char*&>(*buffer) += nbyte;
+    }
+    static void deserialize(void const** buffer, size_t* buffer_size, std::string* value)
+    {
+        size_t nbyte;
+        deserialize_value(buffer, buffer_size, &nbyte);
+        if (nbyte > *buffer_size)
+        {
+            throw std::runtime_error("Deserialization error: string size exceeds available buffer");
+        }
+        value->resize(nbyte);
+        ::memcpy(const_cast<char*>(value->data()), *buffer, nbyte);
+        reinterpret_cast<char const*&>(*buffer) += nbyte;
+        *buffer_size -= nbyte;
+    }
+};
+
+} // namespace
+
+template <typename T>
+inline size_t serialized_size(T const& value)
+{
+    return Serializer<T>::serialized_size(value);
+}
+
+template <typename T>
+inline void serialize_value(void** buffer, T const& value)
+{
+    return Serializer<T>::serialize(buffer, value);
+}
+
+template <typename T>
+inline void deserialize_value(void const** buffer, size_t* buffer_size, T* value)
+{
+    return Serializer<T>::deserialize(buffer, buffer_size, value);
+}

plugin/common/templates.h

@@ -0,0 +1,52 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+#ifndef TRT_PLUGIN_COMMON_TEMPLATES_H_
+#define TRT_PLUGIN_COMMON_TEMPLATES_H_
+namespace nvinfer1
+{
+namespace plugin
+{
+template <typename ToType, typename FromType>
+ToType* toPointer(FromType* ptr)
+{
+    return static_cast<ToType*>(static_cast<void*>(ptr));
+}
+template <typename ToType, typename FromType>
+ToType const* toPointer(FromType const* ptr)
+{
+    return static_cast<ToType const*>(static_cast<void const*>(ptr));
+}
+// Helper function for serializing plugin
+template <typename ValType, typename BufferType>
+void writeToBuffer(BufferType*& buffer, ValType const& val)
+{
+    *toPointer<ValType>(buffer) = val;
+    buffer += sizeof(ValType);
+}
+
+// Helper function for deserializing plugin
+template <typename ValType, typename BufferType>
+ValType readFromBuffer(BufferType const*& buffer)
+{
+    auto val = *toPointer<ValType const>(buffer);
+    buffer += sizeof(ValType);
+    return val;
+}
+
+} // namespace plugin
+} // namespace nvinfer1
+#endif // TRT_PLUGIN_COMMON_TEMPLATES_H_