wehub-resource-sync
933 lines
35 KiB
Plaintext
933 lines
35 KiB
Plaintext
// Copyright (c) 2024 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.
|
|
|
|
#include <math.h>
|
|
|
|
#include <limits>
|
|
#include <string>
|
|
#include <vector>
|
|
#include "paddle/phi/common/memory_utils.h"
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
#include "paddle/phi/backends/dynload/cusparse.h"
|
|
#endif
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/core/utils/data_type.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
namespace phi {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
template <typename T>
|
|
__forceinline__ __device__ T CudaShuffleXorSync(unsigned mask,
|
|
T val,
|
|
int width = warpSize) {
|
|
return __shfl_xor_sync(mask, val, width);
|
|
}
|
|
|
|
template <typename T, int batch_size, int warp_size>
|
|
__device__ __forceinline__ void WarpReduceSum(T* sum) {
|
|
#pragma unroll
|
|
for (int offset = warp_size / 2; offset > 0; offset /= 2) {
|
|
#pragma unroll
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
T sum_val = CudaShuffleXorSync(0xFFFFFFFF, sum[i], offset);
|
|
sum[i] = sum[i] + sum_val;
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int batch_size, int warp_size>
|
|
__device__ __forceinline__ void WarpReduceMax(T* sum) {
|
|
#pragma unroll
|
|
for (int offset = warp_size / 2; offset > 0; offset /= 2) {
|
|
#pragma unroll
|
|
for (int i = 0; i < batch_size; ++i) {
|
|
T max_val = CudaShuffleXorSync(0xFFFFFFFF, sum[i], offset);
|
|
sum[i] = max(sum[i], max_val);
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int BlockSize, int BlockNnzMax>
|
|
__global__ void BlockSparseSoftmaxForward(T* softmax,
|
|
const T* src,
|
|
T scale,
|
|
const T* kp_mask,
|
|
const T* attn_mask,
|
|
const int* layout_rowptr,
|
|
const int* layout_colindex,
|
|
int num_rows) {
|
|
// current thread related info
|
|
const int WarpSize = 32;
|
|
const int cur_row = blockIdx.x * blockDim.y + threadIdx.y;
|
|
if (cur_row < num_rows) {
|
|
const int cur_block_row = cur_row / BlockSize;
|
|
const int cur_block_nnz =
|
|
layout_rowptr[cur_block_row + 1] - layout_rowptr[cur_block_row];
|
|
|
|
T srcdata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize] = {0};
|
|
T attndata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize] = {0};
|
|
|
|
// read tensor data, attn mask
|
|
const int iter = (cur_block_nnz + WarpSize - 1) / WarpSize;
|
|
const T* srcptr = src + layout_rowptr[cur_block_row];
|
|
|
|
const T* attnptr = (attn_mask == nullptr)
|
|
? nullptr
|
|
: (attn_mask + cur_block_row * num_rows);
|
|
// the column start index in current row
|
|
const int* colindex = layout_colindex + layout_rowptr[cur_block_row];
|
|
for (int j = 0; j < iter; j++) {
|
|
int cur_block_col = j * WarpSize + threadIdx.x;
|
|
int cur_reg_index = j;
|
|
if (cur_block_col < cur_block_nnz) {
|
|
// read kp mask
|
|
T cur_kp_mask;
|
|
if ((kp_mask != nullptr) && std::abs(kp_mask[colindex[cur_block_col]]) <
|
|
std::numeric_limits<T>::epsilon()) {
|
|
cur_kp_mask = -std::numeric_limits<T>::infinity();
|
|
} else {
|
|
cur_kp_mask = 0;
|
|
}
|
|
// do mask operation
|
|
if ((attnptr != nullptr) && std::abs(attnptr[colindex[cur_block_col]]) <
|
|
std::numeric_limits<T>::epsilon()) {
|
|
srcdata[cur_reg_index] =
|
|
-std::numeric_limits<T>::infinity() * scale + cur_kp_mask;
|
|
} else {
|
|
srcdata[cur_reg_index] = scale * srcptr[cur_block_col] + cur_kp_mask;
|
|
}
|
|
} else {
|
|
srcdata[cur_reg_index] = -std::numeric_limits<T>::infinity();
|
|
}
|
|
}
|
|
|
|
// max value
|
|
T max_value = srcdata[0];
|
|
const int kIteration =
|
|
(cur_block_nnz * BlockSize + WarpSize - 1) / WarpSize;
|
|
#pragma unroll
|
|
for (int it = 1; it < kIteration; ++it) {
|
|
max_value = (max_value > srcdata[it]) ? max_value : srcdata[it];
|
|
}
|
|
WarpReduceMax<T, 1, WarpSize>(&max_value);
|
|
|
|
// exp sum
|
|
T sum = 0;
|
|
#pragma unroll
|
|
for (int it = 0; it < kIteration; ++it) {
|
|
srcdata[it] = std::exp(srcdata[it] - max_value);
|
|
sum += srcdata[it];
|
|
}
|
|
WarpReduceSum<T, 1, WarpSize>(&sum);
|
|
|
|
// compute softmax and write out
|
|
T* softmaxptr = softmax + layout_rowptr[cur_block_row];
|
|
for (int j = 0; j < iter; j++) {
|
|
int cur_block_col = j * WarpSize + threadIdx.x;
|
|
int cur_reg_index = j;
|
|
if (cur_block_col < cur_block_nnz) {
|
|
softmaxptr[cur_block_col] = srcdata[cur_reg_index] / sum;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
template <typename T, int BlockSize, int BlockNnzMax>
|
|
__global__ void BlockSparseSoftmaxBackward(T* dst,
|
|
const T* grad,
|
|
const T* src,
|
|
T scale,
|
|
const int* layout_rowptr,
|
|
const int* layout_colindex,
|
|
int num_rows) {
|
|
// current thread related info
|
|
const int WarpSize = 32;
|
|
const int cur_row = blockIdx.x * blockDim.y + threadIdx.y;
|
|
if (cur_row < num_rows) {
|
|
const int cur_block_row = cur_row / BlockSize;
|
|
const int cur_block_nnz =
|
|
layout_rowptr[cur_block_row + 1] - layout_rowptr[cur_block_row];
|
|
|
|
T srcdata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize];
|
|
T graddata[(BlockSize * BlockNnzMax + WarpSize - 1) / WarpSize];
|
|
|
|
// read tensor data, attn mask
|
|
const int iter = (cur_block_nnz + WarpSize - 1) / WarpSize;
|
|
const T* srcptr = src + layout_rowptr[cur_block_row];
|
|
const T* gradptr = grad + layout_rowptr[cur_block_row];
|
|
for (int j = 0; j < iter; j++) {
|
|
int cur_block_col = j * WarpSize + threadIdx.x;
|
|
int cur_reg_index = j;
|
|
if (cur_block_col < cur_block_nnz) {
|
|
srcdata[cur_reg_index] = srcptr[cur_block_col];
|
|
graddata[cur_reg_index] = gradptr[cur_block_col];
|
|
} else {
|
|
srcdata[cur_reg_index] = 0;
|
|
graddata[cur_reg_index] = 0;
|
|
}
|
|
}
|
|
|
|
T sum = 0;
|
|
const int kIteration =
|
|
(cur_block_nnz * BlockSize + WarpSize - 1) / WarpSize;
|
|
#pragma unroll
|
|
for (int it = 0; it < kIteration; ++it) {
|
|
sum += srcdata[it] * graddata[it];
|
|
}
|
|
WarpReduceSum<T, 1, WarpSize>(&sum);
|
|
|
|
// compute softmax and write out
|
|
T* dstptr = dst + layout_rowptr[cur_block_row];
|
|
for (int j = 0; j < iter; j++) {
|
|
int cur_block_col = j * WarpSize + threadIdx.x;
|
|
int cur_reg_index = j;
|
|
if (cur_block_col < cur_block_nnz) {
|
|
dstptr[cur_block_col] =
|
|
scale * srcdata[cur_reg_index] * (graddata[cur_reg_index] - sum);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
/*
|
|
input: sparse C in CSR format (num_rows,num_rows)
|
|
output: sparse C after softmax operation
|
|
*/
|
|
template <typename DeviceContext, typename T>
|
|
void SparseSoftmaxForward(const GPUContext& dev_ctx,
|
|
const DenseTensor* offset,
|
|
const DenseTensor* columns,
|
|
DenseTensor* input,
|
|
DenseTensor* output,
|
|
const int blocksize,
|
|
const int num_rows,
|
|
const int num_cols,
|
|
const DenseTensor* key_padding_mask,
|
|
const DenseTensor* attn_mask) {
|
|
const int* offset_data = offset->data<int>();
|
|
const int* columns_data = columns->data<int>();
|
|
T* input_data = input->data<T>();
|
|
T* output_data = output->data<T>();
|
|
// Add mask
|
|
const T* key_padding_mask_data =
|
|
(key_padding_mask != nullptr) ? key_padding_mask->data<T>() : nullptr;
|
|
const T* attn_mask_data =
|
|
(attn_mask != nullptr) ? attn_mask->data<T>() : nullptr;
|
|
|
|
const int block_size = 1;
|
|
dim3 blocks(32, 4, 1);
|
|
int grid = (num_rows * block_size + 3) / 4;
|
|
T scaling = static_cast<T>(1.0) / sqrt(static_cast<T>(num_cols));
|
|
|
|
if (num_cols <= 4) {
|
|
BlockSparseSoftmaxForward<T, block_size, 4>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 4 && num_cols <= 8) {
|
|
BlockSparseSoftmaxForward<T, block_size, 8>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 8 && num_cols <= 16) {
|
|
BlockSparseSoftmaxForward<T, block_size, 16>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 16 && num_cols <= 32) {
|
|
BlockSparseSoftmaxForward<T, block_size, 32>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 32 && num_cols <= 64) {
|
|
BlockSparseSoftmaxForward<T, block_size, 64>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 64 && num_cols <= 128) {
|
|
BlockSparseSoftmaxForward<T, block_size, 128>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 128 && num_cols <= 256) {
|
|
BlockSparseSoftmaxForward<T, block_size, 256>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 256 && num_cols <= 512) {
|
|
BlockSparseSoftmaxForward<T, block_size, 512>
|
|
<<<grid, blocks>>>(output_data,
|
|
input_data,
|
|
scaling,
|
|
key_padding_mask_data,
|
|
attn_mask_data,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"The head_dim of query in sparse_attention op should less or equal "
|
|
"512"));
|
|
}
|
|
}
|
|
|
|
template <typename DeviceContext, typename T>
|
|
void SparseSoftmaxBackward(const GPUContext& dev_ctx,
|
|
const DenseTensor* offset,
|
|
const DenseTensor* columns,
|
|
DenseTensor* dx,
|
|
const DenseTensor* dout,
|
|
const DenseTensor* out,
|
|
const int blocksize,
|
|
const int num_rows,
|
|
const int num_cols) {
|
|
const int* offset_data = offset->data<int>();
|
|
const int* columns_data = columns->data<int>();
|
|
T* dx_data = dx->data<T>();
|
|
const T* dout_data = dout->data<T>();
|
|
const T* out_data = out->data<T>();
|
|
|
|
const int block_size = 1;
|
|
dim3 blocks(32, 4, 1);
|
|
int grid = (num_rows * block_size + 3) / 4;
|
|
T scaling = static_cast<T>(1.0) / sqrt(static_cast<T>(num_cols));
|
|
|
|
if (num_cols <= 4) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 4><<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 4 && num_cols <= 8) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 8><<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 8 && num_cols <= 16) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 16>
|
|
<<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 16 && num_cols <= 32) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 32>
|
|
<<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 32 && num_cols <= 64) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 64>
|
|
<<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 64 && num_cols <= 128) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 128>
|
|
<<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 128 && num_cols <= 256) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 256>
|
|
<<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else if (num_cols > 256 && num_cols <= 512) {
|
|
BlockSparseSoftmaxBackward<T, block_size, 512>
|
|
<<<grid, blocks>>>(dx_data,
|
|
dout_data,
|
|
out_data,
|
|
scaling,
|
|
offset_data,
|
|
columns_data,
|
|
num_rows);
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"The head_dim of query in sparse_attention op should less or equal "
|
|
"512"));
|
|
}
|
|
}
|
|
|
|
inline cudaDataType_t GetGpuType(const DataType data_type) {
|
|
if (data_type == DataType::FLOAT32) {
|
|
return CUDA_R_32F;
|
|
} else if (data_type == DataType::FLOAT64) {
|
|
return CUDA_R_64F;
|
|
} else {
|
|
PADDLE_THROW(common::errors::InvalidArgument(
|
|
"Not support tensor type in sparse_attention OP: %s",
|
|
phi::DataTypeToString(data_type)));
|
|
}
|
|
}
|
|
|
|
inline cusparseOperation_t GetTransposeOperation(const bool transpose) {
|
|
if (transpose) {
|
|
return CUSPARSE_OPERATION_TRANSPOSE;
|
|
} else {
|
|
return CUSPARSE_OPERATION_NON_TRANSPOSE;
|
|
}
|
|
}
|
|
|
|
void CusparseDestroy(cusparseDnMatDescr_t* dn_mat_first,
|
|
cusparseDnMatDescr_t* dn_mat_second,
|
|
cusparseSpMatDescr_t* sp_mat) {
|
|
phi::dynload::cusparseDestroyDnMat(*dn_mat_first);
|
|
phi::dynload::cusparseDestroyDnMat(*dn_mat_second);
|
|
phi::dynload::cusparseDestroySpMat(*sp_mat);
|
|
}
|
|
#endif
|
|
|
|
/*
|
|
input: dense A (num_rows,num_cols), dense B (num_rows,num_cols)
|
|
output: sparse C in CSR format (num_rows,num_rows)
|
|
*/
|
|
template <typename DeviceContext, typename T>
|
|
void DotSdd(const GPUContext& dev_ctx,
|
|
const DenseTensor* a,
|
|
const DenseTensor* b,
|
|
const DenseTensor* c_offset,
|
|
const DenseTensor* c_columns,
|
|
DenseTensor* c_value,
|
|
const int num_rows,
|
|
const int num_cols,
|
|
const bool a_transpose,
|
|
const bool b_transpose) {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
const T* a_data = a->data<T>();
|
|
const T* b_data = b->data<T>();
|
|
const int* c_offset_data = c_offset->data<int>();
|
|
const int* c_columns_data = c_columns->data<int>();
|
|
T* c_value_data = c_value->data<T>();
|
|
|
|
cudaDataType_t gpu_type = GetGpuType(c_value->dtype());
|
|
cusparseHandle_t handle = nullptr;
|
|
cusparseDnMatDescr_t mat_a, mat_b;
|
|
cusparseSpMatDescr_t mat_c;
|
|
phi::dynload::cusparseCreate(&handle);
|
|
|
|
// Create dense matrix A
|
|
phi::dynload::cusparseCreateDnMat(&mat_a,
|
|
num_rows,
|
|
num_cols,
|
|
num_cols,
|
|
const_cast<T*>(a_data),
|
|
gpu_type,
|
|
CUSPARSE_ORDER_ROW);
|
|
// Create dense matrix B
|
|
phi::dynload::cusparseCreateDnMat(&mat_b,
|
|
num_rows,
|
|
num_cols,
|
|
num_cols,
|
|
const_cast<T*>(b_data),
|
|
gpu_type,
|
|
CUSPARSE_ORDER_ROW);
|
|
// Create sparse matrix C in CSR format
|
|
int64_t c_nnz = c_columns->numel();
|
|
phi::dynload::cusparseCreateCsr(&mat_c,
|
|
num_rows,
|
|
num_rows,
|
|
c_nnz,
|
|
const_cast<int*>(c_offset_data),
|
|
const_cast<int*>(c_columns_data),
|
|
c_value_data,
|
|
CUSPARSE_INDEX_32I,
|
|
CUSPARSE_INDEX_32I,
|
|
CUSPARSE_INDEX_BASE_ZERO,
|
|
gpu_type);
|
|
|
|
T alpha = 1;
|
|
T beta = 0;
|
|
|
|
size_t buffer_size = 0;
|
|
phi::dynload::cusparseSDDMM_bufferSize(handle,
|
|
GetTransposeOperation(a_transpose),
|
|
GetTransposeOperation(b_transpose),
|
|
&alpha,
|
|
mat_a,
|
|
mat_b,
|
|
&beta,
|
|
mat_c,
|
|
gpu_type,
|
|
CUSPARSE_SDDMM_ALG_DEFAULT,
|
|
&buffer_size);
|
|
auto d_buffer_ptr = phi::memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
buffer_size,
|
|
phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
|
|
void* d_buffer = static_cast<void*>(d_buffer_ptr->ptr());
|
|
|
|
phi::dynload::cusparseSDDMM(handle,
|
|
GetTransposeOperation(a_transpose),
|
|
GetTransposeOperation(b_transpose),
|
|
&alpha,
|
|
mat_a,
|
|
mat_b,
|
|
&beta,
|
|
mat_c,
|
|
gpu_type,
|
|
CUSPARSE_SDDMM_ALG_DEFAULT,
|
|
d_buffer);
|
|
|
|
CusparseDestroy(&mat_a, &mat_b, &mat_c);
|
|
phi::dynload::cusparseDestroy(handle);
|
|
#else
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"DotSdd use cusparseSDDMM, which is supported "
|
|
"from CUDA 11.3"));
|
|
#endif
|
|
}
|
|
|
|
/*
|
|
input: sparse A in CSR format (num_rows,num_rows), dense B (num_rows,num_cols)
|
|
output: dense C (num_rows,num_cols)
|
|
*/
|
|
template <typename DeviceContext, typename T>
|
|
void DotDsd(const GPUContext& dev_ctx,
|
|
const DenseTensor* a_offset,
|
|
const DenseTensor* a_columns,
|
|
const DenseTensor* a_value,
|
|
const DenseTensor* b,
|
|
DenseTensor* c,
|
|
const int num_rows,
|
|
const int num_cols,
|
|
const bool a_transpose,
|
|
const bool b_transpose) {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
const int* a_offset_data = a_offset->data<int>();
|
|
const int* a_columns_data = a_columns->data<int>();
|
|
const T* a_value_data = a_value->data<T>();
|
|
const T* b_data = b->data<T>();
|
|
T* c_data = c->data<T>();
|
|
|
|
cudaDataType_t gpu_type = GetGpuType(c->dtype());
|
|
cusparseHandle_t handle = nullptr;
|
|
cusparseSpMatDescr_t mat_a;
|
|
cusparseDnMatDescr_t mat_b, mat_c;
|
|
phi::dynload::cusparseCreate(&handle);
|
|
|
|
// Create sparse matrix A in CSR format
|
|
int64_t a_nnz = a_columns->numel();
|
|
phi::dynload::cusparseCreateCsr(&mat_a,
|
|
num_rows,
|
|
num_rows,
|
|
a_nnz,
|
|
const_cast<int*>(a_offset_data),
|
|
const_cast<int*>(a_columns_data),
|
|
const_cast<T*>(a_value_data),
|
|
CUSPARSE_INDEX_32I,
|
|
CUSPARSE_INDEX_32I,
|
|
CUSPARSE_INDEX_BASE_ZERO,
|
|
gpu_type);
|
|
|
|
// Create dense matrix B
|
|
phi::dynload::cusparseCreateDnMat(&mat_b,
|
|
num_rows,
|
|
num_cols,
|
|
num_cols,
|
|
const_cast<T*>(b_data),
|
|
gpu_type,
|
|
CUSPARSE_ORDER_ROW);
|
|
// Create dense matrix C
|
|
phi::dynload::cusparseCreateDnMat(&mat_c,
|
|
num_rows,
|
|
num_cols,
|
|
num_cols,
|
|
c_data,
|
|
gpu_type,
|
|
CUSPARSE_ORDER_ROW);
|
|
|
|
T alpha = 1;
|
|
T beta = 0;
|
|
|
|
size_t buffer_size = 0;
|
|
// allocate an external buffer if needed
|
|
phi::dynload::cusparseSpMM_bufferSize(handle,
|
|
GetTransposeOperation(a_transpose),
|
|
GetTransposeOperation(b_transpose),
|
|
&alpha,
|
|
mat_a,
|
|
mat_b,
|
|
&beta,
|
|
mat_c,
|
|
gpu_type,
|
|
CUSPARSE_SPMM_ALG_DEFAULT,
|
|
&buffer_size);
|
|
auto d_buffer_ptr = phi::memory_utils::Alloc(
|
|
dev_ctx.GetPlace(),
|
|
buffer_size,
|
|
phi::Stream(reinterpret_cast<phi::StreamId>(dev_ctx.stream())));
|
|
void* d_buffer = static_cast<void*>(d_buffer_ptr->ptr());
|
|
|
|
phi::dynload::cusparseSpMM(handle,
|
|
GetTransposeOperation(a_transpose),
|
|
GetTransposeOperation(b_transpose),
|
|
&alpha,
|
|
mat_a,
|
|
mat_b,
|
|
&beta,
|
|
mat_c,
|
|
gpu_type,
|
|
CUSPARSE_SPMM_ALG_DEFAULT,
|
|
d_buffer);
|
|
|
|
CusparseDestroy(&mat_b, &mat_c, &mat_a);
|
|
phi::dynload::cusparseDestroy(handle);
|
|
#else
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"DotDsd use cusparseSpMM, which is supported "
|
|
"from CUDA 11.0"));
|
|
#endif
|
|
}
|
|
|
|
std::vector<DenseTensor> GetSplitTensor(DenseTensor* input) {
|
|
auto dims = input->dims();
|
|
int batch_size = dims[0];
|
|
int num_heads = dims[1];
|
|
std::vector<int> new_dims(dims.size() - 1);
|
|
new_dims[0] = batch_size * num_heads;
|
|
for (int i = 1; i < new_dims.size(); i++) {
|
|
new_dims[i] = dims[i + 1];
|
|
}
|
|
input->Resize(new_dims);
|
|
return input->Split(1, 0);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SparseAttentionCUDAKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& offset,
|
|
const DenseTensor& columns,
|
|
const optional<DenseTensor>& key_padding_mask,
|
|
const optional<DenseTensor>& attn_mask,
|
|
DenseTensor* out,
|
|
DenseTensor* sparse_dot_sdd,
|
|
DenseTensor* softmax) {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
auto query = q;
|
|
auto key = k;
|
|
auto value = v;
|
|
auto output_ptr = out;
|
|
dev_ctx.template Alloc<T>(out);
|
|
auto sparse_dot_sdd_ptr = sparse_dot_sdd;
|
|
dev_ctx.template Alloc<T>(sparse_dot_sdd);
|
|
auto softmax_ptr = softmax;
|
|
dev_ctx.template Alloc<T>(softmax);
|
|
|
|
auto output = *output_ptr;
|
|
auto result_sdd = *sparse_dot_sdd_ptr;
|
|
auto result_softmax = *softmax_ptr;
|
|
|
|
auto query_dims = query.dims();
|
|
int batch_size = query_dims[0];
|
|
int num_heads = query_dims[1];
|
|
int M = query_dims[2];
|
|
int N = query_dims[3];
|
|
|
|
DenseTensor q2 = q;
|
|
DenseTensor k2 = k;
|
|
DenseTensor v2 = v;
|
|
DenseTensor offset2 = offset;
|
|
DenseTensor columns2 = columns;
|
|
std::vector<DenseTensor> query_lists = GetSplitTensor(&q2);
|
|
std::vector<DenseTensor> key_lists = GetSplitTensor(&k2);
|
|
std::vector<DenseTensor> value_lists = GetSplitTensor(&v2);
|
|
std::vector<DenseTensor> offset_lists = GetSplitTensor(&offset2);
|
|
std::vector<DenseTensor> columns_lists = GetSplitTensor(&columns2);
|
|
std::vector<DenseTensor> result_sdd_lists = GetSplitTensor(&result_sdd);
|
|
std::vector<DenseTensor> result_softmax_lists =
|
|
GetSplitTensor(&result_softmax);
|
|
std::vector<DenseTensor> output_lists = GetSplitTensor(&output);
|
|
|
|
const int iter_num = batch_size * num_heads;
|
|
for (int i = 0; i < iter_num; i++) {
|
|
DotSdd<Context, T>(dev_ctx,
|
|
&query_lists[i],
|
|
&key_lists[i],
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&result_sdd_lists[i],
|
|
M,
|
|
N,
|
|
false,
|
|
true);
|
|
|
|
if (key_padding_mask && attn_mask) {
|
|
SparseSoftmaxForward<Context, T>(
|
|
dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&result_sdd_lists[i],
|
|
&result_softmax_lists[i],
|
|
1,
|
|
M,
|
|
N,
|
|
key_padding_mask.get_ptr() + (i / num_heads) * M,
|
|
attn_mask.get_ptr());
|
|
} else if (key_padding_mask && !attn_mask.is_initialized()) {
|
|
SparseSoftmaxForward<Context, T>(
|
|
dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&result_sdd_lists[i],
|
|
&result_softmax_lists[i],
|
|
1,
|
|
M,
|
|
N,
|
|
key_padding_mask.get_ptr() + (i / num_heads) * M,
|
|
nullptr);
|
|
} else if (!key_padding_mask.is_initialized() && attn_mask) {
|
|
SparseSoftmaxForward<Context, T>(dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&result_sdd_lists[i],
|
|
&result_softmax_lists[i],
|
|
1,
|
|
M,
|
|
N,
|
|
nullptr,
|
|
attn_mask.get_ptr());
|
|
} else {
|
|
SparseSoftmaxForward<Context, T>(dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&result_sdd_lists[i],
|
|
&result_softmax_lists[i],
|
|
1,
|
|
M,
|
|
N,
|
|
nullptr,
|
|
nullptr);
|
|
}
|
|
|
|
DotDsd<Context, T>(dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&result_softmax_lists[i],
|
|
&value_lists[i],
|
|
&output_lists[i],
|
|
M,
|
|
N,
|
|
false,
|
|
false);
|
|
}
|
|
#endif
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void SparseAttentionGradCUDAKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& offset,
|
|
const DenseTensor& columns,
|
|
const DenseTensor& sparse_dot_sdd,
|
|
const DenseTensor& softmax,
|
|
const DenseTensor& out_grad,
|
|
DenseTensor* q_grad,
|
|
DenseTensor* k_grad,
|
|
DenseTensor* v_grad) {
|
|
#if defined(PADDLE_WITH_CUDA)
|
|
auto query = q;
|
|
auto key = k;
|
|
auto value = v;
|
|
|
|
auto dout = out_grad;
|
|
auto* dquery_ptr = q_grad;
|
|
auto* dkey_ptr = k_grad;
|
|
auto* dvalue_ptr = v_grad;
|
|
dev_ctx.template Alloc<T>(q_grad);
|
|
dev_ctx.template Alloc<T>(k_grad);
|
|
dev_ctx.template Alloc<T>(v_grad);
|
|
|
|
auto dquery = *dquery_ptr;
|
|
auto dkey = *dkey_ptr;
|
|
auto dvalue = *dvalue_ptr;
|
|
|
|
auto query_dims = query.dims();
|
|
int batch_size = query_dims[0];
|
|
int num_heads = query_dims[1];
|
|
int M = query_dims[2];
|
|
int N = query_dims[3];
|
|
|
|
DenseTensor q2 = q;
|
|
DenseTensor k2 = k;
|
|
DenseTensor v2 = v;
|
|
DenseTensor offset2 = offset;
|
|
DenseTensor columns2 = columns;
|
|
DenseTensor sparse_dot_sdd2 = sparse_dot_sdd;
|
|
DenseTensor softmax2 = softmax;
|
|
DenseTensor dout2 = out_grad;
|
|
std::vector<DenseTensor> query_lists = GetSplitTensor(&q2);
|
|
std::vector<DenseTensor> key_lists = GetSplitTensor(&k2);
|
|
std::vector<DenseTensor> value_lists = GetSplitTensor(&v2);
|
|
std::vector<DenseTensor> offset_lists = GetSplitTensor(&offset2);
|
|
std::vector<DenseTensor> columns_lists = GetSplitTensor(&columns2);
|
|
std::vector<DenseTensor> sparse_dot_sdd_lists =
|
|
GetSplitTensor(&sparse_dot_sdd2);
|
|
std::vector<DenseTensor> softmax_lists = GetSplitTensor(&softmax2);
|
|
std::vector<DenseTensor> dout_lists = GetSplitTensor(&dout2);
|
|
std::vector<DenseTensor> dquery_lists = GetSplitTensor(&dquery);
|
|
std::vector<DenseTensor> dkey_lists = GetSplitTensor(&dkey);
|
|
std::vector<DenseTensor> dvalue_lists = GetSplitTensor(&dvalue);
|
|
|
|
const int iter_num = batch_size * num_heads;
|
|
for (int i = 0; i < iter_num; i++) {
|
|
// dValue = transpose(result_softmax) * dOut
|
|
DotDsd<Context, T>(dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&softmax_lists[i],
|
|
&dout_lists[i],
|
|
&dvalue_lists[i],
|
|
M,
|
|
N,
|
|
true,
|
|
false);
|
|
|
|
// dSoftmax = dOut * transpose(Value)
|
|
int64_t nnz_num = columns_lists[i].numel();
|
|
DenseTensor dsoftmax;
|
|
dsoftmax.Resize({nnz_num});
|
|
dev_ctx.template Alloc<T>(&dsoftmax);
|
|
DotSdd<Context, T>(dev_ctx,
|
|
&dout_lists[i],
|
|
&value_lists[i],
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&dsoftmax,
|
|
M,
|
|
N,
|
|
false,
|
|
true);
|
|
|
|
// dSparseDotSdd = dSoftmax * softmax'(SparseDotSdd)
|
|
DenseTensor dsparse_dot_sdd;
|
|
dsparse_dot_sdd.Resize({nnz_num});
|
|
dev_ctx.template Alloc<T>(&dsparse_dot_sdd);
|
|
SparseSoftmaxBackward<Context, T>(dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&dsparse_dot_sdd,
|
|
&dsoftmax,
|
|
&softmax_lists[i],
|
|
1,
|
|
M,
|
|
N);
|
|
|
|
// dQuery = dSparseDotSdd * Key
|
|
DotDsd<Context, T>(dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&dsparse_dot_sdd,
|
|
&key_lists[i],
|
|
&dquery_lists[i],
|
|
M,
|
|
N,
|
|
false,
|
|
false);
|
|
|
|
// dKey = transpose(dSparseDotSdd) * Query
|
|
DotDsd<Context, T>(dev_ctx,
|
|
&offset_lists[i],
|
|
&columns_lists[i],
|
|
&dsparse_dot_sdd,
|
|
&query_lists[i],
|
|
&dkey_lists[i],
|
|
M,
|
|
N,
|
|
true,
|
|
false);
|
|
}
|
|
#endif
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(sparse_attention,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::SparseAttentionCUDAKernel,
|
|
float,
|
|
double) {
|
|
kernel->InputAt(3).SetDataType(phi::DataType::INT32);
|
|
kernel->InputAt(4).SetDataType(phi::DataType::INT32);
|
|
}
|
|
PD_REGISTER_KERNEL(sparse_attention_grad,
|
|
GPU,
|
|
ALL_LAYOUT,
|
|
phi::SparseAttentionGradCUDAKernel,
|
|
float,
|
|
double) {
|
|
kernel->InputAt(3).SetDataType(phi::DataType::INT32);
|
|
kernel->InputAt(4).SetDataType(phi::DataType::INT32);
|
|
}
|