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