chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled

This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
@@ -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(
multiscaleDeformableAttn.cu
multiscaleDeformableAttn.h
multiscaleDeformableAttnPlugin.cpp
multiscaleDeformableAttnPlugin.h
multiscaleDeformableAttnPluginLegacy.cpp
multiscaleDeformableAttnPluginLegacy.h
multiscaleDeformableIm2ColCuda.cuh
)
@@ -0,0 +1,69 @@
# multiscaleDeformableAttn Plugin
**Table Of Contents**
- [Description](#description)
* [Structure](#structure)
- [Parameters](#parameters)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## 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.
The `multiscaleDeformableAttnPlugin` is used to perform attention computation over a small set of key sampling points around a reference point rather than looking over all possible spatial locations. It makes use of multiscale feature maps to effectively represent objects at different scales. It helps to achieve faster convergence and better performance on small objects.
### Structure
The `multiscaleDeformableAttnPlugin` takes 5 inputs in the following order : `value`, `spatial_shapes`, `level_start_index`, `sampling_locations`, and `atttention_weights`.
`value`
The input feature maps from different scales concatenated to provide the input feature vector. The shape of this tensor is `[N, S, M, D]` where `N` is batch size, `S` is the length of the feature maps, `M` is the number of attentions heads, `D` is hidden_dim/num_heads.
`spatial_shapes`
The shape of each feature map. The shape of this tensor is `[L, 2]` where `L` is the number of feature maps.
`level_start_index`
This tensor is used to find the sampling locations for different feature levels as the input feature tensors are flattened. The shape of this tensor is `[L,]`.
`sampling_locations`
This tensor acts as a pre-filter for prominent key elements out of all the feature map pixels. The shape of this tensor is `[N, Lq, M, L, P, 2]` where `P` is the number of points, `Lq` is the length of feature maps(encoder)/length of queries(decoder).
`attention_weights`
This tensor consists of the attention weights whose shape is `[N, Lq, M, L, P]`.
The `multiscaleDeformableAttnPlugin` generates the attention output of shape `[N, Lq, M, D]`.
## Parameters
`multiscaleDeformableAttnPlugin` has plugin creator class `multiscaleDeformableAttnPluginCreator` and plugin class `multiscaleDeformableAttnPlugin`.
The plugin does not require any parameters to be built and used.
## Additional resources
The following resources provide a deeper understanding of the `multiscaleDeformableAttnPlugin` plugin:
**Networks:**
- [Deformable DETR](https://arxiv.org/pdf/2010.04159.pdf)
## 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
Apr 2025
Added version 2 of the plugin that uses the IPluginV3 interface. The version 1 (using IPluginV2DynamicExt interface) is now deprecated. The version 2 mirrors version 1 in IO and attributes.
Feb 2022
This is the first release of this `README.md` file.
## Known issues
There are no known issues in this plugin.
@@ -0,0 +1,79 @@
/*
* 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.
*/
/*
**************************************************************************
* Modified from Deformable DETR
* Copyright (c) 2020 SenseTime. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
* https://github.com/fundamentalvision/Deformable-DETR/blob/main/LICENSE
**************************************************************************
* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
* Copyright (c) 2018 Microsoft
**************************************************************************
*/
#include <iostream>
#include <vector>
#include <cuda.h>
#include <cuda_fp16.h>
#include <cuda_runtime.h>
#include "multiscaleDeformableIm2ColCuda.cuh"
int32_t ms_deform_attn_cuda_forward(cudaStream_t stream, float const* value, int32_t const* spatialShapes,
int32_t const* levelStartIndex, float const* samplingLoc, float const* attnWeight, float* output, int32_t batch,
int32_t mSpatialSize, int32_t mNumHeads, int32_t mChannels, int32_t mNumLevels, int32_t mNumQuery, int32_t mNumPoint)
{
auto perValueSize = mSpatialSize * mNumHeads * mChannels;
auto perSampleLocSize = mNumQuery * mNumHeads * mNumLevels * mNumPoint * 2;
auto perAttnWeightSize = mNumQuery * mNumHeads * mNumLevels * mNumPoint;
int32_t mIm2colStep = batch;
for (int32_t n = 0; n < batch / mIm2colStep; ++n)
{
auto columns = output + perValueSize * n * mIm2colStep;
ms_deformable_im2col_cuda<float>(stream, value + n * mIm2colStep * perValueSize, spatialShapes, levelStartIndex,
samplingLoc + n * mIm2colStep * perSampleLocSize, attnWeight + n * mIm2colStep * perAttnWeightSize, batch,
mSpatialSize, mNumHeads, mChannels, mNumLevels, mNumQuery, mNumPoint, columns);
}
return 0;
}
int32_t ms_deform_attn_cuda_forward(cudaStream_t stream, __half const* value, int32_t const* spatialShapes,
int32_t const* levelStartIndex, __half const* samplingLoc, __half const* attnWeight, __half* output, int32_t batch,
int32_t mSpatialSize, int32_t mNumHeads, int32_t mChannels, int32_t mNumLevels, int32_t mNumQuery, int32_t mNumPoint)
{
auto perValueSize = mSpatialSize * mNumHeads * mChannels;
auto perSampleLocSize = mNumQuery * mNumHeads * mNumLevels * mNumPoint * 2;
auto perAttnWeightSize = mNumQuery * mNumHeads * mNumLevels * mNumPoint;
int32_t mIm2colStep = batch;
for (int32_t n = 0; n < batch / mIm2colStep; ++n)
{
auto columns = output + perValueSize * n * mIm2colStep;
ms_deformable_im2col_cuda<__half>(stream, value + n * mIm2colStep * perValueSize, spatialShapes,
levelStartIndex, samplingLoc + n * mIm2colStep * perSampleLocSize,
attnWeight + n * mIm2colStep * perAttnWeightSize, batch, mSpatialSize, mNumHeads, mChannels, mNumLevels,
mNumQuery, mNumPoint, columns);
}
return 0;
}
@@ -0,0 +1,45 @@
/*
* 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.
*/
/*
**************************************************************************
* Modified from Deformable DETR
* Copyright (c) 2020-2023 SenseTime. All Rights Reserved.
* Licensed under the Apache License, Version 2.0 [see LICENSE for details]
* https://github.com/fundamentalvision/Deformable-DETR/blob/main/LICENSE
**************************************************************************
* Modified from DCN (https://github.com/msracver/Deformable-ConvNets)
* Copyright (c) 2018-2023 Microsoft
**************************************************************************
* Modified from https://github.com/chengdazhi/Deformable-Convolution-V2-PyTorch/tree/pytorch_1.0.0
**************************************************************************************************
*/
#ifndef TRT_MULTISCALE_DEFORMABLE_ATTN_H
#define TRT_MULTISCALE_DEFORMABLE_ATTN_H
int32_t ms_deform_attn_cuda_forward(cudaStream_t stream, float const* value, int32_t const* spatialShapes,
int32_t const* levelStartIndex, float const* samplingLoc, float const* attnWeight, float* output, int32_t batch,
int32_t mSpatialSize, int32_t mNumHeads, int32_t mChannels, int32_t mNumLevels, int32_t mNumQuery,
int32_t mNumPoint);
int32_t ms_deform_attn_cuda_forward(cudaStream_t stream, __half const* value, int32_t const* spatialShapes,
int32_t const* levelStartIndex, __half const* samplingLoc, __half const* attnWeight, __half* output, int32_t batch,
int32_t mSpatialSize, int32_t mNumHeads, int32_t mChannels, int32_t mNumLevels, int32_t mNumQuery,
int32_t mNumPoint);
#endif
@@ -0,0 +1,418 @@
/*
* 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 "multiscaleDeformableAttnPlugin.h"
#include "multiscaleDeformableAttn.h"
#include <memory>
using namespace nvinfer1;
using namespace nvinfer1::plugin;
namespace
{
static char const* DMHA_VERSION{"2"};
static char const* DMHA_NAME{"MultiscaleDeformableAttnPlugin_TRT"};
} // namespace
namespace nvinfer1::plugin
{
MultiscaleDeformableAttnPlugin::MultiscaleDeformableAttnPlugin() {}
IPluginCapability* MultiscaleDeformableAttnPlugin::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;
}
// IPluginV3OneCore methods
char const* MultiscaleDeformableAttnPlugin::getPluginName() const noexcept
{
return DMHA_NAME;
}
char const* MultiscaleDeformableAttnPlugin::getPluginVersion() const noexcept
{
return DMHA_VERSION;
}
int32_t MultiscaleDeformableAttnPlugin::getNbOutputs() const noexcept
{
return 1;
}
void MultiscaleDeformableAttnPlugin::setPluginNamespace(char const* pluginNamespace) noexcept
{
mNamespace = pluginNamespace;
}
char const* MultiscaleDeformableAttnPlugin::getPluginNamespace() const noexcept
{
return mNamespace.c_str();
}
IPluginV3* MultiscaleDeformableAttnPlugin::clone() noexcept
{
try
{
auto plugin = std::make_unique<MultiscaleDeformableAttnPlugin>();
plugin->setPluginNamespace(mNamespace.c_str());
return plugin.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
// IPluginV3OneBuild methods
int32_t MultiscaleDeformableAttnPlugin::getOutputDataTypes(
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
{
try
{
PLUGIN_VALIDATE(outputTypes != nullptr, "outputTypes pointer is null");
PLUGIN_VALIDATE(nbOutputs > 0, "nbOutputs is not positive");
PLUGIN_VALIDATE(inputTypes != nullptr, "inputTypes pointer is null");
PLUGIN_VALIDATE(nbInputs > 0, "nbInputs is not positive");
// Output type is the same as the first input type
std::fill_n(outputTypes, nbOutputs, inputTypes[0]);
return STATUS_SUCCESS;
}
catch (std::exception const& e)
{
caughtError(e);
}
return STATUS_FAILURE;
}
int32_t MultiscaleDeformableAttnPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
IExprBuilder& exprBuilder) noexcept
{
try
{
PLUGIN_VALIDATE(outputs != nullptr, "outputs pointer is null");
PLUGIN_VALIDATE(nbOutputs > 0, "nbOutputs is not positive");
PLUGIN_VALIDATE(inputs != nullptr, "inputs pointer is null");
PLUGIN_VALIDATE(nbInputs == 5, "Expected 5 inputs");
// Output shape: [N, Lq, M, D]
outputs[0].nbDims = 4;
outputs[0].d[0] = inputs[0].d[0]; // Batch size
outputs[0].d[1] = inputs[3].d[1]; // Lq (query length)
outputs[0].d[2] = inputs[0].d[2]; // Number of heads
outputs[0].d[3] = inputs[0].d[3]; // Hidden dimension per head
return STATUS_SUCCESS;
}
catch (std::exception const& e)
{
caughtError(e);
}
return STATUS_FAILURE;
}
bool MultiscaleDeformableAttnPlugin::supportsFormatCombination(
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
try
{
PLUGIN_VALIDATE(inOut != nullptr, "inOut pointer is null");
PLUGIN_VALIDATE(nbInputs == 5, "Expected 5 inputs");
PLUGIN_VALIDATE(nbOutputs == 1, "Expected 1 output");
// Check format
PluginTensorDesc const& desc = inOut[pos].desc;
if (desc.format != TensorFormat::kLINEAR)
{
return false;
}
// Special handling for spatial_shapes and level_start_index (inputs 1 and 2)
if (pos == 1 || pos == 2)
{
return desc.type == DataType::kINT32;
}
// Other inputs and output must have the same type, either FP32 or FP16
if (pos == 0 || pos == 3 || pos == 4 || pos == nbInputs)
{
// Check that the data type matches input[0]
bool const isFloatType = desc.type == DataType::kFLOAT || desc.type == DataType::kHALF;
if (pos == 0) // First tensor, just check if it's a supported type
{
return isFloatType;
}
// Other tensors must match the first
return desc.type == inOut[0].desc.type && isFloatType;
}
return false;
}
catch (std::exception const& e)
{
caughtError(e);
}
return false;
}
int32_t MultiscaleDeformableAttnPlugin::configurePlugin(
DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
{
try
{
PLUGIN_VALIDATE(in != nullptr, "in pointer is null");
PLUGIN_VALIDATE(out != nullptr, "out pointer is null");
PLUGIN_VALIDATE(nbInputs == 5, "Expected 5 inputs");
PLUGIN_VALIDATE(nbOutputs == 1, "Expected 1 output");
// Check for valid input dimensions
PLUGIN_VALIDATE(in[0].desc.dims.nbDims == 4, "First input must have 4 dimensions");
PLUGIN_VALIDATE(in[1].desc.dims.nbDims == 2, "Second input must have 2 dimensions");
PLUGIN_VALIDATE(in[2].desc.dims.nbDims == 1, "Third input must have 1 dimension");
PLUGIN_VALIDATE(in[3].desc.dims.nbDims == 6, "Fourth input must have 6 dimensions");
PLUGIN_VALIDATE(in[4].desc.dims.nbDims == 5, "Fifth input must have 5 dimensions");
// Check M dimensions consistency
PLUGIN_VALIDATE(in[0].desc.dims.d[2] == in[3].desc.dims.d[2], "Inconsistent dimensions for number of heads");
PLUGIN_VALIDATE(in[0].desc.dims.d[2] == in[4].desc.dims.d[2], "Inconsistent dimensions for number of heads");
// Check L dimensions consistency
PLUGIN_VALIDATE(in[1].desc.dims.d[0] == in[2].desc.dims.d[0], "Inconsistent dimensions for number of levels");
PLUGIN_VALIDATE(in[1].desc.dims.d[0] == in[3].desc.dims.d[3], "Inconsistent dimensions for number of levels");
PLUGIN_VALIDATE(in[1].desc.dims.d[0] == in[4].desc.dims.d[3], "Inconsistent dimensions for number of levels");
// Check P dimensions consistency
PLUGIN_VALIDATE(in[3].desc.dims.d[4] == in[4].desc.dims.d[4], "Inconsistent dimensions for number of points");
// Check Lq dimensions consistency
PLUGIN_VALIDATE(in[3].desc.dims.d[1] == in[4].desc.dims.d[1], "Inconsistent dimensions for query length");
return STATUS_SUCCESS;
}
catch (std::exception const& e)
{
caughtError(e);
}
return STATUS_FAILURE;
}
PluginFieldCollection const* MultiscaleDeformableAttnPlugin::getFieldsToSerialize() noexcept
{
try
{
mDataToSerialize.clear();
// This plugin has no fields to serialize
mFCToSerialize.nbFields = mDataToSerialize.size();
mFCToSerialize.fields = mDataToSerialize.data();
return &mFCToSerialize;
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
// IPluginV3OneRuntime methods
size_t MultiscaleDeformableAttnPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
{
// No workspace needed for this plugin
return 0;
}
int32_t MultiscaleDeformableAttnPlugin::onShapeChange(
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
{
try
{
PLUGIN_VALIDATE(inputs != nullptr, "inputs pointer is null");
PLUGIN_VALIDATE(outputs != nullptr, "outputs pointer is null");
PLUGIN_VALIDATE(nbInputs == 5, "Expected 5 inputs");
PLUGIN_VALIDATE(nbOutputs == 1, "Expected 1 output");
// Check for valid input dimensions
PLUGIN_VALIDATE(inputs[0].dims.nbDims == 4, "First input must have 4 dimensions");
PLUGIN_VALIDATE(inputs[1].dims.nbDims == 2, "Second input must have 2 dimensions");
PLUGIN_VALIDATE(inputs[2].dims.nbDims == 1, "Third input must have 1 dimension");
PLUGIN_VALIDATE(inputs[3].dims.nbDims == 6, "Fourth input must have 6 dimensions");
PLUGIN_VALIDATE(inputs[4].dims.nbDims == 5, "Fifth input must have 5 dimensions");
// Check M dimensions consistency
PLUGIN_VALIDATE(inputs[0].dims.d[2] == inputs[3].dims.d[2], "Inconsistent dimensions for number of heads");
PLUGIN_VALIDATE(inputs[0].dims.d[2] == inputs[4].dims.d[2], "Inconsistent dimensions for number of heads");
// Check L dimensions consistency
PLUGIN_VALIDATE(inputs[1].dims.d[0] == inputs[2].dims.d[0], "Inconsistent dimensions for number of levels");
PLUGIN_VALIDATE(inputs[1].dims.d[0] == inputs[3].dims.d[3], "Inconsistent dimensions for number of levels");
PLUGIN_VALIDATE(inputs[1].dims.d[0] == inputs[4].dims.d[3], "Inconsistent dimensions for number of levels");
// Check P dimensions consistency
PLUGIN_VALIDATE(inputs[3].dims.d[4] == inputs[4].dims.d[4], "Inconsistent dimensions for number of points");
// Check Lq dimensions consistency
PLUGIN_VALIDATE(inputs[3].dims.d[1] == inputs[4].dims.d[1], "Inconsistent dimensions for query length");
return STATUS_SUCCESS;
}
catch (std::exception const& e)
{
caughtError(e);
}
return STATUS_FAILURE;
}
IPluginV3* MultiscaleDeformableAttnPlugin::attachToContext(IPluginResourceContext* context) noexcept
{
try
{
// No resources need to be attached
return clone();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
int32_t MultiscaleDeformableAttnPlugin::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 && inputs != nullptr && outputs != nullptr, "Null pointers found in enqueue");
int32_t const batch = inputDesc[0].dims.d[0];
int32_t spatialSize = inputDesc[0].dims.d[1];
int32_t numHeads = inputDesc[0].dims.d[2];
int32_t channels = inputDesc[0].dims.d[3];
int32_t numLevels = inputDesc[1].dims.d[0];
int32_t numQuery = inputDesc[3].dims.d[1];
int32_t numPoint = inputDesc[3].dims.d[4];
int32_t rc = 0;
if (inputDesc[0].type == DataType::kFLOAT)
{
auto const* value = static_cast<float const*>(inputs[0]);
auto const* spatialShapes = static_cast<int32_t const*>(inputs[1]);
auto const* levelStartIndex = static_cast<int32_t const*>(inputs[2]);
auto const* samplingLoc = static_cast<float const*>(inputs[3]);
auto const* attnWeight = static_cast<float const*>(inputs[4]);
auto* output = static_cast<float*>(outputs[0]);
rc = ms_deform_attn_cuda_forward(stream, value, spatialShapes, levelStartIndex, samplingLoc, attnWeight,
output, batch, spatialSize, numHeads, channels, numLevels, numQuery, numPoint);
}
else if (inputDesc[0].type == DataType::kHALF)
{
auto const* value = static_cast<__half const*>(inputs[0]);
auto const* spatialShapes = static_cast<int32_t const*>(inputs[1]);
auto const* levelStartIndex = static_cast<int32_t const*>(inputs[2]);
auto const* samplingLoc = static_cast<__half const*>(inputs[3]);
auto const* attnWeight = static_cast<__half const*>(inputs[4]);
auto* output = static_cast<__half*>(outputs[0]);
rc = ms_deform_attn_cuda_forward(stream, value, spatialShapes, levelStartIndex, samplingLoc, attnWeight,
output, batch, spatialSize, numHeads, channels, numLevels, numQuery, numPoint);
}
else
{
PLUGIN_VALIDATE(false, "Unsupported data type");
}
return rc;
}
catch (std::exception const& e)
{
caughtError(e);
}
return STATUS_FAILURE;
}
// Plugin Creator Implementation
MultiscaleDeformableAttnPluginCreator::MultiscaleDeformableAttnPluginCreator()
{
mPluginAttributes.clear();
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
char const* MultiscaleDeformableAttnPluginCreator::getPluginName() const noexcept
{
return DMHA_NAME;
}
char const* MultiscaleDeformableAttnPluginCreator::getPluginVersion() const noexcept
{
return DMHA_VERSION;
}
PluginFieldCollection const* MultiscaleDeformableAttnPluginCreator::getFieldNames() noexcept
{
return &mFC;
}
IPluginV3* MultiscaleDeformableAttnPluginCreator::createPlugin(
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
{
try
{
// This plugin doesn't have any configurable parameters
return new MultiscaleDeformableAttnPlugin();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
void MultiscaleDeformableAttnPluginCreator::setPluginNamespace(char const* pluginNamespace) noexcept
{
mNamespace = pluginNamespace;
}
char const* MultiscaleDeformableAttnPluginCreator::getPluginNamespace() const noexcept
{
return mNamespace.c_str();
}
} // namespace nvinfer1::plugin
@@ -0,0 +1,121 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
/*
* V3 version of the plugin using IPluginV3 interfaces.
* This implementation follows TensorRT's plugin V3 API.
*/
#ifndef TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_H
#define TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_H
// Standard library includes
#include <memory>
#include <string>
#include <vector>
#include "NvInferPlugin.h"
// TensorRT includes
#include "common/plugin.h"
namespace nvinfer1
{
namespace plugin
{
// Forward declarations
class MultiscaleDeformableAttnPlugin;
class MultiscaleDeformableAttnPluginCreator;
// V3 Plugin implementation
class MultiscaleDeformableAttnPlugin : public IPluginV3,
public IPluginV3OneCore,
public IPluginV3OneBuild,
public IPluginV3OneRuntime
{
public:
// Constructors/destructors
MultiscaleDeformableAttnPlugin();
~MultiscaleDeformableAttnPlugin() = default;
// IPluginV3 methods
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
// IPluginV3OneCore methods
char const* getPluginName() const noexcept override;
char const* getPluginVersion() const noexcept override;
char const* getPluginNamespace() const noexcept override;
void setPluginNamespace(char const* pluginNamespace) noexcept;
int32_t getNbOutputs() const noexcept override;
IPluginV3* clone() noexcept override;
// IPluginV3OneBuild methods
bool supportsFormatCombination(
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
int32_t getOutputDataTypes(
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
int32_t getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override;
int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out,
int32_t nbOutputs) noexcept override;
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
// IPluginV3OneRuntime methods
size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs,
void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
int32_t onShapeChange(PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs,
int32_t nbOutputs) noexcept override;
private:
// Serialization helpers
std::vector<PluginField> mDataToSerialize;
PluginFieldCollection mFCToSerialize;
// Plugin namespace
std::string mNamespace;
};
// Plugin creator class
class MultiscaleDeformableAttnPluginCreator : public IPluginCreatorV3One
{
public:
// Constructor
MultiscaleDeformableAttnPluginCreator();
// IPluginCreatorV3One methods
char const* getPluginName() const noexcept override;
char const* getPluginVersion() const noexcept override;
PluginFieldCollection const* getFieldNames() noexcept override;
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
void setPluginNamespace(char const* pluginNamespace) noexcept;
char const* getPluginNamespace() const noexcept override;
private:
// Plugin fields and namespace
PluginFieldCollection mFC;
std::vector<PluginField> mPluginAttributes;
std::string mNamespace;
};
} // namespace plugin
} // namespace nvinfer1
#endif // TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_H
@@ -0,0 +1,289 @@
/*
* 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 "multiscaleDeformableAttnPluginLegacy.h"
#include "multiscaleDeformableAttn.h"
#include <memory>
using namespace nvinfer1;
using namespace nvinfer1::plugin;
namespace nvinfer1::plugin
{
namespace
{
static char const* DMHA_VERSION{"1"};
static char const* DMHA_NAME{"MultiscaleDeformableAttnPlugin_TRT"};
} // namespace
// // Register the plugin with TensorRT
// REGISTER_TENSORRT_PLUGIN(MultiscaleDeformableAttnPluginCreatorLegacy);
MultiscaleDeformableAttnPluginLegacy::MultiscaleDeformableAttnPluginLegacy() {}
MultiscaleDeformableAttnPluginLegacy::MultiscaleDeformableAttnPluginLegacy(void const* data, size_t length) {}
nvinfer1::IPluginV2DynamicExt* MultiscaleDeformableAttnPluginLegacy::clone() const noexcept
{
try
{
auto plugin = std::make_unique<MultiscaleDeformableAttnPluginLegacy>();
plugin->setPluginNamespace(getPluginNamespace());
return plugin.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
nvinfer1::DimsExprs MultiscaleDeformableAttnPluginLegacy::getOutputDimensions(int32_t outputIndex,
nvinfer1::DimsExprs const* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
{
nvinfer1::DimsExprs ret;
ret.nbDims = 4;
ret.d[0] = inputs[0].d[0];
ret.d[1] = inputs[3].d[1];
ret.d[2] = inputs[0].d[2];
ret.d[3] = inputs[0].d[3];
return ret;
}
bool MultiscaleDeformableAttnPluginLegacy::supportsFormatCombination(
int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
PLUGIN_ASSERT((nbInputs == 5));
PLUGIN_ASSERT((nbOutputs == 1));
if (inOut[pos].format == nvinfer1::TensorFormat::kLINEAR)
{
if ((pos == 1) || (pos == 2))
{
return (inOut[pos].type == nvinfer1::DataType::kINT32);
}
return ((inOut[pos].type == inOut[0].type)
&& ((inOut[pos].type == nvinfer1::DataType::kFLOAT) || (inOut[pos].type == nvinfer1::DataType::kHALF)));
}
return false;
}
void MultiscaleDeformableAttnPluginLegacy::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* inputs,
int32_t nbInputs, nvinfer1::DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
{
// Check for valid input dimensions
PLUGIN_ASSERT(inputs[0].desc.dims.nbDims == 4);
PLUGIN_ASSERT(inputs[1].desc.dims.nbDims == 2);
PLUGIN_ASSERT(inputs[2].desc.dims.nbDims == 1);
PLUGIN_ASSERT(inputs[3].desc.dims.nbDims == 6);
PLUGIN_ASSERT(inputs[4].desc.dims.nbDims == 5);
// Check M dimensions consistency
PLUGIN_ASSERT(inputs[0].desc.dims.d[2] == inputs[3].desc.dims.d[2]);
PLUGIN_ASSERT(inputs[0].desc.dims.d[2] == inputs[4].desc.dims.d[2]);
// Check L dimensions consistency
PLUGIN_ASSERT(inputs[1].desc.dims.d[0] == inputs[2].desc.dims.d[0]);
PLUGIN_ASSERT(inputs[1].desc.dims.d[0] == inputs[3].desc.dims.d[3]);
PLUGIN_ASSERT(inputs[1].desc.dims.d[0] == inputs[4].desc.dims.d[3]);
// Check P dimensions consistency
PLUGIN_ASSERT(inputs[3].desc.dims.d[4] == inputs[4].desc.dims.d[4]);
// Check Lq dimensions consistency
PLUGIN_ASSERT(inputs[3].desc.dims.d[1] == inputs[4].desc.dims.d[1]);
}
size_t MultiscaleDeformableAttnPluginLegacy::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs,
int32_t nbInputs, nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
{
return 0;
}
int32_t MultiscaleDeformableAttnPluginLegacy::enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
nvinfer1::PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs,
void* /* workSpace */, cudaStream_t stream) noexcept
{
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
int32_t const batch = inputDesc[0].dims.d[0];
int32_t spatial_size = inputDesc[0].dims.d[1];
int32_t num_heads = inputDesc[0].dims.d[2];
int32_t channels = inputDesc[0].dims.d[3];
int32_t num_levels = inputDesc[1].dims.d[0];
int32_t num_query = inputDesc[3].dims.d[1];
int32_t num_point = inputDesc[3].dims.d[4];
int32_t rc = 0;
if (inputDesc[0].type == nvinfer1::DataType::kFLOAT)
{
float const* value = static_cast<float const*>(inputs[0]);
int32_t const* spatialShapes = static_cast<int32_t const*>(inputs[1]);
int32_t const* levelStartIndex = static_cast<int32_t const*>(inputs[2]);
float const* samplingLoc = static_cast<float const*>(inputs[3]);
float const* attnWeight = static_cast<float const*>(inputs[4]);
float* output = static_cast<float*>(outputs[0]);
rc = ms_deform_attn_cuda_forward(stream, value, spatialShapes, levelStartIndex, samplingLoc, attnWeight, output,
batch, spatial_size, num_heads, channels, num_levels, num_query, num_point);
}
else if (inputDesc[0].type == nvinfer1::DataType::kHALF)
{
__half const* value = static_cast<__half const*>(inputs[0]);
int32_t const* spatialShapes = static_cast<int32_t const*>(inputs[1]);
int32_t const* levelStartIndex = static_cast<int32_t const*>(inputs[2]);
__half const* samplingLoc = static_cast<__half const*>(inputs[3]);
__half const* attnWeight = static_cast<__half const*>(inputs[4]);
__half* output = static_cast<__half*>(outputs[0]);
rc = ms_deform_attn_cuda_forward(stream, value, spatialShapes, levelStartIndex, samplingLoc, attnWeight, output,
batch, spatial_size, num_heads, channels, num_levels, num_query, num_point);
}
return rc;
}
void MultiscaleDeformableAttnPluginLegacy::attachToContext(
cudnnContext* cudnnContext, cublasContext* cublasContext, nvinfer1::IGpuAllocator* gpuAllocator) noexcept
{
}
void MultiscaleDeformableAttnPluginLegacy::detachFromContext() noexcept {}
// IPluginV2Ext Methods
nvinfer1::DataType MultiscaleDeformableAttnPluginLegacy::getOutputDataType(
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
{
return inputTypes[0];
}
// IPluginV2 Methods
char const* MultiscaleDeformableAttnPluginLegacy::getPluginType() const noexcept
{
return DMHA_NAME;
}
char const* MultiscaleDeformableAttnPluginLegacy::getPluginVersion() const noexcept
{
return DMHA_VERSION;
}
int32_t MultiscaleDeformableAttnPluginLegacy::getNbOutputs() const noexcept
{
return 1;
}
int32_t MultiscaleDeformableAttnPluginLegacy::initialize() noexcept
{
return 0;
}
void MultiscaleDeformableAttnPluginLegacy::terminate() noexcept {}
size_t MultiscaleDeformableAttnPluginLegacy::getSerializationSize() const noexcept
{
return 0;
}
void MultiscaleDeformableAttnPluginLegacy::serialize(void* buffer) const noexcept {}
void MultiscaleDeformableAttnPluginLegacy::destroy() noexcept
{
delete this;
}
void MultiscaleDeformableAttnPluginLegacy::setPluginNamespace(char const* pluginNamespace) noexcept
{
mNamespace = pluginNamespace;
}
char const* MultiscaleDeformableAttnPluginLegacy::getPluginNamespace() const noexcept
{
return mNamespace.c_str();
}
// Pluginv1 Creator
MultiscaleDeformableAttnPluginCreatorLegacy::MultiscaleDeformableAttnPluginCreatorLegacy()
{
mPluginAttributes.clear();
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
char const* MultiscaleDeformableAttnPluginCreatorLegacy::getPluginName() const noexcept
{
return DMHA_NAME;
}
char const* MultiscaleDeformableAttnPluginCreatorLegacy::getPluginVersion() const noexcept
{
return DMHA_VERSION;
}
nvinfer1::PluginFieldCollection const* MultiscaleDeformableAttnPluginCreatorLegacy::getFieldNames() noexcept
{
return &mFC;
}
IPluginV2* MultiscaleDeformableAttnPluginCreatorLegacy::createPlugin(
char const* name, PluginFieldCollection const* fc) noexcept
{
try
{
auto plugin = std::make_unique<MultiscaleDeformableAttnPluginLegacy>();
return plugin.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
IPluginV2* MultiscaleDeformableAttnPluginCreatorLegacy::deserializePlugin(
char const* name, void const* serialData, size_t serialLength) noexcept
{
try
{
auto plugin = std::make_unique<MultiscaleDeformableAttnPluginLegacy>(serialData, serialLength);
plugin->setPluginNamespace(getPluginNamespace());
return plugin.release();
}
catch (std::exception const& e)
{
caughtError(e);
}
return nullptr;
}
void MultiscaleDeformableAttnPluginCreatorLegacy::setPluginNamespace(char const* pluginNamespace) noexcept
{
mNamespace = pluginNamespace;
}
char const* MultiscaleDeformableAttnPluginCreatorLegacy::getPluginNamespace() const noexcept
{
return mNamespace.c_str();
}
} // namespace nvinfer1::plugin
@@ -0,0 +1,119 @@
/*
* 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.
*/
#ifndef TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_LEGACY_H
#define TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_LEGACY_H
// Standard library includes
#include <memory>
#include <string>
#include <vector>
#include "NvInferPlugin.h"
// TensorRT includes
#include "common/plugin.h"
namespace nvinfer1
{
namespace plugin
{
// Legacy V2 Plugin implementation
class MultiscaleDeformableAttnPluginLegacy : public nvinfer1::IPluginV2DynamicExt
{
public:
// Constructors/destructors
MultiscaleDeformableAttnPluginLegacy();
MultiscaleDeformableAttnPluginLegacy(void const* data, size_t length);
// IPluginV2DynamicExt methods
nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
nvinfer1::DimsExprs getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs, int32_t nbInputs,
nvinfer1::IExprBuilder& exprBuilder) noexcept override;
bool supportsFormatCombination(
int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext,
nvinfer1::IGpuAllocator* gpuAllocator) noexcept override;
void detachFromContext() noexcept override;
// IPluginV2Ext Methods
nvinfer1::DataType getOutputDataType(
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override;
// IPluginV2 Methods
char const* getPluginType() const noexcept override;
char const* getPluginVersion() const noexcept override;
int32_t getNbOutputs() const noexcept override;
int32_t initialize() noexcept override;
void terminate() noexcept override;
size_t getSerializationSize() const noexcept override;
void serialize(void* buffer) const noexcept override;
void destroy() noexcept override;
void setPluginNamespace(char const* pluginNamespace) noexcept override;
char const* getPluginNamespace() const noexcept override;
private:
std::string mNamespace;
#if NV_TENSORRT_MAJOR < 8
using nvinfer1::IPluginV2DynamicExt::configurePlugin;
using nvinfer1::IPluginV2DynamicExt::enqueue;
using nvinfer1::IPluginV2DynamicExt::getOutputDimensions;
using nvinfer1::IPluginV2DynamicExt::getWorkspaceSize;
using nvinfer1::IPluginV2DynamicExt::supportsFormat;
#endif
};
// Legacy creator class
class MultiscaleDeformableAttnPluginCreatorLegacy : public nvinfer1::IPluginCreator
{
public:
// Constructor
MultiscaleDeformableAttnPluginCreatorLegacy();
// IPluginCreator methods
char const* getPluginName() const noexcept override;
char const* getPluginVersion() const noexcept override;
nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
nvinfer1::IPluginV2* deserializePlugin(
char const* name, void const* serialData, size_t serialLength) noexcept override;
void setPluginNamespace(char const* pluginNamespace) noexcept override;
char const* getPluginNamespace() const noexcept override;
private:
nvinfer1::PluginFieldCollection mFC;
std::vector<nvinfer1::PluginField> mPluginAttributes;
std::string mNamespace;
};
} // namespace plugin
} // namespace nvinfer1
#endif // TRT_MULTISCALE_DEFORMABLE_ATTN_PLUGIN_LEGACY_H
File diff suppressed because it is too large Load Diff