chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,932 @@
|
||||
// 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);
|
||||
}
|
||||
Reference in New Issue
Block a user