wehub-resource-sync
132 lines
4.5 KiB
Plaintext
132 lines
4.5 KiB
Plaintext
/*
|
|
* 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
|