chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,30 @@
|
||||
#
|
||||
# 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(
|
||||
embLayerNormKernel.cu
|
||||
embLayerNormPlugin.cpp
|
||||
embLayerNormPlugin.h
|
||||
embLayerNormPluginLegacy.cpp
|
||||
embLayerNormPluginLegacy.h
|
||||
embLayerNormVarSeqlenKernelHFace.cu
|
||||
embLayerNormVarSeqlenKernelMTron.cu
|
||||
embLayerNormVarSeqlenPlugin.cpp
|
||||
embLayerNormVarSeqlenPlugin.h
|
||||
embLayerNormVarSeqlenPluginLegacy.cpp
|
||||
embLayerNormVarSeqlenPluginLegacy.h
|
||||
)
|
||||
|
||||
@@ -0,0 +1,170 @@
|
||||
#
|
||||
# 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: CustomEmbLayerNormPluginDynamic
|
||||
interface: "IPluginV3"
|
||||
versions:
|
||||
"6":
|
||||
inputs:
|
||||
- token_id
|
||||
- segment_id
|
||||
- input_mask
|
||||
outputs:
|
||||
- embedded_input
|
||||
- maskIdx
|
||||
input_dims:
|
||||
token_id: 2
|
||||
segment_id: 2
|
||||
input_mask: 2
|
||||
input_dim_constraints:
|
||||
- "token_id_0 == bert_embeddings_position_embeddings_0"
|
||||
- "segment_id_0 == bert_embeddings_position_embeddings_0"
|
||||
- "segment_id_1 == token_id_1"
|
||||
- "input_mask_0 == bert_embeddings_position_embeddings_0"
|
||||
- "input_mask_1 == token_id_1"
|
||||
input_dim_range:
|
||||
token_id:
|
||||
min: "=1, =1"
|
||||
max: "=pinf, =pinf"
|
||||
segment_id:
|
||||
min: "=1, =1"
|
||||
max: "=pinf, =pinf"
|
||||
input_mask:
|
||||
min: "=1, =1"
|
||||
max: "=pinf, =pinf"
|
||||
supported_input_types:
|
||||
combination1:
|
||||
token_id: int32
|
||||
segment_id: int32
|
||||
input_mask: int32
|
||||
output_dims:
|
||||
embedded_input: "token_id_0, token_id_1, bert_embeddings_layernorm_beta_0, 1, 1"
|
||||
maskIdx: "token_id_1, 1"
|
||||
attributes:
|
||||
- output_fp16
|
||||
- full_mask
|
||||
- mha_type_id
|
||||
- bert_embeddings_layernorm_beta
|
||||
- bert_embeddings_layernorm_gamma
|
||||
- bert_embeddings_word_embeddings
|
||||
- bert_embeddings_token_type_embeddings
|
||||
- bert_embeddings_position_embeddings
|
||||
attribute_types:
|
||||
output_fp16: int32
|
||||
full_mask: int32
|
||||
mha_type_id: int32
|
||||
bert_embeddings_layernorm_beta: float32
|
||||
bert_embeddings_layernorm_gamma: float32
|
||||
bert_embeddings_word_embeddings: float32
|
||||
bert_embeddings_token_type_embeddings: float32
|
||||
bert_embeddings_position_embeddings: float32
|
||||
attribute_dims:
|
||||
output_fp16: 1
|
||||
full_mask: 1
|
||||
mha_type_id: 1
|
||||
bert_embeddings_layernorm_beta: 1
|
||||
bert_embeddings_layernorm_gamma: 1
|
||||
bert_embeddings_word_embeddings: 2
|
||||
bert_embeddings_token_type_embeddings: 2
|
||||
bert_embeddings_position_embeddings: 2
|
||||
attribute_dim_range:
|
||||
output_fp16:
|
||||
- min: "=1"
|
||||
- max: "=1"
|
||||
full_mask:
|
||||
- min: "=1"
|
||||
- max: "=1"
|
||||
mha_type_id:
|
||||
- min: "=1"
|
||||
- max: "=1"
|
||||
bert_embeddings_layernorm_beta:
|
||||
- min: "=1"
|
||||
- max: "=pinf"
|
||||
bert_embeddings_layernorm_gamma:
|
||||
- min: "=1"
|
||||
- max: "=pinf"
|
||||
bert_embeddings_word_embedding:
|
||||
- min: "=1, =1"
|
||||
- max: "=pinf, =pinf"
|
||||
bert_embeddings_token_type_embeddings:
|
||||
- min: "=1, =1"
|
||||
- max: "=pinf, =pinf"
|
||||
bert_embeddings_position_embeddings:
|
||||
- min: "=1, =1"
|
||||
- max: "=pinf, =pinf"
|
||||
attribute_options:
|
||||
output_fp16:
|
||||
- 0
|
||||
- 1
|
||||
full_mask:
|
||||
- 0
|
||||
- 1
|
||||
mha_type_id:
|
||||
- 0
|
||||
- 1
|
||||
- 2
|
||||
bert_embeddings_layernorm_beta:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
bert_embeddings_layernorm_gamma:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
bert_embeddings_word_embeddings:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
bert_embeddings_token_type_embeddings:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
bert_embeddings_position_embeddings:
|
||||
min: "=ninf"
|
||||
max: "=pinf"
|
||||
attributes_required:
|
||||
- bert_embeddings_layernorm_beta
|
||||
- bert_embeddings_layernorm_gamma
|
||||
- bert_embeddings_word_embeddings
|
||||
- bert_embeddings_token_type_embeddings
|
||||
- bert_embeddings_position_embeddings
|
||||
golden_reference_script: "plugin/CustomEmbLayerNormPluginDynamic_PluginReference.py"
|
||||
abs_tol: 1e-5
|
||||
rel_tol: 1e-5
|
||||
configs:
|
||||
config1:
|
||||
input_types:
|
||||
token_id: int32
|
||||
segment_id: int32
|
||||
input_mask: int32
|
||||
attribute_options:
|
||||
output_fp16:
|
||||
value: 0
|
||||
shape: "1"
|
||||
full_mask:
|
||||
value: 0
|
||||
shape: "1"
|
||||
mha_type_id:
|
||||
value: 0
|
||||
shape: "1"
|
||||
bert_embeddings_layernorm_beta:
|
||||
shape: "128"
|
||||
bert_embeddings_layernorm_gamma:
|
||||
shape: "128"
|
||||
bert_embeddings_word_embeddings:
|
||||
shape: "100, 128"
|
||||
bert_embeddings_token_type_embeddings:
|
||||
shape: "2, 128"
|
||||
bert_embeddings_position_embeddings:
|
||||
shape: "20, 128"
|
||||
...
|
||||
@@ -0,0 +1,144 @@
|
||||
# embLayerNormPlugin
|
||||
|
||||
**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-3 of this plugin (using IPluginV2DynamicExt interface) are deprecated and will be removed in a future release. Versions 4-6 (using IPluginV3 interface) are the recommended replacements.
|
||||
|
||||
The plugin performs the following two tasks:
|
||||
1. Embeds an input sequence consisting of token ids and segment ids. This consists of token embedding lookup, segment embedding lookup, adding positional embeddings and finally, layer normalization.
|
||||
|
||||
2. For version 1 of the plugin only, preprocesses input masks, that are used to mark valid input tokens in sequences that are padded to the target sequence length.
|
||||
Assuming contiguous input masks, encodes the masks as a single number denoting the number of valid elements, e.g.:
|
||||
|
||||
```
|
||||
111100 => 4
|
||||
110000 => 2
|
||||
110100: Invalid mask, because it is not contiguous
|
||||
```
|
||||
For subsequent versions (2,3,4,5), the input mask is returned after casting to `half` and reshaping to the shape of the embedded output.
|
||||
|
||||
|
||||
### Structure
|
||||
|
||||
The version 1 `embLayerNormPlugin` takes three inputs; `token_id`, `segment_id`, and `input_mask`.
|
||||
The subsequent versions 2,3,4,5 (variable seqlen) take four inputs; `token_id`, `segment_id`, `cu_seqlen`, and `max_seqlen`.
|
||||
|
||||
### Version 1 & 6
|
||||
Inputs:
|
||||
- `token_id`
|
||||
An input sequence containing token ids. token_id is an `int32` tensor with shape `[S, B,]` where `S` is the sequence length and `B` is the batch size.
|
||||
Tokens typically identify words or word pieces that were obtained by preprocessing the input text.
|
||||
|
||||
- `segment_id`
|
||||
An input sequence containing segment ids. segment_id is an `int32` tensor with shape `[S, B]` where `S` is the sequence length and `B` is the batch size.
|
||||
The segment id is used to distinguish between different parts of the input sequence that might serve different purposes. E.g. in a squad task, the input sequence might consist of a segment representing the knowledge base (i.e. a paragraph of text) and a segment representing the question.
|
||||
|
||||
- `input_mask`
|
||||
input_mask is an `int32` tensor with shape `[S, B]` where `S` is the sequence length and `B` is the batch size.
|
||||
The input mask denotes valid elements in a sequence that was padded to the sequence length `S`.
|
||||
|
||||
Outputs:
|
||||
|
||||
- `embedded_output`
|
||||
embedded_output is a floating point tensor with shape `[S, B, E]` where `S` is sequence length, `B` is batch size, and `E` is hidden size.
|
||||
The final output embedding is the sum of embeddings for the token, the segment and the position in the sequence.
|
||||
|
||||
|
||||
- `maskIdx`
|
||||
The `maskIdx` is a more compact representation of the input mask, consisting of the number of valid elements, assuming that the original mask was contiguous.
|
||||
For fixed sequence length version 1, the `maskIdx` is an `int32` tensor with shape `[B, packSize]` where `B` is batch size, `packSize` is the packed mask size that depends on the sequence length.
|
||||
|
||||
### 6 > Version >= 2
|
||||
|
||||
Inputs:
|
||||
- `token_id`
|
||||
An input sequence containing token ids. token_id is a 1-D, `int32` tensor with shape `[SxB]` where `S` is the sequence length and `B` is the batch size.
|
||||
Tokens typically identify words or word pieces that were obtained by preprocessing the input text.
|
||||
|
||||
- `segment_id`
|
||||
An input sequence containing segment ids. segment_id is also a 1-D, `int32` tensor with shape `[SxB]` where `S` is the sequence length and `B` is the batch size.
|
||||
The segment id is used to distinguish between different parts of the input sequence that might serve different purposes. E.g. in a squad task, the input sequence might consist of a segment representing the knowledge base (i.e. a paragraph of text) and a segment representing the question.
|
||||
|
||||
- `input_mask`
|
||||
input_mask is also a 1-D, `int32` tensor with shape `[SxB]` where `S` is the sequence length and `B` is the batch size.
|
||||
The input mask denotes valid elements in a sequence that was padded to the sequence length `S`.
|
||||
|
||||
- `cu_seqlen` (Version 2,3,4,5 only)
|
||||
An input sequence containing the "cumulative sequence lengths", used to index into the right sequence when sequences have variable lengths. `cu_seqlen` is a 1-D, `int32` tensor with shape `[B+1]` where `B` is the batch size.
|
||||
|
||||
- `max_seqlen` (Version 2,3,4,5 only)
|
||||
Scalar `int32` value that specifies the maximum sequence length.
|
||||
|
||||
Outputs:
|
||||
|
||||
- `embedded_output`
|
||||
embedded_output is a floating point tensor with shape `[SxB, E, 1, 1]` where `S` is sequence length, `B` is batch size, and `E` is hidden size.
|
||||
The final output embedding is the sum of embeddings for the token, the segment and the position in the sequence.
|
||||
|
||||
- `maskIdx`
|
||||
(1) Huggingface variant (versions 2,4): An empty tensor (for backwards compatibility)
|
||||
(2) Megatron variant (versions 3,5): The inputs masks returned as a `half` tensor with the same shape as `embedded_output` - `[SxB, E, 1, 1]`.
|
||||
|
||||
|
||||
|
||||
## Parameters
|
||||
|
||||
`embLayerNormPlugin` has plugin creator class `EmbLayerNormPluginDynamicCreator` and plugin class `CustomEmbLayerNormPluginDynamic`.
|
||||
|
||||
The parameters are defined below and consists of the following attributes:
|
||||
|
||||
| Type | Parameter | Version | Description
|
||||
|----------|----------------------------------------|-------------------|--------------------------------------------------------
|
||||
|`int` |`output_fp16` | 1, 2, 3, 4, 5, 6 |Integer encoding the DataType, set 0 when build FP32 network and set 1 when build FP32/INT8 network (0: FP32, 1: FP16)
|
||||
|`int` |`full_mask` | 1, 6 |Whether to output the full mask that works with the specialized multi-head-attention plugin kernels (this is deprecated, please use mha_type_id)
|
||||
|`int` |`mha_type_id` | 1, 6 |Integer encoding the multi-head-attention plugin DataType (0: FP32, 1: FP16, 2: INT8)
|
||||
|`Weights` |`bert_embeddings_layernorm_beta` | 1, 2, 3, 4, 5, 6 |Beta parameter for layer norm. Shape: `[E,]` where `E` is hidden size
|
||||
|`Weights` |`bert_embeddings_layernorm_gamma` | 1, 2, 3, 4, 5, 6 |Gamma parameter for layer norm. Shape: `[E,]` where `E` is hidden size
|
||||
|`Weights` |`bert_embeddings_word_embeddings` | 1, 2, 3, 4, 5, 6 |Token embedding matrix. Shape: `[word_vocab_size, E]` where `E` is hidden size
|
||||
|`Weights` |`bert_embeddings_token_type_embeddings` | 1, 2, 3, 4, 5, 6 |Token type embedding matrix. Shape: `[type_vocab_size, E]` where `E` is hidden size
|
||||
|`Weights` |`bert_embeddings_position_embeddings` | 1, 2, 3, 4, 5, 6 |Positional embedding matrix. Shape: `[S, E]` where `S` is the maximum sequence length and `E` is hidden size
|
||||
Note: version 1, 2, 3 are deprecated and will be removed in a future release; please use their corresponding updated versions: 6, 4, 5 respectively.
|
||||
|
||||
## Additional resources
|
||||
|
||||
The following resources provide a deeper understanding of the `embLayerNormPlugin` plugin:
|
||||
|
||||
**Networks:**
|
||||
- [BERT](https://arxiv.org/abs/1810.04805)
|
||||
|
||||
|
||||
## 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
|
||||
|
||||
September 2024:
|
||||
Added `EmblayerNormPlugin` version 6 that mirrors version 1 in IO and attributes (but uses underlying `IPluginV3` implementation instead of the deprecated `IPluginV2DynamicExt` interface)
|
||||
|
||||
July 2024:
|
||||
Add `EmbLayerNormPlugin` versions 3 & 4 that duplicate the behavior of v2 and v3 plugins respectively, but implement the `IPluginV3` interface instead of the deprecated `IPluginV2DynamicExt` interface.
|
||||
Update this README with updated description of I/O and structure.
|
||||
|
||||
October 2020:
|
||||
Add V2 plugin that supports 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,263 @@
|
||||
/*
|
||||
* 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
|
||||
|
||||
#include <cassert>
|
||||
#include <cstring>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "common/bertCommon.h"
|
||||
#include "common/common.cuh"
|
||||
#include "common/serialize.hpp"
|
||||
#include "common/cubCcclCompat.h"
|
||||
#include "embLayerNormPlugin.h"
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
__global__ void fillSBSMaskKernel(
|
||||
uint32_t const warps_m, uint32_t const warps_n, uint32_t const S, int32_t const* inputMaskSB, uint32_t* inputMaskX)
|
||||
{
|
||||
|
||||
extern __shared__ int shm_mask[]; // S mask elements of this batch
|
||||
|
||||
size_t const xmmas_n = (S + 16 * warps_n - 1) / (16 * warps_n);
|
||||
uint32_t const threads_per_cta = blockDim.x;
|
||||
uint32_t const xmmas_m = gridDim.x;
|
||||
uint32_t const B = gridDim.y;
|
||||
|
||||
uint32_t const mi = blockIdx.x;
|
||||
uint32_t const bi = blockIdx.y;
|
||||
uint32_t const tidx = threadIdx.x;
|
||||
|
||||
size_t const warp = tidx / 32;
|
||||
size_t const warp_n = warp / warps_m;
|
||||
size_t const lane = tidx % 32;
|
||||
size_t const col = warp_n * 16 + lane % 4 * 2;
|
||||
|
||||
// load the mask corresponding to one batch
|
||||
for (uint32_t si = tidx; si < S; si += threads_per_cta)
|
||||
{
|
||||
// not coalesced to conform to current input format: SxB
|
||||
shm_mask[si] = inputMaskSB[si * B + bi];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
uint32_t mask = 0u;
|
||||
|
||||
for (size_t ni = 0; ni < xmmas_n; ++ni)
|
||||
{
|
||||
int32_t const offset = ni * 16 * warps_n + col;
|
||||
mask |= (shm_mask[offset + 0] == 1.f ? 1u : 0u) << (8 * ni + 0);
|
||||
mask |= (shm_mask[offset + 1] == 1.f ? 1u : 0u) << (8 * ni + 1);
|
||||
mask |= (shm_mask[offset + 0] == 1.f ? 1u : 0u) << (8 * ni + 2);
|
||||
mask |= (shm_mask[offset + 1] == 1.f ? 1u : 0u) << (8 * ni + 3);
|
||||
mask |= (shm_mask[offset + 8] == 1.f ? 1u : 0u) << (8 * ni + 4);
|
||||
mask |= (shm_mask[offset + 9] == 1.f ? 1u : 0u) << (8 * ni + 5);
|
||||
mask |= (shm_mask[offset + 8] == 1.f ? 1u : 0u) << (8 * ni + 6);
|
||||
mask |= (shm_mask[offset + 9] == 1.f ? 1u : 0u) << (8 * ni + 7);
|
||||
}
|
||||
|
||||
inputMaskX[(bi * xmmas_m + mi) * threads_per_cta + tidx] = mask;
|
||||
}
|
||||
|
||||
cudaError_t convertMask(uint32_t const S, uint32_t const B, uint32_t const warps_m, uint32_t const warps_n,
|
||||
uint32_t const warps_k, int32_t const* inputMaskSB, uint32_t* inputMaskX, cudaStream_t stream)
|
||||
{
|
||||
size_t const xmmas_m = (S + 16 * warps_m - 1) / (16 * warps_m);
|
||||
|
||||
size_t const threads_per_cta = warps_m * warps_n * warps_k * 32;
|
||||
dim3 grid(xmmas_m, B);
|
||||
fillSBSMaskKernel<<<grid, threads_per_cta, S * sizeof(int), stream>>>(warps_m, warps_n, S, inputMaskSB, inputMaskX);
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
template <unsigned TPB>
|
||||
__global__ void maskIdxKernelSmall(int ld, int32_t const* mask, int* maskIdx)
|
||||
{
|
||||
|
||||
using BlockReduce = cub::BlockReduce<int32_t, TPB>;
|
||||
__shared__ typename BlockReduce::TempStorage tmpStorage;
|
||||
|
||||
auto min = compat::getCudaMinOp();
|
||||
int threadData(ld); // if the mask admits all values
|
||||
|
||||
if (threadIdx.x < ld)
|
||||
{
|
||||
// mask has input dims {S, B} and gridDims.x is B
|
||||
int32_t const idx = threadIdx.x * gridDim.x + blockIdx.x;
|
||||
|
||||
int32_t const val = mask[idx];
|
||||
if (val == 0) // masked position: report thread idx
|
||||
{
|
||||
threadData = threadIdx.x;
|
||||
}
|
||||
}
|
||||
|
||||
const auto minIdx = BlockReduce(tmpStorage).Reduce(threadData, min);
|
||||
|
||||
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
maskIdx[blockIdx.x] = minIdx;
|
||||
}
|
||||
}
|
||||
|
||||
template <unsigned TPB>
|
||||
__global__ void maskIdxKernel(int ld, int32_t const* mask, int* maskIdx)
|
||||
{
|
||||
|
||||
using BlockReduce = cub::BlockReduce<int32_t, TPB>;
|
||||
__shared__ typename BlockReduce::TempStorage tmpStorage;
|
||||
|
||||
auto min = compat::getCudaMinOp();
|
||||
int threadData(ld); // if the mask admits all values
|
||||
|
||||
for (int i = threadIdx.x; i < ld; i += TPB)
|
||||
{
|
||||
// mask has input dims {S, B} and gridDims.x is B
|
||||
int32_t const idx = i * gridDim.x + blockIdx.x;
|
||||
|
||||
int32_t const val = mask[idx];
|
||||
if (val == 0) // masked position: report thread idx
|
||||
{
|
||||
threadData = min(threadData, i);
|
||||
}
|
||||
}
|
||||
|
||||
const auto minIdx = BlockReduce(tmpStorage).Reduce(threadData, min);
|
||||
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
maskIdx[blockIdx.x] = minIdx;
|
||||
}
|
||||
}
|
||||
|
||||
int computeMaskIdx(cudaStream_t stream, int32_t const S, int32_t const B, int32_t const* mask, int* maskIdx)
|
||||
{
|
||||
// Mask idx is of length B and assumes the valid region is contiguous starting
|
||||
// from the beginning of the sequence
|
||||
|
||||
// Assume n = BxS
|
||||
if (S <= 32)
|
||||
{
|
||||
maskIdxKernelSmall<32><<<B, 32, 0, stream>>>(S, mask, maskIdx);
|
||||
}
|
||||
else if (S <= 128)
|
||||
{
|
||||
maskIdxKernelSmall<128><<<B, 128, 0, stream>>>(S, mask, maskIdx);
|
||||
}
|
||||
else if (S == 384)
|
||||
{
|
||||
maskIdxKernelSmall<384><<<B, 384, 0, stream>>>(S, mask, maskIdx);
|
||||
}
|
||||
else
|
||||
{
|
||||
maskIdxKernel<256><<<B, 256, 0, stream>>>(S, mask, maskIdx);
|
||||
}
|
||||
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
template <typename T, unsigned TPB>
|
||||
__global__ void embLayerNormKernel(int ld, int32_t const* inputIds, int32_t const* tokenIds, float const* beta,
|
||||
float const* gamma, T const* wordEmb, T const* posEmb, T const* tokEmb, int32_t const wordSize,
|
||||
int32_t const tokSize, T* output)
|
||||
{
|
||||
|
||||
// 1. lookup word and token of the block
|
||||
// blockIdx.x = position in the sequence
|
||||
// blockIdx.y = batch
|
||||
// gridDim.x = S
|
||||
// gridDim.y = B
|
||||
__shared__ int wordId;
|
||||
__shared__ int tokenId;
|
||||
|
||||
T const rld = T(1.f) / T(ld);
|
||||
int32_t const seqPos = blockIdx.y + blockIdx.x * gridDim.y;
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
wordId = inputIds[seqPos];
|
||||
tokenId = tokenIds[seqPos];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 2. load pos/tok/word embeddings and add them toghether
|
||||
// offset into embeddings is given by wordId * hidden_size
|
||||
int32_t const poffset = blockIdx.x * ld;
|
||||
int32_t const woffset = wordId * ld;
|
||||
int32_t const toffset = tokenId * ld;
|
||||
// the output offset is given by b * (S*hidden_size) + s * hidden_size
|
||||
int32_t const outOffset = seqPos * ld;
|
||||
|
||||
kvp<T> threadData(0, 0);
|
||||
|
||||
if (wordId >= 0 && wordId < wordSize && tokenId >= 0 && tokenId < tokSize)
|
||||
{
|
||||
for (int it = threadIdx.x; it < ld; it += TPB)
|
||||
{
|
||||
T const w(wordEmb[woffset + it]);
|
||||
T const t(tokEmb[toffset + it]);
|
||||
T const p(posEmb[poffset + it]);
|
||||
T const val = w + t + p;
|
||||
|
||||
output[outOffset + it] = val;
|
||||
T const rldval = rld * val;
|
||||
threadData = threadData + kvp<T>(rldval, rldval * val);
|
||||
}
|
||||
}
|
||||
|
||||
// 3. layer norm on the sum
|
||||
layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
int embSkipLayerNorm(cudaStream_t stream, int ld, int B, int S, int32_t const* inputIds, int32_t const* tokenIds,
|
||||
float const* beta, float const* gamma, T const* wordEmb, T const* posEmb, T const* tokEmb, int32_t const wordSize,
|
||||
int32_t const tokSize, T* output)
|
||||
{
|
||||
|
||||
constexpr int tpb = 256;
|
||||
dim3 const grid(S, B, 1);
|
||||
dim3 const block(tpb, 1, 1);
|
||||
|
||||
embLayerNormKernel<T, tpb><<<grid, block, 0, stream>>>(
|
||||
ld, inputIds, tokenIds, beta, gamma, wordEmb, posEmb, tokEmb, wordSize, tokSize, output);
|
||||
PLUGIN_CHECK(cudaPeekAtLastError());
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
template int embSkipLayerNorm<float>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
|
||||
float const*, float const*, float const*, float const*, float const*, int32_t const, int32_t const, float*);
|
||||
|
||||
template int embSkipLayerNorm<half>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
|
||||
float const*, float const*, half const*, half const*, half const*, int32_t const, int32_t const, half*);
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
@@ -0,0 +1,687 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
#if CUDA_VERSION >= 10010
|
||||
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "embLayerNormPlugin.h"
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::plugin::bert;
|
||||
|
||||
namespace
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
char const* gEmbLayerNormVersion{"6"};
|
||||
char const* gEmbLayerNormName{"CustomEmbLayerNormPluginDynamic"};
|
||||
} // namespace
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(EmbLayerNormPluginDynamicCreator);
|
||||
|
||||
EmbLayerNormPluginDynamic::EmbLayerNormPluginDynamic(std::string const& name, DataType const type,
|
||||
DataType const mhaType, Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb,
|
||||
Weights const& tokEmb, bool const useFullMask)
|
||||
: mLayerName(name)
|
||||
, mLd(beta.count)
|
||||
, mType(type)
|
||||
, mMhaType(mhaType)
|
||||
{
|
||||
// Assuming Weights.count is the number of elements and not bytes
|
||||
PLUGIN_VALIDATE(beta.count == gamma.count);
|
||||
PLUGIN_VALIDATE(mLd > 0U);
|
||||
PLUGIN_VALIDATE(wordEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(posEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(tokEmb.count % mLd == 0);
|
||||
mWordVocabSize = wordEmb.count / mLd;
|
||||
mPosVocabSize = posEmb.count / mLd;
|
||||
mTokVocabSize = tokEmb.count / mLd;
|
||||
mSM = getSmVersion();
|
||||
mOutputFp16 = mType == DataType::kHALF ? 1 : 0;
|
||||
mUseFullMask = static_cast<int32_t>(useFullMask);
|
||||
// NOTE: mS is set during configure
|
||||
mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
|
||||
mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
|
||||
mWordEmb.convertAndCopy(wordEmb, mType);
|
||||
mTokEmb.convertAndCopy(tokEmb, mType);
|
||||
mPosEmb.convertAndCopy(posEmb, mType);
|
||||
|
||||
copyToDevice(mGamma, sizeof(float) * mGamma.count, mGammaDev);
|
||||
copyToDevice(mBeta, sizeof(float) * mBeta.count, mBetaDev);
|
||||
copyToDevice(mWordEmb, getWeightsSize(mWordEmb, mType), mWordEmbDev);
|
||||
copyToDevice(mPosEmb, getWeightsSize(mPosEmb, mType), mPosEmbDev);
|
||||
copyToDevice(mTokEmb, getWeightsSize(mTokEmb, mType), mTokEmbDev);
|
||||
}
|
||||
|
||||
EmbLayerNormPluginDynamic::~EmbLayerNormPluginDynamic()
|
||||
{
|
||||
try
|
||||
{
|
||||
// This gets called when the network containing plugin is destroyed
|
||||
mGammaDev.reset(nullptr);
|
||||
mBetaDev.reset(nullptr);
|
||||
mWordEmbDev.reset(nullptr);
|
||||
mPosEmbDev.reset(nullptr);
|
||||
mTokEmbDev.reset(nullptr);
|
||||
// delete this; TRT or the creator of the plugin will delete this plugin object
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
//////
|
||||
// IPluginV3 method definitions:
|
||||
// - getCapabilityInterface() (Base)
|
||||
// - clone() (HFace, MTron)
|
||||
//////
|
||||
IPluginCapability* EmbLayerNormPluginDynamic::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* EmbLayerNormPluginDynamic::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamic clone.");
|
||||
|
||||
auto p = std::make_unique<EmbLayerNormPluginDynamic>(
|
||||
mLayerName, mType, mMhaType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb, mUseFullMask == 1);
|
||||
p->mS = mS;
|
||||
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()
|
||||
// - onShapeChange()
|
||||
// - attachToContext()
|
||||
// - enqueue()
|
||||
/////
|
||||
|
||||
PluginFieldCollection const* EmbLayerNormPluginDynamic::getFieldsToSerialize() noexcept
|
||||
{
|
||||
mDataToSerialize.clear();
|
||||
mDataToSerialize.emplace_back("output_fp16", &mOutputFp16, PluginFieldType::kINT32, 1);
|
||||
mDataToSerialize.emplace_back("full_mask", &mUseFullMask, PluginFieldType::kINT32, 1);
|
||||
mDataToSerialize.emplace_back("mha_type_id", &mMhaType, PluginFieldType::kINT32, 1);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_layernorm_beta", static_cast<float const*>(mBeta.values),
|
||||
PluginFieldType::kFLOAT32, mBeta.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_layernorm_gamma", static_cast<float const*>(mGamma.values),
|
||||
PluginFieldType::kFLOAT32, mGamma.count);
|
||||
if (mOutputFp16)
|
||||
{
|
||||
mDataToSerialize.emplace_back("bert_embeddings_word_embeddings", static_cast<half const*>(mWordEmb.values),
|
||||
PluginFieldType::kFLOAT16, mWordEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_token_type_embeddings", static_cast<half const*>(mTokEmb.values),
|
||||
PluginFieldType::kFLOAT16, mTokEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_position_embeddings", static_cast<half const*>(mPosEmb.values),
|
||||
PluginFieldType::kFLOAT16, mPosEmb.count);
|
||||
}
|
||||
else
|
||||
{
|
||||
mDataToSerialize.emplace_back("bert_embeddings_word_embeddings", static_cast<float const*>(mWordEmb.values),
|
||||
PluginFieldType::kFLOAT32, mWordEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_token_type_embeddings",
|
||||
static_cast<float const*>(mTokEmb.values), PluginFieldType::kFLOAT32, mTokEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_position_embeddings", static_cast<float const*>(mPosEmb.values),
|
||||
PluginFieldType::kFLOAT32, mPosEmb.count);
|
||||
}
|
||||
mFCToSerialize.nbFields = mDataToSerialize.size();
|
||||
mFCToSerialize.fields = mDataToSerialize.data();
|
||||
return &mFCToSerialize;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamic::onShapeChange(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamic configurePlugin.");
|
||||
try
|
||||
{
|
||||
// Validate input arguments
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
PLUGIN_ASSERT(nbInputs == 3);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].dims.nbDims == 2);
|
||||
int32_t const S = inputs[0].dims.d[SDIM];
|
||||
mS = S;
|
||||
int32_t const B = inputs[0].dims.d[BDIM];
|
||||
TRT_UNUSED B;
|
||||
PLUGIN_ASSERT(mS == static_cast<size_t>(inputs[1].dims.d[SDIM]));
|
||||
PLUGIN_ASSERT(B == inputs[1].dims.d[BDIM]);
|
||||
PLUGIN_ASSERT(mS == static_cast<size_t>(inputs[2].dims.d[SDIM]));
|
||||
PLUGIN_ASSERT(B == inputs[2].dims.d[BDIM]);
|
||||
|
||||
PLUGIN_ASSERT(outputs[0].dims.nbDims == 5);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[SDIM]) == mS);
|
||||
PLUGIN_ASSERT(outputs[0].dims.d[BDIM] == B);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[2]) == mLd);
|
||||
PLUGIN_ASSERT(outputs[0].dims.d[3] == 1);
|
||||
PLUGIN_ASSERT(outputs[0].dims.d[4] == 1);
|
||||
|
||||
if (mUseFullMask)
|
||||
{
|
||||
// user force full_mask
|
||||
PLUGIN_ASSERT(outputs[1].dims.nbDims == 2);
|
||||
PLUGIN_ASSERT(outputs[1].dims.d[0] == B);
|
||||
PLUGIN_ASSERT((outputs[1].dims.d[1] == -1) || (outputs[1].dims.d[1] == packedMaskSize384)
|
||||
|| (outputs[1].dims.d[1] == packedMaskSize128));
|
||||
}
|
||||
else
|
||||
{
|
||||
// auto detect using mhatype
|
||||
if (S != -1 && B != -1)
|
||||
{
|
||||
PLUGIN_ASSERT(outputs[1].dims.nbDims == 2);
|
||||
PLUGIN_ASSERT(outputs[1].dims.d[0] == B);
|
||||
int32_t packedSize = getMHAMaskPackedSize(mSM, mMhaType, S);
|
||||
TRT_UNUSED packedSize;
|
||||
PLUGIN_ASSERT(outputs[1].dims.d[1] == -1 || outputs[1].dims.d[1] == packedSize);
|
||||
}
|
||||
}
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[1].type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[2].type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(outputs[0].type == mType);
|
||||
PLUGIN_ASSERT(outputs[1].type == DataType::kINT32);
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
IPluginV3* EmbLayerNormPluginDynamic::attachToContext(IPluginResourceContext* context) noexcept
|
||||
{
|
||||
return clone();
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamic::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* /* outputDesc */,
|
||||
void const* const* inputs, void* const* outputs, void* /* workspace */, cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
|
||||
|
||||
int32_t const batchSize = inputDesc->dims.d[BDIM];
|
||||
int32_t const S = inputDesc->dims.d[SDIM];
|
||||
int32_t status = STATUS_FAILURE;
|
||||
|
||||
// Our plugin outputs only one tensor
|
||||
auto const inputIds = static_cast<int32_t const*>(inputs[0]);
|
||||
auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
|
||||
auto const inputMask = static_cast<int32_t const*>(inputs[2]);
|
||||
|
||||
float const* beta = mBetaDev.get();
|
||||
float const* gamma = mGammaDev.get();
|
||||
if (mType == DataType::kFLOAT)
|
||||
{
|
||||
auto output = static_cast<float*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
|
||||
status = embSkipLayerNorm<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
|
||||
if (status != cudaSuccess)
|
||||
{
|
||||
return status;
|
||||
}
|
||||
}
|
||||
else if (mType == DataType::kHALF)
|
||||
{
|
||||
auto output = static_cast<half*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
|
||||
status = embSkipLayerNorm<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds, beta,
|
||||
gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
|
||||
if (status != cudaSuccess)
|
||||
{
|
||||
return status;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
|
||||
<< std::endl;
|
||||
|
||||
return STATUS_NOT_SUPPORTED;
|
||||
}
|
||||
|
||||
// check mha use fused kernel
|
||||
if (mUseFullMask || unfusedMaskSize != getMHAMaskPackedSize(mSM, mMhaType, S))
|
||||
{
|
||||
size_t warps_m = 0, warps_n = 0, warps_k = 1;
|
||||
if (S == 64 || S == 96 || S == 128)
|
||||
{
|
||||
warps_m = 2;
|
||||
warps_n = 2;
|
||||
}
|
||||
else if (S == 384)
|
||||
{
|
||||
warps_m = 1;
|
||||
warps_n = 8;
|
||||
}
|
||||
uint32_t* inputMaskX = static_cast<uint32_t*>(outputs[1]);
|
||||
|
||||
status = convertMask(S, batchSize, warps_m, warps_n, warps_k, inputMask, inputMaskX, stream);
|
||||
}
|
||||
else
|
||||
{
|
||||
int32_t* maskIdx = static_cast<int32_t*>(outputs[1]);
|
||||
status = computeMaskIdx(stream, S, batchSize, inputMask, maskIdx);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
// end IPluginV3OneRuntime method definitions
|
||||
|
||||
///////
|
||||
// IPluginV3OneBuild method definitions
|
||||
// - getNbOutputs()
|
||||
// - supportsFormatCombination()
|
||||
// - getOutputShapes
|
||||
// - getOutputDataTypes()
|
||||
// - configurePlugin()
|
||||
// - getWorkSpaceSize()
|
||||
//////
|
||||
|
||||
int32_t EmbLayerNormPluginDynamic::getNbOutputs() const noexcept
|
||||
{
|
||||
return 2;
|
||||
}
|
||||
|
||||
bool EmbLayerNormPluginDynamic::supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// 3 inputs of size BxS
|
||||
PLUGIN_ASSERT(nbInputs == 3);
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PluginTensorDesc const& desc = inOut[pos].desc;
|
||||
if (desc.format != TensorFormat::kLINEAR)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if (pos == 0)
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 2;
|
||||
}
|
||||
|
||||
PluginTensorDesc const& prev = inOut[pos - 1].desc;
|
||||
if (pos == 1 || pos == 2)
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 2 && desc.dims.d[BDIM] == prev.dims.d[BDIM]
|
||||
&& desc.dims.d[SDIM] == prev.dims.d[SDIM];
|
||||
}
|
||||
|
||||
// embedded sequence
|
||||
if (pos == 3)
|
||||
{
|
||||
return desc.type == mType && desc.dims.nbDims == 5 && desc.dims.d[BDIM] == prev.dims.d[BDIM]
|
||||
&& desc.dims.d[SDIM] == prev.dims.d[SDIM] && desc.dims.d[3] == 1 && desc.dims.d[4] == 1;
|
||||
}
|
||||
// mask
|
||||
return desc.type == DataType::kINT32;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamic::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
|
||||
DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
|
||||
IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// Input should be input ids and token ids and the input mask
|
||||
// Output should be the embeddings tensor and mask indices
|
||||
PLUGIN_ASSERT(nbInputs == 3);
|
||||
PLUGIN_ASSERT(inputs != nullptr);
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == 2); // BxS
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[2].nbDims);
|
||||
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
PLUGIN_ASSERT(outputs != nullptr);
|
||||
|
||||
// output 0: embeddings tensor
|
||||
outputs[0].nbDims = 5;
|
||||
outputs[0].d[0] = inputs[0].d[0];
|
||||
outputs[0].d[1] = inputs[0].d[1];
|
||||
outputs[0].d[2] = exprBuilder.constant(mLd);
|
||||
outputs[0].d[3] = exprBuilder.constant(1);
|
||||
outputs[0].d[4] = exprBuilder.constant(1);
|
||||
|
||||
// output 1: mask indices
|
||||
outputs[1].nbDims = 2;
|
||||
outputs[1].d[0] = inputs[0].d[BDIM];
|
||||
auto cms0 = exprBuilder.constant(unfusedMaskSize);
|
||||
|
||||
// this code must match getMHAMaskPackedSize in bertCommon.h
|
||||
bool const isSmOK = elem(mSM, {kSM_75, kSM_80, kSM_86, kSM_87, kSM_89, kSM_90, kSM_100, kSM_120});
|
||||
bool const isPrecisionOK = (mMhaType == nvinfer1::DataType::kHALF || mMhaType == nvinfer1::DataType::kINT8);
|
||||
if (mUseFullMask || (isSmOK && isPrecisionOK))
|
||||
{
|
||||
// support 128, 384 in both int8 and fp16
|
||||
auto cms128 = exprBuilder.constant(packedMaskSize128);
|
||||
auto cms384 = exprBuilder.constant(packedMaskSize384);
|
||||
auto c128 = exprBuilder.constant(128);
|
||||
auto c384 = exprBuilder.constant(384);
|
||||
auto is128 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c128);
|
||||
auto is384 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c384);
|
||||
auto sel128 = exprBuilder.operation(DimensionOperation::kPROD, *is128, *cms128);
|
||||
auto sel384 = exprBuilder.operation(DimensionOperation::kPROD, *is384, *cms384);
|
||||
auto maskSize = exprBuilder.operation(DimensionOperation::kSUM, *sel384, *sel128);
|
||||
|
||||
// support 64, 96 in both int8 and fp16
|
||||
auto cms64 = exprBuilder.constant(packedMaskSize64);
|
||||
auto cms96 = exprBuilder.constant(packedMaskSize96);
|
||||
auto c64 = exprBuilder.constant(64);
|
||||
auto c96 = exprBuilder.constant(96);
|
||||
|
||||
auto is64 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c64);
|
||||
auto is96 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c96);
|
||||
auto sel64 = exprBuilder.operation(DimensionOperation::kPROD, *is64, *cms64);
|
||||
auto sel96 = exprBuilder.operation(DimensionOperation::kPROD, *is96, *cms96);
|
||||
auto maskSize2 = exprBuilder.operation(DimensionOperation::kSUM, *sel64, *sel96);
|
||||
maskSize = exprBuilder.operation(DimensionOperation::kSUM, *maskSize, *maskSize2);
|
||||
|
||||
auto is0 = exprBuilder.operation(DimensionOperation::kEQUAL, *maskSize, *exprBuilder.constant(0));
|
||||
auto sel0 = exprBuilder.operation(DimensionOperation::kPROD, *is0, *cms0);
|
||||
auto combinedMaskSize = exprBuilder.operation(DimensionOperation::kSUM, *maskSize, *sel0);
|
||||
outputs[1].d[1] = combinedMaskSize;
|
||||
}
|
||||
else
|
||||
{
|
||||
outputs[1].d[1] = cms0;
|
||||
}
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamic::getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_ASSERT(outputTypes != nullptr);
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
PLUGIN_ASSERT(inputTypes != nullptr);
|
||||
PLUGIN_ASSERT(nbInputs == 3);
|
||||
PLUGIN_ASSERT(mType == DataType::kHALF || mType == DataType::kFLOAT);
|
||||
outputTypes[0] = mType;
|
||||
outputTypes[1] = DataType::kINT32;
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamic::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
size_t EmbLayerNormPluginDynamic::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()
|
||||
// - getPluginName()
|
||||
// - getPluginNamespace()
|
||||
// - setPluginNamespace()
|
||||
//////
|
||||
char const* EmbLayerNormPluginDynamic::getPluginVersion() const noexcept
|
||||
{
|
||||
return gEmbLayerNormVersion;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamic::getPluginName() const noexcept
|
||||
{
|
||||
return gEmbLayerNormName;
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamic::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamic::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
// End IPluginV3OneCore method definitions
|
||||
|
||||
//////////////////////////// Plugin Creator member definitions /////////////////////////////
|
||||
|
||||
EmbLayerNormPluginDynamicCreator::EmbLayerNormPluginDynamicCreator()
|
||||
{
|
||||
static std::mutex sMutex;
|
||||
std::lock_guard<std::mutex> lock(sMutex);
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_beta"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_gamma"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_word_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_token_type_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_position_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("output_fp16"));
|
||||
mPluginAttributes.emplace_back(PluginField("full_mask"));
|
||||
mPluginAttributes.emplace_back(PluginField("mha_type_id"));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicCreator::getPluginName() const noexcept
|
||||
{
|
||||
return gEmbLayerNormName;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return gEmbLayerNormVersion;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* EmbLayerNormPluginDynamicCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV3* EmbLayerNormPluginDynamicCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamic createPlugin.");
|
||||
|
||||
bool output_fp16 = false;
|
||||
bool useFullMask = false;
|
||||
Weights beta{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights gamma{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights word_emb{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights pos_emb{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights tok_emb{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
int32_t mhaTypeId = 0;
|
||||
std::set<std::string> const requiredAttributes{
|
||||
"bert_embeddings_layernorm_beta",
|
||||
"bert_embeddings_layernorm_gamma",
|
||||
"bert_embeddings_word_embeddings",
|
||||
"bert_embeddings_token_type_embeddings",
|
||||
"bert_embeddings_position_embeddings",
|
||||
};
|
||||
plugin::validateRequiredAttributesExist(requiredAttributes, fc);
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; i++)
|
||||
{
|
||||
std::string_view const field_name = fc->fields[i].name;
|
||||
if (field_name == "bert_embeddings_layernorm_beta"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
|
||||
beta.values = fc->fields[i].data;
|
||||
beta.count = fc->fields[i].length;
|
||||
beta.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_layernorm_gamma"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
|
||||
gamma.values = fc->fields[i].data;
|
||||
gamma.count = fc->fields[i].length;
|
||||
gamma.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_word_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_word_embeddings...");
|
||||
word_emb.values = fc->fields[i].data;
|
||||
word_emb.count = fc->fields[i].length;
|
||||
word_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_token_type_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_token_type_embeddings...");
|
||||
tok_emb.values = fc->fields[i].data;
|
||||
tok_emb.count = fc->fields[i].length;
|
||||
tok_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_position_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_position_embeddings...");
|
||||
pos_emb.values = fc->fields[i].data;
|
||||
pos_emb.count = fc->fields[i].length;
|
||||
pos_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
if (field_name == "output_fp16"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building output_fp16...");
|
||||
PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kINT32);
|
||||
output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
|
||||
}
|
||||
if (field_name == "full_mask"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building full_mask...");
|
||||
PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kINT32);
|
||||
useFullMask = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
|
||||
}
|
||||
if (field_name == "mha_type_id"sv)
|
||||
{
|
||||
mhaTypeId = *static_cast<int32_t const*>(fc->fields[i].data);
|
||||
PLUGIN_VALIDATE(mhaTypeId >= 0 && mhaTypeId <= 3);
|
||||
BERT_DEBUG_VALUE("Building mha typeId: ", mhaTypeId);
|
||||
}
|
||||
}
|
||||
|
||||
BERT_DEBUG_MSG("Building the Plugin...");
|
||||
DataType mhaType = static_cast<DataType>(mhaTypeId);
|
||||
auto p = std::make_unique<EmbLayerNormPluginDynamic>(name, output_fp16 ? DataType::kHALF : DataType::kFLOAT,
|
||||
mhaType, beta, gamma, word_emb, pos_emb, tok_emb, useFullMask);
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamicCreator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
@@ -0,0 +1,174 @@
|
||||
/*
|
||||
* 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_EMB_LAYER_NORM_PLUGIN_H
|
||||
#define TRT_EMB_LAYER_NORM_PLUGIN_H
|
||||
|
||||
#include "NvInferPlugin.h"
|
||||
#include "NvInferRuntime.h"
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
int32_t computeMaskIdx(cudaStream_t stream, int32_t const S, int32_t const B, int32_t const* mask, int32_t* maskIdx);
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNorm(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* token_ids, float const* beta, float const* gamma, T const* wordEmb, T const* posEmb, T const* tokEmb,
|
||||
int32_t const wordSize, int32_t const tokSize, T* output);
|
||||
|
||||
cudaError_t convertMask(uint32_t const S, uint32_t const B, uint32_t const warps_m, uint32_t const warps_n,
|
||||
uint32_t const warps_k, int32_t const* inputMaskSB, uint32_t* inputMaskX, cudaStream_t stream);
|
||||
|
||||
class EmbLayerNormPluginDynamic : public IPluginV3,
|
||||
public IPluginV3OneCore,
|
||||
public IPluginV3OneBuild,
|
||||
public IPluginV3OneRuntime
|
||||
{
|
||||
public:
|
||||
EmbLayerNormPluginDynamic(std::string const& name, nvinfer1::DataType const type, nvinfer1::DataType const mhaType,
|
||||
nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma, nvinfer1::Weights const& word_emb,
|
||||
nvinfer1::Weights const& pos_emb, nvinfer1::Weights const& tok_emb, bool const useFullMask);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormPluginDynamic without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormPluginDynamic() = delete;
|
||||
|
||||
~EmbLayerNormPluginDynamic() 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;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
// end of IPluginV3OneCore Methods
|
||||
|
||||
// IPluginV3Build Methods
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) 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;
|
||||
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
int32_t getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
|
||||
int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override;
|
||||
// end IPluginV3Build Methods
|
||||
|
||||
// IPluginV3Runtime Methods
|
||||
IPluginV3* clone() noexcept override;
|
||||
|
||||
int32_t onShapeChange(
|
||||
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) 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;
|
||||
|
||||
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
|
||||
// end IPluginV3Runtime Methods
|
||||
|
||||
private:
|
||||
// metadata fields
|
||||
std::string const mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
// device-side
|
||||
bert::cuda_unique_ptr<float> mGammaDev;
|
||||
bert::cuda_unique_ptr<float> mBetaDev;
|
||||
bert::cuda_unique_ptr<void> mWordEmbDev;
|
||||
bert::cuda_unique_ptr<void> mTokEmbDev;
|
||||
bert::cuda_unique_ptr<void> mPosEmbDev;
|
||||
size_t mLd; // leading dim = hidden size
|
||||
size_t mS; // sequence length
|
||||
size_t mWordVocabSize;
|
||||
size_t mPosVocabSize;
|
||||
size_t mTokVocabSize;
|
||||
|
||||
// members that partcipate in ser/deserialization
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mWordEmb;
|
||||
bert::WeightsWithOwnership mTokEmb;
|
||||
bert::WeightsWithOwnership mPosEmb;
|
||||
nvinfer1::DataType mType;
|
||||
int32_t mOutputFp16;
|
||||
int32_t mUseFullMask;
|
||||
nvinfer1::DataType mMhaType;
|
||||
int32_t mSM;
|
||||
|
||||
// IPluginV3 serialization related
|
||||
std::vector<nvinfer1::PluginField> mDataToSerialize;
|
||||
nvinfer1::PluginFieldCollection mFCToSerialize;
|
||||
};
|
||||
|
||||
class EmbLayerNormPluginDynamicCreator : public nvinfer1::IPluginCreatorV3One
|
||||
{
|
||||
public:
|
||||
EmbLayerNormPluginDynamicCreator();
|
||||
~EmbLayerNormPluginDynamicCreator() 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* pluginNamespace) 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_EMB_LAYER_NORM_PLUGIN_H
|
||||
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
@@ -0,0 +1,666 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cuda.h>
|
||||
#if CUDA_VERSION >= 10010
|
||||
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "common/serialize.hpp"
|
||||
#include "embLayerNormPluginLegacy.h"
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::plugin::bert;
|
||||
|
||||
namespace
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
char const* gEmbLayerNormVersion{"1"};
|
||||
char const* gEmbLayerNormName{"CustomEmbLayerNormPluginDynamic"};
|
||||
} // namespace
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(EmbLayerNormPluginDynamicLegacyCreator);
|
||||
|
||||
EmbLayerNormPluginDynamicLegacy::EmbLayerNormPluginDynamicLegacy(std::string const& name, DataType const type,
|
||||
DataType const mhaType, Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb,
|
||||
Weights const& tokEmb, bool const useFullMask)
|
||||
: mLayerName(name)
|
||||
, mLd(beta.count)
|
||||
, mType(type)
|
||||
, mUseFullMask(useFullMask)
|
||||
, mMhaType(mhaType)
|
||||
{
|
||||
// Assuming Weights.count is the number of elements and not bytes
|
||||
PLUGIN_VALIDATE(beta.count == gamma.count);
|
||||
PLUGIN_VALIDATE(mLd > 0U);
|
||||
PLUGIN_VALIDATE(wordEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(posEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(tokEmb.count % mLd == 0);
|
||||
mWordVocabSize = wordEmb.count / mLd;
|
||||
mPosVocabSize = posEmb.count / mLd;
|
||||
mTokVocabSize = tokEmb.count / mLd;
|
||||
mSM = getSmVersion();
|
||||
// mS is set during configure
|
||||
|
||||
mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
|
||||
mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
|
||||
mWordEmb.convertAndCopy(wordEmb, mType);
|
||||
mTokEmb.convertAndCopy(tokEmb, mType);
|
||||
mPosEmb.convertAndCopy(posEmb, mType);
|
||||
|
||||
copyToDevice(mGamma, sizeof(float) * mGamma.count, mGammaDev);
|
||||
copyToDevice(mBeta, sizeof(float) * mBeta.count, mBetaDev);
|
||||
copyToDevice(mWordEmb, getWeightsSize(mWordEmb, mType), mWordEmbDev);
|
||||
copyToDevice(mPosEmb, getWeightsSize(mPosEmb, mType), mPosEmbDev);
|
||||
copyToDevice(mTokEmb, getWeightsSize(mTokEmb, mType), mTokEmbDev);
|
||||
}
|
||||
|
||||
EmbLayerNormPluginDynamicLegacy::EmbLayerNormPluginDynamicLegacy(
|
||||
std::string const& name, void const* data, size_t length)
|
||||
: mLayerName(name)
|
||||
, mGammaDev(nullptr)
|
||||
, mBetaDev(nullptr)
|
||||
, mWordEmbDev(nullptr)
|
||||
, mTokEmbDev(nullptr)
|
||||
, mPosEmbDev(nullptr)
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamicLegacy deserialize.");
|
||||
|
||||
// Deserialize in the same order as serialization
|
||||
deserialize_value(&data, &length, &mType);
|
||||
deserialize_value(&data, &length, &mMhaType);
|
||||
deserialize_value(&data, &length, &mLd);
|
||||
deserialize_value(&data, &length, &mS);
|
||||
deserialize_value(&data, &length, &mWordVocabSize);
|
||||
deserialize_value(&data, &length, &mPosVocabSize);
|
||||
deserialize_value(&data, &length, &mTokVocabSize);
|
||||
deserialize_value(&data, &length, &mUseFullMask);
|
||||
deserialize_value(&data, &length, &mSM);
|
||||
|
||||
char const* d = static_cast<char const*>(data);
|
||||
mBeta.convertAndCopy(d, mLd, nvinfer1::DataType::kFLOAT);
|
||||
mGamma.convertAndCopy(d, mLd, nvinfer1::DataType::kFLOAT);
|
||||
mWordEmb.convertAndCopy(d, mLd * mWordVocabSize, mType);
|
||||
mPosEmb.convertAndCopy(d, mLd * mPosVocabSize, mType);
|
||||
mTokEmb.convertAndCopy(d, mLd * mTokVocabSize, mType);
|
||||
|
||||
copyToDevice(mGamma, sizeof(float) * mGamma.count, mGammaDev);
|
||||
copyToDevice(mBeta, sizeof(float) * mBeta.count, mBetaDev);
|
||||
copyToDevice(mWordEmb, getWeightsSize(mWordEmb, mType), mWordEmbDev);
|
||||
copyToDevice(mPosEmb, getWeightsSize(mPosEmb, mType), mPosEmbDev);
|
||||
copyToDevice(mTokEmb, getWeightsSize(mTokEmb, mType), mTokEmbDev);
|
||||
}
|
||||
|
||||
// IPluginV2DynamicExt Methods
|
||||
IPluginV2DynamicExt* EmbLayerNormPluginDynamicLegacy::clone() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamicLegacy clone.");
|
||||
|
||||
auto p = std::make_unique<EmbLayerNormPluginDynamicLegacy>(
|
||||
mLayerName, mType, mMhaType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb, mUseFullMask);
|
||||
p->mS = mS;
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
DimsExprs EmbLayerNormPluginDynamicLegacy::getOutputDimensions(
|
||||
int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// Input should be input ids and token ids and the input mask
|
||||
// Output should be the embeddings tensor and mask indices
|
||||
PLUGIN_ASSERT(nbInputs == 3);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == 2); // BxS
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[2].nbDims);
|
||||
|
||||
PLUGIN_ASSERT(outputIndex == 0 || outputIndex == 1);
|
||||
|
||||
if (outputIndex == 0)
|
||||
{
|
||||
DimsExprs ret;
|
||||
ret.nbDims = 5;
|
||||
ret.d[0] = inputs[0].d[0];
|
||||
ret.d[1] = inputs[0].d[1];
|
||||
ret.d[2] = exprBuilder.constant(mLd);
|
||||
ret.d[3] = exprBuilder.constant(1);
|
||||
ret.d[4] = exprBuilder.constant(1);
|
||||
return ret;
|
||||
}
|
||||
|
||||
DimsExprs ret;
|
||||
ret.nbDims = 2;
|
||||
ret.d[0] = inputs[0].d[BDIM];
|
||||
auto cms0 = exprBuilder.constant(unfusedMaskSize);
|
||||
|
||||
// this code must match getMHAMaskPackedSize in bertCommon.h
|
||||
bool const isSmOK = elem(mSM, {kSM_75, kSM_80, kSM_86, kSM_87, kSM_89, kSM_90, kSM_100, kSM_120});
|
||||
bool const isPrecisionOK = (mMhaType == nvinfer1::DataType::kHALF || mMhaType == nvinfer1::DataType::kINT8);
|
||||
if (mUseFullMask || (isSmOK && isPrecisionOK))
|
||||
{
|
||||
// support 128, 384 in both int8 and fp16
|
||||
auto cms128 = exprBuilder.constant(packedMaskSize128);
|
||||
auto cms384 = exprBuilder.constant(packedMaskSize384);
|
||||
auto c128 = exprBuilder.constant(128);
|
||||
auto c384 = exprBuilder.constant(384);
|
||||
auto is128 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c128);
|
||||
auto is384 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c384);
|
||||
auto sel128 = exprBuilder.operation(DimensionOperation::kPROD, *is128, *cms128);
|
||||
auto sel384 = exprBuilder.operation(DimensionOperation::kPROD, *is384, *cms384);
|
||||
auto maskSize = exprBuilder.operation(DimensionOperation::kSUM, *sel384, *sel128);
|
||||
|
||||
// support 64, 96 in both int8 and fp16
|
||||
auto cms64 = exprBuilder.constant(packedMaskSize64);
|
||||
auto cms96 = exprBuilder.constant(packedMaskSize96);
|
||||
auto c64 = exprBuilder.constant(64);
|
||||
auto c96 = exprBuilder.constant(96);
|
||||
|
||||
auto is64 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c64);
|
||||
auto is96 = exprBuilder.operation(DimensionOperation::kEQUAL, *inputs[0].d[SDIM], *c96);
|
||||
auto sel64 = exprBuilder.operation(DimensionOperation::kPROD, *is64, *cms64);
|
||||
auto sel96 = exprBuilder.operation(DimensionOperation::kPROD, *is96, *cms96);
|
||||
auto maskSize2 = exprBuilder.operation(DimensionOperation::kSUM, *sel64, *sel96);
|
||||
maskSize = exprBuilder.operation(DimensionOperation::kSUM, *maskSize, *maskSize2);
|
||||
|
||||
auto is0 = exprBuilder.operation(DimensionOperation::kEQUAL, *maskSize, *exprBuilder.constant(0));
|
||||
auto sel0 = exprBuilder.operation(DimensionOperation::kPROD, *is0, *cms0);
|
||||
auto combinedMaskSize = exprBuilder.operation(DimensionOperation::kSUM, *maskSize, *sel0);
|
||||
ret.d[1] = combinedMaskSize;
|
||||
}
|
||||
else
|
||||
{
|
||||
ret.d[1] = cms0;
|
||||
}
|
||||
|
||||
return ret;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return DimsExprs{};
|
||||
}
|
||||
|
||||
bool EmbLayerNormPluginDynamicLegacy::supportsFormatCombination(
|
||||
int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// 3 inputs of size BxS
|
||||
PLUGIN_ASSERT(nbInputs == 3);
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PluginTensorDesc const& desc = inOut[pos];
|
||||
if (desc.format != TensorFormat::kLINEAR)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if (pos == 0)
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 2;
|
||||
}
|
||||
|
||||
PluginTensorDesc const& prev = inOut[pos - 1];
|
||||
if (pos == 1 || pos == 2)
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 2 && desc.dims.d[BDIM] == prev.dims.d[BDIM]
|
||||
&& desc.dims.d[SDIM] == prev.dims.d[SDIM];
|
||||
}
|
||||
|
||||
// embedded sequence
|
||||
if (pos == 3)
|
||||
{
|
||||
return desc.type == mType && desc.dims.nbDims == 5 && desc.dims.d[BDIM] == prev.dims.d[BDIM]
|
||||
&& desc.dims.d[SDIM] == prev.dims.d[SDIM] && desc.dims.d[3] == 1 && desc.dims.d[4] == 1;
|
||||
}
|
||||
// mask
|
||||
return desc.type == DataType::kINT32;
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamicLegacy::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamicLegacy configurePlugin.");
|
||||
|
||||
// Validate input arguments
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
PLUGIN_ASSERT(nbInputs == 3);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].desc.dims.nbDims == 2);
|
||||
int32_t const S = inputs[0].desc.dims.d[SDIM];
|
||||
mS = S;
|
||||
int32_t const B = inputs[0].desc.dims.d[BDIM];
|
||||
TRT_UNUSED B;
|
||||
PLUGIN_ASSERT(mS == static_cast<size_t>(inputs[1].desc.dims.d[SDIM]));
|
||||
PLUGIN_ASSERT(B == inputs[1].desc.dims.d[BDIM]);
|
||||
PLUGIN_ASSERT(mS == static_cast<size_t>(inputs[2].desc.dims.d[SDIM]));
|
||||
PLUGIN_ASSERT(B == inputs[2].desc.dims.d[BDIM]);
|
||||
|
||||
PLUGIN_ASSERT(outputs[0].desc.dims.nbDims == 5);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].desc.dims.d[SDIM]) == mS);
|
||||
PLUGIN_ASSERT(outputs[0].desc.dims.d[BDIM] == B);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].desc.dims.d[2]) == mLd);
|
||||
PLUGIN_ASSERT(outputs[0].desc.dims.d[3] == 1);
|
||||
PLUGIN_ASSERT(outputs[0].desc.dims.d[4] == 1);
|
||||
|
||||
if (mUseFullMask)
|
||||
{
|
||||
// user force full_mask
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.nbDims == 2);
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.d[0] == B);
|
||||
PLUGIN_ASSERT((outputs[1].desc.dims.d[1] == -1) || (outputs[1].desc.dims.d[1] == packedMaskSize384)
|
||||
|| (outputs[1].desc.dims.d[1] == packedMaskSize128));
|
||||
}
|
||||
else
|
||||
{
|
||||
// auto detect using mhatype
|
||||
if (S != -1 && B != -1)
|
||||
{
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.nbDims == 2);
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.d[0] == B);
|
||||
int32_t packedSize = getMHAMaskPackedSize(mSM, mMhaType, S);
|
||||
TRT_UNUSED packedSize;
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.d[1] == -1 || outputs[1].desc.dims.d[1] == packedSize);
|
||||
}
|
||||
}
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].desc.type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[1].desc.type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[2].desc.type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(outputs[0].desc.type == mType);
|
||||
PLUGIN_ASSERT(outputs[1].desc.type == DataType::kINT32);
|
||||
}
|
||||
|
||||
size_t EmbLayerNormPluginDynamicLegacy::getWorkspaceSize(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamicLegacy::enqueue(PluginTensorDesc const* inputDesc,
|
||||
PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs, void* /* workspace */,
|
||||
cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
|
||||
|
||||
int32_t const batchSize = inputDesc->dims.d[BDIM];
|
||||
int32_t const S = inputDesc->dims.d[SDIM];
|
||||
int32_t status = STATUS_FAILURE;
|
||||
|
||||
// Our plugin outputs only one tensor
|
||||
auto const inputIds = static_cast<int32_t const*>(inputs[0]);
|
||||
auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
|
||||
auto const inputMask = static_cast<int32_t const*>(inputs[2]);
|
||||
|
||||
float const* beta = mBetaDev.get();
|
||||
float const* gamma = mGammaDev.get();
|
||||
if (mType == DataType::kFLOAT)
|
||||
{
|
||||
auto output = static_cast<float*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
|
||||
status = embSkipLayerNorm<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
|
||||
if (status != cudaSuccess)
|
||||
{
|
||||
return status;
|
||||
}
|
||||
}
|
||||
else if (mType == DataType::kHALF)
|
||||
{
|
||||
auto output = static_cast<half*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
|
||||
status = embSkipLayerNorm<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds, beta,
|
||||
gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
|
||||
if (status != cudaSuccess)
|
||||
{
|
||||
return status;
|
||||
}
|
||||
}
|
||||
else
|
||||
{
|
||||
gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
|
||||
<< std::endl;
|
||||
|
||||
return STATUS_NOT_SUPPORTED;
|
||||
}
|
||||
|
||||
// check mha use fused kernel
|
||||
if (mUseFullMask || unfusedMaskSize != getMHAMaskPackedSize(mSM, mMhaType, S))
|
||||
{
|
||||
size_t warps_m = 0, warps_n = 0, warps_k = 1;
|
||||
if (S == 64 || S == 96 || S == 128)
|
||||
{
|
||||
warps_m = 2;
|
||||
warps_n = 2;
|
||||
}
|
||||
else if (S == 384)
|
||||
{
|
||||
warps_m = 1;
|
||||
warps_n = 8;
|
||||
}
|
||||
uint32_t* inputMaskX = static_cast<uint32_t*>(outputs[1]);
|
||||
|
||||
status = convertMask(S, batchSize, warps_m, warps_n, warps_k, inputMask, inputMaskX, stream);
|
||||
}
|
||||
else
|
||||
{
|
||||
int32_t* maskIdx = static_cast<int32_t*>(outputs[1]);
|
||||
status = computeMaskIdx(stream, S, batchSize, inputMask, maskIdx);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
// IPluginV2Ext Methods
|
||||
DataType EmbLayerNormPluginDynamicLegacy::getOutputDataType(
|
||||
int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
|
||||
PLUGIN_ASSERT(index == 0 || index == 1);
|
||||
if (index == 0)
|
||||
{
|
||||
PLUGIN_ASSERT(mType == DataType::kHALF || mType == DataType::kFLOAT);
|
||||
return mType;
|
||||
}
|
||||
return DataType::kINT32;
|
||||
}
|
||||
|
||||
// IPluginV2 Methods
|
||||
char const* EmbLayerNormPluginDynamicLegacy::getPluginType() const noexcept
|
||||
{
|
||||
return gEmbLayerNormName;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicLegacy::getPluginVersion() const noexcept
|
||||
{
|
||||
return gEmbLayerNormVersion;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamicLegacy::getNbOutputs() const noexcept
|
||||
{
|
||||
return 2;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormPluginDynamicLegacy::initialize() noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamicLegacy::terminate() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamicLegacy terminate.");
|
||||
}
|
||||
|
||||
size_t EmbLayerNormPluginDynamicLegacy::getSerializationSize() const noexcept
|
||||
{
|
||||
size_t const wordSize = getElementSize(mType);
|
||||
return sizeof(mType) // type
|
||||
+ sizeof(mMhaType) // mha plugin datatype
|
||||
+ sizeof(mLd) * 5 // mLd, mS, m*VocabSize
|
||||
+ sizeof(mUseFullMask) // mask type
|
||||
+ sizeof(mSM) // smversion
|
||||
+ 2 * sizeof(float) * mLd // beta + gamma
|
||||
+ wordSize * mLd * mWordVocabSize // word emb
|
||||
+ wordSize * mLd * mPosVocabSize // pos emb
|
||||
+ wordSize * mLd * mTokVocabSize // tok emb
|
||||
;
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamicLegacy::serialize(void* buffer) const noexcept
|
||||
{
|
||||
serialize_value(&buffer, mType);
|
||||
serialize_value(&buffer, mMhaType);
|
||||
serialize_value(&buffer, mLd);
|
||||
serialize_value(&buffer, mS);
|
||||
serialize_value(&buffer, mWordVocabSize);
|
||||
serialize_value(&buffer, mPosVocabSize);
|
||||
serialize_value(&buffer, mTokVocabSize);
|
||||
serialize_value(&buffer, mUseFullMask);
|
||||
serialize_value(&buffer, mSM);
|
||||
|
||||
char* d = static_cast<char*>(buffer);
|
||||
serFromDev(d, mBetaDev.get(), mLd);
|
||||
serFromDev(d, mGammaDev.get(), mLd);
|
||||
size_t const wordSize = getElementSize(mType);
|
||||
serFromDev(d, static_cast<char*>(mWordEmbDev.get()), mLd * mWordVocabSize * wordSize);
|
||||
serFromDev(d, static_cast<char*>(mPosEmbDev.get()), mLd * mPosVocabSize * wordSize);
|
||||
serFromDev(d, static_cast<char*>(mTokEmbDev.get()), mLd * mTokVocabSize * wordSize);
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamicLegacy::destroy() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamicLegacy destroy.");
|
||||
// This gets called when the network containing plugin is destroyed
|
||||
mGammaDev.reset(nullptr);
|
||||
mBetaDev.reset(nullptr);
|
||||
mWordEmbDev.reset(nullptr);
|
||||
mPosEmbDev.reset(nullptr);
|
||||
mTokEmbDev.reset(nullptr);
|
||||
delete this;
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamicLegacy::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicLegacy::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
///////////////////////
|
||||
|
||||
EmbLayerNormPluginDynamicLegacyCreator::EmbLayerNormPluginDynamicLegacyCreator()
|
||||
{
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_beta"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_gamma"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_word_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_token_type_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_position_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("output_fp16"));
|
||||
mPluginAttributes.emplace_back(PluginField("full_mask"));
|
||||
mPluginAttributes.emplace_back(PluginField("mha_type_id"));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicLegacyCreator::getPluginName() const noexcept
|
||||
{
|
||||
return gEmbLayerNormName;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicLegacyCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return gEmbLayerNormVersion;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* EmbLayerNormPluginDynamicLegacyCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV2* EmbLayerNormPluginDynamicLegacyCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormPluginDynamicLegacy createPlugin.");
|
||||
|
||||
bool output_fp16 = false;
|
||||
bool useFullMask = false;
|
||||
Weights beta{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights gamma{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights word_emb{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights pos_emb{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
Weights tok_emb{}; // required attribute - validateRequiredAttributesExist() will verify existence
|
||||
int32_t mhaTypeId = 0;
|
||||
std::set<std::string> const requiredAttributes{
|
||||
"bert_embeddings_layernorm_beta",
|
||||
"bert_embeddings_layernorm_gamma",
|
||||
"bert_embeddings_word_embeddings",
|
||||
"bert_embeddings_token_type_embeddings",
|
||||
"bert_embeddings_position_embeddings",
|
||||
};
|
||||
plugin::validateRequiredAttributesExist(requiredAttributes, fc);
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; i++)
|
||||
{
|
||||
std::string_view const field_name = fc->fields[i].name;
|
||||
if (field_name == "bert_embeddings_layernorm_beta"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
|
||||
beta.values = fc->fields[i].data;
|
||||
beta.count = fc->fields[i].length;
|
||||
beta.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_layernorm_gamma"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
|
||||
gamma.values = fc->fields[i].data;
|
||||
gamma.count = fc->fields[i].length;
|
||||
gamma.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_word_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_word_embeddings...");
|
||||
word_emb.values = fc->fields[i].data;
|
||||
word_emb.count = fc->fields[i].length;
|
||||
word_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_token_type_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_token_type_embeddings...");
|
||||
tok_emb.values = fc->fields[i].data;
|
||||
tok_emb.count = fc->fields[i].length;
|
||||
tok_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_position_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_position_embeddings...");
|
||||
pos_emb.values = fc->fields[i].data;
|
||||
pos_emb.count = fc->fields[i].length;
|
||||
pos_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
if (field_name == "output_fp16"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building output_fp16...");
|
||||
PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kINT32);
|
||||
output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
|
||||
}
|
||||
if (field_name == "full_mask"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building full_mask...");
|
||||
PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kINT32);
|
||||
useFullMask = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
|
||||
}
|
||||
if (field_name == "mha_type_id"sv)
|
||||
{
|
||||
mhaTypeId = *static_cast<int32_t const*>(fc->fields[i].data);
|
||||
PLUGIN_VALIDATE(mhaTypeId >= 0 && mhaTypeId <= 3);
|
||||
BERT_DEBUG_VALUE("Building mha typeId: ", mhaTypeId);
|
||||
}
|
||||
}
|
||||
|
||||
BERT_DEBUG_MSG("Building the Plugin...");
|
||||
DataType mhaType = static_cast<DataType>(mhaTypeId);
|
||||
auto p
|
||||
= std::make_unique<EmbLayerNormPluginDynamicLegacy>(name, output_fp16 ? DataType::kHALF : DataType::kFLOAT,
|
||||
mhaType, beta, gamma, word_emb, pos_emb, tok_emb, useFullMask);
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* EmbLayerNormPluginDynamicLegacyCreator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// This object will be deleted when the network is destroyed, which will
|
||||
// call EmbLayerNormPluginDynamicLegacy::destroy()
|
||||
return new EmbLayerNormPluginDynamicLegacy(name, serialData, serialLength);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void EmbLayerNormPluginDynamicLegacyCreator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* EmbLayerNormPluginDynamicLegacyCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
@@ -0,0 +1,151 @@
|
||||
/*
|
||||
* 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_EMB_LAYER_NORM_PLUGIN_LEGACY_H
|
||||
#define TRT_EMB_LAYER_NORM_PLUGIN_LEGACY_H
|
||||
|
||||
#include "NvInferPlugin.h"
|
||||
#include "NvInferRuntime.h"
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
int32_t computeMaskIdx(cudaStream_t stream, int32_t const S, int32_t const B, int32_t const* mask, int32_t* maskIdx);
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNorm(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* token_ids, float const* beta, float const* gamma, T const* wordEmb, T const* posEmb, T const* tokEmb,
|
||||
int32_t const wordSize, int32_t const tokSize, T* output);
|
||||
|
||||
cudaError_t convertMask(uint32_t const S, uint32_t const B, uint32_t const warps_m, uint32_t const warps_n,
|
||||
uint32_t const warps_k, int32_t const* inputMaskSB, uint32_t* inputMaskX, cudaStream_t stream);
|
||||
|
||||
class EmbLayerNormPluginDynamicLegacy : public nvinfer1::IPluginV2DynamicExt
|
||||
{
|
||||
public:
|
||||
EmbLayerNormPluginDynamicLegacy(std::string const& name, nvinfer1::DataType const type,
|
||||
nvinfer1::DataType const mhaType, nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma,
|
||||
nvinfer1::Weights const& word_emb, nvinfer1::Weights const& pos_emb, nvinfer1::Weights const& tok_emb,
|
||||
bool const useFullMask);
|
||||
|
||||
EmbLayerNormPluginDynamicLegacy(std::string const& name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormPluginDynamicLegacy without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormPluginDynamicLegacy() = 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:
|
||||
std::string const mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
bert::cuda_unique_ptr<float> mGammaDev;
|
||||
bert::cuda_unique_ptr<float> mBetaDev;
|
||||
bert::cuda_unique_ptr<void> mWordEmbDev;
|
||||
bert::cuda_unique_ptr<void> mTokEmbDev;
|
||||
bert::cuda_unique_ptr<void> mPosEmbDev;
|
||||
size_t mLd; // leading dim = hidden size
|
||||
size_t mS; // sequence length
|
||||
size_t mWordVocabSize;
|
||||
size_t mPosVocabSize;
|
||||
size_t mTokVocabSize;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mWordEmb;
|
||||
bert::WeightsWithOwnership mTokEmb;
|
||||
bert::WeightsWithOwnership mPosEmb;
|
||||
nvinfer1::DataType mType;
|
||||
bool mUseFullMask;
|
||||
nvinfer1::DataType mMhaType;
|
||||
int32_t mSM;
|
||||
|
||||
using IPluginV2::getOutputDimensions;
|
||||
using IPluginV2::getWorkspaceSize;
|
||||
using IPluginV2::enqueue;
|
||||
using IPluginV2Ext::configurePlugin;
|
||||
};
|
||||
|
||||
class EmbLayerNormPluginDynamicLegacyCreator : public nvinfer1::IPluginCreator
|
||||
{
|
||||
public:
|
||||
EmbLayerNormPluginDynamicLegacyCreator();
|
||||
|
||||
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_EMB_LAYER_NORM_PLUGIN_LEGACY_H
|
||||
|
||||
#endif // CUDA_VERSION >= 10010
|
||||
@@ -0,0 +1,130 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "common/bertCommon.h"
|
||||
#include "common/common.cuh"
|
||||
#include "common/plugin.h"
|
||||
#include "common/serialize.hpp"
|
||||
|
||||
#include <cassert>
|
||||
#include <cstring>
|
||||
#include <cuda.h>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
template <typename T, unsigned TPB>
|
||||
__global__ void embLayerNormKernelHFace(int32_t ld, int32_t const* inputIds, int32_t const* tokenIds,
|
||||
int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb, T const* posEmb, T const* tokEmb,
|
||||
int32_t const wordSize, int32_t const tokSize, T* output)
|
||||
{
|
||||
// this code currently assumes the input shape is SxB, row-major => seqPos = s * B + b
|
||||
// instead we want BxS, row-major => seqPos = b * S + s
|
||||
|
||||
// 1. lookup word and token of the block
|
||||
// blockIdx.x = position in the sequence
|
||||
// blockIdx.y = batch
|
||||
// gridDim.x = S
|
||||
// gridDim.y = B
|
||||
int32_t const s = blockIdx.x;
|
||||
int32_t const b = blockIdx.y;
|
||||
|
||||
int32_t const sumS = cuSeqlens[b];
|
||||
int32_t const s_b = cuSeqlens[b + 1] - sumS;
|
||||
if (s >= s_b)
|
||||
{
|
||||
return; // This CTA has nothing to do
|
||||
}
|
||||
__shared__ int32_t wordId;
|
||||
__shared__ int32_t tokenId;
|
||||
|
||||
T const rld = T(1.f) / T(ld);
|
||||
// seqPos = b + s * B
|
||||
// int32_t const seqPos = blockIdx.y + blockIdx.x * gridDim.y;
|
||||
|
||||
// int32_t const seqPos = s * B + s;
|
||||
int32_t const seqPos = sumS + s;
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
wordId = inputIds[seqPos];
|
||||
tokenId = tokenIds[seqPos];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 2. load pos/tok/word embeddings and add them toghether
|
||||
// offset into embeddings is given by wordId * hidden_size
|
||||
int32_t const poffset = blockIdx.x * ld;
|
||||
int32_t const woffset = wordId * ld;
|
||||
int32_t const toffset = tokenId * ld;
|
||||
// the output offset is given by b * (S*hidden_size) + s * hidden_size
|
||||
int32_t const outOffset = seqPos * ld;
|
||||
|
||||
kvp<T> threadData(0, 0);
|
||||
|
||||
if (wordId >= 0 && wordId < wordSize && tokenId >= 0 && tokenId < tokSize)
|
||||
{
|
||||
for (int32_t it = threadIdx.x; it < ld; it += TPB)
|
||||
{
|
||||
T const w(wordEmb[woffset + it]);
|
||||
T const t(tokEmb[toffset + it]);
|
||||
T const p(posEmb[poffset + it]);
|
||||
T const val = w + t + p;
|
||||
|
||||
output[outOffset + it] = val;
|
||||
T const rldval = rld * val;
|
||||
threadData = threadData + kvp<T>(rldval, rldval * val);
|
||||
}
|
||||
}
|
||||
|
||||
// 3. layer norm on the sum
|
||||
layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNormHFace(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* tokenIds, int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb,
|
||||
T const* posEmb, T const* tokEmb, int32_t const wordSize, int32_t const tokSize, T* output)
|
||||
{
|
||||
|
||||
constexpr int32_t tpb = 256;
|
||||
dim3 const grid(S, B, 1);
|
||||
dim3 const block(tpb, 1, 1);
|
||||
|
||||
embLayerNormKernelHFace<T, tpb><<<grid, block, 0, stream>>>(
|
||||
ld, inputIds, tokenIds, cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, wordSize, tokSize, output);
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
template int32_t embSkipLayerNormHFace<float>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
|
||||
int32_t const*, float const*, float const*, float const*, float const*, float const*, int32_t const, int32_t const,
|
||||
float*);
|
||||
|
||||
template int32_t embSkipLayerNormHFace<half>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
|
||||
int32_t const*, float const*, float const*, half const*, half const*, half const*, int32_t const, int32_t const,
|
||||
half*);
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
@@ -0,0 +1,131 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "common/bertCommon.h"
|
||||
#include "common/common.cuh"
|
||||
#include "common/plugin.h"
|
||||
#include "common/serialize.hpp"
|
||||
|
||||
#include <cassert>
|
||||
#include <cstring>
|
||||
#include <cuda.h>
|
||||
#include <vector>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
template <typename T, unsigned TPB>
|
||||
__global__ void embLayerNormKernelMTron(int32_t ld, int32_t const* inputIds, int32_t const* tokenIds,
|
||||
int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb, T const* posEmb, T const* tokEmb,
|
||||
int32_t const wordSize, int32_t const tokSize, T* output, T* skip)
|
||||
{
|
||||
// this code currently assumes the input shape is SxB, row-major => seqPos = s * B + b
|
||||
// instead we want BxS, row-major => seqPos = b * S + s
|
||||
|
||||
// 1. lookup word and token of the block
|
||||
// blockIdx.x = position in the sequence
|
||||
// blockIdx.y = batch
|
||||
// gridDim.x = S
|
||||
// gridDim.y = B
|
||||
int32_t const s = blockIdx.x;
|
||||
int32_t const b = blockIdx.y;
|
||||
|
||||
int32_t const sumS = cuSeqlens[b];
|
||||
int32_t const s_b = cuSeqlens[b + 1] - sumS;
|
||||
if (s >= s_b)
|
||||
{
|
||||
return; // This CTA has nothing to do
|
||||
}
|
||||
__shared__ int32_t wordId;
|
||||
__shared__ int32_t tokenId;
|
||||
|
||||
T const rld = T(1.f) / T(ld);
|
||||
// seqPos = b + s * B
|
||||
// int32_t const seqPos = blockIdx.y + blockIdx.x * gridDim.y;
|
||||
|
||||
// int32_t const seqPos = s * B + s;
|
||||
int32_t const seqPos = sumS + s;
|
||||
if (threadIdx.x == 0)
|
||||
{
|
||||
wordId = inputIds[seqPos];
|
||||
tokenId = tokenIds[seqPos];
|
||||
}
|
||||
__syncthreads();
|
||||
|
||||
// 2. load pos/tok/word embeddings and add them toghether
|
||||
// offset into embeddings is given by wordId * hidden_size
|
||||
int32_t const poffset = blockIdx.x * ld;
|
||||
int32_t const woffset = wordId * ld;
|
||||
int32_t const toffset = tokenId * ld;
|
||||
// the output offset is given by b * (S*hidden_size) + s * hidden_size
|
||||
int32_t const outOffset = seqPos * ld;
|
||||
|
||||
kvp<T> threadData(0, 0);
|
||||
|
||||
if (wordId >= 0 && wordId < wordSize && tokenId >= 0 && tokenId < tokSize)
|
||||
{
|
||||
for (int32_t it = threadIdx.x; it < ld; it += TPB)
|
||||
{
|
||||
T const w(wordEmb[woffset + it]);
|
||||
T const t(tokEmb[toffset + it]);
|
||||
T const p(posEmb[poffset + it]);
|
||||
T const val = w + t + p;
|
||||
|
||||
output[outOffset + it] = val;
|
||||
skip[outOffset + it] = val;
|
||||
T const rldval = rld * val;
|
||||
threadData = threadData + kvp<T>(rldval, rldval * val);
|
||||
}
|
||||
}
|
||||
|
||||
// 3. layer norm on the sum
|
||||
layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNormMTron(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* tokenIds, int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb,
|
||||
T const* posEmb, T const* tokEmb, int32_t const wordSize, int32_t const tokSize, T* output, T* skip)
|
||||
{
|
||||
|
||||
constexpr int32_t tpb = 256;
|
||||
dim3 const grid(S, B, 1);
|
||||
dim3 const block(tpb, 1, 1);
|
||||
|
||||
embLayerNormKernelMTron<T, tpb><<<grid, block, 0, stream>>>(
|
||||
ld, inputIds, tokenIds, cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, wordSize, tokSize, output, skip);
|
||||
return cudaPeekAtLastError();
|
||||
}
|
||||
|
||||
template int32_t embSkipLayerNormMTron<float>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
|
||||
int32_t const*, float const*, float const*, float const*, float const*, float const*, int32_t const, int32_t const,
|
||||
float*, float*);
|
||||
|
||||
template int32_t embSkipLayerNormMTron<half>(cudaStream_t, int32_t, int32_t, int32_t, int32_t const*, int32_t const*,
|
||||
int32_t const*, float const*, float const*, half const*, half const*, half const*, int32_t const, int32_t const,
|
||||
half*, half*);
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
@@ -0,0 +1,832 @@
|
||||
/*
|
||||
* 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 <cstring>
|
||||
#include <cuda.h>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "common/serialize.hpp"
|
||||
#include "embLayerNormVarSeqlenPlugin.h"
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::plugin::bert;
|
||||
|
||||
namespace
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE{"4"};
|
||||
constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON{"5"};
|
||||
constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_NAME{"CustomEmbLayerNormPluginDynamic"};
|
||||
|
||||
void checkConfigurationInputs(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// Validate input arguments
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].dims.nbDims == 1);
|
||||
PLUGIN_ASSERT(inputs[1].dims.nbDims == 1);
|
||||
|
||||
PLUGIN_ASSERT(inputs[1].dims.d[0] == inputs[0].dims.d[0]);
|
||||
|
||||
PLUGIN_ASSERT(inputs[2].dims.nbDims == 1);
|
||||
|
||||
PLUGIN_ASSERT(outputs[0].dims.nbDims == 4);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[0]) == static_cast<size_t>(inputs[0].dims.d[0]));
|
||||
PLUGIN_ASSERT(outputs[0].dims.d[2] == 1);
|
||||
PLUGIN_ASSERT(outputs[0].dims.d[3] == 1);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[1].type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[2].type == DataType::kINT32);
|
||||
}
|
||||
|
||||
bool initializeFields(char const* name, PluginFieldCollection const* fc, Weights& beta, Weights& gamma,
|
||||
Weights& word_emb, Weights& pos_emb, Weights& tok_emb)
|
||||
{
|
||||
bool output_fp16 = false;
|
||||
std::set<std::string> const requiredAttributes{
|
||||
"bert_embeddings_layernorm_beta",
|
||||
"bert_embeddings_layernorm_gamma",
|
||||
"bert_embeddings_word_embeddings",
|
||||
"bert_embeddings_token_type_embeddings",
|
||||
"bert_embeddings_position_embeddings",
|
||||
};
|
||||
plugin::validateRequiredAttributesExist(requiredAttributes, fc);
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; i++)
|
||||
{
|
||||
std::string_view const field_name = fc->fields[i].name;
|
||||
if (field_name == "bert_embeddings_layernorm_beta"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
|
||||
beta.values = fc->fields[i].data;
|
||||
beta.count = fc->fields[i].length;
|
||||
beta.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
else if (field_name == "bert_embeddings_layernorm_gamma"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
|
||||
gamma.values = fc->fields[i].data;
|
||||
gamma.count = fc->fields[i].length;
|
||||
gamma.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
else if (field_name == "bert_embeddings_word_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_word_embeddings...");
|
||||
word_emb.values = fc->fields[i].data;
|
||||
word_emb.count = fc->fields[i].length;
|
||||
word_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
else if (field_name == "bert_embeddings_token_type_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_token_type_embeddings...");
|
||||
tok_emb.values = fc->fields[i].data;
|
||||
tok_emb.count = fc->fields[i].length;
|
||||
tok_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
else if (field_name == "bert_embeddings_position_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_position_embeddings...");
|
||||
pos_emb.values = fc->fields[i].data;
|
||||
pos_emb.count = fc->fields[i].length;
|
||||
pos_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
else if (field_name == "output_fp16"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building output_fp16...");
|
||||
PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kINT32);
|
||||
output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
|
||||
}
|
||||
}
|
||||
return output_fp16;
|
||||
}
|
||||
|
||||
} // namespace
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(EmbLayerNormVarSeqlenPluginHFaceCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(EmbLayerNormVarSeqlenPluginMTronCreator);
|
||||
|
||||
EmbLayerNormVarSeqlenPluginBase::EmbLayerNormVarSeqlenPluginBase(std::string const& name, DataType type,
|
||||
Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb, Weights const& tokEmb,
|
||||
DataType maskType)
|
||||
: mLayerName(name)
|
||||
, mLd(beta.count)
|
||||
, mType(type)
|
||||
, mMaskType(maskType)
|
||||
{
|
||||
// Assuming Weights.count is the number of elements and not bytes
|
||||
PLUGIN_VALIDATE(beta.count == gamma.count);
|
||||
PLUGIN_VALIDATE(mLd > 0U);
|
||||
PLUGIN_VALIDATE(wordEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(posEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(tokEmb.count % mLd == 0);
|
||||
mWordVocabSize = wordEmb.count / mLd;
|
||||
mPosVocabSize = posEmb.count / mLd;
|
||||
mTokVocabSize = tokEmb.count / mLd;
|
||||
|
||||
mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
|
||||
mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
|
||||
|
||||
mWordEmb.convertAndCopy(wordEmb, mType);
|
||||
mTokEmb.convertAndCopy(tokEmb, mType);
|
||||
mPosEmb.convertAndCopy(posEmb, mType);
|
||||
|
||||
copyToDevice(mGamma, sizeof(float) * mGamma.count, mGammaDev);
|
||||
copyToDevice(mBeta, sizeof(float) * mBeta.count, mBetaDev);
|
||||
|
||||
copyToDevice(mWordEmb, getWeightsSize(mWordEmb, mType), mWordEmbDev);
|
||||
copyToDevice(mPosEmb, getWeightsSize(mPosEmb, mType), mPosEmbDev);
|
||||
copyToDevice(mTokEmb, getWeightsSize(mTokEmb, mType), mTokEmbDev);
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginHFace::EmbLayerNormVarSeqlenPluginHFace(std::string const& name, DataType const type,
|
||||
Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb, Weights const& tokEmb)
|
||||
: EmbLayerNormVarSeqlenPluginBase(name, type, beta, gamma, wordEmb, posEmb, tokEmb, DataType::kINT32)
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace creation");
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginMTron::EmbLayerNormVarSeqlenPluginMTron(std::string const& name, DataType const type,
|
||||
Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb, Weights const& tokEmb)
|
||||
: EmbLayerNormVarSeqlenPluginBase(name, type, beta, gamma, wordEmb, posEmb, tokEmb, type)
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron creation");
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginBase::~EmbLayerNormVarSeqlenPluginBase()
|
||||
{
|
||||
try
|
||||
{
|
||||
// This gets called when the network containing plugin is destroyed
|
||||
mGammaDev.reset(nullptr);
|
||||
mBetaDev.reset(nullptr);
|
||||
mWordEmbDev.reset(nullptr);
|
||||
mPosEmbDev.reset(nullptr);
|
||||
mTokEmbDev.reset(nullptr);
|
||||
// delete this; (TRT will delete this plugin object)
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginHFace::~EmbLayerNormVarSeqlenPluginHFace()
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace destruction");
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginMTron::~EmbLayerNormVarSeqlenPluginMTron()
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron destruction");
|
||||
}
|
||||
|
||||
//////
|
||||
// IPluginV3 method definitions:
|
||||
// - getCapabilityInterface() (Base)
|
||||
// - clone() (HFace, MTron)
|
||||
//////
|
||||
IPluginCapability* EmbLayerNormVarSeqlenPluginBase::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* EmbLayerNormVarSeqlenPluginHFace::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace clone");
|
||||
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginHFace>(
|
||||
mLayerName, mType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb);
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV3* EmbLayerNormVarSeqlenPluginMTron::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron clone");
|
||||
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginMTron>(
|
||||
mLayerName, mType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb);
|
||||
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() (Base)
|
||||
// - enqueue() (HFace, MTron)
|
||||
/////
|
||||
|
||||
PluginFieldCollection const* EmbLayerNormVarSeqlenPluginBase::getFieldsToSerialize() noexcept
|
||||
{
|
||||
mDataToSerialize.clear();
|
||||
bool output_fp16 = mType == DataType::kHALF;
|
||||
mDataToSerialize.emplace_back("output_fp16", &output_fp16, PluginFieldType::kINT32, 1);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_layernorm_beta", static_cast<float const*>(mBeta.values),
|
||||
PluginFieldType::kFLOAT32, mBeta.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_layernorm_gamma", static_cast<float const*>(mGamma.values),
|
||||
PluginFieldType::kFLOAT32, mGamma.count);
|
||||
if (output_fp16)
|
||||
{
|
||||
mDataToSerialize.emplace_back("bert_embeddings_word_embeddings", static_cast<half const*>(mWordEmb.values),
|
||||
PluginFieldType::kFLOAT16, mWordEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_token_type_embeddings", static_cast<half const*>(mTokEmb.values),
|
||||
PluginFieldType::kFLOAT16, mTokEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_position_embeddings", static_cast<half const*>(mPosEmb.values),
|
||||
PluginFieldType::kFLOAT16, mPosEmb.count);
|
||||
}
|
||||
else
|
||||
{
|
||||
mDataToSerialize.emplace_back("bert_embeddings_word_embeddings", static_cast<float const*>(mWordEmb.values),
|
||||
PluginFieldType::kFLOAT32, mWordEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_token_type_embeddings",
|
||||
static_cast<float const*>(mTokEmb.values), PluginFieldType::kFLOAT32, mTokEmb.count);
|
||||
mDataToSerialize.emplace_back("bert_embeddings_position_embeddings", static_cast<float const*>(mPosEmb.values),
|
||||
PluginFieldType::kFLOAT32, mPosEmb.count);
|
||||
}
|
||||
mFCToSerialize.nbFields = mDataToSerialize.size();
|
||||
mFCToSerialize.fields = mDataToSerialize.data();
|
||||
return &mFCToSerialize;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginHFace::onShapeChange(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace onShapeChange");
|
||||
checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
|
||||
|
||||
// output 0 is the embedding
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[1]) == static_cast<size_t>(mLd));
|
||||
PLUGIN_ASSERT(outputs[0].type == mType);
|
||||
// output 1 is the mask indices (empty for HFace variant)
|
||||
PLUGIN_ASSERT(outputs[1].dims.nbDims == 0);
|
||||
PLUGIN_ASSERT(outputs[1].type == mMaskType);
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginMTron::onShapeChange(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// Validate input arguments
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron onShapeChange");
|
||||
checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[1]) == static_cast<size_t>(mLd));
|
||||
|
||||
PLUGIN_ASSERT(outputs[1].dims.nbDims == 4);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[1].dims.d[0]) == static_cast<size_t>(inputs[0].dims.d[0]));
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[1].dims.d[1]) == static_cast<size_t>(mLd));
|
||||
PLUGIN_ASSERT(outputs[1].dims.d[2] == 1);
|
||||
PLUGIN_ASSERT(outputs[1].dims.d[3] == 1);
|
||||
|
||||
PLUGIN_ASSERT(outputs[0].type == mType);
|
||||
PLUGIN_ASSERT(outputs[1].type == mMaskType);
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
IPluginV3* EmbLayerNormVarSeqlenPluginBase::attachToContext(IPluginResourceContext* context) noexcept
|
||||
{
|
||||
return clone();
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginHFace::enqueue(PluginTensorDesc const* inputDesc,
|
||||
PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs, void* /* workspace */,
|
||||
cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
|
||||
|
||||
int32_t const batchSize = inputDesc[2].dims.d[0] - 1;
|
||||
// read out the maximum sequence length from the dummy input
|
||||
int32_t const maxSeqlen = inputDesc[3].dims.d[0];
|
||||
|
||||
// There are four versions of the kernel which are optimized for sequence lengths 384, 256, 192 and 128.
|
||||
// Find the closest sequence length bigger than the max seq length in this batch.
|
||||
int32_t S = 384;
|
||||
if (maxSeqlen <= 128)
|
||||
{
|
||||
S = 128;
|
||||
}
|
||||
else if (maxSeqlen <= 192)
|
||||
{
|
||||
S = 192;
|
||||
}
|
||||
else if (maxSeqlen <= 256)
|
||||
{
|
||||
S = 256;
|
||||
}
|
||||
|
||||
// Our plugin outputs only one tensor
|
||||
auto const inputIds = static_cast<int32_t const*>(inputs[0]);
|
||||
auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
|
||||
int32_t const* cuSeqlens = static_cast<int32_t const*>(inputs[2]);
|
||||
|
||||
float const* beta = mBetaDev.get();
|
||||
float const* gamma = mGammaDev.get();
|
||||
if (mType == DataType::kFLOAT)
|
||||
{
|
||||
auto output = static_cast<float*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormHFace<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
}
|
||||
if (mType == DataType::kHALF)
|
||||
{
|
||||
auto output = static_cast<half*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormHFace<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
}
|
||||
else
|
||||
{
|
||||
gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
|
||||
<< std::endl;
|
||||
|
||||
return STATUS_NOT_SUPPORTED;
|
||||
}
|
||||
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginMTron::enqueue(PluginTensorDesc const* inputDesc,
|
||||
PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs, void* /* workspace */,
|
||||
cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
|
||||
|
||||
int32_t const batchSize = inputDesc[2].dims.d[0] - 1;
|
||||
// read out the maximum sequence length from the dummy input
|
||||
int32_t const maxSeqlen = inputDesc[3].dims.d[0];
|
||||
|
||||
// There are four versions of the kernel which are optimized for sequence lengths 384, 256, 192 and 128.
|
||||
// Find the closest sequence length bigger than the max seq length in this batch.
|
||||
int32_t S = 384;
|
||||
if (maxSeqlen <= 128)
|
||||
{
|
||||
S = 128;
|
||||
}
|
||||
else if (maxSeqlen <= 192)
|
||||
{
|
||||
S = 192;
|
||||
}
|
||||
else if (maxSeqlen <= 256)
|
||||
{
|
||||
S = 256;
|
||||
}
|
||||
|
||||
// Our plugin outputs only one tensor
|
||||
auto const inputIds = static_cast<int32_t const*>(inputs[0]);
|
||||
auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
|
||||
int32_t const* cuSeqlens = static_cast<int32_t const*>(inputs[2]);
|
||||
|
||||
float const* beta = mBetaDev.get();
|
||||
float const* gamma = mGammaDev.get();
|
||||
if (mType == DataType::kFLOAT)
|
||||
{
|
||||
auto output = static_cast<float*>(outputs[0]);
|
||||
auto skip = static_cast<float*>(outputs[1]);
|
||||
auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormMTron<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output, skip);
|
||||
}
|
||||
if (mType == DataType::kHALF)
|
||||
{
|
||||
auto output = static_cast<half*>(outputs[0]);
|
||||
auto skip = static_cast<half*>(outputs[1]);
|
||||
auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormMTron<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output, skip);
|
||||
}
|
||||
else
|
||||
{
|
||||
gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
|
||||
<< std::endl;
|
||||
|
||||
return STATUS_NOT_SUPPORTED;
|
||||
}
|
||||
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
// end IPluginV3OneRuntime method definitions
|
||||
|
||||
///////
|
||||
// IPluginV3OneBuild method definitions
|
||||
// - getNbOutputs() (Base)
|
||||
// - supportsFormatCombination() (Base)
|
||||
// - getOutputShapes (HFace, MTron)
|
||||
// - getOutputDataTypes() (Base)
|
||||
// - configurePlugin() (Base)
|
||||
// - getWorkSpaceSize() (Base)
|
||||
//////
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginBase::getNbOutputs() const noexcept
|
||||
{
|
||||
return 2;
|
||||
}
|
||||
|
||||
bool EmbLayerNormVarSeqlenPluginBase::supportsFormatCombination(
|
||||
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// The four inputs to this plugin input_ids, segment_ids, cu_seqlens and a dummy input with the
|
||||
// size of the max seq length in that order
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
// The two outputs of the plugin are embedding and the mask
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PluginTensorDesc const& desc = inOut[pos].desc;
|
||||
if (desc.format != TensorFormat::kLINEAR)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if (pos == 0 || pos == 2) // input_ids and cu_seqlens
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 1;
|
||||
}
|
||||
|
||||
PluginTensorDesc const& prev = inOut[pos - 1].desc;
|
||||
if (pos == 1) // segment ids: check it's the same as input_ids
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 1 && desc.dims.d[0] == prev.dims.d[0];
|
||||
}
|
||||
|
||||
if (pos == 3)
|
||||
{
|
||||
return desc.dims.nbDims == 1;
|
||||
}
|
||||
|
||||
// embedded sequence
|
||||
if (pos == nbInputs)
|
||||
{
|
||||
return desc.type == mType && desc.dims.nbDims == 4 && desc.dims.d[0] == inOut[0].desc.dims.d[0]
|
||||
&& desc.dims.d[2] == 1 && desc.dims.d[3] == 1;
|
||||
}
|
||||
// mask
|
||||
return desc.type == mMaskType;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginHFace::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(outputs != nullptr);
|
||||
|
||||
// Input should be input ids and token ids and cumulative seqlens
|
||||
// Output should be the embeddings tensor and mask indices
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == 1); // sum of all s
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
|
||||
|
||||
PLUGIN_ASSERT(inputs[2].nbDims == 1); // B+1
|
||||
|
||||
// output 0 : embedded input
|
||||
outputs[0].nbDims = 4;
|
||||
outputs[0].d[0] = inputs[0].d[0];
|
||||
outputs[0].d[1] = exprBuilder.constant(mLd);
|
||||
outputs[0].d[2] = exprBuilder.constant(1);
|
||||
outputs[0].d[3] = exprBuilder.constant(1);
|
||||
|
||||
// Output 1 : maskIdx
|
||||
// Return empty tensor since this is dummy output, we do not delete it for backward compatibility.
|
||||
outputs[1].nbDims = 0;
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (const std::exception& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginMTron::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(outputs != nullptr);
|
||||
// Input should be input ids and token ids and cumulative seqlens
|
||||
// Output should be the embeddings tensor and mask indices
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == 1); // sum of all s
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
|
||||
PLUGIN_ASSERT(inputs[2].nbDims == 1); // B+1
|
||||
|
||||
// Output 0 : embedded input
|
||||
outputs[0].nbDims = 4;
|
||||
outputs[0].d[0] = inputs[0].d[0];
|
||||
outputs[0].d[1] = exprBuilder.constant(mLd);
|
||||
outputs[0].d[2] = exprBuilder.constant(1);
|
||||
outputs[0].d[3] = exprBuilder.constant(1);
|
||||
|
||||
// Output 1 : maskIdx
|
||||
outputs[1].nbDims = 4;
|
||||
outputs[1].d[0] = inputs[0].d[0];
|
||||
outputs[1].d[1] = exprBuilder.constant(mLd);
|
||||
outputs[1].d[2] = exprBuilder.constant(1);
|
||||
outputs[1].d[3] = exprBuilder.constant(1);
|
||||
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (const std::exception& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginBase::getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_ASSERT(mType == DataType::kHALF || mType == DataType::kFLOAT);
|
||||
outputTypes[0] = mType;
|
||||
outputTypes[1] = mMaskType;
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return pluginStatus_t::STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginBase::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
return pluginStatus_t::STATUS_SUCCESS;
|
||||
}
|
||||
|
||||
size_t EmbLayerNormVarSeqlenPluginBase::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* EmbLayerNormVarSeqlenPluginHFace::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginMTron::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginBase::getPluginName() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_NAME;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginBase::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginBase::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
// End IPluginV3OneCore method definitions
|
||||
|
||||
//////////////////////////// Plugin Creator member definitions /////////////////////////////
|
||||
|
||||
EmbLayerNormVarSeqlenPluginBaseCreator::EmbLayerNormVarSeqlenPluginBaseCreator()
|
||||
{
|
||||
static std::mutex sMutex;
|
||||
std::lock_guard<std::mutex> lock(sMutex);
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("output_fp16", nullptr, PluginFieldType::kINT32, 1));
|
||||
// the length of beta, gamma, word_emb, pos_emb, and tok_emb will only be known at the time of plugin creation
|
||||
// so we set it to 0 here
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_beta", nullptr, PluginFieldType::kFLOAT32, 0));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_gamma", nullptr, PluginFieldType::kFLOAT32, 0));
|
||||
// the embeddings datatype is determined by the output_fp16 attribute known at runtime
|
||||
// so we set it to kUNKNOWN here
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_word_embeddings", nullptr, PluginFieldType::kUNKNOWN, 0));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_token_type_embeddings", nullptr, PluginFieldType::kUNKNOWN, 0));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_position_embeddings", nullptr, PluginFieldType::kUNKNOWN, 0));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginBaseCreator::getPluginName() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_NAME;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginHFaceCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginMTronCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* EmbLayerNormVarSeqlenPluginBaseCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV3* EmbLayerNormVarSeqlenPluginHFaceCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenHFace createPlugin");
|
||||
|
||||
Weights beta{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights gamma{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights word_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights pos_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights tok_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
bool output_fp16 = initializeFields(name, fc, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
|
||||
BERT_DEBUG_MSG("Building the Plugin...");
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginHFace>(
|
||||
name, output_fp16 ? DataType::kHALF : DataType::kFLOAT, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV3* EmbLayerNormVarSeqlenPluginMTronCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenMTron createPlugin");
|
||||
|
||||
Weights beta{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights gamma{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights word_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights pos_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights tok_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
bool output_fp16 = initializeFields(name, fc, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
|
||||
BERT_DEBUG_MSG("Building the Plugin...");
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginMTron>(
|
||||
name, output_fp16 ? DataType::kHALF : DataType::kFLOAT, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginBaseCreator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginBaseCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
@@ -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_EMB_LAYER_NORM_VARSEQ_PLUGIN_H
|
||||
#define TRT_EMB_LAYER_NORM_VARSEQ_PLUGIN_H
|
||||
|
||||
#include <cuda.h>
|
||||
|
||||
#include "NvInferPlugin.h"
|
||||
#include "NvInferRuntime.h"
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNormHFace(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* tokenIds, int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb,
|
||||
T const* posEmb, T const* tokEmb, int32_t const wordSize, int32_t const tokSize, T* output);
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNormMTron(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* tokenIds, int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb,
|
||||
T const* posEmb, T const* tokEmb, int32_t const wordSize, int32_t const tokSize, T* output, T* skip);
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginBase : public IPluginV3,
|
||||
public IPluginV3OneCore,
|
||||
public IPluginV3OneBuild,
|
||||
public IPluginV3OneRuntime
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginBase(std::string const& name, DataType type, Weights const& beta, Weights const& gamma,
|
||||
Weights const& word_emb, Weights const& pos_emb, Weights const& tok_emb, DataType maskType);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormVarSeqlenPlugin without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormVarSeqlenPluginBase() = delete;
|
||||
|
||||
~EmbLayerNormVarSeqlenPluginBase() 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 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;
|
||||
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
// end IPluginV3Build Methods
|
||||
|
||||
// IPluginV3Runtime Methods
|
||||
|
||||
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
|
||||
// end IPluginV3Runtime Methods
|
||||
|
||||
protected:
|
||||
// metadata fields
|
||||
std::string const mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
// device-side
|
||||
bert::cuda_unique_ptr<float> mGammaDev;
|
||||
bert::cuda_unique_ptr<float> mBetaDev;
|
||||
bert::cuda_unique_ptr<void> mWordEmbDev;
|
||||
bert::cuda_unique_ptr<void> mTokEmbDev;
|
||||
bert::cuda_unique_ptr<void> mPosEmbDev;
|
||||
size_t mLd; // leading dim = hidden size
|
||||
size_t mWordVocabSize;
|
||||
size_t mPosVocabSize;
|
||||
size_t mTokVocabSize;
|
||||
|
||||
// members that partcipate in ser/deserialization
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mWordEmb;
|
||||
bert::WeightsWithOwnership mTokEmb;
|
||||
bert::WeightsWithOwnership mPosEmb;
|
||||
DataType mType{};
|
||||
DataType mMaskType{};
|
||||
|
||||
// IPluginV3 serialization related
|
||||
std::vector<nvinfer1::PluginField> mDataToSerialize;
|
||||
nvinfer1::PluginFieldCollection mFCToSerialize;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginHFace : public EmbLayerNormVarSeqlenPluginBase
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginHFace(std::string const& name, nvinfer1::DataType const type,
|
||||
nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma, nvinfer1::Weights const& word_emb,
|
||||
nvinfer1::Weights const& pos_emb, nvinfer1::Weights const& tok_emb);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormVarSeqlenPlugin without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormVarSeqlenPluginHFace() = delete;
|
||||
|
||||
~EmbLayerNormVarSeqlenPluginHFace() override;
|
||||
|
||||
// IPluginV3Runtime overrides
|
||||
IPluginV3* clone() noexcept override;
|
||||
|
||||
int32_t onShapeChange(
|
||||
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) 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 getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
|
||||
int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginMTron : public EmbLayerNormVarSeqlenPluginBase
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginMTron(std::string const& name, nvinfer1::DataType const type,
|
||||
nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma, nvinfer1::Weights const& word_emb,
|
||||
nvinfer1::Weights const& pos_emb, nvinfer1::Weights const& tok_emb);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormVarSeqlenPlugin without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormVarSeqlenPluginMTron() = delete;
|
||||
|
||||
~EmbLayerNormVarSeqlenPluginMTron() override;
|
||||
|
||||
// IPluginV3Runtime overrides
|
||||
IPluginV3* clone() noexcept override;
|
||||
|
||||
int32_t onShapeChange(
|
||||
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) 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 getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
|
||||
int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginBaseCreator : public nvinfer1::IPluginCreatorV3One
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginBaseCreator();
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
|
||||
|
||||
void setPluginNamespace(char const* libNamespace) noexcept;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
protected:
|
||||
nvinfer1::PluginFieldCollection mFC;
|
||||
std::vector<nvinfer1::PluginField> mPluginAttributes;
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginHFaceCreator : public EmbLayerNormVarSeqlenPluginBaseCreator
|
||||
{
|
||||
public:
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginMTronCreator : public EmbLayerNormVarSeqlenPluginBaseCreator
|
||||
{
|
||||
public:
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
};
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // TRT_EMB_LAYER_NORM_VARSEQ_PLUGIN_H
|
||||
@@ -0,0 +1,813 @@
|
||||
/*
|
||||
* 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 <cstring>
|
||||
#include <cuda.h>
|
||||
#include <memory>
|
||||
#include <set>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "common/serialize.hpp"
|
||||
#include "embLayerNormVarSeqlenPluginLegacy.h"
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace nvinfer1::plugin;
|
||||
using namespace nvinfer1::plugin::bert;
|
||||
|
||||
namespace
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE{"2"};
|
||||
constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON{"3"};
|
||||
constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_NAME{"CustomEmbLayerNormPluginDynamic"};
|
||||
} // namespace
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(EmbLayerNormVarSeqlenPluginLegacyHFaceCreator);
|
||||
REGISTER_TENSORRT_PLUGIN(EmbLayerNormVarSeqlenPluginLegacyMTronCreator);
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyBase::EmbLayerNormVarSeqlenPluginLegacyBase(std::string const& name, DataType type,
|
||||
Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb, Weights const& tokEmb,
|
||||
DataType maskType)
|
||||
: mLayerName(name)
|
||||
, mLd(beta.count)
|
||||
, mType(type)
|
||||
, mMaskType(maskType)
|
||||
{
|
||||
// Assuming Weights.count is the number of elements and not bytes
|
||||
PLUGIN_VALIDATE(beta.count == gamma.count);
|
||||
PLUGIN_VALIDATE(mLd > 0U);
|
||||
PLUGIN_VALIDATE(wordEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(posEmb.count % mLd == 0);
|
||||
PLUGIN_VALIDATE(tokEmb.count % mLd == 0);
|
||||
mWordVocabSize = wordEmb.count / mLd;
|
||||
mPosVocabSize = posEmb.count / mLd;
|
||||
mTokVocabSize = tokEmb.count / mLd;
|
||||
|
||||
mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
|
||||
mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
|
||||
|
||||
mWordEmb.convertAndCopy(wordEmb, mType);
|
||||
mTokEmb.convertAndCopy(tokEmb, mType);
|
||||
mPosEmb.convertAndCopy(posEmb, mType);
|
||||
|
||||
copyToDevice(mGamma, sizeof(float) * mGamma.count, mGammaDev);
|
||||
copyToDevice(mBeta, sizeof(float) * mBeta.count, mBetaDev);
|
||||
|
||||
copyToDevice(mWordEmb, getWeightsSize(mWordEmb, mType), mWordEmbDev);
|
||||
copyToDevice(mPosEmb, getWeightsSize(mPosEmb, mType), mPosEmbDev);
|
||||
copyToDevice(mTokEmb, getWeightsSize(mTokEmb, mType), mTokEmbDev);
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyBase::EmbLayerNormVarSeqlenPluginLegacyBase(
|
||||
std::string const& name, void const* data, size_t length)
|
||||
: mLayerName(name)
|
||||
, mGammaDev(nullptr)
|
||||
, mBetaDev(nullptr)
|
||||
, mWordEmbDev(nullptr)
|
||||
, mTokEmbDev(nullptr)
|
||||
, mPosEmbDev(nullptr)
|
||||
{
|
||||
// Deserialize in the same order as serialization
|
||||
deserialize_value(&data, &length, &mType);
|
||||
deserialize_value(&data, &length, &mLd);
|
||||
deserialize_value(&data, &length, &mWordVocabSize);
|
||||
deserialize_value(&data, &length, &mPosVocabSize);
|
||||
deserialize_value(&data, &length, &mTokVocabSize);
|
||||
deserialize_value(&data, &length, &mMaskType);
|
||||
|
||||
char const* d = static_cast<char const*>(data);
|
||||
mBeta.convertAndCopy(d, mLd, nvinfer1::DataType::kFLOAT);
|
||||
mGamma.convertAndCopy(d, mLd, nvinfer1::DataType::kFLOAT);
|
||||
|
||||
mWordEmb.convertAndCopy(d, mLd * mWordVocabSize, mType);
|
||||
mPosEmb.convertAndCopy(d, mLd * mPosVocabSize, mType);
|
||||
mTokEmb.convertAndCopy(d, mLd * mTokVocabSize, mType);
|
||||
|
||||
copyToDevice(mGamma, sizeof(float) * mGamma.count, mGammaDev);
|
||||
copyToDevice(mBeta, sizeof(float) * mBeta.count, mBetaDev);
|
||||
|
||||
copyToDevice(mWordEmb, getWeightsSize(mWordEmb, mType), mWordEmbDev);
|
||||
copyToDevice(mPosEmb, getWeightsSize(mPosEmb, mType), mPosEmbDev);
|
||||
copyToDevice(mTokEmb, getWeightsSize(mTokEmb, mType), mTokEmbDev);
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyHFace::EmbLayerNormVarSeqlenPluginLegacyHFace(std::string const& name,
|
||||
DataType const type, Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb,
|
||||
Weights const& tokEmb)
|
||||
: EmbLayerNormVarSeqlenPluginLegacyBase(name, type, beta, gamma, wordEmb, posEmb, tokEmb, DataType::kINT32)
|
||||
{
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyHFace::EmbLayerNormVarSeqlenPluginLegacyHFace(
|
||||
std::string const& name, void const* data, size_t length)
|
||||
: EmbLayerNormVarSeqlenPluginLegacyBase(name, data, length)
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyHFace deserialize");
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyMTron::EmbLayerNormVarSeqlenPluginLegacyMTron(std::string const& name,
|
||||
DataType const type, Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb,
|
||||
Weights const& tokEmb)
|
||||
: EmbLayerNormVarSeqlenPluginLegacyBase(name, type, beta, gamma, wordEmb, posEmb, tokEmb, type)
|
||||
{
|
||||
}
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyMTron::EmbLayerNormVarSeqlenPluginLegacyMTron(
|
||||
std::string const& name, void const* data, size_t length)
|
||||
: EmbLayerNormVarSeqlenPluginLegacyBase(name, data, length)
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyMTron deserialize");
|
||||
}
|
||||
|
||||
// IPluginV2DynamicExt Methods
|
||||
IPluginV2DynamicExt* EmbLayerNormVarSeqlenPluginLegacyHFace::clone() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyHFace clone");
|
||||
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginLegacyHFace>(
|
||||
mLayerName, mType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb);
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* EmbLayerNormVarSeqlenPluginLegacyMTron::clone() const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyMTron clone");
|
||||
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginLegacyMTron>(
|
||||
mLayerName, mType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb);
|
||||
p->setPluginNamespace(mNamespace.c_str());
|
||||
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
DimsExprs EmbLayerNormVarSeqlenPluginLegacyHFace::getOutputDimensions(
|
||||
int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
// Input should be input ids and token ids and cumulative seqlens
|
||||
// Output should be the embeddings tensor and mask indices
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == 1); // sum of all s
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
|
||||
|
||||
PLUGIN_ASSERT(inputs[2].nbDims == 1); // B+1
|
||||
|
||||
PLUGIN_ASSERT(outputIndex == 0 || outputIndex == 1);
|
||||
|
||||
if (outputIndex == 0)
|
||||
{
|
||||
DimsExprs ret;
|
||||
ret.nbDims = 4;
|
||||
ret.d[0] = inputs[0].d[0];
|
||||
ret.d[1] = exprBuilder.constant(mLd);
|
||||
ret.d[2] = exprBuilder.constant(1);
|
||||
ret.d[3] = exprBuilder.constant(1);
|
||||
return ret;
|
||||
}
|
||||
|
||||
// Return empty tensor since this is dummy output, we do not delete it for backward compatibility.
|
||||
DimsExprs ret{};
|
||||
ret.nbDims = 0;
|
||||
return ret;
|
||||
}
|
||||
|
||||
DimsExprs EmbLayerNormVarSeqlenPluginLegacyMTron::getOutputDimensions(
|
||||
int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
// Input should be input ids and token ids and cumulative seqlens
|
||||
// Output should be the embeddings tensor and mask indices
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == 1); // sum of all s
|
||||
PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
|
||||
|
||||
PLUGIN_ASSERT(inputs[2].nbDims == 1); // B+1
|
||||
|
||||
PLUGIN_ASSERT(outputIndex == 0 || outputIndex == 1);
|
||||
|
||||
DimsExprs ret;
|
||||
ret.nbDims = 4;
|
||||
ret.d[0] = inputs[0].d[0];
|
||||
ret.d[1] = exprBuilder.constant(mLd);
|
||||
ret.d[2] = exprBuilder.constant(1);
|
||||
ret.d[3] = exprBuilder.constant(1);
|
||||
return ret;
|
||||
}
|
||||
|
||||
bool EmbLayerNormVarSeqlenPluginLegacyBase::supportsFormatCombination(
|
||||
int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// The four inputs to this plugin input_ids, segment_ids, cu_seqlens and a dummy input with the
|
||||
// size of the max seq length in that order
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
// The two outputs of the plugin are embedding and the mask
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PluginTensorDesc const& desc = inOut[pos];
|
||||
if (desc.format != TensorFormat::kLINEAR)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
if (pos == 0 || pos == 2) // input_ids and cu_seqlens
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 1;
|
||||
}
|
||||
|
||||
PluginTensorDesc const& prev = inOut[pos - 1];
|
||||
if (pos == 1) // segment ids: check it's the same as input_ids
|
||||
{
|
||||
return desc.type == DataType::kINT32 && desc.dims.nbDims == 1 && desc.dims.d[0] == prev.dims.d[0];
|
||||
}
|
||||
|
||||
if (pos == 3)
|
||||
{
|
||||
return desc.dims.nbDims == 1;
|
||||
}
|
||||
|
||||
// embedded sequence
|
||||
if (pos == nbInputs)
|
||||
{
|
||||
return desc.type == mType && desc.dims.nbDims == 4 && desc.dims.d[0] == inOut[0].dims.d[0]
|
||||
&& desc.dims.d[2] == 1 && desc.dims.d[3] == 1;
|
||||
}
|
||||
// mask
|
||||
return desc.type == mMaskType;
|
||||
}
|
||||
|
||||
void checkConfigurationInputs(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// Validate input arguments
|
||||
PLUGIN_ASSERT(nbInputs == 4);
|
||||
PLUGIN_ASSERT(nbOutputs == 2);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].desc.dims.nbDims == 1);
|
||||
PLUGIN_ASSERT(inputs[1].desc.dims.nbDims == 1);
|
||||
|
||||
PLUGIN_ASSERT(inputs[1].desc.dims.d[0] == inputs[0].desc.dims.d[0]);
|
||||
|
||||
PLUGIN_ASSERT(inputs[2].desc.dims.nbDims == 1);
|
||||
|
||||
PLUGIN_ASSERT(outputs[0].desc.dims.nbDims == 4);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].desc.dims.d[0]) == static_cast<size_t>(inputs[0].desc.dims.d[0]));
|
||||
PLUGIN_ASSERT(outputs[0].desc.dims.d[2] == 1);
|
||||
PLUGIN_ASSERT(outputs[0].desc.dims.d[3] == 1);
|
||||
|
||||
PLUGIN_ASSERT(inputs[0].desc.type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[1].desc.type == DataType::kINT32);
|
||||
PLUGIN_ASSERT(inputs[2].desc.type == DataType::kINT32);
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyHFace::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyHFace configurePlugin");
|
||||
checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].desc.dims.d[1]) == static_cast<size_t>(mLd));
|
||||
|
||||
// check mask
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.nbDims == 0);
|
||||
PLUGIN_ASSERT(outputs[0].desc.type == mType);
|
||||
PLUGIN_ASSERT(outputs[1].desc.type == mMaskType);
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyMTron::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
||||
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyMTron configurePlugin");
|
||||
checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[0].desc.dims.d[1]) == static_cast<size_t>(mLd));
|
||||
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.nbDims == 4);
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[1].desc.dims.d[0]) == static_cast<size_t>(inputs[0].desc.dims.d[0]));
|
||||
PLUGIN_ASSERT(static_cast<size_t>(outputs[1].desc.dims.d[1]) == static_cast<size_t>(mLd));
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.d[2] == 1);
|
||||
PLUGIN_ASSERT(outputs[1].desc.dims.d[3] == 1);
|
||||
|
||||
PLUGIN_ASSERT(outputs[0].desc.type == mType);
|
||||
PLUGIN_ASSERT(outputs[1].desc.type == mMaskType);
|
||||
}
|
||||
|
||||
size_t EmbLayerNormVarSeqlenPluginLegacyBase::getWorkspaceSize(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginLegacyHFace::enqueue(PluginTensorDesc const* inputDesc,
|
||||
PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs, void* /* workspace */,
|
||||
cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
|
||||
|
||||
int32_t const batchSize = inputDesc[2].dims.d[0] - 1;
|
||||
// read out the maximum sequence length from the dummy input
|
||||
int32_t const maxSeqlen = inputDesc[3].dims.d[0];
|
||||
|
||||
// There are four versions of the kernel which are optimized for sequence lengths 384, 256, 192 and 128.
|
||||
// Find the closest sequence length bigger than the max seq length in this batch.
|
||||
int32_t S = 384;
|
||||
if (maxSeqlen <= 128)
|
||||
{
|
||||
S = 128;
|
||||
}
|
||||
else if (maxSeqlen <= 192)
|
||||
{
|
||||
S = 192;
|
||||
}
|
||||
else if (maxSeqlen <= 256)
|
||||
{
|
||||
S = 256;
|
||||
}
|
||||
|
||||
// Our plugin outputs only one tensor
|
||||
auto const inputIds = static_cast<int32_t const*>(inputs[0]);
|
||||
auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
|
||||
int32_t const* cuSeqlens = static_cast<int32_t const*>(inputs[2]);
|
||||
|
||||
float const* beta = mBetaDev.get();
|
||||
float const* gamma = mGammaDev.get();
|
||||
if (mType == DataType::kFLOAT)
|
||||
{
|
||||
auto output = static_cast<float*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormHFace<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
}
|
||||
if (mType == DataType::kHALF)
|
||||
{
|
||||
auto output = static_cast<half*>(outputs[0]);
|
||||
auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormHFace<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
|
||||
}
|
||||
else
|
||||
{
|
||||
gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
|
||||
<< std::endl;
|
||||
|
||||
return STATUS_NOT_SUPPORTED;
|
||||
}
|
||||
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginLegacyMTron::enqueue(PluginTensorDesc const* inputDesc,
|
||||
PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs, void* /* workspace */,
|
||||
cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
|
||||
|
||||
int32_t const batchSize = inputDesc[2].dims.d[0] - 1;
|
||||
// read out the maximum sequence length from the dummy input
|
||||
int32_t const maxSeqlen = inputDesc[3].dims.d[0];
|
||||
|
||||
// There are four versions of the kernel which are optimized for sequence lengths 384, 256, 192 and 128.
|
||||
// Find the closest sequence length bigger than the max seq length in this batch.
|
||||
int32_t S = 384;
|
||||
if (maxSeqlen <= 128)
|
||||
{
|
||||
S = 128;
|
||||
}
|
||||
else if (maxSeqlen <= 192)
|
||||
{
|
||||
S = 192;
|
||||
}
|
||||
else if (maxSeqlen <= 256)
|
||||
{
|
||||
S = 256;
|
||||
}
|
||||
|
||||
// Our plugin outputs only one tensor
|
||||
auto const inputIds = static_cast<int32_t const*>(inputs[0]);
|
||||
auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
|
||||
int32_t const* cuSeqlens = static_cast<int32_t const*>(inputs[2]);
|
||||
|
||||
float const* beta = mBetaDev.get();
|
||||
float const* gamma = mGammaDev.get();
|
||||
if (mType == DataType::kFLOAT)
|
||||
{
|
||||
auto output = static_cast<float*>(outputs[0]);
|
||||
auto skip = static_cast<float*>(outputs[1]);
|
||||
auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormMTron<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output, skip);
|
||||
}
|
||||
if (mType == DataType::kHALF)
|
||||
{
|
||||
auto output = static_cast<half*>(outputs[0]);
|
||||
auto skip = static_cast<half*>(outputs[1]);
|
||||
auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
|
||||
auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
|
||||
auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
|
||||
|
||||
return embSkipLayerNormMTron<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
|
||||
cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output, skip);
|
||||
}
|
||||
else
|
||||
{
|
||||
gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
|
||||
<< std::endl;
|
||||
|
||||
return STATUS_NOT_SUPPORTED;
|
||||
}
|
||||
|
||||
return STATUS_SUCCESS;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return STATUS_FAILURE;
|
||||
}
|
||||
|
||||
// IPluginV2Ext Methods
|
||||
DataType EmbLayerNormVarSeqlenPluginLegacyBase::getOutputDataType(
|
||||
int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
PLUGIN_ASSERT(index == 0 || index == 1);
|
||||
PLUGIN_ASSERT(mType == DataType::kHALF || mType == DataType::kFLOAT);
|
||||
return index == 0 ? mType : mMaskType;
|
||||
}
|
||||
|
||||
// IPluginV2 Methods
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyBase::getPluginType() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_NAME;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyHFace::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyMTron::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginLegacyBase::getNbOutputs() const noexcept
|
||||
{
|
||||
return 2;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginLegacyHFace::initialize() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyHFace initialize");
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t EmbLayerNormVarSeqlenPluginLegacyMTron::initialize() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyMTron initialize");
|
||||
return 0;
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyHFace::terminate() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyHFace terminate");
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyMTron::terminate() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyMTron terminate");
|
||||
}
|
||||
|
||||
size_t EmbLayerNormVarSeqlenPluginLegacyBase::getSerializationSize() const noexcept
|
||||
{
|
||||
size_t const wordSize = getElementSize(mType);
|
||||
return 2 * sizeof(float) * mLd // beta + gamma
|
||||
+ sizeof(mType) //
|
||||
+ sizeof(mLd) //
|
||||
+ sizeof(mWordVocabSize) //
|
||||
+ sizeof(mPosVocabSize) //
|
||||
+ sizeof(mTokVocabSize) //
|
||||
+ wordSize * mLd * mWordVocabSize // word emb
|
||||
+ wordSize * mLd * mPosVocabSize // pos emb
|
||||
+ wordSize * mLd * mTokVocabSize // tok emb
|
||||
+ sizeof(mMaskType) // mask type
|
||||
;
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyBase::serialize(void* buffer) const noexcept
|
||||
{
|
||||
serialize_value(&buffer, mType);
|
||||
serialize_value(&buffer, mLd);
|
||||
serialize_value(&buffer, mWordVocabSize);
|
||||
serialize_value(&buffer, mPosVocabSize);
|
||||
serialize_value(&buffer, mTokVocabSize);
|
||||
serialize_value(&buffer, mMaskType);
|
||||
|
||||
char* d = static_cast<char*>(buffer);
|
||||
size_t const wordSize = getElementSize(mType);
|
||||
|
||||
serFromDev(d, mBetaDev.get(), mLd);
|
||||
serFromDev(d, mGammaDev.get(), mLd);
|
||||
serFromDev(d, static_cast<char*>(mWordEmbDev.get()), mLd * mWordVocabSize * wordSize);
|
||||
serFromDev(d, static_cast<char*>(mPosEmbDev.get()), mLd * mPosVocabSize * wordSize);
|
||||
serFromDev(d, static_cast<char*>(mTokEmbDev.get()), mLd * mTokVocabSize * wordSize);
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyBase::destroy() noexcept
|
||||
{
|
||||
// This gets called when the network containing plugin is destroyed
|
||||
mGammaDev.reset(nullptr);
|
||||
mBetaDev.reset(nullptr);
|
||||
mWordEmbDev.reset(nullptr);
|
||||
mPosEmbDev.reset(nullptr);
|
||||
mTokEmbDev.reset(nullptr);
|
||||
delete this;
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyHFace::destroy() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyHFace destroy");
|
||||
EmbLayerNormVarSeqlenPluginLegacyBase::destroy();
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyMTron::destroy() noexcept
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginLegacyMTron destroy");
|
||||
EmbLayerNormVarSeqlenPluginLegacyBase::destroy();
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyBase::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyBase::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
///////////////////////
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyBaseCreator::EmbLayerNormVarSeqlenPluginLegacyBaseCreator()
|
||||
{
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_beta"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_gamma"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_word_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_token_type_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("bert_embeddings_position_embeddings"));
|
||||
mPluginAttributes.emplace_back(PluginField("output_fp16"));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyBaseCreator::getPluginName() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_NAME;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyHFaceCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyMTronCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* EmbLayerNormVarSeqlenPluginLegacyBaseCreator::getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
bool initializeFields(char const* name, PluginFieldCollection const* fc, Weights& beta, Weights& gamma,
|
||||
Weights& word_emb, Weights& pos_emb, Weights& tok_emb)
|
||||
{
|
||||
bool output_fp16 = false;
|
||||
std::set<std::string> const requiredAttributes{
|
||||
"bert_embeddings_layernorm_beta",
|
||||
"bert_embeddings_layernorm_gamma",
|
||||
"bert_embeddings_word_embeddings",
|
||||
"bert_embeddings_token_type_embeddings",
|
||||
"bert_embeddings_position_embeddings",
|
||||
};
|
||||
plugin::validateRequiredAttributesExist(requiredAttributes, fc);
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; i++)
|
||||
{
|
||||
std::string_view const field_name = fc->fields[i].name;
|
||||
if (field_name == "bert_embeddings_layernorm_beta"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
|
||||
beta.values = fc->fields[i].data;
|
||||
beta.count = fc->fields[i].length;
|
||||
beta.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_layernorm_gamma"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
|
||||
gamma.values = fc->fields[i].data;
|
||||
gamma.count = fc->fields[i].length;
|
||||
gamma.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_word_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_word_embeddings...");
|
||||
word_emb.values = fc->fields[i].data;
|
||||
word_emb.count = fc->fields[i].length;
|
||||
word_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_token_type_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_token_type_embeddings...");
|
||||
tok_emb.values = fc->fields[i].data;
|
||||
tok_emb.count = fc->fields[i].length;
|
||||
tok_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
|
||||
if (field_name == "bert_embeddings_position_embeddings"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building bert_embeddings_position_embeddings...");
|
||||
pos_emb.values = fc->fields[i].data;
|
||||
pos_emb.count = fc->fields[i].length;
|
||||
pos_emb.type = fieldTypeToDataType(fc->fields[i].type);
|
||||
}
|
||||
if (field_name == "output_fp16"sv)
|
||||
{
|
||||
BERT_DEBUG_MSG("Building output_fp16...");
|
||||
PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kINT32);
|
||||
output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
|
||||
}
|
||||
}
|
||||
return output_fp16;
|
||||
}
|
||||
|
||||
IPluginV2* EmbLayerNormVarSeqlenPluginLegacyHFaceCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenHFace createPlugin");
|
||||
|
||||
Weights beta{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights gamma{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights word_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights pos_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights tok_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
bool output_fp16 = initializeFields(name, fc, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
|
||||
BERT_DEBUG_MSG("Building the Plugin...");
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginLegacyHFace>(
|
||||
name, output_fp16 ? DataType::kHALF : DataType::kFLOAT, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* EmbLayerNormVarSeqlenPluginLegacyMTronCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
BERT_DEBUG_MSG("EmbLayerNormVarSeqlenMTron createPlugin");
|
||||
|
||||
Weights beta{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights gamma{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
|
||||
// existence
|
||||
Weights word_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights pos_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
Weights tok_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
|
||||
// verify existence
|
||||
bool output_fp16 = initializeFields(name, fc, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
|
||||
BERT_DEBUG_MSG("Building the Plugin...");
|
||||
auto p = std::make_unique<EmbLayerNormVarSeqlenPluginLegacyMTron>(
|
||||
name, output_fp16 ? DataType::kHALF : DataType::kFLOAT, beta, gamma, word_emb, pos_emb, tok_emb);
|
||||
return p.release();
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* EmbLayerNormVarSeqlenPluginLegacyHFaceCreator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// This object will be deleted when the network is destroyed, which will
|
||||
// call EmbLayerNormVarSeqlen::destroy()
|
||||
return new EmbLayerNormVarSeqlenPluginLegacyHFace(name, serialData, serialLength);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* EmbLayerNormVarSeqlenPluginLegacyMTronCreator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
// This object will be deleted when the network is destroyed, which will
|
||||
// call EmbLayerNormVarSeqlen::destroy()
|
||||
return new EmbLayerNormVarSeqlenPluginLegacyMTron(name, serialData, serialLength);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void EmbLayerNormVarSeqlenPluginLegacyBaseCreator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
char const* EmbLayerNormVarSeqlenPluginLegacyBaseCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
@@ -0,0 +1,198 @@
|
||||
/*
|
||||
* 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_EMB_LAYER_NORM_VARSEQ_PLUGIN_LEGACY_H
|
||||
#define TRT_EMB_LAYER_NORM_VARSEQ_PLUGIN_LEGACY_H
|
||||
|
||||
#include <cuda.h>
|
||||
|
||||
#include "NvInferPlugin.h"
|
||||
#include "NvInferRuntime.h"
|
||||
|
||||
#include "common/bertCommon.h"
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
namespace bert
|
||||
{
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNormHFace(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* tokenIds, int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb,
|
||||
T const* posEmb, T const* tokEmb, int32_t const wordSize, int32_t const tokSize, T* output);
|
||||
|
||||
template <typename T>
|
||||
int32_t embSkipLayerNormMTron(cudaStream_t stream, int32_t ld, int32_t B, int32_t S, int32_t const* inputIds,
|
||||
int32_t const* tokenIds, int32_t const* cuSeqlens, float const* beta, float const* gamma, T const* wordEmb,
|
||||
T const* posEmb, T const* tokEmb, int32_t const wordSize, int32_t const tokSize, T* output, T* skip);
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginLegacyBase : public nvinfer1::IPluginV2DynamicExt
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginLegacyBase(std::string const& name, DataType type, Weights const& beta,
|
||||
Weights const& gamma, Weights const& word_emb, Weights const& pos_emb, Weights const& tok_emb,
|
||||
DataType maskType);
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyBase(std::string const& name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormVarSeqlenPluginLegacy without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormVarSeqlenPluginLegacyBase() = delete;
|
||||
|
||||
// IPluginV2DynamicExt Methods
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, 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;
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
size_t getSerializationSize() const noexcept override;
|
||||
void serialize(void* buffer) const noexcept override;
|
||||
void destroy() noexcept override;
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
void setPluginNamespace(char const* pluginNamespace) noexcept override;
|
||||
|
||||
protected:
|
||||
std::string const mLayerName;
|
||||
std::string mNamespace;
|
||||
|
||||
bert::cuda_unique_ptr<float> mGammaDev;
|
||||
bert::cuda_unique_ptr<float> mBetaDev;
|
||||
bert::cuda_unique_ptr<void> mWordEmbDev;
|
||||
bert::cuda_unique_ptr<void> mTokEmbDev;
|
||||
bert::cuda_unique_ptr<void> mPosEmbDev;
|
||||
size_t mLd; // leading dim = hidden size
|
||||
size_t mWordVocabSize;
|
||||
size_t mPosVocabSize;
|
||||
size_t mTokVocabSize;
|
||||
bert::WeightsWithOwnership mBeta;
|
||||
bert::WeightsWithOwnership mGamma;
|
||||
bert::WeightsWithOwnership mWordEmb;
|
||||
bert::WeightsWithOwnership mTokEmb;
|
||||
bert::WeightsWithOwnership mPosEmb;
|
||||
DataType mType{};
|
||||
DataType mMaskType{};
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginLegacyHFace : public EmbLayerNormVarSeqlenPluginLegacyBase
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginLegacyHFace(std::string const& name, nvinfer1::DataType const type,
|
||||
nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma, nvinfer1::Weights const& word_emb,
|
||||
nvinfer1::Weights const& pos_emb, nvinfer1::Weights const& tok_emb);
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyHFace(std::string const& name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormVarSeqlenPluginLegacy without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormVarSeqlenPluginLegacyHFace() = 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;
|
||||
void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
|
||||
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) 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;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginLegacyMTron : public EmbLayerNormVarSeqlenPluginLegacyBase
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginLegacyMTron(std::string const& name, nvinfer1::DataType const type,
|
||||
nvinfer1::Weights const& beta, nvinfer1::Weights const& gamma, nvinfer1::Weights const& word_emb,
|
||||
nvinfer1::Weights const& pos_emb, nvinfer1::Weights const& tok_emb);
|
||||
|
||||
EmbLayerNormVarSeqlenPluginLegacyMTron(std::string const& name, void const* data, size_t length);
|
||||
|
||||
// It doesn't make sense to make EmbLayerNormVarSeqlenPluginLegacy without arguments, so we
|
||||
// delete default constructor.
|
||||
EmbLayerNormVarSeqlenPluginLegacyMTron() = 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;
|
||||
void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
|
||||
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) 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;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginLegacyBaseCreator : public nvinfer1::IPluginCreator
|
||||
{
|
||||
public:
|
||||
EmbLayerNormVarSeqlenPluginLegacyBaseCreator();
|
||||
|
||||
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;
|
||||
|
||||
protected:
|
||||
nvinfer1::PluginFieldCollection mFC;
|
||||
std::vector<nvinfer1::PluginField> mPluginAttributes;
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginLegacyHFaceCreator : public EmbLayerNormVarSeqlenPluginLegacyBaseCreator
|
||||
{
|
||||
public:
|
||||
nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
nvinfer1::IPluginV2* deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept override;
|
||||
};
|
||||
|
||||
class EmbLayerNormVarSeqlenPluginLegacyMTronCreator : public EmbLayerNormVarSeqlenPluginLegacyBaseCreator
|
||||
{
|
||||
public:
|
||||
nvinfer1::IPluginV2* createPlugin(char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept override;
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
nvinfer1::IPluginV2* deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept override;
|
||||
};
|
||||
|
||||
} // namespace bert
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
#endif // TRT_EMB_LAYER_NORM_VARSEQ_PLUGIN_LEGACY_H
|
||||
Reference in New Issue
Block a user