chore: import upstream snapshot with attribution

plugin/disentangledAttentionPlugin/CMakeLists.txt

@@ -0,0 +1,26 @@
+#
+# 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.
+#
+
+add_plugin_source(
+    disentangledAttentionCommon.h
+    disentangledAttentionPlugin.cpp
+    disentangledAttentionPlugin.h
+    disentangledAttentionPluginLegacy.cpp
+    disentangledAttentionPluginLegacy.h
+    disentangledKernel.cu
+)
+

plugin/disentangledAttentionPlugin/DisentangledAttentionPlugin_PluginConfig.yaml

@@ -0,0 +1,170 @@
+#
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+---
+name: "DisentangledAttention_TRT"
+versions:
+  "2":  # Current version (v3 implementation)
+    interface: "IPluginV3"  # Interface type for the new version
+    inputs:
+      - data0
+      - data1
+      - data2
+    outputs:
+      - output
+    attributes:
+      - span
+      - factor
+    attribute_types:
+      span: int32
+      factor: float32
+    attribute_length:
+      span: 1
+      factor: 1
+    attribute_options:
+      span:
+        min: "1"
+        max: "=pinf"
+      factor:
+        min: "0.0"
+        max: "=pinf"
+    attributes_required:
+      - span
+      - factor
+    golden_io_path: "plugin/DisentangledAttention_PluginGoldenIO.json"
+    abs_tol: 1e-5
+    rel_tol: 1e-5
+    fp16_atol: 1e-2
+    fp16_rtol: 1e-2
+    configs:
+      config1:
+        input_types:
+          data0: float32
+          data1: float32
+          data2: float32
+        attribute_options:
+          "span":
+            value: 4
+            shape: "1"
+          "factor":
+            value: 0.1
+            shape: "1"
+        output_types:
+          output: float32
+      config2:
+        input_types:
+          data0: float32
+          data1: float32
+          data2: float32
+        attribute_options:
+          "span":
+            value: 8
+            shape: "1"
+          "factor":
+            value: 0.05
+            shape: "1"
+        output_types:
+          output: float32
+      config_fp16:
+        input_types:
+          data0: float16
+          data1: float16
+          data2: float16
+        attribute_options:
+          "span":
+            value: 4
+            shape: "1"
+          "factor":
+            value: 0.1
+            shape: "1"
+        output_types:
+          output: float16
+  "1":  # Legacy version (v2 implementation)
+    interface: "IPluginV2DynamicExt"  # Interface type for the new version
+    inputs:
+      - data0
+      - data1
+      - data2
+    outputs:
+      - output
+    attributes:
+      - span
+      - factor
+    attribute_types:
+      span: int32
+      factor: float32
+    attribute_length:
+      span: 1
+      factor: 1
+    attribute_options:
+      span:
+        min: "1"
+        max: "=pinf"
+      factor:
+        min: "0.0"
+        max: "=pinf"
+    attributes_required:
+      - span
+      - factor
+    golden_io_path: "plugin/DisentangledAttention_PluginGoldenIO.json"
+    abs_tol: 1e-5
+    rel_tol: 1e-5
+    fp16_atol: 1e-2
+    fp16_rtol: 1e-2
+    configs:
+      config1:
+        input_types:
+          data0: float32
+          data1: float32
+          data2: float32
+        attribute_options:
+          "span":
+            value: 4
+            shape: "1"
+          "factor":
+            value: 0.1
+            shape: "1"
+        output_types:
+          output: float32
+      config2:
+        input_types:
+          data0: float32
+          data1: float32
+          data2: float32
+        attribute_options:
+          "span":
+            value: 8
+            shape: "1"
+          "factor":
+            value: 0.05
+            shape: "1"
+        output_types:
+          output: float32
+      config_fp16:
+        input_types:
+          data0: float16
+          data1: float16
+          data2: float16
+        attribute_options:
+          "span":
+            value: 4
+            shape: "1"
+          "factor":
+            value: 0.1
+            shape: "1"
+        output_types:
+          output: float16

plugin/disentangledAttentionPlugin/README.md

@@ -0,0 +1,83 @@
+# Disentangled Attention Plugin
+
+- [Disentangled Attention Plugin](#disentangled-attention-plugin)
+  - [Description](#description)
+  - [Structure](#structure)
+    - [Input(s)](#inputs)
+    - [Output(s)](#outputs)
+    - [Parameters](#parameters)
+  - [Additional Resources](#additional-resources)
+  - [License](#license)
+  - [Changelog](#changelog)
+
+## Description
+
+> NOTE: Version 1 of this plugin (using IPluginV2DynamicExt interface) is deprecated since TensorRT 10.11. Version 2 (using IPluginV3 interface) is the recommended replacement.
+
+This TensorRT plugin implements an efficient algorithm to perform the calculation of disentangled attention matrices for DeBERTa-variant types of Transformers.
+
+Unlike [BERT](https://arxiv.org/abs/1810.04805) where each word is represented by one vector that sums the content embedding and position embedding, [DeBERTa](https://arxiv.org/abs/2006.03654) design first proposed the concept of disentangled attention, which uses two vectors to encode content and position respectively and forms attention weights by summing disentangled matrices. Performance gap has been identified between the new attention scheme and the original self-attention, mainly due to extra indexing and gather operations. Major optimizations implemented in this plugin includes: (i) fusion of gather and pointwise operations (ii) utilizing the pattern of relative position matrix and shortcuting out-of-boundary index calculation (iii) parallel index calculation (iv) log tables for relative position index calculation (used for DeBERTa-V2, enabling capture of long-range dependencies without significantly increasing the number of position embeddings).
+
+This TensorRT plugin is primarily intended to be used together with DeBERTa network (with HuggingFace [DeBERTa](https://huggingface.co/docs/transformers/model_doc/deberta) and [DeBERTa-V2](https://huggingface.co/docs/transformers/model_doc/deberta-v2) implementation), but also applies to generic architectures that adopt disentangled attention.
+
+## Structure
+This plugin works for network with graph node named `DisentangledAttention_TRT`. The corresponding graph modification script can be found under the `demo/DeBERTa` folder of TensorRT OSS.
+
+### Input(s)
+This plugin takes three inputs:
+
+* `data0`: Content-to-content ("c2c") Attention Matrix
+
+  > **Input Shape:** `[batch_size*number_heads, sequence_length, sequence_length]`
+  >
+  > **Data Type:** `float32` or `float16` or `int8`
+
+  This is the content-to-content attention, Q<sub>c</sub>K<sub>c</sub><sup>T</sup>, which is essentially the BERT self-attention.
+
+* `data1`: Content-to-position ("c2p") Attention Matrix
+
+  > **Input Shape:** `[batch_size*number_heads, sequence_length, relative_distance*2]`
+  >
+  > **Data Type:** `float32` or `float16` or `int8`
+
+  This is the content-to-position attention, Q<sub>c</sub>K<sub>r</sub><sup>T</sup>.
+
+* `data2`: Position-to-content ("p2c") Attention Matrix
+
+  > **Input Shape:** `[batch_size*number_heads, sequence_length,  relative_distance*2]`
+  >
+  > **Data Type:** `float32` or `float16` or `int8`
+
+   This is the position-to-content attention, K<sub>c</sub>Q<sub>r</sub><sup>T</sup>. Relative distance is the distance span `k` for disentangled attention.
+
+### Output(s)
+This plugin generates one output.
+
+* `result`: Disentangled Attention Matrix
+
+  > **Input Shape:** `[batch_size*number_heads, sequence_length, sequence_length]`
+  >
+  > **Data Type:** `float32` or `float16` or `int8`
+
+  This is the disentangled attention matrix after applying the scaling factor.
+
+### Parameters
+| Type     | Parameter                | Description
+|----------|--------------------------|--------------------------------------------------------
+|`int`   |`span`      | The distance span `k` for relative position. Scalar.
+|`float`   |`factor`           | The scaling factor (multiply) to be applied to the attention weights, `1/sqrt(3d)`, where `d` is hidden size per head `H/N`. `H` is hidden size, `N` is number of heads. Scalar (less than 1).
+
+## Additional Resources
+- [BERT](https://arxiv.org/abs/1810.04805)
+- [DeBERTa](https://arxiv.org/abs/2006.03654)
+- [DeBERTa HuggingFace Implementation](https://github.com/huggingface/transformers/tree/main/src/transformers/models/deberta)
+- [DeBERTa-V2 HuggingFace Implementation](https://github.com/huggingface/transformers/tree/main/src/transformers/models/deberta_v2)
+
+## License
+For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html)
+documentation.
+
+## Changelog
+- 2024.03: Migrated to IPluginV3 interface. The legacy plugin (version 1) using IPluginV2DynamicExt interface is maintained for backward compatibility.
+- 2022.07: Added log bucket for the relative position index calculation (since DeBERTa V2).
+- 2022.04: This is the first release of this `README` file.

plugin/disentangledAttentionPlugin/disentangledAttentionCommon.h

@@ -0,0 +1,48 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef TRT_DISENTANGLED_ATTENTION_COMMON_H
+#define TRT_DISENTANGLED_ATTENTION_COMMON_H
+
+#include "NvInferPlugin.h"
+#include <cstdint>
+
+namespace nvinfer1
+{
+namespace plugin
+{
+
+// Version 1: regular relative position index
+// Version 2: log bucket relative position index
+#define kDISENTANGLED_VERSION 2
+#if kDISENTANGLED_VERSION == 1
+constexpr int32_t kDISENTANGLED_TILESIZE = 32;
+constexpr int32_t kDISENTANGLED_BLOCKDIMY = 8;
+#elif kDISENTANGLED_VERSION == 2
+constexpr int32_t kDISENTANGLED_TILESIZE = 64;
+constexpr int32_t kDISENTANGLED_BLOCKDIMY = 4;
+#endif
+
+template <typename TDataType, int32_t tTileSize, int32_t tBlockDimY>
+void disentangled_kernel_wrapper(TDataType const* data0, TDataType const* data1, TDataType const* data2,
+    TDataType* result, dim3 dimData0, dim3 dimData1, dim3 dimData2, dim3 dimResult, TDataType factor, int32_t span,
+    dim3 block, dim3 grid, cudaStream_t stream);
+
+} // namespace plugin
+} // namespace nvinfer1
+
+#endif // TRT_DISENTANGLED_ATTENTION_COMMON_H

plugin/disentangledAttentionPlugin/disentangledAttentionPlugin.cpp

@@ -0,0 +1,422 @@
+/*
+ * 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 "disentangledAttentionPlugin.h"
+#include "NvInferPlugin.h"
+#include <cuda_fp16.h>
+#include <memory>
+#include <numeric>
+#include <optional>
+#include <stdexcept>
+#include <string_view>
+
+using namespace nvinfer1;
+using nvinfer1::plugin::DisentangledAttentionPlugin;
+using nvinfer1::plugin::DisentangledAttentionPluginCreator;
+
+REGISTER_TENSORRT_PLUGIN(DisentangledAttentionPluginCreator);
+
+namespace
+{
+constexpr char const* kDEBERTA_PLUGIN_NAME{"DisentangledAttention_TRT"};
+constexpr char const* kDEBERTA_PLUGIN_VERSION{"2"};
+} // namespace
+
+DisentangledAttentionPlugin::DisentangledAttentionPlugin()
+    : mSpan(0)
+    , mFactor(0.0f)
+{
+}
+
+DisentangledAttentionPlugin::DisentangledAttentionPlugin(int32_t span, float factor)
+    : mSpan(span)
+    , mFactor(factor)
+{
+}
+
+// IPluginV3OneCore methods
+
+int32_t DisentangledAttentionPlugin::getNbOutputs() const noexcept
+{
+    return 1;
+}
+
+char const* DisentangledAttentionPlugin::getPluginName() const noexcept
+{
+    return kDEBERTA_PLUGIN_NAME;
+}
+
+char const* DisentangledAttentionPlugin::getPluginVersion() const noexcept
+{
+    return kDEBERTA_PLUGIN_VERSION;
+}
+
+IPluginV3* DisentangledAttentionPlugin::clone() noexcept
+{
+    try
+    {
+        auto plugin = std::make_unique<DisentangledAttentionPlugin>(mSpan, mFactor);
+        plugin->setPluginNamespace(mNamespace.c_str());
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+void DisentangledAttentionPlugin::setPluginNamespace(char const* pluginNamespace) noexcept
+{
+    try
+    {
+        mNamespace = pluginNamespace;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+char const* DisentangledAttentionPlugin::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+IPluginCapability* DisentangledAttentionPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept
+{
+    try
+    {
+        if (type == PluginCapabilityType::kBUILD)
+        {
+            return static_cast<IPluginV3OneBuild*>(this);
+        }
+        if (type == PluginCapabilityType::kRUNTIME)
+        {
+            return static_cast<IPluginV3OneRuntime*>(this);
+        }
+        PLUGIN_ASSERT(type == PluginCapabilityType::kCORE);
+        return static_cast<IPluginV3OneCore*>(this);
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+PluginFieldCollection const* DisentangledAttentionPlugin::getFieldsToSerialize() noexcept
+{
+    try
+    {
+        mDataToSerialize.clear();
+
+        mDataToSerialize.emplace_back("span", &mSpan, PluginFieldType::kINT32, 1);
+        mDataToSerialize.emplace_back("factor", &mFactor, PluginFieldType::kFLOAT32, 1);
+
+        mFCToSerialize.nbFields = mDataToSerialize.size();
+        mFCToSerialize.fields = mDataToSerialize.data();
+
+        return &mFCToSerialize;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+// IPluginV3OneBuild methods
+
+int32_t DisentangledAttentionPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
+    DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
+    IExprBuilder& exprBuilder) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputs != nullptr);
+        PLUGIN_VALIDATE(nbInputs == 3);
+        PLUGIN_VALIDATE(outputs != nullptr);
+        PLUGIN_VALIDATE(nbOutputs == 1);
+
+        // Output has the same shape as the first input
+        outputs[0] = inputs[0];
+
+        return STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return STATUS_FAILURE;
+}
+
+int32_t DisentangledAttentionPlugin::configurePlugin(
+    DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(in != nullptr && out != nullptr && nbInputs == 3 && nbOutputs == 1);
+
+        // Validate input and output shapes
+        for (int32_t i = 0; i < nbInputs; i++)
+        {
+            PLUGIN_VALIDATE(in[i].desc.dims.nbDims == in[0].desc.dims.nbDims);
+        }
+
+        // Check data types are consistent
+        PLUGIN_VALIDATE(in[0].desc.type == in[1].desc.type && in[0].desc.type == in[2].desc.type);
+        PLUGIN_VALIDATE(out[0].desc.type == in[0].desc.type);
+
+        return STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return STATUS_FAILURE;
+}
+
+int32_t DisentangledAttentionPlugin::getOutputDataTypes(
+    DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputTypes != nullptr && outputTypes != nullptr);
+        PLUGIN_VALIDATE(nbInputs == 3 && nbOutputs == 1);
+
+        // Output has the same data type as the first input
+        outputTypes[0] = inputTypes[0];
+
+        return STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return STATUS_FAILURE;
+}
+
+bool DisentangledAttentionPlugin::supportsFormatCombination(
+    int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
+{
+    try
+    {
+        PLUGIN_ASSERT(inOut && pos < (nbInputs + nbOutputs));
+
+        // All inputs and outputs should have the same precision type
+        bool const consistentFloatPrecision = (inOut[pos].desc.type == inOut[0].desc.type);
+
+        return (inOut[pos].desc.type == DataType::kINT8 || inOut[pos].desc.type == DataType::kHALF
+                   || inOut[pos].desc.type == DataType::kFLOAT)
+            && inOut[pos].desc.format == PluginFormat::kLINEAR && consistentFloatPrecision;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return false;
+}
+
+// IPluginV3OneRuntime methods
+
+template <typename TDataType>
+void DisentangledAttentionPlugin::enqueueType(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
+    void const* const* inputs, void* const* outputs, cudaStream_t stream, TDataType factor)
+{
+    Dims dims0 = inputDesc[0].dims;
+    Dims dims1 = inputDesc[1].dims;
+    Dims dims2 = inputDesc[2].dims;
+    dim3 dimData0(dims0.d[0], dims0.d[1], dims0.d[2]);
+    dim3 dimData1(dims1.d[0], dims1.d[1], dims1.d[2]);
+    dim3 dimData2(dims2.d[0], dims2.d[1], dims2.d[2]);
+    dim3 dimResult(dimData0);
+
+    dim3 blockOptimized(kDISENTANGLED_TILESIZE, kDISENTANGLED_BLOCKDIMY);
+    dim3 gridOptimized(
+        (dimResult.z - 1) / kDISENTANGLED_TILESIZE + 1, (dimResult.y - 1) / kDISENTANGLED_TILESIZE + 1, dimResult.x);
+
+    auto const* data0 = static_cast<TDataType const*>(inputs[0]);
+    auto const* data1 = static_cast<TDataType const*>(inputs[1]);
+    auto const* data2 = static_cast<TDataType const*>(inputs[2]);
+    auto* result = static_cast<TDataType*>(outputs[0]);
+    disentangled_kernel_wrapper<TDataType, kDISENTANGLED_TILESIZE, kDISENTANGLED_BLOCKDIMY>(data0, data1, data2, result,
+        dimData0, dimData1, dimData2, dimResult, factor, mSpan, blockOptimized, gridOptimized, stream);
+}
+
+int32_t DisentangledAttentionPlugin::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
+    void const* const* inputs, void* const* outputs, void* /* workspace */, cudaStream_t stream) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr);
+
+        switch (inputDesc[0].type)
+        {
+        case DataType::kFLOAT: enqueueType<float>(inputDesc, outputDesc, inputs, outputs, stream, mFactor); break;
+        case DataType::kHALF:
+            enqueueType<__half>(inputDesc, outputDesc, inputs, outputs, stream, __float2half(mFactor));
+            break;
+        case DataType::kINT8:
+            enqueueType<int8_t>(inputDesc, outputDesc, inputs, outputs, stream, static_cast<int8_t>(mFactor));
+            break;
+        default: PLUGIN_VALIDATE(false, "Unsupported Datatype"); break;
+        }
+        return cudaPeekAtLastError();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+        return STATUS_FAILURE;
+    }
+}
+
+size_t DisentangledAttentionPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
+{
+    return 0;
+}
+
+int32_t DisentangledAttentionPlugin::onShapeChange(
+    PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputs != nullptr && outputs != nullptr);
+        PLUGIN_VALIDATE(nbInputs == 3 && nbOutputs == 1);
+
+        // Check that all inputs have the same data type
+        DataType dataType = inputs[0].type;
+        PLUGIN_VALIDATE(inputs[1].type == dataType && inputs[2].type == dataType);
+
+        // Check that output has the same data type
+        PLUGIN_VALIDATE(outputs[0].type == dataType);
+
+        // Validate dimensions
+        PLUGIN_VALIDATE(inputs[0].dims.nbDims == inputs[1].dims.nbDims);
+        PLUGIN_VALIDATE(inputs[0].dims.nbDims == inputs[2].dims.nbDims);
+        PLUGIN_VALIDATE(outputs[0].dims.nbDims == inputs[0].dims.nbDims);
+
+        return STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return STATUS_FAILURE;
+}
+
+IPluginV3* DisentangledAttentionPlugin::attachToContext(IPluginResourceContext* context) noexcept
+{
+    try
+    {
+        return this->clone();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+// -------------------- Creator class Implementation --------------------
+
+DisentangledAttentionPluginCreator::DisentangledAttentionPluginCreator()
+{
+    mPluginAttributes.clear();
+    mPluginAttributes.emplace_back(PluginField("span", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("factor", nullptr, PluginFieldType::kFLOAT32, 1));
+
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+char const* DisentangledAttentionPluginCreator::getPluginName() const noexcept
+{
+    return kDEBERTA_PLUGIN_NAME;
+}
+
+char const* DisentangledAttentionPluginCreator::getPluginVersion() const noexcept
+{
+    return kDEBERTA_PLUGIN_VERSION;
+}
+
+PluginFieldCollection const* DisentangledAttentionPluginCreator::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+IPluginV3* DisentangledAttentionPluginCreator::createPlugin(
+    char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
+{
+    using namespace std::string_view_literals;
+    try
+    {
+        PLUGIN_VALIDATE(fc != nullptr);
+        PluginField const* fields = fc->fields;
+        std::optional<int32_t> span;
+        std::optional<float> factor;
+
+        for (int32_t i = 0; i < fc->nbFields; ++i)
+        {
+            std::string_view const attrName = fields[i].name;
+            if (attrName == "span"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kINT32);
+                span = *static_cast<int32_t const*>(fields[i].data);
+            }
+            else if (attrName == "factor"sv)
+            {
+                PLUGIN_VALIDATE(fields[i].type == PluginFieldType::kFLOAT32);
+                factor = *static_cast<float const*>(fields[i].data);
+            }
+        }
+
+        // Validate that all required fields were found
+        PLUGIN_VALIDATE(span.has_value(), "Required attribute 'span' not found");
+        PLUGIN_VALIDATE(factor.has_value(), "Required attribute 'factor' not found");
+        PLUGIN_VALIDATE(span.value() >= 0);
+        PLUGIN_VALIDATE(
+            factor.value() > 0.F && factor.value() < 1.F); // factor is 1/sqrt(3d), therefore must less than 1
+
+        auto plugin = std::make_unique<DisentangledAttentionPlugin>(span.value(), factor.value());
+        plugin->setPluginNamespace(mNamespace.c_str());
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+void DisentangledAttentionPluginCreator::setPluginNamespace(char const* pluginNamespace) noexcept
+{
+    try
+    {
+        mNamespace = pluginNamespace;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+char const* DisentangledAttentionPluginCreator::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}

plugin/disentangledAttentionPlugin/disentangledAttentionPlugin.h

@@ -0,0 +1,131 @@
+/*
+ * 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_DISENTANGLED_ATTENTION_PLUGIN_H
+#define TRT_DISENTANGLED_ATTENTION_PLUGIN_H
+
+#include "NvInferPlugin.h"
+#include "common/plugin.h"
+#include "common/serialize.hpp"
+#include "disentangledAttentionCommon.h"
+#include <cstdint>
+#include <iostream>
+#include <string>
+#include <vector>
+
+// One of the preferred ways of making TensorRT to be able to see
+// our custom layer requires extending IPluginV2 and IPluginCreator classes.
+// For requirements for overriden functions, check TensorRT API docs.
+namespace nvinfer1
+{
+namespace plugin
+{
+
+// using namespace nvinfer1;
+
+class DisentangledAttentionPlugin : public nvinfer1::IPluginV3,
+                                    public nvinfer1::IPluginV3OneCore,
+                                    public nvinfer1::IPluginV3OneBuild,
+                                    public nvinfer1::IPluginV3OneRuntime
+{
+public:
+    DisentangledAttentionPlugin();
+
+    DisentangledAttentionPlugin(int32_t span, float factor);
+
+    DisentangledAttentionPlugin(void const* serialData, size_t serialLength);
+
+    // Destructor
+    virtual ~DisentangledAttentionPlugin() {}
+
+    // IPluginV3OneCore methods
+    int32_t getNbOutputs() const noexcept override;
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+    char const* getPluginNamespace() const noexcept override;
+    char const* getPluginName() const noexcept override;
+    char const* getPluginVersion() const noexcept override;
+    nvinfer1::IPluginV3* clone() noexcept override;
+    nvinfer1::PluginFieldCollection const* getFieldsToSerialize() noexcept override;
+    nvinfer1::IPluginCapability* getCapabilityInterface(nvinfer1::PluginCapabilityType type) noexcept override;
+
+    // IPluginV3OneBuild methods
+    int32_t getOutputShapes(nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::DimsExprs const* shapeInputs,
+        int32_t nbShapeInputs, nvinfer1::DimsExprs* outputs, int32_t nbOutputs,
+        nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+    int32_t configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
+        nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
+    int32_t getOutputDataTypes(nvinfer1::DataType* outputTypes, int32_t nbOutputs, nvinfer1::DataType const* inputTypes,
+        int32_t nbInputs) const noexcept override;
+    bool supportsFormatCombination(int32_t pos, nvinfer1::DynamicPluginTensorDesc const* inOut, int32_t nbInputs,
+        int32_t nbOutputs) noexcept override;
+
+    // IPluginV3OneRuntime methods
+    int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
+        void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
+    size_t getWorkspaceSize(nvinfer1::DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+        nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
+    int32_t onShapeChange(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
+        nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept override;
+    nvinfer1::IPluginV3* attachToContext(nvinfer1::IPluginResourceContext* context) noexcept override;
+
+private:
+    // Helper method for enqueue()
+    template <typename TDataType>
+    void enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
+        void const* const* inputs, void* const* outputs, cudaStream_t stream, TDataType factor);
+
+    std::string mNamespace;
+
+    // attributes
+    int32_t mSpan;
+    float mFactor;
+
+    // Field serialization storage
+    std::vector<nvinfer1::PluginField> mDataToSerialize;
+    nvinfer1::PluginFieldCollection mFCToSerialize;
+};
+
+class DisentangledAttentionPluginCreator : public nvinfer1::IPluginCreatorV3One
+{
+public:
+    DisentangledAttentionPluginCreator();
+
+    ~DisentangledAttentionPluginCreator() override = default;
+
+    char const* getPluginName() const noexcept override;
+
+    char const* getPluginVersion() const noexcept override;
+
+    nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
+
+    nvinfer1::IPluginV3* createPlugin(
+        char const* name, nvinfer1::PluginFieldCollection const* fc, nvinfer1::TensorRTPhase phase) noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept;
+
+    char const* getPluginNamespace() const noexcept override;
+
+private:
+    nvinfer1::PluginFieldCollection mFC;
+    std::vector<nvinfer1::PluginField> mPluginAttributes;
+    std::string mNamespace;
+};
+
+} // namespace plugin
+} // namespace nvinfer1
+
+#endif // TRT_DISENTANGLED_ATTENTION_PLUGIN_H

plugin/disentangledAttentionPlugin/disentangledAttentionPluginLegacy.cpp

@@ -0,0 +1,375 @@
+/*
+ * 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.
+ */
+
+/*
+ * Legacy version of the plugin maintained for backward compatibility.
+ * This implementation is based on IPluginV2 interfaces.
+ */
+#include "disentangledAttentionPluginLegacy.h"
+#include "NvInferPlugin.h"
+#include <cuda_fp16.h>
+#include <memory>
+#include <numeric>
+#include <stdexcept>
+#include <string_view>
+
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+
+REGISTER_TENSORRT_PLUGIN(DisentangledAttentionPluginCreatorLegacy);
+
+namespace
+{
+using namespace std::string_view_literals;
+constexpr char const* kDEBERTA_PLUGIN_NAME{"DisentangledAttention_TRT"};
+constexpr char const* kDEBERTA_PLUGIN_VERSION{"1"};
+} // namespace
+
+DisentangledAttentionPluginLegacy::DisentangledAttentionPluginLegacy() {}
+
+DisentangledAttentionPluginLegacy::DisentangledAttentionPluginLegacy(int32_t span, float factor)
+    : mSpan(span)
+    , mFactor(factor)
+{
+}
+
+DisentangledAttentionPluginLegacy::DisentangledAttentionPluginLegacy(void const* serialData, size_t serialLength)
+{
+    // Deserialize in the same order as serialization
+    deserialize_value(&serialData, &serialLength, &mSpan);
+    deserialize_value(&serialData, &serialLength, &mFactor);
+}
+
+int32_t DisentangledAttentionPluginLegacy::getNbOutputs() const noexcept
+{
+    return 1;
+}
+
+int32_t DisentangledAttentionPluginLegacy::initialize() noexcept
+{
+    return 0;
+}
+
+char const* DisentangledAttentionPluginLegacy::getPluginType() const noexcept
+{
+    return kDEBERTA_PLUGIN_NAME;
+}
+
+char const* DisentangledAttentionPluginLegacy::getPluginVersion() const noexcept
+{
+    return kDEBERTA_PLUGIN_VERSION;
+}
+
+// IPluginV2DynamicExt Methods
+nvinfer1::DimsExprs DisentangledAttentionPluginLegacy::getOutputDimensions(
+    int32_t index, nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputs != nullptr);
+        PLUGIN_VALIDATE(index == 0); // Only one output
+        return inputs[0];
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nvinfer1::DimsExprs{};
+}
+
+template <typename TDataType>
+void DisentangledAttentionPluginLegacy::enqueueType(nvinfer1::PluginTensorDesc const* inputDesc,
+    nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, cudaStream_t stream,
+    TDataType factor)
+{
+    nvinfer1::Dims dims0 = inputDesc[0].dims;
+    nvinfer1::Dims dims1 = inputDesc[1].dims;
+    nvinfer1::Dims dims2 = inputDesc[2].dims;
+    dim3 dimData0(dims0.d[0], dims0.d[1], dims0.d[2]);
+    dim3 dimData1(dims1.d[0], dims1.d[1], dims1.d[2]);
+    dim3 dimData2(dims2.d[0], dims2.d[1], dims2.d[2]);
+    dim3 dimResult(dimData0);
+
+    dim3 blockOptimized(kDISENTANGLED_TILESIZE, kDISENTANGLED_BLOCKDIMY);
+    dim3 gridOptimized(
+        (dimResult.z - 1) / kDISENTANGLED_TILESIZE + 1, (dimResult.y - 1) / kDISENTANGLED_TILESIZE + 1, dimResult.x);
+
+    auto const* data0 = static_cast<TDataType const*>(inputs[0]);
+    auto const* data1 = static_cast<TDataType const*>(inputs[1]);
+    auto const* data2 = static_cast<TDataType const*>(inputs[2]);
+    auto* result = static_cast<TDataType*>(outputs[0]);
+    disentangled_kernel_wrapper<TDataType, kDISENTANGLED_TILESIZE, kDISENTANGLED_BLOCKDIMY>(data0, data1, data2, result,
+        dimData0, dimData1, dimData2, dimResult, factor, mSpan, blockOptimized, gridOptimized, stream);
+}
+
+int32_t DisentangledAttentionPluginLegacy::enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
+    nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs,
+    void* /* workspace */, cudaStream_t stream) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr);
+
+        switch (inputDesc[0].type)
+        {
+        case nvinfer1::DataType::kFLOAT:
+            enqueueType<float>(inputDesc, outputDesc, inputs, outputs, stream, mFactor);
+            break;
+        case nvinfer1::DataType::kHALF:
+            enqueueType<__half>(inputDesc, outputDesc, inputs, outputs, stream, __float2half(mFactor));
+            break;
+        case nvinfer1::DataType::kINT8:
+            enqueueType<int8_t>(inputDesc, outputDesc, inputs, outputs, stream, static_cast<int8_t>(mFactor));
+            break;
+        default: PLUGIN_VALIDATE(false, "Unsupported Datatype"); break;
+        }
+        return cudaPeekAtLastError();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+        return STATUS_FAILURE;
+    }
+}
+
+size_t DisentangledAttentionPluginLegacy::getSerializationSize() const noexcept
+{
+    return sizeof(mSpan) + sizeof(mFactor);
+}
+
+void DisentangledAttentionPluginLegacy::serialize(void* buffer) const noexcept
+{
+    serialize_value(&buffer, mSpan);
+    serialize_value(&buffer, mFactor);
+}
+
+bool DisentangledAttentionPluginLegacy::supportsFormatCombination(
+    int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
+{
+
+    PLUGIN_ASSERT(inOut && pos < (nbInputs + nbOutputs));
+
+    bool const consistentFloatPrecision
+        = (inOut[pos].type == inOut[0].type); // all inputs & outputs should have the same precision type
+
+    return (inOut[pos].type == nvinfer1::DataType::kINT8 || inOut[pos].type == nvinfer1::DataType::kHALF
+               || inOut[pos].type == nvinfer1::DataType::kFLOAT)
+        && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR && consistentFloatPrecision;
+}
+
+void DisentangledAttentionPluginLegacy::terminate() noexcept {}
+
+void DisentangledAttentionPluginLegacy::destroy() noexcept
+{
+    // This gets called when the network containing plugin is destroyed
+    delete this;
+}
+
+IPluginV2DynamicExt* DisentangledAttentionPluginLegacy::clone() const noexcept
+{
+    try
+    {
+        auto plugin = std::make_unique<DisentangledAttentionPluginLegacy>(mSpan, mFactor);
+        plugin->setPluginNamespace(mNamespace.c_str());
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+void DisentangledAttentionPluginLegacy::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
+    nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
+{
+    try
+    {
+        // inputs
+        PLUGIN_VALIDATE(nbInputs == 3); // 3 inputs
+
+        // check for valid input dimensions
+        PLUGIN_VALIDATE(in[0].desc.dims.nbDims == 3);
+        PLUGIN_VALIDATE(in[1].desc.dims.nbDims == 3);
+        PLUGIN_VALIDATE(in[2].desc.dims.nbDims == 3);
+
+        // check BN (batch_size * num_heads) dimension consistency
+        PLUGIN_VALIDATE(in[0].desc.dims.d[0] == in[1].desc.dims.d[0]);
+        PLUGIN_VALIDATE(in[0].desc.dims.d[0] == in[2].desc.dims.d[0]);
+
+        // check S (sequence_length) dimension consistency
+        PLUGIN_VALIDATE(in[0].desc.dims.d[1] == in[1].desc.dims.d[1]);
+        PLUGIN_VALIDATE(in[0].desc.dims.d[1] == in[2].desc.dims.d[1]);
+        PLUGIN_VALIDATE(in[0].desc.dims.d[1] == in[0].desc.dims.d[2]);
+
+        // check K (2 * span) dimension consistency for in[1] and in[2]
+        PLUGIN_VALIDATE(in[1].desc.dims.d[2] == 2 * mSpan);
+        PLUGIN_VALIDATE(in[2].desc.dims.d[2] == 2 * mSpan);
+
+        // Outputs (same dimension as in[0])
+        PLUGIN_VALIDATE(nbOutputs == 1);
+        PLUGIN_VALIDATE(out[0].desc.dims.nbDims == 3);
+        PLUGIN_VALIDATE(in[0].desc.dims.d[0] == out[0].desc.dims.d[0]);
+        PLUGIN_VALIDATE(in[0].desc.dims.d[1] == out[0].desc.dims.d[1]);
+        PLUGIN_VALIDATE(in[0].desc.dims.d[2] == out[0].desc.dims.d[2]);
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+nvinfer1::DataType DisentangledAttentionPluginLegacy::getOutputDataType(
+    int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputTypes != nullptr);
+        PLUGIN_VALIDATE(nbInputs > 0);
+        PLUGIN_VALIDATE(index == 0);
+        return inputTypes[0]; // version 1, same as data1; version 2, same as data0
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nvinfer1::DataType{};
+}
+
+size_t DisentangledAttentionPluginLegacy::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
+    nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
+{
+    return 0;
+}
+
+void DisentangledAttentionPluginLegacy::setPluginNamespace(char const* libNamespace) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(libNamespace != nullptr);
+        mNamespace = libNamespace;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+char const* DisentangledAttentionPluginLegacy::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+DisentangledAttentionPluginCreatorLegacy::DisentangledAttentionPluginCreatorLegacy()
+{
+    mPluginAttributes.clear();
+
+    // consistent with the ONNX model attr fields
+    mPluginAttributes.emplace_back(PluginField("span", nullptr, PluginFieldType::kINT32, 1));
+    mPluginAttributes.emplace_back(PluginField("factor", nullptr, PluginFieldType::kFLOAT32, 1));
+
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+char const* DisentangledAttentionPluginCreatorLegacy::getPluginName() const noexcept
+{
+    return kDEBERTA_PLUGIN_NAME;
+}
+
+char const* DisentangledAttentionPluginCreatorLegacy::getPluginVersion() const noexcept
+{
+    return kDEBERTA_PLUGIN_VERSION;
+}
+
+PluginFieldCollection const* DisentangledAttentionPluginCreatorLegacy::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+char const* DisentangledAttentionPluginCreatorLegacy::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+void DisentangledAttentionPluginCreatorLegacy::setPluginNamespace(char const* libNamespace) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(libNamespace != nullptr);
+        mNamespace = libNamespace;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+IPluginV2DynamicExt* DisentangledAttentionPluginCreatorLegacy::createPlugin(
+    char const* /*name*/, PluginFieldCollection const* fc) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(fc != nullptr);
+
+        // Set default invalid values (for assert in case when attributes are missing)
+        int32_t span = 0;
+        float factor = 0.F;
+        for (int32_t i = 0; i < fc->nbFields; i++)
+        {
+            std::string_view const fieldName = fc->fields[i].name;
+            if (fieldName == "span"sv)
+            {
+                span = *static_cast<int32_t const*>(fc->fields[i].data);
+            }
+            if (fieldName == "factor"sv)
+            {
+                factor = *static_cast<float const*>(fc->fields[i].data);
+            }
+        }
+
+        PLUGIN_VALIDATE(span >= 0);
+        PLUGIN_VALIDATE(factor > 0.F && factor < 1.F); // factor is 1/sqrt(3d), therefore must less than 1
+
+        auto plugin = std::make_unique<DisentangledAttentionPluginLegacy>(span, factor);
+        plugin->setPluginNamespace(mNamespace.c_str());
+
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+IPluginV2DynamicExt* DisentangledAttentionPluginCreatorLegacy::deserializePlugin(
+    char const* /*name*/, void const* serialData, size_t serialLength) noexcept
+{
+    try
+    {
+        auto plugin = std::make_unique<DisentangledAttentionPluginLegacy>(serialData, serialLength);
+        plugin->setPluginNamespace(mNamespace.c_str());
+
+        return plugin.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}

plugin/disentangledAttentionPlugin/disentangledAttentionPluginLegacy.h

@@ -0,0 +1,144 @@
+/*
+ * 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 DISENTANGLEDATTENTIONPLUGIN_LEGACY_PLUGIN_H
+#define DISENTANGLEDATTENTIONPLUGIN_LEGACY_PLUGIN_H
+
+/*
+ * Legacy version of the plugin maintained for backward compatibility.
+ * This implementation is based on IPluginV2 interfaces.
+ */
+#include "NvInferPlugin.h"
+#include "common/plugin.h"
+#include "common/serialize.hpp"
+#include "disentangledAttentionCommon.h"
+#include <cstdint>
+#include <iostream>
+#include <string>
+#include <vector>
+
+// One of the preferred ways of making TensorRT to be able to see
+// our custom layer requires extending IPluginV2 and IPluginCreator classes.
+// For requirements for overriden functions, check TensorRT API docs.
+namespace nvinfer1
+{
+namespace plugin
+{
+
+// using namespace nvinfer1;
+
+class DisentangledAttentionPluginLegacy : public nvinfer1::IPluginV2DynamicExt
+{
+public:
+    DisentangledAttentionPluginLegacy();
+
+    DisentangledAttentionPluginLegacy(int32_t span, float factor);
+
+    DisentangledAttentionPluginLegacy(void const* serialData, size_t serialLength);
+
+    int32_t getNbOutputs() const noexcept override;
+
+    // DynamicExt plugins returns DimsExprs class instead of Dims
+    nvinfer1::DimsExprs getOutputDimensions(int32_t index, nvinfer1::DimsExprs const* inputs, int32_t nbInputDims,
+        nvinfer1::IExprBuilder& exprBuilder) noexcept override; // determine output dims based on input info
+
+    int32_t initialize() noexcept override;
+
+    void terminate() noexcept override;
+
+    size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
+        nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
+
+    // This is where the plugin work is done.
+    int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
+        void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
+
+    size_t getSerializationSize() const noexcept override;
+
+    void serialize(void* buffer) const noexcept override;
+
+    bool supportsFormatCombination(
+        int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
+
+    char const* getPluginType() const noexcept override;
+
+    char const* getPluginVersion() const noexcept override;
+
+    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
+
+    void destroy() noexcept override;
+
+    nvinfer1::DataType getOutputDataType(
+        int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+    char const* getPluginNamespace() const noexcept override;
+
+    void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
+        nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
+
+private:
+    // Helper method for enqueue()
+    template <typename TDataType>
+    void enqueueType(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
+        void const* const* inputs, void* const* outputs, cudaStream_t stream, TDataType factor);
+
+    std::string mNamespace;
+
+    // attributes
+    int32_t mSpan;
+    float mFactor;
+
+    using IPluginV2::getOutputDimensions;
+    using IPluginV2::getWorkspaceSize;
+    using IPluginV2::enqueue;
+    using IPluginV2Ext::configurePlugin;
+};
+
+class DisentangledAttentionPluginCreatorLegacy : public nvinfer1::IPluginCreator
+{
+public:
+    DisentangledAttentionPluginCreatorLegacy();
+
+    ~DisentangledAttentionPluginCreatorLegacy() override = default;
+
+    char const* getPluginName() const noexcept override;
+
+    char const* getPluginVersion() const noexcept override;
+
+    nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
+
+    nvinfer1::IPluginV2DynamicExt* createPlugin(
+        char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
+
+    nvinfer1::IPluginV2DynamicExt* deserializePlugin(
+        char const* name, void const* serialData, size_t serialLength) noexcept override;
+
+    void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+    char const* getPluginNamespace() const noexcept override;
+
+private:
+    nvinfer1::PluginFieldCollection mFC;
+    std::vector<nvinfer1::PluginField> mPluginAttributes;
+    std::string mNamespace;
+};
+} // namespace plugin
+} // namespace nvinfer1
+
+#endif // DISENTANGLEDATTENTIONPLUGIN_LEGACY_PLUGIN_H

plugin/disentangledAttentionPlugin/disentangledKernel.cu

@@ -0,0 +1,279 @@
+/*
+ * 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 "disentangledAttentionCommon.h"
+#include <cuda_fp16.h>
+#include <stdio.h>
+#include <assert.h>
+
+#define IND(i, j, k, dim)                                                                                              \
+    ((i) *dim.y * dim.z + (j) *dim.z + (k)) // caveat: must use brackets around var name! otherwise IND(i,j+3,k,dim) =
+                                            // (i*dim.y*dim.z + j+3*dim.z + k)...
+
+namespace nvinfer1
+{
+namespace plugin
+{
+
+using namespace nvinfer1;
+
+// template specialization for double/float
+template <typename TDataType,
+    std::enable_if_t<std::is_same_v<std::decay_t<TDataType>, double>
+            || std::is_same_v<std::decay_t<TDataType>, float>,
+        TDataType>* dummy
+    = nullptr>
+__forceinline__ __device__ void compute_attention(
+    TDataType& res, const TDataType& res0, const TDataType& res1, const TDataType& res2, const TDataType& factor)
+{
+    res = (res0 + res1 + res2) * factor;
+}
+
+// template specialization for half
+template <typename TDataType,
+    std::enable_if_t<std::is_same_v<std::decay_t<TDataType>, __half>
+            || std::is_same_v<std::decay_t<TDataType>, half>,
+        TDataType>* dummy
+    = nullptr>
+__forceinline__ __device__ void compute_attention(
+    TDataType& res, const TDataType& res0, const TDataType& res1, const TDataType& res2, const TDataType& factor)
+{
+#if __CUDA_ARCH__ >= 530
+    // __hmul only supported >= sm_53
+    res = __hmul(__hadd(res0, __hadd(res1, res2)), factor);
+#else
+    // for < sm_53, workaround/fallback is convert to float and downconvert
+    res = __float2half((__half2float(res0) + __half2float(res1) + __half2float(res2)) * __half2float(factor));
+#endif
+}
+
+// template specialization for int8
+template <typename TDataType,
+    std::enable_if_t<std::is_same_v<std::decay_t<TDataType>, int8_t>
+            || std::is_same_v<std::decay_t<TDataType>, uint8_t>,
+        TDataType>* dummy
+    = nullptr>
+__forceinline__ __device__ void compute_attention(
+    TDataType& res, const TDataType& res0, const TDataType& res1, const TDataType& res2, const TDataType& factor)
+{
+    res = (res0 + res1 + res2) * factor;
+}
+
+/**
+ * Fused kernel for Disentangled Attention design (first proposed in Microsoft DeBERTa), Version 2.
+ *
+ * @tparam TDataType type of the input data
+ * @tparam tTileSize dimension of the shared memory tile (square) and also the BlockDimX
+ * @tparam tBlockDimY 2D thread block is (tTileSize, tBlockDimY)
+ * @param data0 content-to-content ("c2c") attention QcKc^T
+ * @param data1 content-to-position ("c2p") attention QcKr^T
+ * @param data2 position-to-content ("p2c") attention KcQr^T
+ * @param result attention result
+ * @param dimData0, dimData1, dimData2, dimResult dimension of the tensors
+ * @param factor scaling factor applied on attention for stabilizing model training, 1/sqrt(3d), d is hidden size per
+ * head = H/N. H is hidden size, N is number of heads
+ * @param span relative distance hyper-parameter, k, in Disentangled attention
+ * @note C++ 17 and above due to constexpr if
+ */
+template <typename TDataType = __half, int32_t tTileSize = 32, int32_t tBlockDimY = 8>
+__global__ void GatherAddGatherTransposeAddMul_fused(TDataType const* data0, TDataType const* data1,
+    TDataType const* data2, TDataType* result, dim3 dimData0, dim3 dimData1, dim3 dimData2, dim3 dimResult,
+    TDataType factor, int32_t span)
+{
+    // Tile size should be a multiple of number of block rows
+    assert(tBlockDimY * (tTileSize / tBlockDimY) == tTileSize);
+
+    // map block to the output (result)
+    int32_t i;
+    int32_t j;
+    int32_t k;
+    int32_t c;
+    int32_t ty;
+    TDataType res0;
+    TDataType res1;
+    TDataType res2;
+    TDataType res;
+
+#if kDISENTANGLED_VERSION == 2
+    int32_t bucket;
+    int32_t mid = span / 2;
+    int32_t index;
+
+    // tmp values are precomputed for re-use; must be at least float to ensure accuracy
+    float tmp1 = logf(mid);
+
+    // Multiply by (1 - epsilon) to ensure that taking the ceil of approximately an integer
+    // results in that integer when computing the bucket later on.
+    // This corrects for the mathematical imprecision from using float.
+    constexpr float kEPSILON = 1e-7;
+    float tmp = (mid - 1) / (logf(dimData1.z - 1) - tmp1) * (1 - kEPSILON);
+#endif
+
+    __shared__ TDataType T[tTileSize][tTileSize + 1]; // +1 to avoid bank conflict
+
+    // (i,j,k) location of data2 (transposed)
+    i = blockIdx.z;
+    j = blockIdx.x * tTileSize + threadIdx.y;
+    k = blockIdx.y * tTileSize + threadIdx.x;
+
+// gather data2
+#pragma unroll
+    for (c = 0, ty = 0; c < tTileSize / tBlockDimY; c++, ty += tBlockDimY)
+    {
+#if kDISENTANGLED_VERSION == 1
+        // relative position -- version 1
+        if (k - (j + ty) >= span)
+        {
+            res2 = data2[IND(i, j + ty, 2 * span - 1, dimData2)];
+        }
+        else if (k - (j + ty) <= -span)
+        {
+            res2 = data2[IND(i, j + ty, 0, dimData2)];
+        }
+        else
+        {
+            res2 = data2[IND(i, j + ty, k - (j + ty) + span, dimData2)]; // compute index on the fly
+        }
+        T[ty + threadIdx.y][threadIdx.x] = res2;
+#elif kDISENTANGLED_VERSION == 2
+        // relative position w/ log bucket -- version 2
+        if (k - (j + ty) >= -mid && k - (j + ty) <= mid)
+        {
+            // preserved region, (i - j) + span
+            bucket = k - (j + ty);
+        }
+        else
+        {
+            // log bucket region, bucket(i,j) + span
+            bucket = ceilf((logf(fabsf(k - (j + ty))) - tmp1) * tmp) + mid;
+            bucket = k - (j + ty) < 0 ? -bucket : bucket;
+        }
+        // clamp [0,2k]. Although this is guaranteed by equation, but numerically the floating precision can still break
+        // boundary
+        index = bucket + span;
+        index = min(max(0, index), 2 * span - 1);
+        res2 = data2[IND(i, j + ty, index, dimData2)];
+        T[ty + threadIdx.y][threadIdx.x] = res2;
+#endif
+    }
+
+    __syncthreads();
+
+    // (i,j,k) location of data1 (non-transposed) and output. i unchanged
+    j = blockIdx.y * tTileSize + threadIdx.y;
+    k = blockIdx.x * tTileSize + threadIdx.x;
+
+// read data0 + gather data1 + add all + write
+#pragma unroll
+    for (c = 0, ty = 0; c < tTileSize / tBlockDimY; c++, ty += tBlockDimY)
+    {
+#if kDISENTANGLED_VERSION == 1
+        // relative position -- version 1
+        // for non-transposed matrix 1, just fetch element at the transposed location & add to the result)
+        if (j + ty - k <= -span)
+        {
+            res1 = data1[IND(i, j + ty, 0, dimData1)];
+        }
+        else if (j + ty - k >= span)
+        {
+            res1 = data1[IND(i, j + ty, 2 * span - 1, dimData1)];
+        }
+        else
+        {
+            res1 = data1[IND(i, j + ty, j + ty - k + span, dimData1)]; // compute index on the fly
+        }
+#elif kDISENTANGLED_VERSION == 2
+        // relative position w/ log bucket -- version 2
+        if (j + ty - k >= -mid && j + ty - k <= mid)
+        {
+            // preserved region, (i - j) + span
+            bucket = j + ty - k;
+        }
+        else
+        {
+            // log bucket region, bucket(i,j) + span
+            bucket = ceilf((logf(fabsf((j + ty) - k)) - tmp1) * tmp) + mid;
+            bucket = (j + ty) - k < 0 ? -bucket : bucket;
+        }
+        // clamp [0,2k]. Although this is guaranteed by equation, but numerically the floating precision can still break
+        // boundary
+        index = bucket + span;
+        index = min(max(0, index), 2 * span - 1);
+        res1 = data1[IND(i, j + ty, index, dimData1)];
+#endif
+
+        // for non-tranposed matrix 0, same as matrix 1
+        res0 = data0[IND(i, j + ty, k, dimData0)];
+
+        // (res0 + res1 + res2) / sqrt(3d), d is the hidden states size per head
+#if __cplusplus >= 201703L
+        // C++ 17 has more convenient `if constexpr` for conditional implementation at compile time; before C++ 17,
+        // switch to template specialization
+        if constexpr (std::is_same_v<TDataType, double> || std::is_same_v<TDataType, float>)
+        {
+            // double, float32
+            res = (res0 + res1 + T[threadIdx.x][ty + threadIdx.y]) * factor;
+        }
+        else if constexpr (std::is_same_v<TDataType, __half> || std::is_same_v<TDataType, half>)
+        {
+            // fp16
+#if __CUDA_ARCH__ >= 530
+            // __hmul only supported >= sm_53
+            res = __hmul(__hadd(res0, __hadd(res1, T[threadIdx.x][ty + threadIdx.y])), factor);
+#else
+            // for < sm_53, workaround/fallback is convert to float and downconvert
+            res = __float2half(
+                (__half2float(res0) + __half2float(res1) + __half2float(T[threadIdx.x][ty + threadIdx.y]))
+                * __half2float(factor));
+#endif
+        }
+        else if constexpr (std::is_same_v<TDataType, int8_t> || std::is_same_v<TDataType, uint8_t>)
+        {
+            // int8_t
+            res = (res0 + res1 + T[threadIdx.x][ty + threadIdx.y]) * factor;
+        }
+#else
+        // before C++ 17, use template specialization
+        compute_attention<TDataType>(res, res0, res1, T[threadIdx.x][ty + threadIdx.y], factor);
+#endif
+        // write
+        result[IND(i, j + ty, k, dimResult)] = res;
+    }
+}
+
+template <typename TDataType, int32_t tTileSize, int32_t tBlockDimY>
+void disentangled_kernel_wrapper(TDataType const* data0, TDataType const* data1, TDataType const* data2,
+    TDataType* result, dim3 dimData0, dim3 dimData1, dim3 dimData2, dim3 dimResult, TDataType factor, int32_t span,
+    dim3 block, dim3 grid, cudaStream_t stream)
+{
+    GatherAddGatherTransposeAddMul_fused<TDataType, tTileSize, tBlockDimY><<<grid, block, 0, stream>>>(
+        data0, data1, data2, result, dimData0, dimData1, dimData2, dimResult, factor, span);
+}
+
+template void disentangled_kernel_wrapper<float, kDISENTANGLED_TILESIZE, kDISENTANGLED_BLOCKDIMY>(
+    float const*, float const*, float const*, float*, dim3, dim3, dim3, dim3, float, int32_t, dim3, dim3, cudaStream_t);
+
+template void disentangled_kernel_wrapper<__half, kDISENTANGLED_TILESIZE, kDISENTANGLED_BLOCKDIMY>(__half const*,
+    __half const*, __half const*, __half*, dim3, dim3, dim3, dim3, __half, int32_t, dim3, dim3, cudaStream_t);
+
+template void disentangled_kernel_wrapper<int8_t, kDISENTANGLED_TILESIZE, kDISENTANGLED_BLOCKDIMY>(int8_t const*,
+    int8_t const*, int8_t const*, int8_t*, dim3, dim3, dim3, dim3, int8_t, int32_t, dim3, dim3, cudaStream_t);
+
+#undef IND
+
+} // namespace plugin
+} // namespace nvinfer1