wehub-resource-sync
749 lines
27 KiB
Plaintext
749 lines
27 KiB
Plaintext
// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
|
|
//
|
|
// 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.
|
|
|
|
#ifdef PADDLE_WITH_CUDA
|
|
#include <cuda.h>
|
|
#include <cuda_runtime.h>
|
|
#endif
|
|
#ifdef PADDLE_WITH_HIP
|
|
#include <hip/hip_runtime.h>
|
|
#endif
|
|
|
|
#include "paddle/phi/kernels/funcs/blas/blas.h"
|
|
#include "paddle/phi/kernels/funcs/cub.h"
|
|
#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
|
|
#include "paddle/phi/kernels/funcs/multihead_matmul_functor.h"
|
|
|
|
namespace phi {
|
|
namespace funcs {
|
|
|
|
template <typename T>
|
|
struct CUDATypeTraits;
|
|
|
|
template <>
|
|
struct CUDATypeTraits<half> {
|
|
typedef phi::float16 TYPE;
|
|
};
|
|
|
|
template <>
|
|
struct CUDATypeTraits<float> {
|
|
typedef float TYPE;
|
|
};
|
|
|
|
using funcs::operator+;
|
|
|
|
template <typename T>
|
|
__global__ void SoftmaxKernelWithEltadd(T *qk_buf_,
|
|
const T *bias_qk_,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
float tmp = threadIdx.x < seq_len
|
|
? static_cast<float>(qk_buf_[threadIdx.x + qk_offset] +
|
|
bias_qk_[threadIdx.x + qk_offset])
|
|
: -1e20f;
|
|
float max_val = funcs::BlockReduceMax<float>(tmp, mask);
|
|
|
|
float qk_tmp = threadIdx.x < seq_len ? __expf(tmp - max_val) : 0.0f;
|
|
float sum_val = funcs::BlockReduceSum<float>(qk_tmp, mask);
|
|
|
|
if (threadIdx.x < seq_len)
|
|
qk_buf_[threadIdx.x + qk_offset] = (T)(qk_tmp / sum_val);
|
|
}
|
|
|
|
template <>
|
|
__global__ void SoftmaxKernelWithEltadd<half>(half *qk_buf_,
|
|
const half *bias_qk_,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
float tmp = threadIdx.x < seq_len
|
|
? static_cast<float>(qk_buf_[threadIdx.x + qk_offset] +
|
|
bias_qk_[threadIdx.x + qk_offset])
|
|
: -1e20f;
|
|
float max_val = funcs::BlockReduceMax<float>(tmp, mask);
|
|
|
|
float qk_tmp = threadIdx.x < seq_len ? __expf(tmp - max_val) : 0.0f;
|
|
float sum_val = funcs::BlockReduceSum<float>(qk_tmp, mask);
|
|
|
|
if (threadIdx.x < seq_len)
|
|
qk_buf_[threadIdx.x + qk_offset] = (half)(qk_tmp / sum_val);
|
|
#endif
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void SoftmaxKernelWithEltadd2(T *qk_buf_,
|
|
const T *bias_qk_,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
int idx = threadIdx.x;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
float2 tmp = idx < seq_len ? funcs::ToFloat2<T>(qk_buf_[idx + qk_offset] +
|
|
bias_qk_[idx + qk_offset])
|
|
: make_float2(-1e20f, -1e20f);
|
|
float max_val = funcs::BlockReduceMax<float>(max(tmp.x, tmp.y), mask);
|
|
float2 qk_tmp = idx < seq_len ? make_float2(__expf(tmp.x - max_val),
|
|
__expf(tmp.y - max_val))
|
|
: make_float2(0.f, 0.f);
|
|
float sum_val =
|
|
funcs::BlockReduceSum<float>(qk_tmp.x + qk_tmp.y, mask) + 1e-6f;
|
|
|
|
if (idx < seq_len) {
|
|
qk_buf_[idx + qk_offset] =
|
|
funcs::FloatsToPair<T>(qk_tmp.x / sum_val, qk_tmp.y / sum_val);
|
|
}
|
|
}
|
|
|
|
template <>
|
|
__global__ void SoftmaxKernelWithEltadd2<half2>(half2 *qk_buf_,
|
|
const half2 *bias_qk_,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
// operator "+" of half only supported after cuda version 10.0
|
|
// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
int idx = threadIdx.x;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
float2 tmp = idx < seq_len ? funcs::ToFloat2<half2>(qk_buf_[idx + qk_offset] +
|
|
bias_qk_[idx + qk_offset])
|
|
: make_float2(-1e20f, -1e20f);
|
|
float max_val = funcs::BlockReduceMax<float>(max(tmp.x, tmp.y), mask);
|
|
float2 qk_tmp = idx < seq_len ? make_float2(__expf(tmp.x - max_val),
|
|
__expf(tmp.y - max_val))
|
|
: make_float2(0.f, 0.f);
|
|
float sum_val =
|
|
funcs::BlockReduceSum<float>(qk_tmp.x + qk_tmp.y, mask) + 1e-6f;
|
|
|
|
if (idx < seq_len) {
|
|
qk_buf_[idx + qk_offset] =
|
|
funcs::FloatsToPair<half2>(qk_tmp.x / sum_val, qk_tmp.y / sum_val);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void SoftmaxKernelWithEltaddForLarge(T *qk_buf,
|
|
const T *bias_qk,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
T stride_max = -1e20f;
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
stride_max = qk_buf[threadIdx.x + i + qk_offset] +
|
|
bias_qk[threadIdx.x + i + qk_offset] >
|
|
stride_max
|
|
? qk_buf[threadIdx.x + i + qk_offset] +
|
|
bias_qk[threadIdx.x + i + qk_offset]
|
|
: stride_max;
|
|
}
|
|
T max_val = funcs::BlockReduceMax<T>(stride_max, mask);
|
|
|
|
T stride_sum = 0.f;
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
stride_sum += __expf(qk_buf[threadIdx.x + i + qk_offset] +
|
|
bias_qk[threadIdx.x + i + qk_offset] - max_val);
|
|
}
|
|
T sum_val = funcs::BlockReduceSum<T>(stride_sum, mask);
|
|
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
qk_buf[threadIdx.x + i + qk_offset] =
|
|
(T)(__expf(qk_buf[threadIdx.x + i + qk_offset] +
|
|
bias_qk[threadIdx.x + i + qk_offset] - max_val) /
|
|
sum_val);
|
|
}
|
|
}
|
|
|
|
template <>
|
|
__global__ void SoftmaxKernelWithEltaddForLarge(half *qk_buf,
|
|
const half *bias_qk,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
float stride_max = -1e20f;
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float tmp = static_cast<float>(qk_buf[threadIdx.x + i + qk_offset] +
|
|
bias_qk[threadIdx.x + i + qk_offset]);
|
|
stride_max = tmp > stride_max ? tmp : stride_max;
|
|
}
|
|
float max_val = funcs::BlockReduceMax<float>(stride_max, mask);
|
|
|
|
float stride_sum = 0.f;
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float tmp = static_cast<float>(qk_buf[threadIdx.x + i + qk_offset] +
|
|
bias_qk[threadIdx.x + i + qk_offset]);
|
|
stride_sum += __expf(tmp - max_val);
|
|
}
|
|
float sum_val = funcs::BlockReduceSum<float>(stride_sum, mask);
|
|
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float tmp =
|
|
__expf(static_cast<float>(qk_buf[threadIdx.x + i + qk_offset] +
|
|
bias_qk[threadIdx.x + i + qk_offset]) -
|
|
max_val);
|
|
qk_buf[threadIdx.x + i + qk_offset] = (half)(tmp / sum_val);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
template <typename T>
|
|
__global__ void SoftmaxKernelWithEltaddForLarge2(
|
|
T *qk_buf_,
|
|
const T *bias_qk_,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
float2 stride_max = make_float2(-1e20f, -1e20f);
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float2 cur = funcs::ToFloat2<T>(qk_buf_[threadIdx.x + i + qk_offset] +
|
|
bias_qk_[threadIdx.x + i + qk_offset]);
|
|
stride_max.x = max(stride_max.x, cur.x);
|
|
stride_max.y = max(stride_max.y, cur.y);
|
|
}
|
|
float max_val =
|
|
funcs::BlockReduceMax<float>(max(stride_max.x, stride_max.y), mask);
|
|
|
|
float2 stride_sum = make_float2(0.f, 0.f);
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float2 cur = funcs::ToFloat2<T>(qk_buf_[threadIdx.x + i + qk_offset] +
|
|
bias_qk_[threadIdx.x + i + qk_offset]);
|
|
stride_sum.x += __expf(cur.x - max_val);
|
|
stride_sum.y += __expf(cur.y - max_val);
|
|
}
|
|
|
|
float sum_val =
|
|
funcs::BlockReduceSum<float>(stride_sum.x + stride_sum.y, mask) + 1e-6f;
|
|
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float2 cur = funcs::ToFloat2<T>(qk_buf_[threadIdx.x + i + qk_offset] +
|
|
bias_qk_[threadIdx.x + i + qk_offset]);
|
|
qk_buf_[threadIdx.x + i + qk_offset] = funcs::FloatsToPair<T>(
|
|
__expf(cur.x - max_val) / sum_val, __expf(cur.y - max_val) / sum_val);
|
|
}
|
|
}
|
|
|
|
template <>
|
|
__global__ void SoftmaxKernelWithEltaddForLarge2(
|
|
half2 *qk_buf_,
|
|
const half2 *bias_qk_,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len,
|
|
const funcs::warp_mask_t mask) {
|
|
// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
|
|
int64_t qk_offset = static_cast<int64_t>(blockIdx.x) * seq_len;
|
|
assert(blockDim.x % WARP_SIZE == 0);
|
|
|
|
float2 stride_max = make_float2(-1e20f, -1e20f);
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float2 cur = funcs::ToFloat2<half2>(qk_buf_[threadIdx.x + i + qk_offset] +
|
|
bias_qk_[threadIdx.x + i + qk_offset]);
|
|
stride_max.x = max(stride_max.x, cur.x);
|
|
stride_max.y = max(stride_max.y, cur.y);
|
|
}
|
|
float max_val =
|
|
funcs::BlockReduceMax<float>(max(stride_max.x, stride_max.y), mask);
|
|
|
|
float2 stride_sum = make_float2(0.f, 0.f);
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float2 cur = funcs::ToFloat2<half2>(qk_buf_[threadIdx.x + i + qk_offset] +
|
|
bias_qk_[threadIdx.x + i + qk_offset]);
|
|
stride_sum.x += __expf(cur.x - max_val);
|
|
stride_sum.y += __expf(cur.y - max_val);
|
|
}
|
|
|
|
float sum_val =
|
|
funcs::BlockReduceSum<float>(stride_sum.x + stride_sum.y, mask) + 1e-6f;
|
|
|
|
for (int i = 0; threadIdx.x + i < seq_len; i += blockDim.x) {
|
|
float2 cur = funcs::ToFloat2<half2>(qk_buf_[threadIdx.x + i + qk_offset] +
|
|
bias_qk_[threadIdx.x + i + qk_offset]);
|
|
qk_buf_[threadIdx.x + i + qk_offset] = funcs::FloatsToPair<half2>(
|
|
__expf(cur.x - max_val) / sum_val, __expf(cur.y - max_val) / sum_val);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
template <typename T>
|
|
inline __device__ T ldg(const T *val) {
|
|
return __ldg(val);
|
|
}
|
|
|
|
template <typename T>
|
|
inline __device__ T hexp2(T a) {
|
|
return h2exp(a);
|
|
}
|
|
|
|
template <typename T_IN, typename T_OUT>
|
|
inline __device__ T_OUT type2type2(T_IN a);
|
|
|
|
template <>
|
|
inline __device__ half2 type2type2(half a) {
|
|
return __half2half2(a);
|
|
}
|
|
|
|
template <typename T>
|
|
inline __device__ T float2type2(float a);
|
|
|
|
template <>
|
|
inline __device__ half2 float2type2(float a) {
|
|
return __float2half2_rn(a);
|
|
}
|
|
|
|
template <typename T>
|
|
inline __device__ T hmul2(T a, T b) {
|
|
return __hmul2(a, b);
|
|
}
|
|
|
|
template <typename T>
|
|
inline __device__ T hsub2(T a, T b) {
|
|
return __hsub2(a, b);
|
|
}
|
|
|
|
template <typename T>
|
|
inline __device__ T hadd2(T a, T b) {
|
|
return __hadd2(a, b);
|
|
}
|
|
|
|
template <typename T, int ITEMS_PER_THREAD, int NUM>
|
|
__global__ void softmax_kernel_with_mask(T *qk_buf_,
|
|
const T *attr_mask,
|
|
const int batch_size,
|
|
const int head_num,
|
|
const int seq_len) {
|
|
using T2 = half2;
|
|
T2 *qk_buf_half2 = reinterpret_cast<T2 *>(qk_buf_);
|
|
const T2 *attr_mask_half2 = (const T2 *)attr_mask;
|
|
|
|
for (int64_t seq_id = static_cast<int64_t>(blockIdx.x); seq_id < seq_len;
|
|
seq_id += static_cast<int64_t>(gridDim.x) * NUM) {
|
|
T2 data[NUM][ITEMS_PER_THREAD];
|
|
|
|
int64_t qk_offset[NUM];
|
|
|
|
__shared__ float s_sum[NUM], s_max[NUM];
|
|
float local_max[NUM];
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
local_max[j] = -1e20f;
|
|
}
|
|
|
|
for (int i = 0;
|
|
static_cast<int64_t>(blockDim.x) * i + threadIdx.x < (seq_len / 2) &&
|
|
i < ITEMS_PER_THREAD;
|
|
i++) {
|
|
int64_t mask_offset[NUM];
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
qk_offset[j] =
|
|
((static_cast<int64_t>(blockIdx.y) * head_num +
|
|
static_cast<int64_t>(blockIdx.z)) *
|
|
static_cast<int64_t>(seq_len) +
|
|
seq_id +
|
|
static_cast<int64_t>(j) * static_cast<int64_t>(gridDim.x)) *
|
|
static_cast<int64_t>(seq_len / 2) +
|
|
static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
|
|
mask_offset[j] =
|
|
(static_cast<int64_t>(blockIdx.y) * seq_len + seq_id +
|
|
static_cast<int64_t>(j) * static_cast<int64_t>(gridDim.x)) *
|
|
static_cast<int64_t>(seq_len / 2) +
|
|
static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
|
|
}
|
|
|
|
T2 mask_val[NUM];
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
mask_val[j] = ldg(&attr_mask_half2[mask_offset[j]]);
|
|
}
|
|
|
|
T2 qk[NUM];
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
qk[j] = qk_buf_half2[qk_offset[j]];
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
mask_val[j] = hmul2<T2>(hsub2<T2>(float2type2<T2>(1.0f), mask_val[j]),
|
|
float2type2<T2>(-10000.0f));
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
data[j][i] = hadd2<T2>(qk[j], mask_val[j]);
|
|
local_max[j] = fmax(local_max[j],
|
|
fmax(static_cast<float>(data[j][i].x),
|
|
static_cast<float>(data[j][i].y)));
|
|
}
|
|
}
|
|
|
|
if (blockDim.x <= WARP_SIZE) {
|
|
funcs::WarpReduceMaxV2<float, NUM>(local_max);
|
|
} else {
|
|
funcs::BlockReduceMaxV2<float, NUM>(local_max);
|
|
}
|
|
|
|
if (threadIdx.x == 0) {
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
s_max[j] = local_max[j];
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
float local_sum[NUM];
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
local_sum[j] = {0.f};
|
|
}
|
|
|
|
for (int i = 0;
|
|
static_cast<int64_t>(blockDim.x) * i + threadIdx.x < (seq_len / 2) &&
|
|
i < ITEMS_PER_THREAD;
|
|
i++) {
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
data[j][i] =
|
|
hexp2<T2>(hsub2<T2>(data[j][i], float2type2<T2>(s_max[j])));
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
local_sum[j] += static_cast<float>(data[j][i].x + data[j][i].y);
|
|
}
|
|
}
|
|
|
|
if (blockDim.x <= WARP_SIZE) {
|
|
funcs::WarpReduceSumV2<float, NUM>(local_sum);
|
|
} else {
|
|
funcs::BlockReduceSumV2<float, NUM>(local_sum);
|
|
}
|
|
|
|
if (threadIdx.x == 0) {
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
s_sum[j] = __fdividef(1.0f, local_sum[j] + 1e-6f);
|
|
}
|
|
}
|
|
__syncthreads();
|
|
|
|
for (int i = 0;
|
|
static_cast<int64_t>(blockDim.x) * i + threadIdx.x < (seq_len / 2) &&
|
|
i < ITEMS_PER_THREAD;
|
|
i++) {
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
qk_offset[j] =
|
|
((static_cast<int64_t>(blockIdx.y) * head_num +
|
|
static_cast<int64_t>(blockIdx.z)) *
|
|
static_cast<int64_t>(seq_len) +
|
|
seq_id +
|
|
static_cast<int64_t>(j) * static_cast<int64_t>(gridDim.x)) *
|
|
static_cast<int64_t>(seq_len / 2) +
|
|
static_cast<int64_t>(blockDim.x) * i + threadIdx.x;
|
|
}
|
|
|
|
#pragma unroll
|
|
for (int j = 0; j < NUM; j++) {
|
|
qk_buf_half2[qk_offset[j]] =
|
|
hmul2<T2>(data[j][i], float2type2<T2>(s_sum[j]));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
#define SOFTMAX_KERNEL_WITH_MASK(REPEAT_THREAD) \
|
|
do { \
|
|
block.x /= REPEAT_THREAD; \
|
|
grid.x /= 4; \
|
|
constexpr int NUM = 4; \
|
|
softmax_kernel_with_mask<half, REPEAT_THREAD, NUM> \
|
|
<<<grid, block, 0, stream>>>(reinterpret_cast<half *>(qk_buf_), \
|
|
(const half *)bias_qk, \
|
|
batch_size, \
|
|
head_num, \
|
|
seq_len); \
|
|
} while (0)
|
|
|
|
template <typename T>
|
|
inline void MatmulWithHeadQK(const GPUContext &dev_ctx,
|
|
int head_num,
|
|
int seq_len,
|
|
int size_per_head,
|
|
int batch_size,
|
|
bool q_trans,
|
|
bool k_trans,
|
|
T *q_buf_,
|
|
T *k_buf_,
|
|
T *qk_buf_,
|
|
const T *bias_qk,
|
|
bool bias_is_mask,
|
|
T alpha,
|
|
T beta) {
|
|
CBLAS_TRANSPOSE transA = !q_trans ? CblasNoTrans : CblasTrans;
|
|
CBLAS_TRANSPOSE transB = !k_trans ? CblasNoTrans : CblasTrans;
|
|
|
|
typedef typename CUDATypeTraits<T>::TYPE run_type;
|
|
auto blas = funcs::GetBlas<GPUContext, run_type>(dev_ctx);
|
|
auto stream = dev_ctx.stream();
|
|
|
|
blas.BatchedGEMM(transA,
|
|
transB,
|
|
seq_len,
|
|
seq_len,
|
|
size_per_head,
|
|
static_cast<run_type>(alpha),
|
|
reinterpret_cast<run_type *>(q_buf_),
|
|
reinterpret_cast<run_type *>(k_buf_),
|
|
static_cast<run_type>(beta),
|
|
reinterpret_cast<run_type *>(qk_buf_),
|
|
static_cast<int64_t>(batch_size) * head_num,
|
|
static_cast<int64_t>(seq_len) * size_per_head,
|
|
static_cast<int64_t>(seq_len) * size_per_head);
|
|
|
|
if (seq_len <= 1024) {
|
|
int64_t grid_size64 = static_cast<int64_t>(batch_size) * head_num * seq_len;
|
|
PADDLE_ENFORCE_LE_INT_MAX(
|
|
grid_size64, "CUDA launch grid batch_size * head_num * seq_len");
|
|
int grid_size = static_cast<int>(grid_size64);
|
|
int block = seq_len;
|
|
|
|
// Align block to 32, also limit seq_len to max block size.
|
|
if (seq_len % 2 == 0) {
|
|
block =
|
|
(seq_len <= (2 * WARP_SIZE))
|
|
? WARP_SIZE
|
|
: ((seq_len + (2 * WARP_SIZE - 1)) / (2 * WARP_SIZE)) * WARP_SIZE;
|
|
if (std::is_same<T, float>::value) {
|
|
SoftmaxKernelWithEltadd2<float2><<<grid_size, block, 0, stream>>>(
|
|
reinterpret_cast<float2 *>(qk_buf_),
|
|
reinterpret_cast<const float2 *>(bias_qk),
|
|
batch_size,
|
|
head_num,
|
|
seq_len / 2,
|
|
FINAL_MASK);
|
|
} else {
|
|
if (bias_is_mask) {
|
|
#if defined(__HIPCC__) || (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700)
|
|
PADDLE_ENFORCE_EQ(bias_is_mask,
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"QK_bias is mask can't be supported on rocm or "
|
|
"cuda_arch<700"));
|
|
#else
|
|
dim3 grid(seq_len, batch_size, head_num);
|
|
dim3 block((seq_len / 2 + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE);
|
|
SOFTMAX_KERNEL_WITH_MASK(1);
|
|
#endif
|
|
} else {
|
|
SoftmaxKernelWithEltadd2<__half2><<<grid_size, block, 0, stream>>>(
|
|
reinterpret_cast<__half2 *>(qk_buf_),
|
|
reinterpret_cast<const __half2 *>(bias_qk),
|
|
batch_size,
|
|
head_num,
|
|
seq_len / 2,
|
|
FINAL_MASK);
|
|
}
|
|
}
|
|
} else {
|
|
block = (seq_len <= WARP_SIZE)
|
|
? WARP_SIZE
|
|
: ((seq_len + WARP_SIZE - 1) / WARP_SIZE) * WARP_SIZE;
|
|
SoftmaxKernelWithEltadd<T><<<grid_size, block, 0, stream>>>(
|
|
qk_buf_, bias_qk, batch_size, head_num, seq_len, FINAL_MASK);
|
|
}
|
|
} else {
|
|
int64_t grid_size64 = static_cast<int64_t>(batch_size) * head_num * seq_len;
|
|
PADDLE_ENFORCE_LE_INT_MAX(
|
|
grid_size64, "CUDA launch grid batch_size * head_num * seq_len");
|
|
int grid_size = static_cast<int>(grid_size64);
|
|
int block = 512;
|
|
if (seq_len % 2 == 0) {
|
|
if (std::is_same<T, float>::value) {
|
|
SoftmaxKernelWithEltaddForLarge2<float2>
|
|
<<<grid_size, block, 0, stream>>>(
|
|
reinterpret_cast<float2 *>(qk_buf_),
|
|
reinterpret_cast<const float2 *>(bias_qk),
|
|
batch_size,
|
|
head_num,
|
|
seq_len / 2,
|
|
FINAL_MASK);
|
|
} else {
|
|
if (bias_is_mask) {
|
|
#if defined(__HIPCC__) || (defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 700)
|
|
PADDLE_ENFORCE_EQ(bias_is_mask,
|
|
false,
|
|
common::errors::InvalidArgument(
|
|
"QK_bias is mask can't be supported on rocm or "
|
|
"cuda_arch<700"));
|
|
#else
|
|
dim3 grid(seq_len, batch_size, head_num);
|
|
dim3 block((seq_len / 2 + WARP_SIZE - 1) / WARP_SIZE * WARP_SIZE);
|
|
if (block.x > 0 && block.x <= 1024) {
|
|
SOFTMAX_KERNEL_WITH_MASK(1);
|
|
} else if (block.x <= 2048) {
|
|
SOFTMAX_KERNEL_WITH_MASK(2);
|
|
} else if (block.x <= 4096) {
|
|
SOFTMAX_KERNEL_WITH_MASK(4);
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"Cannot support the length of attention > 8192."));
|
|
}
|
|
#endif
|
|
} else {
|
|
SoftmaxKernelWithEltaddForLarge2<__half2>
|
|
<<<grid_size, block, 0, stream>>>(
|
|
reinterpret_cast<__half2 *>(qk_buf_),
|
|
reinterpret_cast<const __half2 *>(bias_qk),
|
|
batch_size,
|
|
head_num,
|
|
seq_len / 2,
|
|
FINAL_MASK);
|
|
}
|
|
}
|
|
} else {
|
|
SoftmaxKernelWithEltaddForLarge<T><<<grid_size, block, 0, stream>>>(
|
|
qk_buf_, bias_qk, batch_size, head_num, seq_len, FINAL_MASK);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T>
|
|
inline void MatmulWithHeadQKV(const GPUContext &dev_ctx,
|
|
int head_num,
|
|
int seq_len,
|
|
int size_per_head,
|
|
int batch_size,
|
|
bool qk_trans,
|
|
bool v_trans,
|
|
T *v_buf_,
|
|
const T *qk_buf_,
|
|
T *dst,
|
|
T alpha,
|
|
T beta) {
|
|
int64_t m = static_cast<int64_t>(batch_size) * seq_len;
|
|
int64_t k = static_cast<int64_t>(head_num) * size_per_head;
|
|
|
|
typedef typename CUDATypeTraits<T>::TYPE run_type;
|
|
auto blas = funcs::GetBlas<GPUContext, run_type>(dev_ctx);
|
|
auto stream = dev_ctx.stream();
|
|
CBLAS_TRANSPOSE transA = !qk_trans ? CblasNoTrans : CblasTrans;
|
|
CBLAS_TRANSPOSE transB = !v_trans ? CblasNoTrans : CblasTrans;
|
|
|
|
blas.BatchedGEMM(transA,
|
|
transB,
|
|
seq_len,
|
|
size_per_head,
|
|
seq_len,
|
|
static_cast<run_type>(alpha),
|
|
reinterpret_cast<const run_type *>(qk_buf_),
|
|
reinterpret_cast<run_type *>(v_buf_),
|
|
static_cast<run_type>(beta),
|
|
reinterpret_cast<run_type *>(dst),
|
|
static_cast<int64_t>(batch_size) * head_num,
|
|
static_cast<int64_t>(seq_len) * seq_len,
|
|
static_cast<int64_t>(seq_len) * size_per_head);
|
|
}
|
|
|
|
template <typename T>
|
|
void MultiheadGPUComputeFunctor<T>::operator()(const GPUContext &dev_ctx,
|
|
int batch,
|
|
int seq_len,
|
|
int head_num,
|
|
int head_size,
|
|
T *qkptr,
|
|
const T *bias_qk_ptr,
|
|
bool bias_is_mask,
|
|
T *tptr,
|
|
T alpha,
|
|
T beta) {
|
|
auto stream = dev_ctx.stream();
|
|
const int64_t tsize =
|
|
static_cast<int64_t>(batch) * head_num * seq_len * head_size;
|
|
|
|
T *qptr = tptr;
|
|
T *kptr = qptr + tsize;
|
|
T *vptr = kptr + tsize;
|
|
// batch gemm stride, softmaxwithscale.
|
|
MatmulWithHeadQK<T>(dev_ctx,
|
|
head_num,
|
|
seq_len,
|
|
head_size,
|
|
batch,
|
|
false,
|
|
true,
|
|
qptr,
|
|
kptr,
|
|
qkptr,
|
|
bias_qk_ptr,
|
|
bias_is_mask,
|
|
alpha,
|
|
beta);
|
|
// batch gemm stride, transpose.
|
|
MatmulWithHeadQKV<T>(dev_ctx,
|
|
head_num,
|
|
seq_len,
|
|
head_size,
|
|
batch,
|
|
false,
|
|
false,
|
|
vptr,
|
|
qkptr,
|
|
tptr,
|
|
T(1.0),
|
|
beta);
|
|
}
|
|
|
|
template class PADDLE_API MultiheadGPUComputeFunctor<float>;
|
|
|
|
// device function 'operator()' is not supported until cuda 10.0
|
|
// HIP defined __HIP_NO_HALF_CONVERSIONS__ in hip.cmake
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
template class PADDLE_API MultiheadGPUComputeFunctor<half>;
|
|
#endif
|
|
|
|
} // namespace funcs
|
|
} // namespace phi
|