wehub-resource-sync
419 lines
14 KiB
C++
419 lines
14 KiB
C++
// 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.
|
|
|
|
#pragma once
|
|
|
|
#include "paddle/phi/backends/gpu/gpu_context.h"
|
|
#include "paddle/phi/backends/gpu/gpu_launch_config.h"
|
|
#include "paddle/phi/core/enforce.h"
|
|
#include "paddle/phi/kernels/empty_kernel.h"
|
|
#include "paddle/phi/kernels/full_kernel.h"
|
|
|
|
#ifdef PADDLE_WITH_FLASHATTN
|
|
#include "paddle/phi/backends/dynload/flashattn.h"
|
|
#endif
|
|
|
|
#ifdef PADDLE_WITH_FLASHATTN_V3
|
|
#include "paddle/phi/backends/dynload/flashattnv3.h"
|
|
#endif
|
|
|
|
namespace phi {
|
|
|
|
#ifdef PADDLE_WITH_FLASHATTN
|
|
static std::pair<uint64_t, uint64_t> GenerateRNGState(
|
|
const GPUContext& dev_ctx,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const std::string& rng_name,
|
|
const int64_t batch_size,
|
|
const int64_t num_heads) {
|
|
if (fixed_seed_offset.get_ptr()) {
|
|
const int64_t* fixed_seed_offset_data =
|
|
fixed_seed_offset.get_ptr()->data<int64_t>();
|
|
uint64_t seed = static_cast<uint64_t>(fixed_seed_offset_data[0]);
|
|
uint64_t offset = static_cast<uint64_t>(fixed_seed_offset_data[1]);
|
|
return std::make_pair(seed, offset);
|
|
} else {
|
|
uint64_t inc = batch_size * num_heads * 32;
|
|
std::pair<uint64_t, uint64_t> seed_offset_pair;
|
|
if (rng_name != "") {
|
|
auto gen = phi::GetRandomSeedGenerator(rng_name);
|
|
seed_offset_pair = gen->IncrementOffset(inc);
|
|
} else {
|
|
auto* gen = dev_ctx.GetGenerator();
|
|
seed_offset_pair = gen->IncrementOffset(inc);
|
|
}
|
|
return seed_offset_pair;
|
|
}
|
|
}
|
|
|
|
static std::vector<int64_t> GetAttnMaskDims(const DenseTensor* attn_mask) {
|
|
std::vector<int64_t> mask_dim_4d;
|
|
if (attn_mask) {
|
|
const auto& origin_dims = attn_mask->dims();
|
|
auto rank = origin_dims.size();
|
|
PADDLE_ENFORCE_GE(
|
|
rank,
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"The number of dimensions of attn_mask is expected to be greater "
|
|
"or equal to 4, but received %d. The shape of attn_mask is {%s}",
|
|
rank,
|
|
origin_dims));
|
|
|
|
int64_t first_dim = 1;
|
|
for (int i = 0; i < rank - 3; i++) {
|
|
first_dim *= origin_dims[i];
|
|
}
|
|
mask_dim_4d = {first_dim,
|
|
origin_dims[rank - 3],
|
|
origin_dims[rank - 2],
|
|
origin_dims[rank - 1]};
|
|
}
|
|
return mask_dim_4d;
|
|
}
|
|
|
|
static std::vector<int64_t> GetAttnSparseMaskDims(
|
|
const DenseTensor* startend_row_indices, int64_t max_seqlen_q) {
|
|
std::vector<int64_t> mask_dim_4d;
|
|
if (startend_row_indices) {
|
|
const auto& dtype = startend_row_indices->dtype();
|
|
const auto& origin_dims = startend_row_indices->dims();
|
|
auto rank = origin_dims.size();
|
|
PADDLE_ENFORCE_EQ(
|
|
dtype,
|
|
DataType::INT32,
|
|
common::errors::InvalidArgument("dtype of startend_row_indices must be "
|
|
"int32, but received %d",
|
|
dtype));
|
|
PADDLE_ENFORCE_GE(
|
|
rank,
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"The number of dimensions of startend_row_indices is expected to "
|
|
"be greater or equal to 4, but received %d. The shape of "
|
|
"startend_row_indices is [%s]",
|
|
rank,
|
|
origin_dims));
|
|
|
|
int64_t first_dim = 1;
|
|
for (int i = 0; i < rank - 3; i++) {
|
|
first_dim *= origin_dims[i];
|
|
}
|
|
mask_dim_4d = {first_dim,
|
|
origin_dims[rank - 3],
|
|
origin_dims[rank - 2],
|
|
origin_dims[rank - 1]};
|
|
}
|
|
|
|
return mask_dim_4d;
|
|
}
|
|
|
|
struct FlashAttnParamsBase {
|
|
int version;
|
|
bool is_fwd;
|
|
|
|
int kBlockM;
|
|
int batch_size;
|
|
// for padded kernel, max_seqlen_q and seqlen_q is the same.
|
|
int64_t max_seqlen_q;
|
|
// for padded kernel, max_seqlen_k and seqlen_k is the same.
|
|
int64_t max_seqlen_k;
|
|
int num_heads;
|
|
int num_heads_k;
|
|
int head_size;
|
|
|
|
int seqlen_q_rounded;
|
|
int seqlen_k_rounded;
|
|
int head_size_rounded;
|
|
|
|
bool is_bf16;
|
|
bool is_fp8;
|
|
float softmax_scale;
|
|
std::vector<int64_t> softmax_lse_dims;
|
|
std::vector<int64_t> dpsum_dims;
|
|
|
|
bool causal;
|
|
std::vector<int64_t> mask_dims;
|
|
const DenseTensor* attn_mask_tensor;
|
|
|
|
const DenseTensor* startend_row_indices;
|
|
std::vector<int64_t> startend_row_indices_dims;
|
|
|
|
bool unpadded_lse;
|
|
int total_q;
|
|
|
|
FlashAttnParamsBase(const int _version,
|
|
const int _is_fwd,
|
|
const int _batch_size,
|
|
const int64_t _max_seqlen_q,
|
|
const int64_t _max_seqlen_k,
|
|
const int _num_heads,
|
|
const int _num_heads_k,
|
|
const int _head_size,
|
|
const float _scale,
|
|
const bool _causal,
|
|
const DataType q_dtype,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const optional<DenseTensor>& startend_row_indices,
|
|
const bool _unpadded_lse,
|
|
const int _total_q)
|
|
: version(_version),
|
|
is_fwd(_is_fwd),
|
|
batch_size(_batch_size),
|
|
max_seqlen_q(_max_seqlen_q),
|
|
max_seqlen_k(_max_seqlen_k),
|
|
num_heads(_num_heads),
|
|
num_heads_k(_num_heads_k),
|
|
head_size(_head_size),
|
|
softmax_scale(_scale),
|
|
causal(_causal),
|
|
attn_mask_tensor(attn_mask.get_ptr()),
|
|
startend_row_indices(startend_row_indices.get_ptr()),
|
|
unpadded_lse(_unpadded_lse),
|
|
total_q(_total_q) {
|
|
is_bf16 = q_dtype == DataType::BFLOAT16;
|
|
|
|
// TODO(GuoxiaWang): check q, k, v dtype
|
|
|
|
auto round_multiple = [](int64_t x, int m) {
|
|
int64_t rounded = (x + m - 1) / m * static_cast<int64_t>(m);
|
|
PADDLE_ENFORCE_LE_INT_MAX(rounded, "flash_attn rounded sequence length");
|
|
return static_cast<int>(rounded);
|
|
};
|
|
// FLAGS_flash_attn_version
|
|
if (_version == 3 && !_is_fwd) {
|
|
kBlockM = head_size <= 64 ? 128 : (head_size < 256 ? 64 : 32);
|
|
head_size_rounded = head_size <= 64 ? 64 : round_multiple(head_size, 32);
|
|
} else {
|
|
kBlockM = 128;
|
|
head_size_rounded = round_multiple(head_size, 32);
|
|
}
|
|
|
|
seqlen_q_rounded = round_multiple(max_seqlen_q, kBlockM);
|
|
seqlen_k_rounded = round_multiple(max_seqlen_k, 128);
|
|
|
|
softmax_lse_dims = unpadded_lse ? std::vector<int64_t>{num_heads, total_q}
|
|
: std::vector<int64_t>{
|
|
batch_size, num_heads, max_seqlen_q};
|
|
|
|
dpsum_dims = std::vector<int64_t>{batch_size, num_heads, seqlen_q_rounded};
|
|
|
|
if (attn_mask_tensor) {
|
|
PADDLE_ENFORCE_EQ(
|
|
attn_mask->dtype(),
|
|
q_dtype,
|
|
common::errors::InvalidArgument(
|
|
"attn_mask is expected to have the same data type with q."));
|
|
|
|
mask_dims = GetAttnMaskDims(attn_mask_tensor);
|
|
}
|
|
|
|
startend_row_indices_dims = GetAttnSparseMaskDims(
|
|
startend_row_indices ? startend_row_indices.get_ptr() : nullptr,
|
|
max_seqlen_q);
|
|
|
|
if (startend_row_indices.is_initialized()) {
|
|
PADDLE_ENFORCE_EQ(
|
|
attn_mask_tensor,
|
|
nullptr,
|
|
common::errors::InvalidArgument(
|
|
"attn_mask and attn_mask_start_row_indices cannot be "
|
|
"set at same time."));
|
|
}
|
|
}
|
|
};
|
|
|
|
template <typename T>
|
|
struct FlashAttnFwdParamsV2 : public FlashAttnParamsBase {
|
|
float dropout;
|
|
bool return_softmax;
|
|
uint64_t seed;
|
|
uint64_t offset;
|
|
DenseTensor rng_state;
|
|
DenseTensor* softmax;
|
|
DenseTensor* softmax_lse;
|
|
DenseTensor* seed_offset;
|
|
DenseTensor tile_count_semaphore;
|
|
|
|
FlashAttnFwdParamsV2(const GPUContext& dev_ctx,
|
|
const int _version,
|
|
const int _batch_size,
|
|
const int64_t _max_seqlen_q,
|
|
const int64_t _max_seqlen_k,
|
|
const int _num_heads,
|
|
const int _num_heads_k,
|
|
const int _head_size,
|
|
const float _dropout,
|
|
const float _scale,
|
|
const bool _causal,
|
|
const bool _return_softmax,
|
|
const DataType q_dtype,
|
|
const bool is_test,
|
|
const std::string& rng_name,
|
|
const optional<DenseTensor>& fixed_seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const optional<DenseTensor>& startend_row_indices,
|
|
DenseTensor* _softmax,
|
|
DenseTensor* _softmax_lse,
|
|
DenseTensor* _seed_offset,
|
|
const bool _unpadded_lse,
|
|
const int _total_q)
|
|
: FlashAttnParamsBase(_version,
|
|
/*is_fwd=*/true,
|
|
_batch_size,
|
|
_max_seqlen_q,
|
|
_max_seqlen_k,
|
|
_num_heads,
|
|
_num_heads_k,
|
|
_head_size,
|
|
_scale,
|
|
_causal,
|
|
q_dtype,
|
|
attn_mask,
|
|
startend_row_indices,
|
|
_unpadded_lse,
|
|
_total_q),
|
|
dropout(_dropout),
|
|
return_softmax(_return_softmax),
|
|
softmax(_softmax),
|
|
softmax_lse(_softmax_lse),
|
|
seed_offset(_seed_offset) {
|
|
dropout = is_test ? 0.0f : _dropout;
|
|
|
|
// (umiswing): There is no suitable kernel for uint64_t, allocate in int64_t
|
|
// with the same size.
|
|
rng_state = Empty<int64_t>(dev_ctx, {2});
|
|
|
|
if (_dropout > 0.0f) {
|
|
auto seed_offset_pair = GenerateRNGState(
|
|
dev_ctx, fixed_seed_offset, rng_name, batch_size, num_heads);
|
|
seed = seed_offset_pair.first;
|
|
offset = seed_offset_pair.second;
|
|
} else {
|
|
seed = 0;
|
|
offset = 0;
|
|
}
|
|
|
|
seed_offset->Resize({2});
|
|
int64_t* seed_offset_data =
|
|
dev_ctx.template HostAlloc<int64_t>(seed_offset);
|
|
seed_offset_data[0] = static_cast<int64_t>(seed);
|
|
seed_offset_data[1] = static_cast<int64_t>(offset);
|
|
|
|
softmax_lse->Resize(softmax_lse_dims);
|
|
dev_ctx.template Alloc<float>(softmax_lse);
|
|
|
|
if (_version == 3) {
|
|
tile_count_semaphore = Full<int>(dev_ctx, {1}, static_cast<int>(0));
|
|
}
|
|
|
|
if (return_softmax) {
|
|
PADDLE_ENFORCE_EQ(
|
|
dropout > 0.0f,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"return_softmax is only supported when dropout > 0.0"));
|
|
|
|
softmax->Resize(
|
|
{batch_size, num_heads, seqlen_q_rounded, seqlen_k_rounded});
|
|
dev_ctx.template Alloc<T>(softmax);
|
|
}
|
|
}
|
|
};
|
|
|
|
struct FlashAttnBwdParamsV2 : public FlashAttnParamsBase {
|
|
float dropout;
|
|
uint64_t seed;
|
|
uint64_t offset;
|
|
DenseTensor softmax_d;
|
|
DenseTensor dq_accum;
|
|
DenseTensor rng_state;
|
|
|
|
DenseTensor softmax_lse_log2;
|
|
DenseTensor dq_semaphore;
|
|
|
|
FlashAttnBwdParamsV2(const GPUContext& dev_ctx,
|
|
const int _version,
|
|
const int _batch_size,
|
|
const int64_t _max_seqlen_q,
|
|
const int64_t _max_seqlen_k,
|
|
const int _num_heads,
|
|
const int _num_heads_k,
|
|
const int _head_size,
|
|
const float _dropout,
|
|
const float _scale,
|
|
const bool _causal,
|
|
const DataType q_dtype,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const optional<DenseTensor>& startend_row_indices,
|
|
const int64_t* seed_offset_data,
|
|
const bool _unpadded_lse,
|
|
const int _total_q)
|
|
: FlashAttnParamsBase(_version,
|
|
/*is_fwd=*/false,
|
|
_batch_size,
|
|
_max_seqlen_q,
|
|
_max_seqlen_k,
|
|
_num_heads,
|
|
_num_heads_k,
|
|
_head_size,
|
|
_scale,
|
|
_causal,
|
|
q_dtype,
|
|
attn_mask,
|
|
startend_row_indices,
|
|
_unpadded_lse,
|
|
_total_q),
|
|
dropout(_dropout) {
|
|
seed = static_cast<uint64_t>(seed_offset_data[0]);
|
|
offset = static_cast<uint64_t>(seed_offset_data[1]);
|
|
|
|
// (umiswing): There is no suitable kernel for uint64_t, allocate in int64_t
|
|
// with the same size.
|
|
rng_state = Empty<int64_t>(dev_ctx, {2});
|
|
|
|
// gradient of softmax_lse
|
|
softmax_d = Empty<float>(dev_ctx, dpsum_dims);
|
|
|
|
if (_version == 3) {
|
|
softmax_lse_log2 = Empty<float>(dev_ctx, softmax_lse_dims);
|
|
dq_semaphore = Empty<int>(
|
|
dev_ctx,
|
|
{(max_seqlen_q + kBlockM - 1) / kBlockM, batch_size, num_heads});
|
|
}
|
|
|
|
// an internal gradient of q, which will be further accumulated.
|
|
dq_accum = Empty<float>(
|
|
dev_ctx, {batch_size, num_heads, seqlen_q_rounded, head_size_rounded});
|
|
}
|
|
};
|
|
|
|
static void CheckFlashAttnStatus(const bool status) {
|
|
PADDLE_ENFORCE_EQ(status,
|
|
true,
|
|
common::errors::External(
|
|
"Error in Flash-Attention, detail information is: %s",
|
|
dynload::flash_attn_error()));
|
|
}
|
|
#endif
|
|
|
|
static void RaiseNotSupportedError(int version = 2) {
|
|
PADDLE_THROW(common::errors::Unimplemented(
|
|
"FlashAttention %d is unsupported, please check "
|
|
"the GPU compatibility and CUDA Version.",
|
|
version));
|
|
}
|
|
|
|
} // namespace phi
|