Skip to content

Commit cf2db1f

Browse files
zhwesky2010zhwesky2010
committed
[Sparse] optimize sparse attention
1 parent 16506d8 commit cf2db1f

File tree

6 files changed

+54
-101
lines changed

6 files changed

+54
-101
lines changed

‎paddle/fluid/platform/dynload/cusparse.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -56,6 +56,7 @@ CUSPARSE_ROUTINE_EACH(PLATFORM_DECLARE_DYNAMIC_LOAD_CUSPARSE_WRAP)
5656

5757
#if CUDA_VERSION >= 11030
5858
#define CUSPARSE_ROUTINE_EACH_R2(__macro) \
59+
__macro(cusparseSpMM_preprocess); \
5960
__macro(cusparseSDDMM_bufferSize); \
6061
__macro(cusparseSDDMM_preprocess); \
6162
__macro(cusparseSDDMM);

‎paddle/phi/backends/dynload/cusparse.h‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ CUSPARSE_ROUTINE_EACH(DECLARE_DYNAMIC_LOAD_CUSPARSE_WRAP)
6868

6969
#if CUDA_VERSION >= 11030
7070
#define CUSPARSE_ROUTINE_EACH_R2(__macro) \
71+
__macro(cusparseSpMM_preprocess); \
7172
__macro(cusparseSDDMM_bufferSize); \
7273
__macro(cusparseSDDMM_preprocess); \
7374
__macro(cusparseSDDMM);

‎paddle/phi/kernels/funcs/sparse/sparse_blas_impl.cu.h‎

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,15 @@ inline cusparseOperation_t GetTransposeOperation(const bool trans) {
4848
}
4949
}
5050

51+
inline cusparseSpMMAlg_t GetSpMMAlgorithm(const SparseCsrTensor& x) {
52+
// TODO(zhouwei): will change to 'CUSPARSE_SPMM_CSR_ALG2' when support batch
53+
return CUSPARSE_SPMM_CSR_ALG2;
54+
}
55+
56+
inline cusparseSpMMAlg_t GetSpMMAlgorithm(const SparseCooTensor& x) {
57+
return CUSPARSE_SPMM_ALG_DEFAULT;
58+
}
59+
5160
/************* SPARSE MATRIX DESCRIPTOR (COO/CSR) ************/
5261

5362
template <typename T, typename IntT>
@@ -324,7 +333,7 @@ void SparseBlas<phi::GPUContext>::SPMM(bool transa,
324333
&beta,
325334
out_descriptor.descriptor(),
326335
gpu_type,
327-
CUSPARSE_SPMM_ALG_DEFAULT,
336+
GetSpMMAlgorithm(mat_a),
328337
&buffer_size);
329338
});
330339

@@ -341,7 +350,7 @@ void SparseBlas<phi::GPUContext>::SPMM(bool transa,
341350
&beta,
342351
out_descriptor.descriptor(),
343352
gpu_type,
344-
CUSPARSE_SPMM_ALG_DEFAULT,
353+
GetSpMMAlgorithm(mat_a),
345354
tmp_buffer_ptr);
346355
});
347356
}

‎paddle/phi/kernels/sparse/gpu/fused_attention_grad_kernel.cu‎

Lines changed: 8 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -43,21 +43,14 @@ __global__ void AttnSoftmaxGpuGradKernel(const int64_t* out_crows,
4343
int row_nnz = static_cast<int>(out_crows[crow_idx + 1] - out_crows[crow_idx]);
4444
if (row_nnz == 0) return;
4545

46-
int kIteration = (row_nnz + WARP_SIZE - 1) / WARP_SIZE;
47-
T mul_result = 0;
48-
for (int i = 0; i < kIteration; ++i) {
49-
int idx = threadIdx.x + i * WARP_SIZE;
50-
if (idx >= row_nnz) break;
51-
52-
mul_result += out_values[row_first + idx] * dout_values[row_first + idx];
46+
T mul = 0;
47+
for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
48+
mul += out_values[row_first + idx] * dout_values[row_first + idx];
5349
}
54-
T sum = phi::funcs::warpReduceSum<T>(mul_result, 0xFFFFFFFF);
55-
56-
for (int i = 0; i < kIteration; ++i) {
57-
int idx = threadIdx.x + i * WARP_SIZE;
58-
if (idx >= row_nnz) break;
50+
T mul_sum = phi::funcs::warpReduceSum<T>(mul, 0xFFFFFFFF);
5951

60-
dx_values[row_first + idx] = (dout_values[row_first + idx] - sum) *
52+
for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
53+
dx_values[row_first + idx] = (dout_values[row_first + idx] - mul_sum) *
6154
out_values[row_first + idx] / scale;
6255
}
6356
}
@@ -96,8 +89,8 @@ void FusedAttentionCsrGradKernel(const Context& dev_ctx,
9689
int N = q_dim[q_rank - 1];
9790
int batch_nnz = softmax.nnz() / batch_num;
9891

99-
dim3 grid((total_row_num + 3) / 4);
100-
dim3 block(WARP_SIZE, 4);
92+
dim3 grid((total_row_num + 7) / 8);
93+
dim3 block(WARP_SIZE, 8);
10194

10295
AttnSoftmaxGpuGradKernel<T><<<grid, block, 0, dev_ctx.stream()>>>(
10396
softmax.non_zero_crows().data<int64_t>(),

‎paddle/phi/kernels/sparse/gpu/fused_attention_kernel.cu‎

Lines changed: 32 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -26,30 +26,7 @@ limitations under the License. */
2626
namespace phi {
2727
namespace sparse {
2828

29-
#define PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, size, HINT, ...) \
30-
case size: { \
31-
constexpr int HINT = size; \
32-
__VA_ARGS__(); \
33-
break; \
34-
}
35-
36-
#define VISIT_ATTN_SFOTMAX(SIZE, NAME, ...) \
37-
[&] { \
38-
const auto& __size__ = SIZE; \
39-
switch (__size__) { \
40-
PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, 1, KBufferSize, __VA_ARGS__) \
41-
PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, 2, KBufferSize, __VA_ARGS__) \
42-
PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, 3, KBufferSize, __VA_ARGS__) \
43-
PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, 4, KBufferSize, __VA_ARGS__) \
44-
PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, 8, KBufferSize, __VA_ARGS__) \
45-
PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, 12, KBufferSize, __VA_ARGS__) \
46-
PRIVATE_CASE_VISIT_ATTN_SOFTMAX(NAME, 16, KBufferSize, __VA_ARGS__) \
47-
default: \
48-
PD_THROW("function " #NAME " is not implemented for columns>512 "); \
49-
} \
50-
}()
51-
52-
template <typename T, int BufferSize>
29+
template <typename T>
5330
__global__ void AttnSoftmaxGpuKernel(const int64_t* x_crows,
5431
const int64_t* x_cols,
5532
const T* x_values,
@@ -58,7 +35,6 @@ __global__ void AttnSoftmaxGpuKernel(const int64_t* x_crows,
5835
T* out_values,
5936
int M,
6037
int total_row_num,
61-
float scale,
6238
int num_heads,
6339
int batch_nnz) {
6440
// out = exp(x-x_max) / sum(exp(x-x_max))
@@ -72,17 +48,10 @@ __global__ void AttnSoftmaxGpuKernel(const int64_t* x_crows,
7248
int row_nnz = static_cast<int>(x_crows[crow_idx + 1] - x_crows[crow_idx]);
7349
if (row_nnz == 0) return;
7450

75-
T buffer[BufferSize] = {0};
76-
int kIteration = (row_nnz + WARP_SIZE - 1) / WARP_SIZE;
77-
7851
T max_val = -std::numeric_limits<T>::infinity();
79-
for (int i = 0; i < kIteration; ++i) {
52+
for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
8053
bool mask = false;
81-
int idx = threadIdx.x + i * WARP_SIZE;
82-
if (idx >= row_nnz) break;
83-
8454
int col_idx = static_cast<int>(x_cols[row_first + idx]);
85-
8655
if (kp_mask != nullptr &&
8756
kp_mask[(cur_batch / num_heads) * M + col_idx] == 0) {
8857
mask = true;
@@ -92,37 +61,30 @@ __global__ void AttnSoftmaxGpuKernel(const int64_t* x_crows,
9261
}
9362

9463
if (!mask) {
95-
buffer[i] = x_values[row_first + idx] / scale;
96-
if (buffer[i] > max_val) {
97-
max_val = buffer[i];
64+
T val = x_values[row_first + idx];
65+
if (val > max_val) {
66+
max_val = val;
9867
}
68+
out_values[row_first + idx] = val;
69+
} else {
70+
// Note corner case: when all elements of the row are masked, result
71+
// may be wrong because of exp('-inf' - '-inf'), just ignore now.
72+
out_values[row_first + idx] = -std::numeric_limits<T>::infinity();
9973
}
10074
}
10175
T row_max_val = phi::funcs::warpReduceMax<T>(max_val, 0xFFFFFFFF);
10276

103-
auto functor = phi::funcs::CudaExpFunctor<T>();
10477
T exp_sum = 0;
105-
for (int i = 0; i < kIteration; ++i) {
106-
int idx = threadIdx.x + i * WARP_SIZE;
107-
if (idx >= row_nnz) break;
108-
109-
if (buffer[i]) {
110-
T exp = functor(buffer[i] - row_max_val);
111-
exp_sum += exp;
112-
buffer[i] = exp;
113-
}
78+
for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
79+
auto functor = phi::funcs::CudaExpFunctor<T>();
80+
T exp = functor(out_values[row_first + idx] - row_max_val);
81+
exp_sum += exp;
82+
out_values[row_first + idx] = exp;
11483
}
11584
T row_exp_sum = phi::funcs::warpReduceSum<T>(exp_sum, 0xFFFFFFFF);
11685

117-
for (int i = 0; i < kIteration; ++i) {
118-
int idx = threadIdx.x + i * WARP_SIZE;
119-
if (idx >= row_nnz) break;
120-
121-
if (buffer[i]) {
122-
out_values[row_first + idx] = buffer[i] / row_exp_sum;
123-
} else {
124-
out_values[row_first + idx] = static_cast<T>(0);
125-
}
86+
for (int idx = threadIdx.x; idx < row_nnz; idx += blockDim.x) {
87+
out_values[row_first + idx] = out_values[row_first + idx] / row_exp_sum;
12688
}
12789
}
12890

@@ -219,49 +181,36 @@ void FusedAttentionCsrKernel(
219181
"shape of 'attn_mask' must be [seq_len, seq_len]"));
220182
}
221183

222-
/* Step1: SDD Matmul, reuse */
184+
/* Step1: SDD Matmul, reuse matmul */
223185
SparseCsrTensor sdd_result;
224186
EmptyLikeCsrKernel<T, Context>(dev_ctx, sparse_mask, &sdd_result);
225187
auto sparse_blas = phi::funcs::sparse::GetSparseBlas<Context, T>(dev_ctx);
226188
sparse_blas.SDDMM(false,
227189
true,
228-
static_cast<T>(1),
190+
static_cast<T>(1 / std::sqrt(N)),
229191
query,
230192
key,
231193
static_cast<T>(0),
232194
&sdd_result);
233195

234-
/* Step2: Softmax with kp_mask/attn_mask, manualy not reuse */
235196
EmptyLikeCsrKernel<T, Context>(dev_ctx, sdd_result, softmax);
236197

237-
int buffer_size;
238-
if (M < 128) {
239-
buffer_size = (M + 32 - 1) / 32;
240-
} else {
241-
buffer_size = ((M + 128 - 1) / 128) * 4;
242-
}
243-
244-
dim3 grid((total_row_num + 3) / 4);
245-
dim3 block(WARP_SIZE, 4);
198+
dim3 grid((total_row_num + 7) / 8);
199+
dim3 block(WARP_SIZE, 8);
246200

247201
int batch_nnz = sdd_result.nnz() / batch_num;
202+
AttnSoftmaxGpuKernel<T><<<grid, block, 0, dev_ctx.stream()>>>(
203+
sdd_result.non_zero_crows().data<int64_t>(),
204+
sdd_result.non_zero_cols().data<int64_t>(),
205+
sdd_result.non_zero_elements().data<T>(),
206+
kp_mask_ptr ? kp_mask_ptr->data<T>() : nullptr,
207+
attn_mask_ptr ? attn_mask_ptr->data<T>() : nullptr,
208+
softmax->mutable_non_zero_elements()->data<T>(),
209+
M,
210+
total_row_num,
211+
q_dim[1],
212+
batch_nnz);
248213

249-
VISIT_ATTN_SFOTMAX(buffer_size, "AttnSoftmaxGpuKernel", [&] {
250-
AttnSoftmaxGpuKernel<T, KBufferSize><<<grid, block, 0, dev_ctx.stream()>>>(
251-
sdd_result.non_zero_crows().data<int64_t>(),
252-
sdd_result.non_zero_cols().data<int64_t>(),
253-
sdd_result.non_zero_elements().data<T>(),
254-
kp_mask_ptr ? kp_mask_ptr->data<T>() : nullptr,
255-
attn_mask_ptr ? attn_mask_ptr->data<T>() : nullptr,
256-
softmax->mutable_non_zero_elements()->data<T>(),
257-
M,
258-
total_row_num,
259-
std::sqrt(N),
260-
q_dim[1],
261-
batch_nnz);
262-
});
263-
264-
/* Step3: DSD Matmul, reuse */
265214
softmax->set_dims(phi::make_ddim({q_dim[0], q_dim[1], q_dim[2], q_dim[2]}));
266215
MatmulCsrDenseKernel<T, Context>(dev_ctx, *softmax, value, out);
267216
#else

‎python/paddle/fluid/tests/unittests/test_sparse_fused_attention_op.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,7 @@ def get_cuda_version():
3737

3838
@unittest.skipIf(
3939
not core.is_compiled_with_cuda() or get_cuda_version() < 11070,
40-
"core is not compiled with CUDA and cuda version need larger than or equal to 11.3"
40+
"core is not compiled with CUDA and cuda version need larger than or equal to 11.7"
4141
)
4242
class TestSparseAttentionAPI1(unittest.TestCase):
4343

0 commit comments

Comments
 (0)