chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,31 @@
|
||||
#
|
||||
# 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(
|
||||
skipLayerNormInt8InterleavedKernelHFace.cu
|
||||
skipLayerNormInt8InterleavedKernelMTron.cu
|
||||
skipLayerNormInt8InterleavedPlugin.cpp
|
||||
skipLayerNormInt8InterleavedPlugin.h
|
||||
skipLayerNormInt8InterleavedPluginLegacy.cpp
|
||||
skipLayerNormInt8InterleavedPluginLegacy.h
|
||||
skipLayerNormKernel.cu
|
||||
skipLayerNormPlugin.cpp
|
||||
skipLayerNormPlugin.h
|
||||
skipLayerNormPluginLegacy.cpp
|
||||
skipLayerNormPluginLegacy.h
|
||||
)
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
---
|
||||
name: CustomSkipLayerNormPluginDynamic
|
||||
interface: "IPluginV3"
|
||||
versions:
|
||||
"5": # SkipLayerNormPluginV3
|
||||
inputs:
|
||||
- input
|
||||
- skip
|
||||
outputs:
|
||||
- output
|
||||
input_dims:
|
||||
input: 5
|
||||
skip: 5
|
||||
input_dim_constraints:
|
||||
- "input_2 == bias_2"
|
||||
- "skip_0 == input_0"
|
||||
- "skip_1 == input_1"
|
||||
- "skip_2 == input_2"
|
||||
input_dim_range:
|
||||
input:
|
||||
min: "=1, =1, =1, =1, =1"
|
||||
max: "=pinf, =pinf, =pinf, =1, =1"
|
||||
skip:
|
||||
min: "=1, =1, =1, =1, =1"
|
||||
max: "=pinf, =pinf, =pinf, =1, =1"
|
||||
supported_input_types:
|
||||
- combination1:
|
||||
input: float32
|
||||
skip: float32
|
||||
- combination2:
|
||||
input: float16
|
||||
skip: float16
|
||||
output_dims:
|
||||
output: "input_0, input_1, input_2, input_3, input_4"
|
||||
attributes:
|
||||
- type_id
|
||||
- ld
|
||||
- beta
|
||||
- gamma
|
||||
- bias
|
||||
attribute_types:
|
||||
type_id: int32
|
||||
ld: int32
|
||||
beta: float32
|
||||
gamma: float32
|
||||
bias: float32
|
||||
attribute_dims:
|
||||
type_id: 1
|
||||
ld: 1
|
||||
beta: 3
|
||||
gamma: 3
|
||||
bias: 3
|
||||
attribute_dim_range:
|
||||
type_id:
|
||||
min: "=1"
|
||||
max: "=1"
|
||||
ld:
|
||||
min: "=1"
|
||||
max: "=1"
|
||||
beta:
|
||||
min: "=1, =1, =1"
|
||||
max: "=1, =1, =pinf"
|
||||
gamma:
|
||||
min: "=1, =1, =1"
|
||||
max: "=1, =1, =pinf"
|
||||
bias:
|
||||
min: "=1, =1, =1"
|
||||
max: "=1, =1, =pinf"
|
||||
attribute_options:
|
||||
type_id:
|
||||
- 0
|
||||
- 1
|
||||
- 2
|
||||
ld:
|
||||
min: "=1"
|
||||
max: "=pinf"
|
||||
beta:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
gamma:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
bias:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
attributes_required:
|
||||
- type_id
|
||||
- ld
|
||||
- beta
|
||||
- gamma
|
||||
golden_reference_script: "plugin/CustomSkipLayerNormPluginDynamic_PluginReference.py"
|
||||
abs_tol: 1e-2
|
||||
rel_tol: 1e-2
|
||||
configs:
|
||||
config1:
|
||||
input_types:
|
||||
input: float32
|
||||
skip: float32
|
||||
attribute_options:
|
||||
type_id:
|
||||
value: 0
|
||||
ld:
|
||||
value: 128
|
||||
beta:
|
||||
shape: "1, 1, 128"
|
||||
gamma:
|
||||
shape: "1, 1, 128"
|
||||
bias:
|
||||
shape: "1, 1, 128"
|
||||
config2:
|
||||
input_types:
|
||||
input: float16
|
||||
skip: float16
|
||||
attribute_options:
|
||||
type_id:
|
||||
value: 1
|
||||
ld:
|
||||
value: 768
|
||||
beta:
|
||||
shape: "1, 1, 768"
|
||||
gamma:
|
||||
shape: "1, 1, 768"
|
||||
bias:
|
||||
shape: "1, 1, 768"
|
||||
"6": # SkipLayerNormVarSeqlenPluginV3
|
||||
inputs:
|
||||
- input
|
||||
- skip
|
||||
outputs:
|
||||
- output
|
||||
input_dims:
|
||||
input: 5
|
||||
skip: 5
|
||||
input_dim_constraints:
|
||||
- "input_2 == bias_2"
|
||||
- "skip_0 == input_0"
|
||||
- "skip_1 == input_1"
|
||||
- "skip_2 == input_2"
|
||||
input_dim_range:
|
||||
input:
|
||||
min: "=1, =1, =1, =1, =1"
|
||||
max: "=pinf, =pinf, =pinf, =1, =1"
|
||||
skip:
|
||||
min: "=1, =1, =1, =1, =1"
|
||||
max: "=pinf, =pinf, =pinf, =1, =1"
|
||||
supported_input_types:
|
||||
- combination1:
|
||||
input: float32
|
||||
skip: float32
|
||||
- combination2:
|
||||
input: float16
|
||||
skip: float16
|
||||
output_dims:
|
||||
output: "input_0, input_1, input_2, input_3, input_4"
|
||||
attributes:
|
||||
- type_id
|
||||
- beta
|
||||
- gamma
|
||||
- bias
|
||||
attribute_types:
|
||||
type_id: int32
|
||||
beta: float32
|
||||
gamma: float32
|
||||
bias: float32
|
||||
attribute_dims:
|
||||
type_id: 1
|
||||
beta: 3
|
||||
gamma: 3
|
||||
bias: 3
|
||||
attribute_dim_range:
|
||||
type_id:
|
||||
min: "=1"
|
||||
max: "=1"
|
||||
beta:
|
||||
min: "=1, =1, =1"
|
||||
max: "=1, =1, =pinf"
|
||||
gamma:
|
||||
min: "=1, =1, =1"
|
||||
max: "=1, =1, =pinf"
|
||||
bias:
|
||||
min: "=1, =1, =1"
|
||||
max: "=1, =1, =pinf"
|
||||
attribute_options:
|
||||
type_id:
|
||||
- 0
|
||||
- 1
|
||||
- 2
|
||||
beta:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
gamma:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
bias:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
attributes_required:
|
||||
- type_id
|
||||
- beta
|
||||
- gamma
|
||||
golden_reference_script: "plugin/CustomSkipLayerNormPluginDynamic_PluginReference.py"
|
||||
abs_tol: 1e-2
|
||||
rel_tol: 1e-2
|
||||
configs:
|
||||
config1:
|
||||
input_types:
|
||||
input: float32
|
||||
skip: float32
|
||||
attribute_options:
|
||||
type_id:
|
||||
value: 0
|
||||
beta:
|
||||
shape: "1, 1, 128"
|
||||
gamma:
|
||||
shape: "1, 1, 128"
|
||||
bias:
|
||||
shape: "1, 1, 128"
|
||||
config2:
|
||||
input_types:
|
||||
input: float16
|
||||
skip: float16
|
||||
attribute_options:
|
||||
type_id:
|
||||
value: 1
|
||||
beta:
|
||||
shape: "1, 1, 768"
|
||||
gamma:
|
||||
shape: "1, 1, 768"
|
||||
bias:
|
||||
shape: "1, 1, 768"
|
||||
...
|
||||
@@ -0,0 +1,83 @@
|
||||
# skipLayerNormPlugin
|
||||
|
||||
**Table Of Contents**
|
||||
- [Description](#description)
|
||||
* [Structure](#structure)
|
||||
- [Parameters](#parameters)
|
||||
- [Additional resources](#additional-resources)
|
||||
- [License](#license)
|
||||
- [Changelog](#changelog)
|
||||
- [Known issues](#known-issues)
|
||||
|
||||
|
||||
## Description
|
||||
|
||||
> NOTE: Versions 1-4 of this plugin (using IPluginV2DynamicExt interface) are deprecated since TensorRT 10.4. Versions 5-8 (using IPluginV3 interface) are the recommended replacements.
|
||||
|
||||
Adds a residual tensor, applies layer normalization, i.e., transforms the mean and standard deviation to beta and gamma respectively.
|
||||
Optionally, adds a bias vector before layer-normalization.
|
||||
|
||||
|
||||
### Structure
|
||||
|
||||
The `skipLayerNormPlugin` takes two inputs; `input` and `skip`.
|
||||
|
||||
`input`
|
||||
For V1, V2, V5, V6, input is a tensor with shape `[S, B, E, 1, 1]` where `S` is the sequence length, `B` is the batch size, `E` is the hidden size, and the last two dimensions are of size 1.
|
||||
For V3 and V4, input is a tensor with shape `[1, E, S', 1]` where `S'` is the accumulated sequence length, `E` is the hidden size, and the first and last dimensions are of size 1.
|
||||
|
||||
`skip`
|
||||
skip has the same input dimensions as the input.
|
||||
The purpose of this input is to introduce skip (aka. residual) connections to previously computed tensors.
|
||||
|
||||
|
||||
The `skipLayerNormPlugin` generates the following output:
|
||||
|
||||
`output`
|
||||
output is a tensor with the same shape as the input.
|
||||
|
||||
|
||||
## Parameters
|
||||
|
||||
`skipLayerNormPlugin` has plugin creator class `SkipLayerNormPluginDynamicCreator` and plugin class `CustomSkipLayerNormPluginDynamic`.
|
||||
|
||||
The parameters are defined below and consists of the following attributes:
|
||||
|
||||
| Type | Parameter | Version | Description
|
||||
|----------|-----------------------------------------|-------------------------|-------------------------------------------------------------------
|
||||
|`int` |`type_id` | 1, 2, 5, 6 |Integer encoding the DataType (0: FP32, 1: FP16, 2: INT8)
|
||||
|`int` |`ld` | 1, 5 |The leading dimension of the input tensor, corresponding to the hidden size, denoted by `E` above.
|
||||
|`Weights` |`beta` | 1, 2, 3, 4, 5, 6, 7, 8 |The mean to normalize to. Shape: `[1, 1, E]`
|
||||
|`Weights` |`gamma` | 1, 2, 3, 4, 5, 6, 7, 8 |The standard deviation to normalize to. Shape: `[1, 1, E]`
|
||||
|`Weights` |`bias` | 1, 2, 5, 6 |An optional bias vector to add before normalization. Shape: `[1, 1, E]`
|
||||
|
||||
|
||||
## Additional resources
|
||||
|
||||
- [LayerNorm](https://arxiv.org/abs/1607.06450)
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
July 2024
|
||||
Add v5, v6, v7 and v8 plugins that duplicate the behavior of v1, v3, v3 and v4 plugins respectively, but implement the `IPluginV3` interface instead of the deprecated `IPluginV2DynamicExt` interface.
|
||||
|
||||
February 2024
|
||||
Add epsilon to avoid divide by zero.
|
||||
|
||||
October 2020
|
||||
Add V2 plugin that supports variable sequence length.
|
||||
Add v3 plugin that supports int8 interleaved variable sequence length.
|
||||
|
||||
November 2019
|
||||
This is the first release of this `README.md` file.
|
||||
|
||||
## Known issues
|
||||
|
||||
This plugin only supports GPUs with compute capability >= 7.0. For more information see the [CUDA GPU Compute Capability Support Matrix](https://developer.nvidia.com/cuda-gpus#compute)
|
||||
@@ -0,0 +1,363 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include "NvInfer.h"
|
||||
#include "common/bertCommon.h"
|
||||
#include "common/common.cuh"
|
||||
#include "common/cubCcclCompat.h"
|
||||
#include <cassert>
|
||||
#include <cstring>
|
||||
#include <cuda.h>
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
inline __device__ void resAdd(
|
||||
float (&hdata)[4], const uint32_t idata, const uint32_t ires, float const dqData, float const dqRes)
|
||||
{
|
||||
char4 ires4 = reinterpret_cast<char4 const&>(ires);
|
||||
char4 idata4 = reinterpret_cast<char4 const&>(idata);
|
||||
hdata[0] = float(idata4.x) * dqData + float(ires4.x) * dqRes;
|
||||
hdata[1] = float(idata4.y) * dqData + float(ires4.y) * dqRes;
|
||||
hdata[2] = float(idata4.z) * dqData + float(ires4.z) * dqRes;
|
||||
hdata[3] = float(idata4.w) * dqData + float(ires4.w) * dqRes;
|
||||
}
|
||||
|
||||
template <int32_t tWARPS, int32_t tHEADS, int32_t tTHREADS_PER_ROW>
|
||||
__global__ void skipln_vec32_hface(int8_t const* input, int8_t const* skip, int8_t* output, half const* beta,
|
||||
half const* gamma, float const dqScaleIn, float const dqScaleSkip, float const qScale, int32_t const total)
|
||||
{
|
||||
|
||||
// clang-format off
|
||||
enum { kHEAD_SIZE = 64 };
|
||||
enum { kBYTES_PER_LDG = 16 };
|
||||
enum { kTHREADS_PER_CTA = tWARPS * 32 };
|
||||
enum { kROWS_PER_LDG = kTHREADS_PER_CTA / tTHREADS_PER_ROW };
|
||||
enum { kVECS_PER_CTA = tTHREADS_PER_ROW / 2 };
|
||||
enum { kPARAM_BYTES = tHEADS * kHEAD_SIZE * 2 };
|
||||
enum { kPARAM_LDGS = kPARAM_BYTES / (kTHREADS_PER_CTA * kBYTES_PER_LDG) };
|
||||
enum { kLDGS = tHEADS * 2 / kROWS_PER_LDG };
|
||||
// clang-format on
|
||||
static_assert(kVECS_PER_CTA == 4, "");
|
||||
static_assert(kPARAM_LDGS == 1, "");
|
||||
static_assert(kROWS_PER_LDG == tHEADS, "");
|
||||
static_assert(kLDGS == 2, "");
|
||||
static_assert(kLDGS * kROWS_PER_LDG == tHEADS * 2, "");
|
||||
static_assert(kTHREADS_PER_CTA * kBYTES_PER_LDG == kPARAM_BYTES, "");
|
||||
static_assert(kPARAM_LDGS == 1, "");
|
||||
|
||||
extern __shared__ char smem_[];
|
||||
|
||||
// space for CTA-wide reduction
|
||||
__shared__ half2 smemRed[kVECS_PER_CTA][tWARPS];
|
||||
|
||||
constexpr float rld = 1.F / (float(tHEADS) * float(kHEAD_SIZE));
|
||||
int32_t const bidx = blockIdx.x;
|
||||
int32_t const tidx = threadIdx.x;
|
||||
int32_t const row = tidx / tTHREADS_PER_ROW;
|
||||
int32_t const col = tidx % tTHREADS_PER_ROW;
|
||||
int32_t const lane = tidx % 32;
|
||||
int32_t const warp = tidx / 32;
|
||||
|
||||
bool const isWarpLead = (lane < tTHREADS_PER_ROW) && ((lane & 1) == 0);
|
||||
bool const isCtaLead = (tidx < tTHREADS_PER_ROW) && ((tidx & 1) == 0);
|
||||
|
||||
// token position: every two threads load together the 32B at one token
|
||||
// position
|
||||
int32_t const pos = col / 2;
|
||||
|
||||
int32_t const posOffset = bidx * kVECS_PER_CTA + pos; // for token positions per block, disabling 2 threads per pos
|
||||
bool const myPred = posOffset < total;
|
||||
int32_t const rowStrideBytes = total * 32;
|
||||
|
||||
uint4 inData[kLDGS];
|
||||
uint4 inSkip[kLDGS];
|
||||
float hdata[kLDGS * 4][4];
|
||||
int32_t const gmemOffset = row * rowStrideBytes + (bidx * tTHREADS_PER_ROW + col) * kBYTES_PER_LDG;
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
inData[ii] = {0, 0, 0, 0};
|
||||
inSkip[ii] = {0, 0, 0, 0};
|
||||
if (myPred)
|
||||
{
|
||||
ldg(input + gmemOffset + ii * kROWS_PER_LDG * rowStrideBytes, inData[ii]);
|
||||
ldg(skip + gmemOffset + ii * kROWS_PER_LDG * rowStrideBytes, inSkip[ii]);
|
||||
}
|
||||
}
|
||||
|
||||
uint4* smemB = reinterpret_cast<uint4*>(&smem_[0]) + tidx;
|
||||
uint4* smemG = reinterpret_cast<uint4*>(&smem_[kPARAM_BYTES]) + tidx;
|
||||
|
||||
int8_t const* betaPtr = reinterpret_cast<int8_t const*>(beta) + tidx * kBYTES_PER_LDG;
|
||||
int8_t const* gammaPtr = reinterpret_cast<int8_t const*>(gamma) + tidx * kBYTES_PER_LDG;
|
||||
ldg(betaPtr, *smemB);
|
||||
ldg(gammaPtr, *smemG);
|
||||
|
||||
half* b = reinterpret_cast<half*>(&smem_[0]);
|
||||
half* g = reinterpret_cast<half*>(&smem_[kPARAM_BYTES]);
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
resAdd(hdata[ii * 4 + 0], inData[ii].x, inSkip[ii].x, dqScaleIn, dqScaleSkip);
|
||||
resAdd(hdata[ii * 4 + 1], inData[ii].y, inSkip[ii].y, dqScaleIn, dqScaleSkip);
|
||||
resAdd(hdata[ii * 4 + 2], inData[ii].z, inSkip[ii].z, dqScaleIn, dqScaleSkip);
|
||||
resAdd(hdata[ii * 4 + 3], inData[ii].w, inSkip[ii].w, dqScaleIn, dqScaleSkip);
|
||||
}
|
||||
|
||||
half2 statsLocal = {0, 0};
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS * 4; ii++)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t jj = 0; jj < 4; jj++)
|
||||
{
|
||||
float const tmp = hdata[ii][jj] * (rld);
|
||||
statsLocal = statsLocal + __floats2half2_rn(tmp, tmp * hdata[ii][jj]);
|
||||
}
|
||||
}
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 1);
|
||||
__syncwarp();
|
||||
|
||||
if (kVECS_PER_CTA == 1)
|
||||
{
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 2);
|
||||
__syncwarp();
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 4);
|
||||
__syncwarp();
|
||||
}
|
||||
else if (kVECS_PER_CTA == 2)
|
||||
{
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 4);
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 8);
|
||||
__syncwarp();
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 16);
|
||||
__syncwarp();
|
||||
|
||||
if (isWarpLead)
|
||||
{
|
||||
smemRed[pos][warp] = statsLocal;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (isCtaLead)
|
||||
{
|
||||
for (int32_t ii = 1; ii < tWARPS; ii++)
|
||||
{
|
||||
statsLocal = statsLocal + smemRed[pos][ii];
|
||||
}
|
||||
|
||||
float mu = __low2float(statsLocal);
|
||||
float sos = __high2float(statsLocal);
|
||||
float rsigma = rsqrtf(sos - mu * mu + std::numeric_limits<float>::epsilon());
|
||||
|
||||
smemRed[pos][0] = __floats2half2_rn(mu, rsigma);
|
||||
}
|
||||
__syncthreads();
|
||||
// load params into smem: 2x Headsx32x2x2B
|
||||
const float2 statsf = __half22float2(smemRed[pos][0]);
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t jj = 0; jj < 4; jj++)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t kk = 0; kk < 4; kk++)
|
||||
{
|
||||
int32_t const paramIdx = (ii * kROWS_PER_LDG + row) * 32 + (jj * 4 + kk) + (tidx & 1) * 16;
|
||||
float const bb = b[paramIdx];
|
||||
float const gg = g[paramIdx];
|
||||
hdata[ii * 4 + jj][kk] = gg * statsf.y * (hdata[ii * 4 + jj][kk] - statsf.x) + bb;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
inData[ii].x = pack4(hdata[ii * 4 + 0], qScale);
|
||||
inData[ii].y = pack4(hdata[ii * 4 + 1], qScale);
|
||||
inData[ii].z = pack4(hdata[ii * 4 + 2], qScale);
|
||||
inData[ii].w = pack4(hdata[ii * 4 + 3], qScale);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
if (myPred)
|
||||
{
|
||||
stg(output + gmemOffset + ii * kROWS_PER_LDG * rowStrideBytes, inData[ii]);
|
||||
}
|
||||
}
|
||||
// store
|
||||
}
|
||||
|
||||
int32_t launch_large_hface(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale)
|
||||
{
|
||||
if (ld == 1024)
|
||||
{
|
||||
constexpr int32_t tWARPS = 4;
|
||||
constexpr int32_t tTHREADS_PER_ROW = 8;
|
||||
constexpr int32_t tHEADS = 16;
|
||||
constexpr int32_t kPARAM_BYTES = tHEADS * 64 * 2 * sizeof(half);
|
||||
constexpr int32_t kVECS_PER_CTA = tTHREADS_PER_ROW / 2;
|
||||
int32_t const blocks = (total + kVECS_PER_CTA - 1) / kVECS_PER_CTA;
|
||||
|
||||
skipln_vec32_hface<tWARPS, tHEADS, tTHREADS_PER_ROW><<<blocks, tWARPS * 32, kPARAM_BYTES, stream>>>(
|
||||
input, skip, output, beta, gamma, dqScaleIn, dqScaleSkip, qScale, total);
|
||||
}
|
||||
else if (ld == 768)
|
||||
{
|
||||
constexpr int32_t tWARPS = 3;
|
||||
constexpr int32_t tTHREADS_PER_ROW = 8;
|
||||
constexpr int32_t tHEADS = 12;
|
||||
constexpr int32_t kPARAM_BYTES = tHEADS * 64 * 2 * sizeof(half);
|
||||
constexpr int32_t kVECS_PER_CTA = tTHREADS_PER_ROW / 2;
|
||||
int32_t const blocks = (total + kVECS_PER_CTA - 1) / kVECS_PER_CTA;
|
||||
|
||||
skipln_vec32_hface<tWARPS, tHEADS, tTHREADS_PER_ROW><<<blocks, tWARPS * 32, kPARAM_BYTES, stream>>>(
|
||||
input, skip, output, beta, gamma, dqScaleIn, dqScaleSkip, qScale, total);
|
||||
}
|
||||
else
|
||||
{
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
// naive kernel that only changes the addressing seems to be faster for small
|
||||
// problem sizes
|
||||
template <int32_t TPB, int32_t VPT>
|
||||
__global__ void skiplnDQQ_vec3(int32_t const ld, int8_t const* input, int8_t const* skip, int8_t* output,
|
||||
half const* beta, half const* gamma, float const dqScaleIn, float const dqScaleSkip, float const qScale,
|
||||
int32_t const total)
|
||||
{
|
||||
int32_t const hinner = threadIdx.x % 4;
|
||||
int32_t const houter = threadIdx.x / 4;
|
||||
|
||||
int32_t const tidx = threadIdx.x;
|
||||
int32_t const bidx = blockIdx.x;
|
||||
int32_t const idx = houter * total * 32 + bidx * 32 + hinner * VPT;
|
||||
// 4 * 1024 * 4 * 2 Bytes = 16KB per block
|
||||
int8_t inLocal[VPT];
|
||||
int8_t skipLocal[VPT];
|
||||
|
||||
half inLocalDQ[VPT]; // dequantized input + skip
|
||||
half betaLocal[VPT];
|
||||
half gammaLocal[VPT];
|
||||
|
||||
// load input tensors
|
||||
copy<sizeof(int8_t) * VPT>(&input[idx], inLocal);
|
||||
copy<sizeof(int8_t) * VPT>(&skip[idx], skipLocal);
|
||||
|
||||
// load parameters
|
||||
copy<sizeof(half) * VPT>(&beta[tidx * VPT], betaLocal);
|
||||
copy<sizeof(half) * VPT>(&gamma[tidx * VPT], gammaLocal);
|
||||
|
||||
half2 statsLocal = __floats2half2_rn(0.F, 0.F); // accumulator
|
||||
|
||||
half const rld = half(1.F) / half(ld);
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT; it++)
|
||||
{
|
||||
// DQ input and skip
|
||||
float const tmpIn = inLocal[it];
|
||||
float const tmpSkip = skipLocal[it];
|
||||
inLocalDQ[it] = dqScaleIn * tmpIn + dqScaleSkip * tmpSkip;
|
||||
|
||||
half const tmp = rld * inLocalDQ[it];
|
||||
half2 const tmp2 = __halves2half2(tmp, tmp * inLocalDQ[it]);
|
||||
statsLocal = statsLocal + tmp2;
|
||||
}
|
||||
|
||||
using BlockReduce = cub::BlockReduce<half2, TPB>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
__shared__ half mu; // mean
|
||||
__shared__ half rsigma; // 1 / std.dev.
|
||||
|
||||
half2 const sum2 = BlockReduce(tempStorage).Reduce(statsLocal, compat::getCudaSumOp());
|
||||
|
||||
if (tidx == 0)
|
||||
{
|
||||
mu = __low2half(sum2);
|
||||
rsigma = rsqrtf(__high2half(sum2) - mu * mu + std::numeric_limits<half>::epsilon());
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
static_assert(VPT % 4 == 0, "");
|
||||
uint32_t outLocal[VPT / 4];
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT / 4; it++)
|
||||
{
|
||||
float const tmp0 = gammaLocal[it * 4 + 0] * (inLocalDQ[it * 4 + 0] - mu) * rsigma + betaLocal[it * 4 + 0];
|
||||
float const tmp1 = gammaLocal[it * 4 + 1] * (inLocalDQ[it * 4 + 1] - mu) * rsigma + betaLocal[it * 4 + 1];
|
||||
float const tmp2 = gammaLocal[it * 4 + 2] * (inLocalDQ[it * 4 + 2] - mu) * rsigma + betaLocal[it * 4 + 2];
|
||||
float const tmp3 = gammaLocal[it * 4 + 3] * (inLocalDQ[it * 4 + 3] - mu) * rsigma + betaLocal[it * 4 + 3];
|
||||
outLocal[it] = float4_to_char4(tmp0 * qScale, tmp1 * qScale, tmp2 * qScale, tmp3 * qScale);
|
||||
}
|
||||
|
||||
copy<sizeof(int8_t) * VPT>(outLocal, &output[idx]);
|
||||
}
|
||||
|
||||
int32_t launch_small_hface(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale)
|
||||
{
|
||||
int32_t const gridSize = total;
|
||||
// we align reads with the number of parameters, i.e. 8-wide instead of 16
|
||||
constexpr int32_t VPT = 16 / sizeof(half); // 8
|
||||
if (ld == 768)
|
||||
{
|
||||
constexpr int32_t TPB = 768 / VPT;
|
||||
skiplnDQQ_vec3<TPB, VPT>
|
||||
<<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, dqScaleIn, dqScaleSkip, qScale, total);
|
||||
}
|
||||
else if (ld == 1024)
|
||||
{
|
||||
constexpr int32_t TPB = 1024 / VPT; // 128
|
||||
skiplnDQQ_vec3<TPB, VPT>
|
||||
<<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, dqScaleIn, dqScaleSkip, qScale, total);
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "SkipLayerNormDQQ - FATAL: unsupported hidden layer size: " << ld << std::endl;
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
@@ -0,0 +1,415 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include "NvInfer.h"
|
||||
#include "common/bertCommon.h"
|
||||
#include "common/common.cuh"
|
||||
#include "common/cubCcclCompat.h"
|
||||
#include <cassert>
|
||||
#include <cstring>
|
||||
#include <cuda.h>
|
||||
#include <type_traits>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
inline __device__ void res_add(
|
||||
float (&hdata)[4], const uint32_t idata, const uint32_t ires, float const dqData, float const dqRes)
|
||||
{
|
||||
char4 ires4 = reinterpret_cast<char4 const&>(ires);
|
||||
char4 idata4 = reinterpret_cast<char4 const&>(idata);
|
||||
hdata[0] = float(idata4.x) * dqData + float(ires4.x) * dqRes;
|
||||
hdata[1] = float(idata4.y) * dqData + float(ires4.y) * dqRes;
|
||||
hdata[2] = float(idata4.z) * dqData + float(ires4.z) * dqRes;
|
||||
hdata[3] = float(idata4.w) * dqData + float(ires4.w) * dqRes;
|
||||
}
|
||||
|
||||
template <int32_t tWARPS, int32_t tHEADS, int32_t tTHREADS_PER_ROW>
|
||||
__global__ void skipln_vec32_mtron(int8_t const* input, int8_t const* skip, int8_t* output, int8_t* preln,
|
||||
half const* beta, half const* gamma, float const dqScaleIn, float const dqScaleSkip, float const qScale,
|
||||
float const qSkipScale, int32_t const total)
|
||||
{
|
||||
|
||||
// clang-format off
|
||||
enum { kHEAD_SIZE = 64 };
|
||||
enum { kBYTES_PER_LDG = 16 };
|
||||
enum { kTHREADS_PER_CTA = tWARPS * 32 };
|
||||
enum { kROWS_PER_LDG = kTHREADS_PER_CTA / tTHREADS_PER_ROW };
|
||||
enum { kVECS_PER_CTA = tTHREADS_PER_ROW / 2 };
|
||||
enum { kPARAM_BYTES = tHEADS * kHEAD_SIZE * 2 };
|
||||
enum { kPARAM_LDGS = kPARAM_BYTES / (kTHREADS_PER_CTA * kBYTES_PER_LDG) };
|
||||
enum { kLDGS = tHEADS * 2 / kROWS_PER_LDG };
|
||||
// clang-format on
|
||||
static_assert(kVECS_PER_CTA == 4, "");
|
||||
static_assert(kPARAM_LDGS == 1, "");
|
||||
static_assert(kROWS_PER_LDG == tHEADS, "");
|
||||
static_assert(kLDGS == 2, "");
|
||||
static_assert(kLDGS * kROWS_PER_LDG == tHEADS * 2, "");
|
||||
static_assert(kTHREADS_PER_CTA * kBYTES_PER_LDG == kPARAM_BYTES, "");
|
||||
static_assert(kPARAM_LDGS == 1, "");
|
||||
|
||||
extern __shared__ char smem_[];
|
||||
|
||||
// space for CTA-wide reduction
|
||||
__shared__ half2 smemRed[kVECS_PER_CTA][tWARPS];
|
||||
|
||||
constexpr float rld = 1.F / (float(tHEADS) * float(kHEAD_SIZE));
|
||||
int32_t const bidx = blockIdx.x;
|
||||
int32_t const tidx = threadIdx.x;
|
||||
int32_t const row = tidx / tTHREADS_PER_ROW;
|
||||
int32_t const col = tidx % tTHREADS_PER_ROW;
|
||||
int32_t const lane = tidx % 32;
|
||||
int32_t const warp = tidx / 32;
|
||||
|
||||
bool const isWarpLead = (lane < tTHREADS_PER_ROW) && ((lane & 1) == 0);
|
||||
bool const isCtaLead = (tidx < tTHREADS_PER_ROW) && ((tidx & 1) == 0);
|
||||
|
||||
// token position: every two threads load together the 32B at one token
|
||||
// position
|
||||
int32_t const pos = col / 2;
|
||||
|
||||
int32_t const posOffset = bidx * kVECS_PER_CTA + pos; // for token positions per block, disabling 2 threads per pos
|
||||
bool const myPred = posOffset < total;
|
||||
int32_t const rowStrideBytes = total * 32;
|
||||
|
||||
uint4 inData[kLDGS];
|
||||
uint4 in_skip[kLDGS];
|
||||
float hdata[kLDGS * 4][4];
|
||||
int32_t const gmemOffset = row * rowStrideBytes + (bidx * tTHREADS_PER_ROW + col) * kBYTES_PER_LDG;
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
inData[ii] = {0, 0, 0, 0};
|
||||
in_skip[ii] = {0, 0, 0, 0};
|
||||
if (myPred)
|
||||
{
|
||||
ldg(input + gmemOffset + ii * kROWS_PER_LDG * rowStrideBytes, inData[ii]);
|
||||
ldg(skip + gmemOffset + ii * kROWS_PER_LDG * rowStrideBytes, in_skip[ii]);
|
||||
}
|
||||
}
|
||||
|
||||
uint4* smemB = reinterpret_cast<uint4*>(&smem_[0]) + tidx;
|
||||
uint4* smemG = reinterpret_cast<uint4*>(&smem_[kPARAM_BYTES]) + tidx;
|
||||
|
||||
int8_t const* betaPtr = reinterpret_cast<int8_t const*>(beta) + tidx * kBYTES_PER_LDG;
|
||||
int8_t const* gammaPtr = reinterpret_cast<int8_t const*>(gamma) + tidx * kBYTES_PER_LDG;
|
||||
ldg(betaPtr, *smemB);
|
||||
ldg(gammaPtr, *smemG);
|
||||
|
||||
half* b = reinterpret_cast<half*>(&smem_[0]);
|
||||
half* g = reinterpret_cast<half*>(&smem_[kPARAM_BYTES]);
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
res_add(hdata[ii * 4 + 0], inData[ii].x, in_skip[ii].x, dqScaleIn, dqScaleSkip);
|
||||
res_add(hdata[ii * 4 + 1], inData[ii].y, in_skip[ii].y, dqScaleIn, dqScaleSkip);
|
||||
res_add(hdata[ii * 4 + 2], inData[ii].z, in_skip[ii].z, dqScaleIn, dqScaleSkip);
|
||||
res_add(hdata[ii * 4 + 3], inData[ii].w, in_skip[ii].w, dqScaleIn, dqScaleSkip);
|
||||
}
|
||||
|
||||
half2 statsLocal = {0, 0};
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS * 4; ii++)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t jj = 0; jj < 4; jj++)
|
||||
{
|
||||
float const tmp = hdata[ii][jj] * (rld);
|
||||
statsLocal = statsLocal + __floats2half2_rn(tmp, tmp * hdata[ii][jj]);
|
||||
}
|
||||
}
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 1);
|
||||
__syncwarp();
|
||||
|
||||
if (kVECS_PER_CTA == 1)
|
||||
{
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 2);
|
||||
__syncwarp();
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 4);
|
||||
__syncwarp();
|
||||
}
|
||||
else if (kVECS_PER_CTA == 2)
|
||||
{
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 4);
|
||||
__syncwarp();
|
||||
}
|
||||
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 8);
|
||||
__syncwarp();
|
||||
statsLocal = statsLocal + __shfl_xor_sync(uint32_t(-1), statsLocal, 16);
|
||||
__syncwarp();
|
||||
|
||||
if (isWarpLead)
|
||||
{
|
||||
smemRed[pos][warp] = statsLocal;
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
if (isCtaLead)
|
||||
{
|
||||
for (int32_t ii = 1; ii < tWARPS; ii++)
|
||||
{
|
||||
statsLocal = statsLocal + smemRed[pos][ii];
|
||||
}
|
||||
|
||||
float mu = __low2float(statsLocal);
|
||||
float sos = __high2float(statsLocal);
|
||||
float rsigma = rsqrtf(sos - mu * mu + std::numeric_limits<float>::epsilon());
|
||||
|
||||
smemRed[pos][0] = __floats2half2_rn(mu, rsigma);
|
||||
}
|
||||
__syncthreads();
|
||||
// load params into smem: 2x Headsx32x2x2B
|
||||
const float2 statsf = __half22float2(smemRed[pos][0]);
|
||||
|
||||
// Copy skip connection output before Layer Norm
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
inData[ii].x = pack4(hdata[ii * 4 + 0], qSkipScale);
|
||||
inData[ii].y = pack4(hdata[ii * 4 + 1], qSkipScale);
|
||||
inData[ii].z = pack4(hdata[ii * 4 + 2], qSkipScale);
|
||||
inData[ii].w = pack4(hdata[ii * 4 + 3], qSkipScale);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
if (myPred)
|
||||
{
|
||||
stg(preln + gmemOffset + ii * kROWS_PER_LDG * rowStrideBytes, inData[ii]);
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t jj = 0; jj < 4; jj++)
|
||||
{
|
||||
#pragma unroll
|
||||
for (int32_t kk = 0; kk < 4; kk++)
|
||||
{
|
||||
int32_t const paramIdx = (ii * kROWS_PER_LDG + row) * 32 + (jj * 4 + kk) + (tidx & 1) * 16;
|
||||
float const bb = b[paramIdx];
|
||||
float const gg = g[paramIdx];
|
||||
hdata[ii * 4 + jj][kk] = gg * statsf.y * (hdata[ii * 4 + jj][kk] - statsf.x) + bb;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
inData[ii].x = pack4(hdata[ii * 4 + 0], qScale);
|
||||
inData[ii].y = pack4(hdata[ii * 4 + 1], qScale);
|
||||
inData[ii].z = pack4(hdata[ii * 4 + 2], qScale);
|
||||
inData[ii].w = pack4(hdata[ii * 4 + 3], qScale);
|
||||
}
|
||||
|
||||
#pragma unroll
|
||||
for (int32_t ii = 0; ii < kLDGS; ii++)
|
||||
{
|
||||
if (myPred)
|
||||
{
|
||||
stg(output + gmemOffset + ii * kROWS_PER_LDG * rowStrideBytes, inData[ii]);
|
||||
}
|
||||
}
|
||||
// store
|
||||
}
|
||||
|
||||
int32_t launch_large_mtron(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, int8_t* preln, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale, float const qSkipScale)
|
||||
{
|
||||
if (ld == 1024)
|
||||
{
|
||||
constexpr int32_t tWARPS = 4;
|
||||
constexpr int32_t tTHREADS_PER_ROW = 8;
|
||||
constexpr int32_t tHEADS = 16;
|
||||
constexpr int32_t kPARAM_BYTES = tHEADS * 64 * 2 * sizeof(half);
|
||||
constexpr int32_t kVECS_PER_CTA = tTHREADS_PER_ROW / 2;
|
||||
int32_t const blocks = (total + kVECS_PER_CTA - 1) / kVECS_PER_CTA;
|
||||
|
||||
skipln_vec32_mtron<tWARPS, tHEADS, tTHREADS_PER_ROW><<<blocks, tWARPS * 32, kPARAM_BYTES, stream>>>(
|
||||
input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else if (ld == 768)
|
||||
{
|
||||
constexpr int32_t tWARPS = 3;
|
||||
constexpr int32_t tTHREADS_PER_ROW = 8;
|
||||
constexpr int32_t tHEADS = 12;
|
||||
constexpr int32_t kPARAM_BYTES = tHEADS * 64 * 2 * sizeof(half);
|
||||
constexpr int32_t kVECS_PER_CTA = tTHREADS_PER_ROW / 2;
|
||||
int32_t const blocks = (total + kVECS_PER_CTA - 1) / kVECS_PER_CTA;
|
||||
|
||||
skipln_vec32_mtron<tWARPS, tHEADS, tTHREADS_PER_ROW><<<blocks, tWARPS * 32, kPARAM_BYTES, stream>>>(
|
||||
input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else
|
||||
{
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
// naive kernel that only changes the addressing seems to be faster for small
|
||||
// problem sizes
|
||||
template <int32_t TPB, int32_t VPT>
|
||||
__global__ void skiplnDQQ_vec4(int32_t const ld, int8_t const* input, int8_t const* skip, int8_t* output, int8_t* preln,
|
||||
half const* beta, half const* gamma, float const dqScaleIn, float const dqScaleSkip, float const qScale,
|
||||
float const qSkipScale, int32_t const total)
|
||||
{
|
||||
int32_t const hinner = threadIdx.x % 4;
|
||||
int32_t const houter = threadIdx.x / 4;
|
||||
|
||||
int32_t const tidx = threadIdx.x;
|
||||
int32_t const bidx = blockIdx.x;
|
||||
int32_t const idx = houter * total * 32 + bidx * 32 + hinner * VPT;
|
||||
// 4 * 1024 * 4 * 2 Bytes = 16KB per block
|
||||
int8_t inLocal[VPT];
|
||||
int8_t skipLocal[VPT];
|
||||
|
||||
half inLocalDQ[VPT]; // dequantized input + skip
|
||||
half betaLocal[VPT];
|
||||
half gammaLocal[VPT];
|
||||
|
||||
// load input tensors
|
||||
copy<sizeof(int8_t) * VPT>(&input[idx], inLocal);
|
||||
copy<sizeof(int8_t) * VPT>(&skip[idx], skipLocal);
|
||||
|
||||
// load parameters
|
||||
copy<sizeof(half) * VPT>(&beta[tidx * VPT], betaLocal);
|
||||
copy<sizeof(half) * VPT>(&gamma[tidx * VPT], gammaLocal);
|
||||
|
||||
half2 statsLocal = __floats2half2_rn(0.F, 0.F); // accumulator
|
||||
|
||||
half const rld = half(1.F) / half(ld);
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT; it++)
|
||||
{
|
||||
// DQ input and skip
|
||||
float const tmpIn = inLocal[it];
|
||||
float const tmpSkip = skipLocal[it];
|
||||
inLocalDQ[it] = dqScaleIn * tmpIn + dqScaleSkip * tmpSkip;
|
||||
|
||||
half const tmp = rld * inLocalDQ[it];
|
||||
half2 const tmp2 = __halves2half2(tmp, tmp * inLocalDQ[it]);
|
||||
statsLocal = statsLocal + tmp2;
|
||||
}
|
||||
|
||||
using BlockReduce = cub::BlockReduce<half2, TPB>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
__shared__ half mu; // mean
|
||||
__shared__ half rsigma; // 1 / std.dev.
|
||||
|
||||
half2 const sum2 = BlockReduce(tempStorage).Reduce(statsLocal, compat::getCudaSumOp());
|
||||
|
||||
// Copy skip connection output before Layer Norm
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT; it++)
|
||||
{
|
||||
inLocal[it] = quantize(inLocalDQ[it], qSkipScale);
|
||||
}
|
||||
copy<sizeof(int8_t) * VPT>(inLocal, &preln[idx]);
|
||||
|
||||
if (tidx == 0)
|
||||
{
|
||||
mu = __low2half(sum2);
|
||||
rsigma = rsqrtf(__high2half(sum2) - mu * mu + std::numeric_limits<half>::epsilon());
|
||||
}
|
||||
|
||||
__syncthreads();
|
||||
|
||||
static_assert(VPT % 4 == 0, "");
|
||||
uint32_t outLocal[VPT / 4U];
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT / 4U; it++)
|
||||
{
|
||||
float const tmp0 = gammaLocal[it * 4 + 0] * (inLocalDQ[it * 4 + 0] - mu) * rsigma + betaLocal[it * 4 + 0];
|
||||
float const tmp1 = gammaLocal[it * 4 + 1] * (inLocalDQ[it * 4 + 1] - mu) * rsigma + betaLocal[it * 4 + 1];
|
||||
float const tmp2 = gammaLocal[it * 4 + 2] * (inLocalDQ[it * 4 + 2] - mu) * rsigma + betaLocal[it * 4 + 2];
|
||||
float const tmp3 = gammaLocal[it * 4 + 3] * (inLocalDQ[it * 4 + 3] - mu) * rsigma + betaLocal[it * 4 + 3];
|
||||
outLocal[it] = float4_to_char4(tmp0 * qScale, tmp1 * qScale, tmp2 * qScale, tmp3 * qScale);
|
||||
}
|
||||
|
||||
copy<sizeof(int8_t) * VPT>(outLocal, &output[idx]);
|
||||
}
|
||||
|
||||
int32_t launch_small_mtron(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, int8_t* preln, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale, float const qSkipScale)
|
||||
{
|
||||
int32_t const gridSize = total;
|
||||
// we align reads with the number of parameters, i.e. 8-wide instead of 16
|
||||
int32_t constexpr VPT = 16 / sizeof(half); // 8
|
||||
if (ld == 768)
|
||||
{
|
||||
int32_t constexpr TPB = 768 / VPT;
|
||||
skiplnDQQ_vec4<TPB, VPT><<<gridSize, TPB, 0, stream>>>(
|
||||
ld, input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else if (ld == 1024)
|
||||
{
|
||||
int32_t constexpr TPB = 1024 / VPT; // 128
|
||||
skiplnDQQ_vec4<TPB, VPT><<<gridSize, TPB, 0, stream>>>(
|
||||
ld, input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else if (ld == 1536)
|
||||
{
|
||||
int32_t constexpr TPB = 1536 / VPT; // 192
|
||||
skiplnDQQ_vec4<TPB, VPT><<<gridSize, TPB, 0, stream>>>(
|
||||
ld, input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else if (ld == 2048)
|
||||
{
|
||||
int32_t constexpr TPB = 2048 / VPT; // 256
|
||||
skiplnDQQ_vec4<TPB, VPT><<<gridSize, TPB, 0, stream>>>(
|
||||
ld, input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else if (ld == 3072)
|
||||
{
|
||||
int32_t constexpr TPB = 3072 / VPT; // 384
|
||||
skiplnDQQ_vec4<TPB, VPT><<<gridSize, TPB, 0, stream>>>(
|
||||
ld, input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else if (ld == 4096)
|
||||
{
|
||||
int32_t constexpr TPB = 4096 / VPT; // 512
|
||||
skiplnDQQ_vec4<TPB, VPT><<<gridSize, TPB, 0, stream>>>(
|
||||
ld, input, skip, output, preln, beta, gamma, dqScaleIn, dqScaleSkip, qScale, qSkipScale, total);
|
||||
}
|
||||
else
|
||||
{
|
||||
std::cout << "SkipLayerNormDQQ - FATAL: unsupported hidden layer size: " << ld << std::endl;
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
@@ -0,0 +1,584 @@
|
||||
/*
|
||||
* 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 "skipLayerNormInt8InterleavedPlugin.h"
|
||||
#include "NvInfer.h"
|
||||
#include "common/serialize.hpp"
|
||||
|
||||
#include <cuda.h>
|
||||
|
||||
#include <cstring>
|
||||
#include <memory>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::plugin::bert;
|
||||
|
||||
// Clip plugin specific constants
|
||||
namespace
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
constexpr char const* kSKIP_LAYER_NORM_INTERLEAVED_VERSION_HFACE{"7"};
|
||||
constexpr char const* kSKIP_LAYER_NORM_INTERLEAVED_VERSION_MTRON{"8"};
|
||||
constexpr char const* kSKIP_LAYER_NORM_INTERLEAVED_NAME{"CustomSkipLayerNormPluginDynamic"};
|
||||
|
||||
void checkDescs(PluginTensorDesc const& iDesc, PluginTensorDesc const& sDesc, PluginTensorDesc const& oDesc)
|
||||
{
|
||||
PLUGIN_VALIDATE(iDesc.dims.nbDims == 4);
|
||||
PLUGIN_VALIDATE(iDesc.dims.nbDims == sDesc.dims.nbDims);
|
||||
PLUGIN_VALIDATE(std::equal(iDesc.dims.d, iDesc.dims.d + iDesc.dims.nbDims, sDesc.dims.d));
|
||||
PLUGIN_VALIDATE(std::equal(iDesc.dims.d, iDesc.dims.d + iDesc.dims.nbDims, oDesc.dims.d));
|
||||
PLUGIN_VALIDATE(iDesc.dims.d[0] == 1);
|
||||
PLUGIN_VALIDATE(iDesc.dims.d[3] == 1);
|
||||
PLUGIN_VALIDATE(iDesc.format == TensorFormat::kCHW32);
|
||||
PLUGIN_VALIDATE(iDesc.type == DataType::kINT8);
|
||||
PLUGIN_VALIDATE(iDesc.format == sDesc.format);
|
||||
PLUGIN_VALIDATE(iDesc.format == oDesc.format);
|
||||
PLUGIN_VALIDATE(iDesc.type == sDesc.type);
|
||||
PLUGIN_VALIDATE(iDesc.type == oDesc.type);
|
||||
}
|
||||
|
||||
void buildBetaAndGamma(PluginFieldCollection const* fc, Weights& beta, Weights& gamma)
|
||||
{
|
||||
PLUGIN_VALIDATE(fc != nullptr, "SkipLayerNorm: Plugin Field collection is null");
|
||||
PLUGIN_VALIDATE(fc->fields != nullptr, "SkipLayerNorm: Plugin Fields are null");
|
||||
plugin::validateRequiredAttributesExist({"beta", "gamma"}, fc);
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; i++)
|
||||
{
|
||||
std::string_view const fieldName = fc->fields[i].name;
|
||||
|
||||
if (fieldName == "beta"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building beta...");
|
||||
beta.values = fc->fields[i].data;
|
||||
beta.count = fc->fields[i].length;
|
||||
beta.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (fieldName == "gamma"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building gamma...");
|
||||
gamma.values = fc->fields[i].data;
|
||||
gamma.count = fc->fields[i].length;
|
||||
gamma.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
}
|
||||
|
||||
PLUGIN_VALIDATE(beta.values != nullptr, "SkipLayerNorm: invalid beta");
|
||||
PLUGIN_VALIDATE(beta.count > 0, "SkipLayerNorm: invalid beta");
|
||||
|
||||
PLUGIN_VALIDATE(gamma.values != nullptr, "SkipLayerNorm: invalid gamma");
|
||||
PLUGIN_VALIDATE(gamma.count > 0, "SkipLayerNorm: invalid gamma");
|
||||
}
|
||||
} // namespace
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(SkipLayerNormInterleavedPluginHFaceCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(SkipLayerNormInterleavedPluginMTronCreator);
|
||||
|
||||
constexpr auto kPARAM_TYPE = DataType::kHALF;
|
||||
|
||||
SkipLayerNormInterleavedPluginBase::SkipLayerNormInterleavedPluginBase(
|
||||
std::string const& name, Weights const& beta, Weights const& gamma)
|
||||
: mLayerName(name)
|
||||
, mGammaDev(nullptr)
|
||||
, mBetaDev(nullptr)
|
||||
, mLd(beta.count)
|
||||
, mParamsOnDevice(false)
|
||||
{
|
||||
PLUGIN_VALIDATE(mLd > 0);
|
||||
PLUGIN_VALIDATE(beta.count == gamma.count);
|
||||
// dataType for beta, gamma weights is always fp16
|
||||
|
||||
mParamWordsize = getElementSize(kPARAM_TYPE);
|
||||
|
||||
mBeta.convertAndCopy(beta, kPARAM_TYPE);
|
||||
mGamma.convertAndCopy(gamma, kPARAM_TYPE);
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginHFace::SkipLayerNormInterleavedPluginHFace(
|
||||
std::string const& name, Weights const& beta, Weights const& gamma)
|
||||
: SkipLayerNormInterleavedPluginBase(name, beta, gamma)
|
||||
{
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginMTron::SkipLayerNormInterleavedPluginMTron(
|
||||
std::string const& name, Weights const& beta, Weights const& gamma)
|
||||
: SkipLayerNormInterleavedPluginBase(name, beta, gamma)
|
||||
{
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginBase::~SkipLayerNormInterleavedPluginBase()
|
||||
{
|
||||
try
|
||||
{
|
||||
mGammaDev.reset(nullptr);
|
||||
mBetaDev.reset(nullptr);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginHFace::~SkipLayerNormInterleavedPluginHFace()
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFace destructor");
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginMTron::~SkipLayerNormInterleavedPluginMTron()
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTron destructor");
|
||||
}
|
||||
|
||||
//////
|
||||
// IPluginV3 method definitions:
|
||||
// - getCapabilityInterface() (Base)
|
||||
// - clone() (HFace, MTron)
|
||||
//////
|
||||
IPluginCapability* SkipLayerNormInterleavedPluginBase::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;
|
||||
}
|
||||
|
||||
IPluginV3* SkipLayerNormInterleavedPluginHFace::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFace clone");
|
||||
|
||||
auto p = std::make_unique<SkipLayerNormInterleavedPluginHFace>(mLayerName, mBeta, mGamma);
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV3* SkipLayerNormInterleavedPluginMTron::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTron clone");
|
||||
|
||||
auto p = std::make_unique<SkipLayerNormInterleavedPluginMTron>(mLayerName, mBeta, mGamma);
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
// End IPluginV3 method definitions
|
||||
|
||||
//////
|
||||
// IPluginV3OneRuntime method definitions:
|
||||
// - getFieldsToSerialize() (Base)
|
||||
// - onShapeChange() (Base)
|
||||
// - attachToContext() (HFace, MTron)
|
||||
// - execute() (HFace, MTron)
|
||||
/////
|
||||
PluginFieldCollection const* SkipLayerNormInterleavedPluginBase::getFieldsToSerialize() noexcept
|
||||
{
|
||||
mDataToSerialize.clear();
|
||||
mDataToSerialize.emplace_back(
|
||||
"beta", static_cast<half const*>(mBeta.values), PluginFieldType::kFLOAT16, mBeta.count);
|
||||
PLUGIN_ASSERT(mBeta.type == kPARAM_TYPE);
|
||||
mDataToSerialize.emplace_back(
|
||||
"gamma", static_cast<half const*>(mGamma.values), PluginFieldType::kFLOAT16, mGamma.count);
|
||||
PLUGIN_ASSERT(mGamma.type == kPARAM_TYPE);
|
||||
mFCToSerialize.nbFields = mDataToSerialize.size();
|
||||
mFCToSerialize.fields = mDataToSerialize.data();
|
||||
return &mFCToSerialize;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginBase::onShapeChange(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// Validate input arguments
|
||||
PLUGIN_VALIDATE(inputs != nullptr);
|
||||
PLUGIN_VALIDATE(outputs != nullptr);
|
||||
PLUGIN_VALIDATE(nbOutputs == getNbOutputs());
|
||||
PLUGIN_VALIDATE(nbInputs == 2);
|
||||
PLUGIN_VALIDATE(DataType::kINT8 == inputs[0].type);
|
||||
PLUGIN_VALIDATE(DataType::kINT8 == inputs[1].type);
|
||||
|
||||
auto const& inDims0 = inputs[0].dims;
|
||||
auto const& inDims1 = inputs[1].dims;
|
||||
TRT_UNUSED inDims1;
|
||||
|
||||
PLUGIN_VALIDATE(inDims0.nbDims == inDims1.nbDims);
|
||||
PLUGIN_VALIDATE(std::equal(inDims0.d, inDims0.d + inDims0.nbDims, inDims1.d));
|
||||
|
||||
mParamWordsize = getElementSize(kPARAM_TYPE);
|
||||
|
||||
if (!mParamsOnDevice)
|
||||
{
|
||||
copyToDevice(mGamma, getWeightsSize(mGamma, kPARAM_TYPE), mGammaDev);
|
||||
copyToDevice(mBeta, getWeightsSize(mBeta, kPARAM_TYPE), mBetaDev);
|
||||
mParamsOnDevice = true;
|
||||
}
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
IPluginV3* SkipLayerNormInterleavedPluginBase::attachToContext(IPluginResourceContext* context) noexcept
|
||||
{
|
||||
return clone();
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginHFace::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);
|
||||
|
||||
// Input shape: 1x(hxd)xtotalx1
|
||||
auto const iDesc = inputDesc[0];
|
||||
auto const sDesc = inputDesc[1];
|
||||
auto const oDesc = outputDesc[0];
|
||||
checkDescs(iDesc, sDesc, oDesc);
|
||||
|
||||
const int32_t ld = iDesc.dims.d[1];
|
||||
const int32_t total = iDesc.dims.d[2];
|
||||
float const dqScaleIn = iDesc.scale;
|
||||
float const dqScaleSkip = sDesc.scale;
|
||||
float const qScale = 1.F / oDesc.scale;
|
||||
int8_t const* input = static_cast<int8_t const*>(inputs[0]);
|
||||
int8_t const* skip = static_cast<int8_t const*>(inputs[1]);
|
||||
int8_t* output = static_cast<int8_t*>(outputs[0]);
|
||||
half const* gamma = static_cast<half const*>(mGammaDev.get());
|
||||
half const* beta = static_cast<half const*>(mBetaDev.get());
|
||||
|
||||
if (total < 4096)
|
||||
{
|
||||
return launch_small_hface(
|
||||
stream, ld, total, input, skip, beta, gamma, output, dqScaleIn, dqScaleSkip, qScale);
|
||||
}
|
||||
|
||||
return launch_large_hface(stream, ld, total, input, skip, beta, gamma, output, dqScaleIn, dqScaleSkip, qScale);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginMTron::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);
|
||||
|
||||
// Input shape: 1x(hxd)xtotalx1
|
||||
auto const iDesc = inputDesc[0];
|
||||
auto const sDesc = inputDesc[1];
|
||||
auto const oDesc = outputDesc[0];
|
||||
auto const pDesc = outputDesc[1];
|
||||
checkDescs(iDesc, sDesc, oDesc);
|
||||
PLUGIN_VALIDATE(std::equal(iDesc.dims.d, iDesc.dims.d + iDesc.dims.nbDims, pDesc.dims.d));
|
||||
|
||||
const int32_t ld = iDesc.dims.d[1];
|
||||
const int32_t total = iDesc.dims.d[2];
|
||||
float const dqScaleIn = iDesc.scale;
|
||||
float const dqScaleSkip = sDesc.scale;
|
||||
float const qScale = 1.F / oDesc.scale;
|
||||
float const qSkipScale = 1.F / pDesc.scale;
|
||||
int8_t const* input = static_cast<int8_t const*>(inputs[0]);
|
||||
int8_t const* skip = static_cast<int8_t const*>(inputs[1]);
|
||||
int8_t* output = static_cast<int8_t*>(outputs[0]);
|
||||
int8_t* preln = static_cast<int8_t*>(outputs[1]);
|
||||
half const* gamma = static_cast<half const*>(mGammaDev.get());
|
||||
half const* beta = static_cast<half const*>(mBetaDev.get());
|
||||
|
||||
if (total < 4096)
|
||||
{
|
||||
return launch_small_mtron(
|
||||
stream, ld, total, input, skip, beta, gamma, output, preln, dqScaleIn, dqScaleSkip, qScale, qSkipScale);
|
||||
}
|
||||
|
||||
return launch_large_mtron(
|
||||
stream, ld, total, input, skip, beta, gamma, output, preln, dqScaleIn, dqScaleSkip, qScale, qSkipScale);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
// end IPluginV3OneRuntime method definitions
|
||||
|
||||
///////
|
||||
// IPluginV3OneBuild method definitions
|
||||
// - getNbOutputs() (MTron, HFace)
|
||||
// - supportsFormatCombination() (Base)
|
||||
// - getOutputShapes (Base)
|
||||
// - getOutputDataType() (Base)
|
||||
// - configurePlugin() (Base)
|
||||
// - getWorkSpaceSize() (Base)
|
||||
//////
|
||||
int32_t SkipLayerNormInterleavedPluginHFace::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginMTron::getNbOutputs() const noexcept
|
||||
{
|
||||
return 2;
|
||||
}
|
||||
|
||||
bool SkipLayerNormInterleavedPluginBase::supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inOut != nullptr);
|
||||
PLUGIN_VALIDATE(nbInputs == 2);
|
||||
PLUGIN_VALIDATE(nbOutputs == getNbOutputs());
|
||||
PLUGIN_VALIDATE(pos >= 0 && pos < (nbInputs + nbOutputs));
|
||||
PluginTensorDesc const& desc = inOut[pos].desc;
|
||||
return desc.type == DataType::kINT8 && desc.format == TensorFormat::kCHW32;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginBase::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 == 2);
|
||||
PLUGIN_VALIDATE(nbOutputs == getNbOutputs());
|
||||
PLUGIN_VALIDATE(inputs[0].nbDims == inputs[1].nbDims);
|
||||
for (int32_t i = 0; i < nbOutputs; ++i)
|
||||
{
|
||||
outputs[i] = inputs[0];
|
||||
}
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginBase::getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputTypes != nullptr);
|
||||
PLUGIN_VALIDATE(nbOutputs == getNbOutputs());
|
||||
PLUGIN_VALIDATE(nbInputs == 2);
|
||||
for (int32_t i = 0; i < nbOutputs; ++i)
|
||||
{
|
||||
outputTypes[i] = inputTypes[0];
|
||||
}
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginBase::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
size_t SkipLayerNormInterleavedPluginBase::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
// End IPluginV3OneBuild method definitions
|
||||
|
||||
//////
|
||||
// IPluginV3OneCore method definitions
|
||||
// - getPluginVersion() (MTron, HFace)
|
||||
// - getPluginName() (Base)
|
||||
// - getPluginNamespace() (Base)
|
||||
// - setPluginNamespace() (Base)
|
||||
//////
|
||||
char const* SkipLayerNormInterleavedPluginHFace::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_HFACE;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginMTron::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_MTRON;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginBase::getPluginName() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_NAME;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginBase::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginBase::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
// End IPluginV3OneCore method definitions
|
||||
|
||||
//////////////////////////// Plugin Creator member definitions /////////////////////////////
|
||||
|
||||
SkipLayerNormInterleavedPluginBaseCreator::SkipLayerNormInterleavedPluginBaseCreator()
|
||||
{
|
||||
static std::mutex sMutex;
|
||||
std::lock_guard<std::mutex> lock(sMutex);
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("beta"));
|
||||
mPluginAttributes.emplace_back(PluginField("gamma"));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginHFaceCreator::SkipLayerNormInterleavedPluginHFaceCreator()
|
||||
: SkipLayerNormInterleavedPluginBaseCreator()
|
||||
{
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginMTronCreator::SkipLayerNormInterleavedPluginMTronCreator()
|
||||
: SkipLayerNormInterleavedPluginBaseCreator()
|
||||
{
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginBaseCreator::getPluginName() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_NAME;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginHFaceCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_HFACE;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginMTronCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_MTRON;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* SkipLayerNormInterleavedPluginBaseCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV3* SkipLayerNormInterleavedPluginHFaceCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceCreator createPlugin");
|
||||
|
||||
Weights beta{DataType::kFLOAT, nullptr, 0};
|
||||
Weights gamma{DataType::kFLOAT, nullptr, 0};
|
||||
buildBetaAndGamma(fc, beta, gamma);
|
||||
|
||||
return new SkipLayerNormInterleavedPluginHFace(name, beta, gamma);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV3* SkipLayerNormInterleavedPluginMTronCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronCreator createPlugin");
|
||||
|
||||
PLUGIN_VALIDATE(fc != nullptr);
|
||||
|
||||
Weights beta{DataType::kFLOAT, nullptr, 0};
|
||||
Weights gamma{DataType::kFLOAT, nullptr, 0};
|
||||
buildBetaAndGamma(fc, beta, gamma);
|
||||
|
||||
return new SkipLayerNormInterleavedPluginMTron(name, beta, gamma);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginBaseCreator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginBaseCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
// End Plugin Creator member definitions
|
||||
@@ -0,0 +1,223 @@
|
||||
/*
|
||||
* 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_SKIP_LAYER_NORM_INTERLEAVED_PLUGIN_H
|
||||
#define TRT_SKIP_LAYER_NORM_INTERLEAVED_PLUGIN_H
|
||||
#include "NvInferPlugin.h"
|
||||
#include <cuda.h>
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
int32_t launch_small_hface(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale);
|
||||
|
||||
int32_t launch_large_hface(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale);
|
||||
|
||||
int32_t launch_small_mtron(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, int8_t* preln, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale, float const qSkipScale);
|
||||
|
||||
int32_t launch_large_mtron(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, int8_t* preln, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale, float const qSkipScale);
|
||||
|
||||
class SkipLayerNormInterleavedPluginBase : public IPluginV3,
|
||||
public IPluginV3OneCore,
|
||||
public IPluginV3OneBuild,
|
||||
public IPluginV3OneRuntime
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginBase(
|
||||
std::string const& name, nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormInterleavedPlugin without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormInterleavedPluginBase() = delete;
|
||||
|
||||
~SkipLayerNormInterleavedPluginBase() override;
|
||||
|
||||
// IPluginV3 Methods
|
||||
// NOTE: since this is itself is an abstract class, the rest of virtual methods defined in its children classes
|
||||
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
|
||||
// end of IPluginV3 Methods
|
||||
|
||||
// IPluginV3OneCore Methods
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* pluginNamespace) noexcept;
|
||||
// end of IPluginV3OneCore Methods
|
||||
|
||||
// IPluginV3Build Methods
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) 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;
|
||||
|
||||
size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
|
||||
|
||||
int32_t getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
// end IPluginV3Build Methods
|
||||
|
||||
// IPluginV3Runtime Methods
|
||||
int32_t onShapeChange(
|
||||
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
|
||||
|
||||
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
|
||||
// end IPluginV3Runtime Methods
|
||||
|
||||
protected:
|
||||
// metadata fields
|
||||
std::string const& mLayerName;
|
||||
std::string mNamespace;
|
||||
std::vector<nvinfer1::PluginField> mDataToSerialize;
|
||||
nvinfer1::PluginFieldCollection mFCToSerialize;
|
||||
|
||||
// members that participate in ser/deserialization
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
|
||||
// device-side
|
||||
bert::cuda_unique_ptr<void> mGammaDev;
|
||||
bert::cuda_unique_ptr<void> mBetaDev;
|
||||
|
||||
// derived members
|
||||
size_t mLd{}; // leading dim
|
||||
size_t mParamWordsize{};
|
||||
bool mParamsOnDevice{};
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginHFace : public SkipLayerNormInterleavedPluginBase
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginHFace(
|
||||
std::string const& name, nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormInterleavedPlugin without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormInterleavedPluginHFace() = delete;
|
||||
|
||||
~SkipLayerNormInterleavedPluginHFace() override;
|
||||
|
||||
// IPluginV3Runtime overrides
|
||||
IPluginV3* clone() 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;
|
||||
|
||||
// IPluginV3OneCore override
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
// IPluginV3OneBuild override
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginMTron : public SkipLayerNormInterleavedPluginBase
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginMTron(
|
||||
std::string const& name, nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormInterleavedPlugin without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormInterleavedPluginMTron() = delete;
|
||||
|
||||
~SkipLayerNormInterleavedPluginMTron() override;
|
||||
|
||||
// IPluginV3Runtime overrides
|
||||
IPluginV3* clone() 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;
|
||||
|
||||
// IPluginV3OneCore override
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
// IPluginV3OneBuild override
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginBaseCreator : public nvinfer1::IPluginCreatorV3One
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginBaseCreator();
|
||||
~SkipLayerNormInterleavedPluginBaseCreator() override = default;
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
nvinfer1::PluginFieldCollection const* getFieldNames() 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;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginHFaceCreator : public SkipLayerNormInterleavedPluginBaseCreator
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginHFaceCreator();
|
||||
|
||||
~SkipLayerNormInterleavedPluginHFaceCreator() override = default;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginMTronCreator : public SkipLayerNormInterleavedPluginBaseCreator
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginMTronCreator();
|
||||
|
||||
~SkipLayerNormInterleavedPluginMTronCreator() override = default;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
};
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // TRT_SKIP_LAYER_NORM_INTERLEAVED_PLUGIN_H
|
||||
@@ -0,0 +1,604 @@
|
||||
/*
|
||||
* 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 "skipLayerNormInt8InterleavedPluginLegacy.h"
|
||||
#include "NvInfer.h"
|
||||
#include "common/serialize.hpp"
|
||||
#include <cuda.h>
|
||||
|
||||
#include <memory>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::plugin::bert;
|
||||
|
||||
// Clip plugin specific constants
|
||||
namespace
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
constexpr char const* kSKIP_LAYER_NORM_INTERLEAVED_VERSION_HFACE_LEGACY{"3"};
|
||||
constexpr char const* kSKIP_LAYER_NORM_INTERLEAVED_VERSION_MTRON_LEGACY{"4"};
|
||||
constexpr char const* kSKIP_LAYER_NORM_INTERLEAVED_NAME{"CustomSkipLayerNormPluginDynamic"};
|
||||
|
||||
void buildBetaAndGamma(PluginFieldCollection const* fc, Weights& beta, Weights& gamma)
|
||||
{
|
||||
plugin::validateRequiredAttributesExist({"beta", "gamma"}, fc);
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; i++)
|
||||
{
|
||||
std::string_view const field_name = fc->fields[i].name;
|
||||
|
||||
if (field_name == "beta"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building beta...");
|
||||
beta.values = fc->fields[i].data;
|
||||
beta.count = fc->fields[i].length;
|
||||
beta.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "gamma"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building gamma...");
|
||||
gamma.values = fc->fields[i].data;
|
||||
gamma.count = fc->fields[i].length;
|
||||
gamma.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
}
|
||||
|
||||
PLUGIN_VALIDATE(beta.values != nullptr, "SkipLayerNorm: invalid beta");
|
||||
PLUGIN_VALIDATE(beta.count > 0, "SkipLayerNorm: invalid beta");
|
||||
|
||||
PLUGIN_VALIDATE(gamma.values != nullptr, "SkipLayerNorm: invalid gamma");
|
||||
PLUGIN_VALIDATE(gamma.count > 0, "SkipLayerNorm: invalid gamma");
|
||||
}
|
||||
|
||||
void checkDescs(PluginTensorDesc const& iDesc, PluginTensorDesc const& sDesc, PluginTensorDesc const& oDesc)
|
||||
{
|
||||
PLUGIN_VALIDATE(iDesc.dims.nbDims == 4);
|
||||
PLUGIN_VALIDATE(iDesc.dims.nbDims == sDesc.dims.nbDims);
|
||||
PLUGIN_VALIDATE(std::equal(iDesc.dims.d, iDesc.dims.d + iDesc.dims.nbDims, sDesc.dims.d));
|
||||
PLUGIN_VALIDATE(std::equal(iDesc.dims.d, iDesc.dims.d + iDesc.dims.nbDims, oDesc.dims.d));
|
||||
PLUGIN_VALIDATE(iDesc.dims.d[0] == 1);
|
||||
PLUGIN_VALIDATE(iDesc.dims.d[3] == 1);
|
||||
PLUGIN_VALIDATE(iDesc.format == TensorFormat::kCHW32);
|
||||
PLUGIN_VALIDATE(iDesc.type == DataType::kINT8);
|
||||
PLUGIN_VALIDATE(iDesc.format == sDesc.format);
|
||||
PLUGIN_VALIDATE(iDesc.format == oDesc.format);
|
||||
PLUGIN_VALIDATE(iDesc.type == sDesc.type);
|
||||
PLUGIN_VALIDATE(iDesc.type == oDesc.type);
|
||||
}
|
||||
} // namespace
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(SkipLayerNormInterleavedPluginHFaceLegacyCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(SkipLayerNormInterleavedPluginMTronLegacyCreator);
|
||||
|
||||
constexpr auto kPARAM_TYPE = DataType::kHALF;
|
||||
|
||||
SkipLayerNormInterleavedPluginBaseLegacy::SkipLayerNormInterleavedPluginBaseLegacy(
|
||||
std::string const& name, Weights const& beta, Weights const& gamma)
|
||||
: mLayerName(name)
|
||||
, mGammaDev(nullptr)
|
||||
, mBetaDev(nullptr)
|
||||
, mLd(beta.count)
|
||||
, mParamsOnDevice(false)
|
||||
{
|
||||
PLUGIN_VALIDATE(mLd > 0);
|
||||
PLUGIN_VALIDATE(beta.count == gamma.count);
|
||||
// dataType for beta, gamma weights is always fp16
|
||||
|
||||
mParamWordsize = getElementSize(kPARAM_TYPE);
|
||||
|
||||
mBeta.convertAndCopy(beta, kPARAM_TYPE);
|
||||
mGamma.convertAndCopy(gamma, kPARAM_TYPE);
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginHFaceLegacy::SkipLayerNormInterleavedPluginHFaceLegacy(
|
||||
std::string const& name, Weights const& beta, Weights const& gamma)
|
||||
: SkipLayerNormInterleavedPluginBaseLegacy(name, beta, gamma)
|
||||
{
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginMTronLegacy::SkipLayerNormInterleavedPluginMTronLegacy(
|
||||
std::string const& name, Weights const& beta, Weights const& gamma)
|
||||
: SkipLayerNormInterleavedPluginBaseLegacy(name, beta, gamma)
|
||||
{
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginBaseLegacy::SkipLayerNormInterleavedPluginBaseLegacy(
|
||||
std::string const& name, void const* data, size_t length)
|
||||
: mLayerName(name)
|
||||
, mGammaDev(nullptr)
|
||||
, mBetaDev(nullptr)
|
||||
, mParamsOnDevice(false)
|
||||
{
|
||||
// Deserialize in the same order as serialization
|
||||
deserialize_value(&data, &length, &mLd);
|
||||
|
||||
mParamWordsize = getElementSize(kPARAM_TYPE);
|
||||
|
||||
char const* d = static_cast<char const*>(data);
|
||||
mBeta.convertAndCopy(d, mLd, kPARAM_TYPE);
|
||||
mGamma.convertAndCopy(d, mLd, kPARAM_TYPE);
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginHFaceLegacy::SkipLayerNormInterleavedPluginHFaceLegacy(
|
||||
std::string const& name, void const* data, size_t length)
|
||||
: SkipLayerNormInterleavedPluginBaseLegacy(name, data, length)
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceLegacy deserialize");
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginMTronLegacy::SkipLayerNormInterleavedPluginMTronLegacy(
|
||||
std::string const& name, void const* data, size_t length)
|
||||
: SkipLayerNormInterleavedPluginBaseLegacy(name, data, length)
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronLegacy deserialize");
|
||||
}
|
||||
|
||||
// IPluginV2DynamicExt Methods
|
||||
IPluginV2DynamicExt* SkipLayerNormInterleavedPluginHFaceLegacy::clone() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceLegacy clone");
|
||||
|
||||
auto p = std::make_unique<SkipLayerNormInterleavedPluginHFaceLegacy>(mLayerName, mBeta, mGamma);
|
||||
p->initialize();
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* SkipLayerNormInterleavedPluginMTronLegacy::clone() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronLegacy clone");
|
||||
|
||||
auto p = std::make_unique<SkipLayerNormInterleavedPluginMTronLegacy>(mLayerName, mBeta, mGamma);
|
||||
p->initialize();
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
DimsExprs SkipLayerNormInterleavedPluginBaseLegacy::getOutputDimensions(
|
||||
int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputs != nullptr);
|
||||
PLUGIN_VALIDATE(nbInputs == 2);
|
||||
PLUGIN_VALIDATE(outputIndex >= 0 && outputIndex < getNbOutputs());
|
||||
PLUGIN_VALIDATE(inputs[0].nbDims == inputs[1].nbDims);
|
||||
return inputs[0];
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return DimsExprs{};
|
||||
}
|
||||
|
||||
bool SkipLayerNormInterleavedPluginBaseLegacy::supportsFormatCombination(
|
||||
int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inOut != nullptr);
|
||||
PLUGIN_VALIDATE(nbInputs == 2);
|
||||
PLUGIN_VALIDATE(nbOutputs == getNbOutputs());
|
||||
PLUGIN_VALIDATE(pos >= 0 && pos < (nbInputs + nbOutputs));
|
||||
|
||||
PluginTensorDesc const& desc = inOut[pos];
|
||||
return desc.type == DataType::kINT8 && desc.format == TensorFormat::kCHW32;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginBaseLegacy::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// Validate input arguments
|
||||
PLUGIN_VALIDATE(inputs != nullptr);
|
||||
PLUGIN_VALIDATE(outputs != nullptr);
|
||||
PLUGIN_VALIDATE(nbOutputs == getNbOutputs());
|
||||
PLUGIN_VALIDATE(nbInputs == 2);
|
||||
PLUGIN_VALIDATE(DataType::kINT8 == inputs[0].desc.type);
|
||||
PLUGIN_VALIDATE(DataType::kINT8 == inputs[1].desc.type);
|
||||
|
||||
auto const& inDims0 = inputs[0].desc.dims;
|
||||
auto const& inDims1 = inputs[1].desc.dims;
|
||||
TRT_UNUSED inDims1;
|
||||
|
||||
PLUGIN_VALIDATE(inDims0.nbDims == inDims1.nbDims);
|
||||
PLUGIN_VALIDATE(std::equal(inDims0.d, inDims0.d + inDims0.nbDims, inDims1.d));
|
||||
|
||||
mParamWordsize = getElementSize(kPARAM_TYPE);
|
||||
|
||||
if (!mParamsOnDevice)
|
||||
{
|
||||
copyToDevice(mGamma, getWeightsSize(mGamma, kPARAM_TYPE), mGammaDev);
|
||||
copyToDevice(mBeta, getWeightsSize(mBeta, kPARAM_TYPE), mBetaDev);
|
||||
mParamsOnDevice = true;
|
||||
}
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
size_t SkipLayerNormInterleavedPluginBaseLegacy::getWorkspaceSize(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginHFaceLegacy::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);
|
||||
|
||||
// Input shape: 1x(hxd)xtotalx1
|
||||
auto const iDesc = inputDesc[0];
|
||||
auto const sDesc = inputDesc[1];
|
||||
auto const oDesc = outputDesc[0];
|
||||
checkDescs(iDesc, sDesc, oDesc);
|
||||
|
||||
int32_t const ld = iDesc.dims.d[1];
|
||||
int32_t const total = iDesc.dims.d[2];
|
||||
float const dqScaleIn = iDesc.scale;
|
||||
float const dqScaleSkip = sDesc.scale;
|
||||
float const qScale = 1.F / oDesc.scale;
|
||||
int8_t const* input = static_cast<int8_t const*>(inputs[0]);
|
||||
int8_t const* skip = static_cast<int8_t const*>(inputs[1]);
|
||||
int8_t* output = static_cast<int8_t*>(outputs[0]);
|
||||
half const* gamma = static_cast<half const*>(mGammaDev.get());
|
||||
half const* beta = static_cast<half const*>(mBetaDev.get());
|
||||
|
||||
if (total < 4096)
|
||||
{
|
||||
return launch_small_hface(
|
||||
stream, ld, total, input, skip, beta, gamma, output, dqScaleIn, dqScaleSkip, qScale);
|
||||
}
|
||||
|
||||
return launch_large_hface(stream, ld, total, input, skip, beta, gamma, output, dqScaleIn, dqScaleSkip, qScale);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginMTronLegacy::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);
|
||||
|
||||
// Input shape: 1x(hxd)xtotalx1
|
||||
auto const iDesc = inputDesc[0];
|
||||
auto const sDesc = inputDesc[1];
|
||||
auto const oDesc = outputDesc[0];
|
||||
auto const pDesc = outputDesc[1];
|
||||
checkDescs(iDesc, sDesc, oDesc);
|
||||
PLUGIN_VALIDATE(std::equal(iDesc.dims.d, iDesc.dims.d + iDesc.dims.nbDims, pDesc.dims.d));
|
||||
|
||||
int32_t const ld = iDesc.dims.d[1];
|
||||
int32_t const total = iDesc.dims.d[2];
|
||||
float const dqScaleIn = iDesc.scale;
|
||||
float const dqScaleSkip = sDesc.scale;
|
||||
float const qScale = 1.F / oDesc.scale;
|
||||
float const qSkipScale = 1.F / pDesc.scale;
|
||||
int8_t const* input = static_cast<int8_t const*>(inputs[0]);
|
||||
int8_t const* skip = static_cast<int8_t const*>(inputs[1]);
|
||||
int8_t* output = static_cast<int8_t*>(outputs[0]);
|
||||
int8_t* preln = static_cast<int8_t*>(outputs[1]);
|
||||
half const* gamma = static_cast<half const*>(mGammaDev.get());
|
||||
half const* beta = static_cast<half const*>(mBetaDev.get());
|
||||
|
||||
if (total < 4096)
|
||||
{
|
||||
return launch_small_mtron(
|
||||
stream, ld, total, input, skip, beta, gamma, output, preln, dqScaleIn, dqScaleSkip, qScale, qSkipScale);
|
||||
}
|
||||
|
||||
return launch_large_mtron(
|
||||
stream, ld, total, input, skip, beta, gamma, output, preln, dqScaleIn, dqScaleSkip, qScale, qSkipScale);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
// IPluginV2Ext Methods
|
||||
DataType SkipLayerNormInterleavedPluginBaseLegacy::getOutputDataType(
|
||||
int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputTypes != nullptr);
|
||||
PLUGIN_VALIDATE(index >= 0 && index < getNbOutputs());
|
||||
PLUGIN_VALIDATE(nbInputs == 2);
|
||||
return inputTypes[0];
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return DataType{};
|
||||
}
|
||||
|
||||
// IPluginV2 Methods
|
||||
char const* SkipLayerNormInterleavedPluginBaseLegacy::getPluginType() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_NAME;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginHFaceLegacy::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_HFACE_LEGACY;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginMTronLegacy::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_MTRON_LEGACY;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginHFaceLegacy::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginMTronLegacy::getNbOutputs() const noexcept
|
||||
{
|
||||
return 2;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginHFaceLegacy::initialize() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceLegacy initialize");
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t SkipLayerNormInterleavedPluginMTronLegacy::initialize() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronLegacy initialize");
|
||||
return 0;
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginHFaceLegacy::terminate() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceLegacy terminate");
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginMTronLegacy::terminate() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronLegacy terminate");
|
||||
}
|
||||
|
||||
size_t SkipLayerNormInterleavedPluginBaseLegacy::getSerializationSize() const noexcept
|
||||
{
|
||||
return 2 * mParamWordsize * mLd + sizeof(mLd);
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginBaseLegacy::serialize(void* buffer) const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
serialize_value(&buffer, mLd);
|
||||
|
||||
char* d = static_cast<char*>(buffer);
|
||||
serFromDev(d, static_cast<char*>(mBetaDev.get()), mLd * mParamWordsize);
|
||||
serFromDev(d, static_cast<char*>(mGammaDev.get()), mLd * mParamWordsize);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginBaseLegacy::destroy() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// This gets called when the network containing plugin is destroyed
|
||||
mGammaDev.reset(nullptr);
|
||||
mBetaDev.reset(nullptr);
|
||||
delete this;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginHFaceLegacy::destroy() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceLegacy destroy");
|
||||
SkipLayerNormInterleavedPluginBaseLegacy::destroy();
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginMTronLegacy::destroy() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronLegacy destroy");
|
||||
SkipLayerNormInterleavedPluginBaseLegacy::destroy();
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginBaseLegacy::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginBaseLegacy::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
/////////////////////////////////////////////////////////
|
||||
|
||||
SkipLayerNormInterleavedPluginBaseLegacyCreator::SkipLayerNormInterleavedPluginBaseLegacyCreator()
|
||||
{
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("beta"));
|
||||
mPluginAttributes.emplace_back(PluginField("gamma"));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginHFaceLegacyCreator::SkipLayerNormInterleavedPluginHFaceLegacyCreator()
|
||||
: SkipLayerNormInterleavedPluginBaseLegacyCreator()
|
||||
{
|
||||
}
|
||||
|
||||
SkipLayerNormInterleavedPluginMTronLegacyCreator::SkipLayerNormInterleavedPluginMTronLegacyCreator()
|
||||
: SkipLayerNormInterleavedPluginBaseLegacyCreator()
|
||||
{
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginBaseLegacyCreator::getPluginName() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_NAME;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginHFaceLegacyCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_HFACE_LEGACY;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginMTronLegacyCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSKIP_LAYER_NORM_INTERLEAVED_VERSION_MTRON_LEGACY;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* SkipLayerNormInterleavedPluginBaseLegacyCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV2* SkipLayerNormInterleavedPluginHFaceLegacyCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceLegacyCreator createPlugin");
|
||||
|
||||
Weights beta{DataType::kFLOAT, nullptr, 0};
|
||||
Weights gamma{DataType::kFLOAT, nullptr, 0};
|
||||
buildBetaAndGamma(fc, beta, gamma);
|
||||
|
||||
return new SkipLayerNormInterleavedPluginHFaceLegacy(name, beta, gamma);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* SkipLayerNormInterleavedPluginMTronLegacyCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronLegacyCreator createPlugin");
|
||||
|
||||
PLUGIN_VALIDATE(fc != nullptr);
|
||||
|
||||
Weights beta{DataType::kFLOAT, nullptr, 0};
|
||||
Weights gamma{DataType::kFLOAT, nullptr, 0};
|
||||
buildBetaAndGamma(fc, beta, gamma);
|
||||
|
||||
return new SkipLayerNormInterleavedPluginMTronLegacy(name, beta, gamma);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* SkipLayerNormInterleavedPluginHFaceLegacyCreator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
// This object will be deleted when the network is destroyed, which will
|
||||
// call SkipLayerNormInterleavedPlugin::destroy()
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginHFaceLegacyCreator deserializePlugin");
|
||||
return new SkipLayerNormInterleavedPluginHFaceLegacy(name, serialData, serialLength);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* SkipLayerNormInterleavedPluginMTronLegacyCreator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
// This object will be deleted when the network is destroyed, which will
|
||||
// call SkipLayerNormInterleavedPlugin::destroy()
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("SkipLayerNormInterleavedPluginMTronLegacyCreator deserializePlugin");
|
||||
return new SkipLayerNormInterleavedPluginMTronLegacy(name, serialData, serialLength);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void SkipLayerNormInterleavedPluginBaseLegacyCreator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
|
||||
char const* SkipLayerNormInterleavedPluginBaseLegacyCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
@@ -0,0 +1,195 @@
|
||||
/*
|
||||
* 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_SKIP_LAYER_NORM_INTERLEAVED_PLUGIN_LEGACY_H
|
||||
#define TRT_SKIP_LAYER_NORM_INTERLEAVED_PLUGIN_LEGACY_H
|
||||
#include "NvInferPlugin.h"
|
||||
#include <cuda.h>
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
int32_t launch_small_hface(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale);
|
||||
|
||||
int32_t launch_large_hface(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale);
|
||||
|
||||
int32_t launch_small_mtron(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, int8_t* preln, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale, float const qSkipScale);
|
||||
|
||||
int32_t launch_large_mtron(cudaStream_t stream, int32_t const ld, int32_t const total, int8_t const* input,
|
||||
int8_t const* skip, half const* beta, half const* gamma, int8_t* output, int8_t* preln, float const dqScaleIn,
|
||||
float const dqScaleSkip, float const qScale, float const qSkipScale);
|
||||
|
||||
class SkipLayerNormInterleavedPluginBaseLegacy : public nvinfer1::IPluginV2DynamicExt
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginBaseLegacy(
|
||||
std::string const& name, nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma);
|
||||
|
||||
SkipLayerNormInterleavedPluginBaseLegacy(std::string const& name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormInterleavedPlugin without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormInterleavedPluginBaseLegacy() = delete;
|
||||
|
||||
// IPluginV2DynamicExt Methods
|
||||
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;
|
||||
|
||||
// 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;
|
||||
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;
|
||||
|
||||
protected:
|
||||
std::string const& mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
bert::cuda_unique_ptr<void> mGammaDev;
|
||||
bert::cuda_unique_ptr<void> mBetaDev;
|
||||
size_t mLd{}; // leading dim
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
|
||||
size_t mParamWordsize{};
|
||||
bool mParamsOnDevice{};
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginHFaceLegacy : public SkipLayerNormInterleavedPluginBaseLegacy
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginHFaceLegacy(
|
||||
std::string const& name, nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma);
|
||||
|
||||
SkipLayerNormInterleavedPluginHFaceLegacy(std::string const& name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormInterleavedPlugin without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormInterleavedPluginHFaceLegacy() = delete;
|
||||
|
||||
// IPluginV2DynamicExt Methods
|
||||
nvinfer1::IPluginV2DynamicExt* clone() 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;
|
||||
|
||||
// IPluginV2 Methods
|
||||
int32_t initialize() noexcept override;
|
||||
void terminate() noexcept override;
|
||||
void destroy() noexcept override;
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginMTronLegacy : public SkipLayerNormInterleavedPluginBaseLegacy
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginMTronLegacy(
|
||||
std::string const& name, nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma);
|
||||
|
||||
SkipLayerNormInterleavedPluginMTronLegacy(std::string const& name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormInterleavedPlugin without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormInterleavedPluginMTronLegacy() = delete;
|
||||
|
||||
// IPluginV2DynamicExt Methods
|
||||
nvinfer1::IPluginV2DynamicExt* clone() 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;
|
||||
|
||||
// IPluginV2 Methods
|
||||
int32_t initialize() noexcept override;
|
||||
void terminate() noexcept override;
|
||||
void destroy() noexcept override;
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginBaseLegacyCreator : public nvinfer1::IPluginCreator
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginBaseLegacyCreator();
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
nvinfer1::PluginFieldCollection const* getFieldNames() 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;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginHFaceLegacyCreator : public SkipLayerNormInterleavedPluginBaseLegacyCreator
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginHFaceLegacyCreator();
|
||||
|
||||
char const* getPluginVersion() const 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;
|
||||
};
|
||||
|
||||
class SkipLayerNormInterleavedPluginMTronLegacyCreator : public SkipLayerNormInterleavedPluginBaseLegacyCreator
|
||||
{
|
||||
public:
|
||||
SkipLayerNormInterleavedPluginMTronLegacyCreator();
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
|
||||
nvinfer1::IPluginV2* deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept override;
|
||||
};
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // TRT_SKIP_LAYER_NORM_INTERLEAVED_PLUGIN_LEGACY_H
|
||||
@@ -0,0 +1,307 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
#if CUDA_VERSION >= 10010
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "common/bertCommon.h"
|
||||
#include "common/common.cuh"
|
||||
#include "common/serialize.hpp"
|
||||
#include "skipLayerNormPlugin.h"
|
||||
#include "skipLayerNormPluginLegacy.h"
|
||||
|
||||
#include <cassert>
|
||||
#include <cstring>
|
||||
#include <limits>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
template <int32_t TPB, int32_t VPT, bool hasBias>
|
||||
__global__ void skiplnDQQ(int32_t const ld, int8_t const* input, int8_t const* skip, int8_t* output, __half const* beta,
|
||||
__half const* gamma, __half const* bias, float const dqScaleIn, float const dqScaleSkip, float const qScale)
|
||||
{
|
||||
int32_t const idx = ld * blockIdx.x + threadIdx.x * VPT;
|
||||
// 4 * 1024 * 4 * 2 Bytes = 16KB per block
|
||||
int8_t inLocal[VPT];
|
||||
int8_t skipLocal[VPT];
|
||||
|
||||
__half inLocalDQ[VPT]; // dequantized input + skip + bias
|
||||
__half biasLocal[VPT]; // bias and beta
|
||||
__half gammaLocal[VPT];
|
||||
copy<sizeof(int8_t) * VPT>(&input[idx], inLocal);
|
||||
copy<sizeof(int8_t) * VPT>(&skip[idx], skipLocal);
|
||||
copy<sizeof(__half) * VPT>(&bias[threadIdx.x * VPT], biasLocal);
|
||||
__half2 loc = __floats2half2_rn(0.f, 0.f); // accumulator
|
||||
|
||||
const __half rld = __half(1) / __half(ld);
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT; it++)
|
||||
{
|
||||
// DQ input and skip
|
||||
float const tmpIn = inLocal[it];
|
||||
float const tmpSkip = skipLocal[it];
|
||||
inLocalDQ[it] = dqScaleIn * tmpIn + dqScaleSkip * tmpSkip;
|
||||
|
||||
if (hasBias)
|
||||
inLocalDQ[it] += biasLocal[it];
|
||||
const __half tmp = rld * inLocalDQ[it];
|
||||
const __half2 tmp2 = __halves2half2(tmp, tmp * inLocalDQ[it]);
|
||||
loc = loc + tmp2;
|
||||
}
|
||||
// load parameters
|
||||
copy<sizeof(__half) * VPT>(&beta[threadIdx.x * VPT], biasLocal);
|
||||
copy<sizeof(__half) * VPT>(&gamma[threadIdx.x * VPT], gammaLocal);
|
||||
|
||||
using BlockReduce = cub::BlockReduce<__half2, TPB>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
__shared__ __half mu; // mean
|
||||
__shared__ __half rsigma; // 1 / std.dev.
|
||||
|
||||
const __half2 sum2 = BlockReduce(tempStorage).Reduce(loc, [](auto const& lhs, auto const& rhs){return lhs + rhs;});
|
||||
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
mu = __low2half(sum2);
|
||||
rsigma = rsqrt(__high2half(sum2) - mu * mu + std::numeric_limits<half>::epsilon());
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
static_assert(VPT % 4 == 0, "");
|
||||
uint32_t outLocal[VPT / 4U];
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT / 4U; it++)
|
||||
{
|
||||
float const tmp0 = gammaLocal[it * 4 + 0] * (inLocalDQ[it * 4 + 0] - mu) * rsigma + biasLocal[it * 4 + 0];
|
||||
float const tmp1 = gammaLocal[it * 4 + 1] * (inLocalDQ[it * 4 + 1] - mu) * rsigma + biasLocal[it * 4 + 1];
|
||||
float const tmp2 = gammaLocal[it * 4 + 2] * (inLocalDQ[it * 4 + 2] - mu) * rsigma + biasLocal[it * 4 + 2];
|
||||
float const tmp3 = gammaLocal[it * 4 + 3] * (inLocalDQ[it * 4 + 3] - mu) * rsigma + biasLocal[it * 4 + 3];
|
||||
outLocal[it] = float4_to_char4(tmp0 * qScale, tmp1 * qScale, tmp2 * qScale, tmp3 * qScale);
|
||||
}
|
||||
|
||||
copy<sizeof(int8_t) * VPT>(outLocal, &output[idx]);
|
||||
}
|
||||
|
||||
template <typename T, int32_t TPB, int32_t VPT, bool hasBias>
|
||||
__global__ void skipln_vec(
|
||||
int32_t const ld, const T* input, const T* skip, T* output, const T* beta, const T* gamma, const T* bias)
|
||||
{
|
||||
int32_t const idx = ld * blockIdx.x + threadIdx.x * VPT;
|
||||
// 4 * 1024 * 4 * 2 Bytes = 16KB per block
|
||||
T inLocal[VPT];
|
||||
T skipLocal[VPT];
|
||||
T biasLocal[VPT];
|
||||
// T gammaLocal[VPT];
|
||||
copy<sizeof(T) * VPT>(&input[idx], inLocal);
|
||||
copy<sizeof(T) * VPT>(&skip[idx], skipLocal);
|
||||
copy<sizeof(T) * VPT>(&bias[threadIdx.x * VPT], biasLocal);
|
||||
T local = 0.f;
|
||||
T local2 = 0.f;
|
||||
|
||||
const T rld = T(1) / T(ld);
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT; it++)
|
||||
{
|
||||
inLocal[it] += skipLocal[it];
|
||||
if (hasBias)
|
||||
inLocal[it] += biasLocal[it];
|
||||
const T tmp = rld * inLocal[it];
|
||||
local += tmp;
|
||||
local2 += tmp * inLocal[it];
|
||||
}
|
||||
|
||||
copy<sizeof(T) * VPT>(&beta[threadIdx.x * VPT], biasLocal);
|
||||
copy<sizeof(T) * VPT>(&gamma[threadIdx.x * VPT], skipLocal);
|
||||
|
||||
using BlockReduce = cub::BlockReduce<kvp<T>, TPB>;
|
||||
__shared__ typename BlockReduce::TempStorage tempStorage;
|
||||
__shared__ T mu; // mean
|
||||
__shared__ T rsigma; // 1 / std.dev.
|
||||
|
||||
auto const sumKV = BlockReduce(tempStorage).Reduce(kvp<T>(local, local2), [](auto const& lhs, auto const& rhs){return lhs + rhs;});
|
||||
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
mu = sumKV.key;
|
||||
rsigma = rsqrt(sumKV.value - mu * mu + std::numeric_limits<T>::epsilon());
|
||||
}
|
||||
__syncthreads();
|
||||
///*
|
||||
#pragma unroll
|
||||
for (int32_t it = 0; it < VPT; it++)
|
||||
{
|
||||
inLocal[it] = skipLocal[it] * (inLocal[it] - mu) * rsigma + biasLocal[it];
|
||||
}
|
||||
/* */
|
||||
|
||||
copy<sizeof(T) * VPT>(inLocal, &output[idx]);
|
||||
}
|
||||
|
||||
template <typename T, unsigned TPB, bool hasBias>
|
||||
__global__ void skipLayerNormKernelSmall(
|
||||
int32_t const ld, const T* input, const T* skip, const T* beta, const T* gamma, T* output, const T* bias)
|
||||
{
|
||||
|
||||
const T rld = T(1) / T(ld);
|
||||
int32_t const offset = blockIdx.x * ld;
|
||||
|
||||
// reduce x and x^2
|
||||
kvp<T> threadData(0, 0);
|
||||
int32_t const idx = offset + threadIdx.x;
|
||||
T val = 0;
|
||||
|
||||
if (threadIdx.x < ld)
|
||||
{
|
||||
|
||||
val = input[idx] + skip[idx];
|
||||
if (hasBias)
|
||||
{
|
||||
val += bias[threadIdx.x];
|
||||
}
|
||||
|
||||
const T rldval = rld * val;
|
||||
threadData = threadData + kvp<T>(rldval, rldval * val);
|
||||
}
|
||||
|
||||
layerNormSmall<T, T, TPB>(val, threadData, ld, idx, beta, gamma, output);
|
||||
}
|
||||
|
||||
template <typename T, unsigned TPB, bool hasBias>
|
||||
__global__ void skipLayerNormKernel(
|
||||
int32_t const ld, const T* input, const T* skip, const T* beta, const T* gamma, T* output, const T* bias)
|
||||
{
|
||||
const T rld = T(1) / T(ld);
|
||||
int32_t const offset = blockIdx.x * ld;
|
||||
|
||||
// reduce x and x^2
|
||||
kvp<T> threadData(0, 0);
|
||||
|
||||
for (int32_t i = threadIdx.x; i < ld; i += TPB)
|
||||
{
|
||||
int32_t const idx = offset + i;
|
||||
T val = T(input[idx]) + T(skip[idx]);
|
||||
|
||||
if (hasBias)
|
||||
{
|
||||
val += T(bias[i]);
|
||||
}
|
||||
const T rldval = rld * val;
|
||||
threadData = threadData + kvp<T>(rldval, rldval * val);
|
||||
output[idx] = val;
|
||||
}
|
||||
|
||||
layerNorm<T, T, T, TPB>(threadData, ld, offset, beta, gamma, output);
|
||||
}
|
||||
|
||||
template <bool hasBias>
|
||||
int32_t computeSkipLayerNormDQQ(cudaStream_t stream, int32_t const ld, int32_t const n, int8_t const* input,
|
||||
int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output, __half const* bias,
|
||||
float const dqScaleIn, float const dqScaleSkip, float const qScale)
|
||||
{
|
||||
// this must be true because n is the total size of the tensor
|
||||
PLUGIN_VALIDATE(n % ld == 0);
|
||||
|
||||
int32_t const gridSize = n / ld;
|
||||
// we're limited by the size of the parameters, i.e. 8-wide instead of 16
|
||||
constexpr int32_t VPT = 16 / sizeof(__half);
|
||||
if (ld == 768)
|
||||
{
|
||||
constexpr int32_t TPB = 768 / VPT;
|
||||
skiplnDQQ<TPB, VPT, hasBias>
|
||||
<<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias, dqScaleIn, dqScaleSkip, qScale);
|
||||
}
|
||||
else if (ld == 1024)
|
||||
{
|
||||
constexpr int32_t TPB = 1024 / VPT;
|
||||
skiplnDQQ<TPB, VPT, hasBias>
|
||||
<<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias, dqScaleIn, dqScaleSkip, qScale);
|
||||
}
|
||||
else
|
||||
{
|
||||
// TODO need to implement this
|
||||
PLUGIN_ERROR(("SkipLayerNormDQQ - FATAL: unsupported hidden layer size: " + std::to_string(ld)).c_str());
|
||||
}
|
||||
PLUGIN_CHECK(cudaPeekAtLastError());
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
template <typename T, bool hasBias>
|
||||
int32_t computeSkipLayerNorm(cudaStream_t stream, int32_t const ld, int32_t const n, const T* input, const T* skip,
|
||||
const T* beta, const T* gamma, T* output, const T* bias)
|
||||
{
|
||||
|
||||
// this must be true because n is the total size of the tensor
|
||||
PLUGIN_VALIDATE(n % ld == 0);
|
||||
int32_t const gridSize = n / ld;
|
||||
constexpr int32_t VPT = 16 / sizeof(T);
|
||||
if (ld <= 32)
|
||||
{
|
||||
constexpr int32_t blockSize = 32;
|
||||
skipLayerNormKernelSmall<T, blockSize, hasBias>
|
||||
<<<gridSize, blockSize, 0, stream>>>(ld, input, skip, beta, gamma, output, bias);
|
||||
}
|
||||
else if (ld == 768)
|
||||
{
|
||||
constexpr int32_t TPB = 768 / VPT;
|
||||
skipln_vec<T, TPB, VPT, hasBias><<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias);
|
||||
}
|
||||
else if (ld == 1024)
|
||||
{
|
||||
constexpr int32_t TPB = 1024 / VPT;
|
||||
skipln_vec<T, TPB, VPT, hasBias><<<gridSize, TPB, 0, stream>>>(ld, input, skip, output, beta, gamma, bias);
|
||||
}
|
||||
else
|
||||
{
|
||||
constexpr int32_t blockSize = 256;
|
||||
skipLayerNormKernel<T, blockSize, hasBias>
|
||||
<<<gridSize, blockSize, 0, stream>>>(ld, input, skip, beta, gamma, output, bias);
|
||||
}
|
||||
PLUGIN_CHECK(cudaPeekAtLastError());
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
template int32_t computeSkipLayerNormDQQ<true>(cudaStream_t stream, int32_t const ld, int32_t const n,
|
||||
int8_t const* input, int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output,
|
||||
__half const* bias, float const dqScaleIn, float const dqScaleSkip, float const qScale);
|
||||
template int32_t computeSkipLayerNormDQQ<false>(cudaStream_t stream, int32_t const ld, int32_t const n,
|
||||
int8_t const* input, int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output,
|
||||
__half const* bias, float const dqScaleIn, float const dqScaleSkip, float const qScale);
|
||||
|
||||
template int32_t computeSkipLayerNorm<float, true>(cudaStream_t, int32_t const, int32_t const, float const*,
|
||||
float const*, float const*, float const*, float*, float const*);
|
||||
template int32_t computeSkipLayerNorm<float, false>(cudaStream_t, int32_t const, int32_t const, float const*,
|
||||
float const*, float const*, float const*, float*, float const*);
|
||||
template int32_t computeSkipLayerNorm<half, true>(
|
||||
cudaStream_t, int32_t const, int32_t const, half const*, half const*, half const*, half const*, half*, half const*);
|
||||
template int32_t computeSkipLayerNorm<half, false>(
|
||||
cudaStream_t, int32_t const, int32_t const, half const*, half const*, half const*, half const*, half*, half const*);
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,274 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
#if CUDA_VERSION >= 10010
|
||||
|
||||
#ifndef TRT_SKIP_LAYER_NORM_PLUGIN_H
|
||||
#define TRT_SKIP_LAYER_NORM_PLUGIN_H
|
||||
|
||||
#include "NvInferPlugin.h"
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
template <bool hasBias>
|
||||
int32_t computeSkipLayerNormDQQ(cudaStream_t stream, int32_t const ld, int32_t const n, int8_t const* input,
|
||||
int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output, __half const* bias,
|
||||
float const dqScaleIn, float const dqScaleSkip, float const qScale);
|
||||
|
||||
template <typename T, bool hasBias>
|
||||
int32_t computeSkipLayerNorm(cudaStream_t stream, int32_t const ld, int32_t const n, T const* input, T const* skip,
|
||||
T const* beta, T const* gamma, T* output, T const* bias);
|
||||
|
||||
class SkipLayerNormPluginV3 : public IPluginV3,
|
||||
public IPluginV3OneCore,
|
||||
public IPluginV3OneBuild,
|
||||
public IPluginV3OneRuntime
|
||||
{
|
||||
public:
|
||||
SkipLayerNormPluginV3(const std::string name, const nvinfer1::DataType type, int32_t const ld,
|
||||
nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma, nvinfer1::Weights const& bias);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormPluginV3 without arguments,
|
||||
// so we delete default constructor.
|
||||
SkipLayerNormPluginV3() = delete;
|
||||
|
||||
~SkipLayerNormPluginV3() override;
|
||||
|
||||
// IPluginV3 Methods
|
||||
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
|
||||
|
||||
IPluginV3* clone() noexcept override;
|
||||
// end of IPluginV3 Methods
|
||||
|
||||
// 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;
|
||||
// end of IPluginV3OneCore Methods
|
||||
|
||||
// IPluginV3Build Methods
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) 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;
|
||||
|
||||
size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
|
||||
|
||||
int32_t getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
// end IPluginV3Build Methods
|
||||
|
||||
// IPluginV3Runtime 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;
|
||||
|
||||
int32_t onShapeChange(
|
||||
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
|
||||
|
||||
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
|
||||
// end IPluginV3Runtime Methods
|
||||
|
||||
private:
|
||||
// metadata
|
||||
const std::string mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
// members that participate in ser/deserialization
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
bert::WeightsWithOwnership mBias;
|
||||
nvinfer1::DataType mType;
|
||||
nvinfer1::DataType mCfgType;
|
||||
int32_t mLd{}; // leading dim
|
||||
bool mHasBias{};
|
||||
|
||||
// device-side
|
||||
bert::cuda_unique_ptr<void> mGammaDev;
|
||||
bert::cuda_unique_ptr<void> mBetaDev;
|
||||
bert::cuda_unique_ptr<void> mBiasDev;
|
||||
|
||||
// derived member from mCfgType
|
||||
size_t mParamWordsize{};
|
||||
|
||||
// serialization data structures
|
||||
std::vector<nvinfer1::PluginField> mDataToSerialize;
|
||||
nvinfer1::PluginFieldCollection mFCToSerialize;
|
||||
};
|
||||
|
||||
class SkipLayerNormPluginV3Creator : public nvinfer1::IPluginCreatorV3One
|
||||
{
|
||||
public:
|
||||
SkipLayerNormPluginV3Creator();
|
||||
~SkipLayerNormPluginV3Creator() override = default;
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
|
||||
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* libNamespace) noexcept;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
private:
|
||||
PluginFieldCollection mFC;
|
||||
std::vector<PluginField> mPluginAttributes;
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
class SkipLayerNormVarSeqlenPluginV3 : public IPluginV3,
|
||||
public IPluginV3OneCore,
|
||||
public IPluginV3OneBuild,
|
||||
public IPluginV3OneRuntime
|
||||
{
|
||||
public:
|
||||
SkipLayerNormVarSeqlenPluginV3(const std::string name, const nvinfer1::DataType type, nvinfer1::Weights const& beta,
|
||||
nvinfer1::Weights const& gamma, nvinfer1::Weights const& bias);
|
||||
|
||||
SkipLayerNormVarSeqlenPluginV3(const std::string name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormVarSeqlenPluginV3 without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormVarSeqlenPluginV3() = delete;
|
||||
|
||||
~SkipLayerNormVarSeqlenPluginV3() override;
|
||||
|
||||
// IPluginV3 Methods
|
||||
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
|
||||
|
||||
IPluginV3* clone() noexcept override;
|
||||
// end of IPluginV3 Methods
|
||||
|
||||
// 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;
|
||||
// end of IPluginV3OneCore Methods
|
||||
|
||||
// IPluginV3Build Methods
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) 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;
|
||||
|
||||
size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override;
|
||||
|
||||
int32_t getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
// end IPluginV3Build Methods
|
||||
|
||||
// IPluginV3Runtime 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;
|
||||
|
||||
int32_t onShapeChange(
|
||||
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override;
|
||||
|
||||
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
|
||||
// end IPluginV3Runtime Methods
|
||||
|
||||
private:
|
||||
const std::string mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
bert::cuda_unique_ptr<void> mGammaDev;
|
||||
bert::cuda_unique_ptr<void> mBetaDev;
|
||||
int32_t mLd{}; // leading dim
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
nvinfer1::DataType mType;
|
||||
nvinfer1::DataType mCfgType;
|
||||
|
||||
bool mHasBias{};
|
||||
bert::cuda_unique_ptr<void> mBiasDev;
|
||||
bert::WeightsWithOwnership mBias;
|
||||
|
||||
size_t mParamWordsize{};
|
||||
|
||||
std::vector<nvinfer1::PluginField> mDataToSerialize;
|
||||
nvinfer1::PluginFieldCollection mFCToSerialize;
|
||||
};
|
||||
|
||||
class SkipLayerNormVarSeqlenPluginV3Creator : public nvinfer1::IPluginCreatorV3One
|
||||
{
|
||||
public:
|
||||
SkipLayerNormVarSeqlenPluginV3Creator();
|
||||
~SkipLayerNormVarSeqlenPluginV3Creator() override = default;
|
||||
|
||||
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* libNamespace) noexcept;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
private:
|
||||
nvinfer1::PluginFieldCollection mFC;
|
||||
std::vector<nvinfer1::PluginField> mPluginAttributes;
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // TRT_SKIP_LAYER_NORM_PLUGIN_H
|
||||
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,233 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION &
|
||||
* AFFILIATES. All rights reserved. SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
#if CUDA_VERSION >= 10010
|
||||
|
||||
#ifndef TRT_SKIP_LAYER_NORM_PLUGIN_H
|
||||
#define TRT_SKIP_LAYER_NORM_PLUGIN_H
|
||||
|
||||
#include "NvInferPlugin.h"
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
template <bool hasBias>
|
||||
int32_t computeSkipLayerNormDQQ(cudaStream_t stream, int32_t const ld, int32_t const n, int8_t const* input,
|
||||
int8_t const* skip, __half const* beta, __half const* gamma, int8_t* output, __half const* bias,
|
||||
float const dqScaleIn, float const dqScaleSkip, float const qScale);
|
||||
|
||||
template <typename T, bool hasBias>
|
||||
int32_t computeSkipLayerNorm(cudaStream_t stream, int32_t const ld, int32_t const n, T const* input, T const* skip,
|
||||
T const* beta, T const* gamma, T* output, T const* bias);
|
||||
|
||||
class SkipLayerNormPluginDynamic : public nvinfer1::IPluginV2DynamicExt
|
||||
{
|
||||
public:
|
||||
SkipLayerNormPluginDynamic(const std::string name, const nvinfer1::DataType type, int32_t const ld,
|
||||
nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma, nvinfer1::Weights const& bias);
|
||||
|
||||
SkipLayerNormPluginDynamic(const std::string name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormPluginDynamic without arguments,
|
||||
// so we delete default constructor.
|
||||
SkipLayerNormPluginDynamic() = delete;
|
||||
|
||||
// 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;
|
||||
|
||||
// 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:
|
||||
const std::string mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
bert::cuda_unique_ptr<void> mGammaDev;
|
||||
bert::cuda_unique_ptr<void> mBetaDev;
|
||||
size_t mLd{}; // leading dim
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
nvinfer1::DataType mType;
|
||||
nvinfer1::DataType mCfgType;
|
||||
|
||||
bool mHasBias{};
|
||||
bert::cuda_unique_ptr<void> mBiasDev;
|
||||
bert::WeightsWithOwnership mBias;
|
||||
|
||||
size_t mParamWordsize{};
|
||||
|
||||
using IPluginV2::enqueue;
|
||||
using IPluginV2::getOutputDimensions;
|
||||
using IPluginV2::getWorkspaceSize;
|
||||
using IPluginV2Ext::configurePlugin;
|
||||
};
|
||||
|
||||
class SkipLayerNormPluginDynamicCreator : public nvinfer1::IPluginCreator
|
||||
{
|
||||
public:
|
||||
SkipLayerNormPluginDynamicCreator();
|
||||
|
||||
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;
|
||||
};
|
||||
|
||||
class SkipLayerNormVarSeqlenPlugin : public nvinfer1::IPluginV2DynamicExt
|
||||
{
|
||||
public:
|
||||
SkipLayerNormVarSeqlenPlugin(const std::string name, const nvinfer1::DataType type, nvinfer1::Weights const& beta,
|
||||
nvinfer1::Weights const& gamma, nvinfer1::Weights const& bias);
|
||||
|
||||
SkipLayerNormVarSeqlenPlugin(const std::string name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make SkipLayerNormVarSeqlenPlugin without
|
||||
// arguments, so we delete default constructor.
|
||||
SkipLayerNormVarSeqlenPlugin() = delete;
|
||||
|
||||
// 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;
|
||||
|
||||
// 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:
|
||||
const std::string mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
bert::cuda_unique_ptr<void> mGammaDev;
|
||||
bert::cuda_unique_ptr<void> mBetaDev;
|
||||
size_t mLd{}; // leading dim
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
nvinfer1::DataType mType;
|
||||
nvinfer1::DataType mCfgType;
|
||||
|
||||
bool mHasBias{};
|
||||
bert::cuda_unique_ptr<void> mBiasDev;
|
||||
bert::WeightsWithOwnership mBias;
|
||||
|
||||
size_t mParamWordsize{};
|
||||
|
||||
using IPluginV2::enqueue;
|
||||
using IPluginV2::getOutputDimensions;
|
||||
using IPluginV2::getWorkspaceSize;
|
||||
using IPluginV2Ext::configurePlugin;
|
||||
};
|
||||
|
||||
class SkipLayerNormVarSeqlenPluginCreator : public nvinfer1::IPluginCreator
|
||||
{
|
||||
public:
|
||||
SkipLayerNormVarSeqlenPluginCreator();
|
||||
|
||||
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 bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // TRT_SKIP_LAYER_NORM_PLUGIN_H
|
||||
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
Reference in New Issue
Block a user