wehub-resource-sync
555 lines
25 KiB
C++
555 lines
25 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.
|
|
|
|
#include "paddle/phi/kernels/flash_attn_grad_kernel.h"
|
|
|
|
#include "paddle/phi/backends/xpu/enforce_xpu.h"
|
|
#include "paddle/phi/core/kernel_registry.h"
|
|
#include "paddle/phi/kernels/slice_kernel.h"
|
|
#include "paddle/phi/kernels/xpu/flash_attn_utils.h"
|
|
#include "xfa/flash_api.h"
|
|
namespace phi {
|
|
#define MHA_VARLEN_BWD_TYPES_AND_ARGS(T, TACCUM, TGEMM, TID) \
|
|
xpu::Context *ctx, const T *dout, const T *q, const T *k, const T *v, \
|
|
const T *out, const TACCUM *softmax_lse, T *dq, T *dk, T *dv, \
|
|
const xpu::VectorParam<TID>&lod_seqlens_q, \
|
|
const xpu::VectorParam<TID>&lod_seqlens_k, int64_t max_seqlen_q, \
|
|
int64_t max_seqlen_k, int64_t head_num, int64_t head_num_k, \
|
|
int64_t head_dim, const float softmax_scale, const float p_dropout, \
|
|
int seed, const bool is_causal, const TACCUM *attn_mask, \
|
|
const TACCUM *bias, const float *q_maxptr, const float *k_maxptr, \
|
|
const float *v_maxptr, const float *o_maxptr, float *dq_maxptr, \
|
|
float *dk_maxptr, float *dv_maxptr, const float *do_maxptr, \
|
|
const bool is_qkv_fusion, const bool is_dqkv_fusion, \
|
|
const int64_t qkv_layout, const float *alibi_slopes, \
|
|
const std::vector<int64_t>&alibi_slopes_shape, int window_size_left, \
|
|
int window_size_right, int64_t v_head_dim, \
|
|
const int *downstart_row_indices_data, \
|
|
const int *downend_row_indices_data, \
|
|
const int *upstart_row_indices_data, const int *upend_row_indices_data, \
|
|
const int flash_mask_head_num, int *flashmask_maxmin, \
|
|
XPUStream side_stream
|
|
|
|
#define MHA_VARLEN_BWD_ARGS \
|
|
ctx, dout, q, k, v, out, softmax_lse, dq, dk, dv, lod_seqlens_q, \
|
|
lod_seqlens_k, max_seqlen_q, max_seqlen_k, head_num, head_num_k, \
|
|
head_dim, softmax_scale, p_dropout, seed, is_causal, attn_mask, bias, \
|
|
q_maxptr, k_maxptr, v_maxptr, o_maxptr, dq_maxptr, dk_maxptr, dv_maxptr, \
|
|
do_maxptr, is_qkv_fusion, is_dqkv_fusion, qkv_layout, alibi_slopes, \
|
|
alibi_slopes_shape, window_size_left, window_size_right, v_head_dim, \
|
|
downstart_row_indices_data, downend_row_indices_data, \
|
|
upstart_row_indices_data, upend_row_indices_data, flash_mask_head_num, \
|
|
flashmask_maxmin, side_stream
|
|
|
|
template <typename T, typename TACCUM, typename TGEMM, typename TID>
|
|
int mha_varlen_bwd_wrapper(
|
|
MHA_VARLEN_BWD_TYPES_AND_ARGS(T, TACCUM, TGEMM, TID)) {
|
|
PADDLE_THROW(
|
|
"Unsupported template params combination for mha_varlen_bwd, should not "
|
|
"reach here.");
|
|
}
|
|
|
|
#define DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(T, TACCUM, TGEMM, TID) \
|
|
template <> \
|
|
int mha_varlen_bwd_wrapper<T, TACCUM, TGEMM, TID>( \
|
|
MHA_VARLEN_BWD_TYPES_AND_ARGS(T, TACCUM, TGEMM, TID)) { \
|
|
return baidu::xpu::xfa::mha_varlen_bwd<T, TACCUM, TGEMM, TID>( \
|
|
MHA_VARLEN_BWD_ARGS); \
|
|
}
|
|
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(float, float, tfloat32, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(float, float, float, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, float, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, XPUTypeFP16, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, tfloat32, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16,
|
|
XPUTypeFP16,
|
|
XPUTypeFP16,
|
|
int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16, float, tfloat32, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16, float, float, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16, float, XPUTypeFP16, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16,
|
|
XPUTypeFP16,
|
|
XPUTypeFP16,
|
|
int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, XPUTypeBF16, int);
|
|
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16,
|
|
XPUTypeFP16,
|
|
XPUTypeBF16,
|
|
int);
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnGradKernelBase(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const api::VectorParam<int>& lod_seqlen_q,
|
|
const api::VectorParam<int>& lod_seqlen_k,
|
|
const DenseTensor& out,
|
|
const DenseTensor& softmax_lse,
|
|
const DenseTensor& seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const optional<DenseTensor>& startend_row_indices,
|
|
const DenseTensor& dout,
|
|
const int64_t batch_size,
|
|
const Scalar& max_seqlen_q_,
|
|
const Scalar& max_seqlen_k_,
|
|
const int64_t num_heads,
|
|
const int64_t num_heads_k,
|
|
const int64_t head_size,
|
|
const int64_t head_size_v,
|
|
float scale,
|
|
float dropout,
|
|
bool causal,
|
|
DenseTensor* dq,
|
|
DenseTensor* dk,
|
|
DenseTensor* dv) {
|
|
xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());
|
|
|
|
using XPUType = typename XPUTypeTrait<T>::Type;
|
|
using XPUTypeFP16 = typename XPUTypeTrait<phi::float16>::Type;
|
|
const XPUType* q_data = reinterpret_cast<const XPUType*>(q.data<T>());
|
|
const XPUType* k_data = reinterpret_cast<const XPUType*>(k.data<T>());
|
|
const XPUType* v_data = reinterpret_cast<const XPUType*>(v.data<T>());
|
|
const XPUType* out_data = reinterpret_cast<const XPUType*>(out.data<T>());
|
|
const float* softmax_lse_data = softmax_lse.data<float>();
|
|
const XPUType* dout_data = reinterpret_cast<const XPUType*>(dout.data<T>());
|
|
|
|
float real_scale = scale == 0.0f ? 1.0f / std::sqrt(head_size) : scale;
|
|
|
|
const float* bias_data = nullptr;
|
|
int64_t fa_layout = AttnQKVLayout_t::ATTN_BLHD;
|
|
DenseTensor downstart_row_indices, upend_row_indices, downend_row_indices,
|
|
upstart_row_indices;
|
|
void *downstart_row_indices_data = nullptr, *upend_row_indices_data = nullptr,
|
|
*downend_row_indices_data = nullptr, *upstart_row_indices_data = nullptr;
|
|
bool is_flashmask = startend_row_indices.get_ptr() != nullptr;
|
|
XPUStream flashmask_stream;
|
|
if (is_flashmask) {
|
|
xpu_stream_create(&flashmask_stream);
|
|
PADDLE_ENFORCE_EQ(
|
|
startend_row_indices->dims().size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"flashmask_attention receive startend_row_indices with dim "
|
|
"[batch_size, num_heads,seq_len, mask_bounds]"));
|
|
PADDLE_ENFORCE_EQ(startend_row_indices->dims()[3] == 1 ||
|
|
startend_row_indices->dims()[3] == 2 ||
|
|
startend_row_indices->dims()[3] == 4,
|
|
true,
|
|
common::errors::InvalidArgument(
|
|
"flashmask_attention startend_row_indices "
|
|
"mask_bounds must in [1,2,4]"));
|
|
downstart_row_indices =
|
|
phi::Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {0}, {1});
|
|
downstart_row_indices_data = downstart_row_indices.data();
|
|
if (startend_row_indices->dims()[3] == 2) {
|
|
if (!causal) {
|
|
upend_row_indices = phi::Slice<int32_t>(
|
|
dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
|
|
upend_row_indices_data = upend_row_indices.data();
|
|
} else {
|
|
downend_row_indices = phi::Slice<int32_t>(
|
|
dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
|
|
downend_row_indices_data = downend_row_indices.data();
|
|
}
|
|
} else if (startend_row_indices->dims()[3] == 4) {
|
|
upend_row_indices = phi::Slice<int32_t>(
|
|
dev_ctx, startend_row_indices.get(), {3}, {3}, {4});
|
|
upend_row_indices_data = upend_row_indices.data();
|
|
downend_row_indices = phi::Slice<int32_t>(
|
|
dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
|
|
downend_row_indices_data = downend_row_indices.data();
|
|
upstart_row_indices = phi::Slice<int32_t>(
|
|
dev_ctx, startend_row_indices.get(), {3}, {2}, {3});
|
|
upstart_row_indices_data = upstart_row_indices.data();
|
|
}
|
|
} else if (attn_mask.get_ptr() != nullptr) {
|
|
const auto& mask_dims = attn_mask->dims();
|
|
if (mask_dims.size() == 3 || (mask_dims[1] == 1 && mask_dims.size() == 4)) {
|
|
fa_layout |= AttnQKVLayout_t::BIAS_BLL;
|
|
} else {
|
|
PADDLE_ENFORCE_EQ(mask_dims.size(),
|
|
4,
|
|
common::errors::InvalidArgument(
|
|
"flash_attn_bwd requires mask's shape "
|
|
"like [b,l,l] or [b, h, l, l]"));
|
|
}
|
|
if (attn_mask->dtype() == DataType::FLOAT32) {
|
|
bias_data = attn_mask->data<float>();
|
|
} else if (attn_mask->dtype() == DataType::FLOAT16 ||
|
|
attn_mask->dtype() == DataType::BFLOAT16) {
|
|
float* bias_tmp = RAII_GUARD.alloc_l3_or_gm<float>(attn_mask->numel());
|
|
int r = xpu::cast<XPUType, float>(
|
|
dev_ctx.x_context(),
|
|
reinterpret_cast<const XPUType*>(attn_mask->data<T>()),
|
|
bias_tmp,
|
|
attn_mask->numel());
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
|
|
bias_data = bias_tmp;
|
|
} else {
|
|
errors::Unimplemented(
|
|
"Unsupported dtype for attention_mask in xpu flash attention, only "
|
|
"float32, float16 and "
|
|
"bfloat16 are supported.");
|
|
}
|
|
}
|
|
// output
|
|
XPUType* dq_data = reinterpret_cast<XPUType*>(dq->data<T>());
|
|
XPUType* dk_data = reinterpret_cast<XPUType*>(dk->data<T>());
|
|
XPUType* dv_data = reinterpret_cast<XPUType*>(dv->data<T>());
|
|
|
|
int64_t max_seqlen_q = max_seqlen_q_.to<int64_t>();
|
|
int64_t max_seqlen_k = max_seqlen_k_.to<int64_t>();
|
|
|
|
// get seed offset
|
|
const int64_t* seed_offset_data = seed_offset.data<int64_t>();
|
|
int fa_tgemm = get_flash_attn_tgemm<XPUType>();
|
|
auto flash_attention_grad_kernel =
|
|
mha_varlen_bwd_wrapper<XPUType, float, tfloat32, int>;
|
|
if (fa_tgemm == XPU_FA_DTYPE::FA_FLOAT) {
|
|
flash_attention_grad_kernel =
|
|
mha_varlen_bwd_wrapper<XPUType, float, float, int>;
|
|
} else if (fa_tgemm == XPU_FA_DTYPE::FA_FLOAT16) {
|
|
flash_attention_grad_kernel =
|
|
mha_varlen_bwd_wrapper<XPUType, float, XPUTypeFP16, int>;
|
|
}
|
|
// template<typename T, typename TACCUM, typename TGEMM, typename TID = int>
|
|
// int mha_varlen_bwd(xdnn::Context* xpu_ctx, const T* dout, const T* q, const
|
|
// T* k, const T* v, const T* out, const TACCUM* softmax_lse, T* dq, T* dk, T*
|
|
// dv, const xdnn::VectorParam<TID>& lod_seqlens_q, const
|
|
// xdnn::VectorParam<TID>& lod_seqlens_k, int64_t max_seqlen_q, int64_t
|
|
// max_seqlen_k, int64_t head_num, int64_t head_num_k, int64_t head_dim, const
|
|
// float softmax_scale = 0.0f, const float p_dropout = 0.0f, int seed =
|
|
// 0x45678901, const bool is_causal = true, const TACCUM* attn_mask = nullptr,
|
|
// const TACCUM* bias = nullptr, const float* q_maxptr = nullptr, const float*
|
|
// k_maxptr = nullptr, const float* v_maxptr = nullptr, const float* o_maxptr
|
|
// = nullptr, float* dq_maxptr = nullptr, float* dk_maxptr = nullptr, float*
|
|
// dv_maxptr = nullptr, const float* do_maxptr = nullptr, const bool
|
|
// is_qkv_fusion = false, const bool is_dqkv_fusion = false, const int64_t
|
|
// qkv_layout = AttnQKVLayout_t::ATTN_BLHD, const float* alibi_slopes =
|
|
// nullptr, const std::vector<int64_t>& alibi_slopes_shape = {}, int
|
|
// window_size_left = -1, int window_size_right = -1, int64_t v_head_dim =
|
|
// -1, const int* downstart_row_indices_data = nullptr,
|
|
// const int* downend_row_indices_data = nullptr,
|
|
// const int* upstart_row_indices_data = nullptr,
|
|
// const int* upend_row_indices_data = nullptr,
|
|
// const int flash_mask_head_num = 0,
|
|
// int* flashmask_maxmin = nullptr,
|
|
// XPUStream side_stream = nullptr);
|
|
int r = flash_attention_grad_kernel(
|
|
dev_ctx.x_context(),
|
|
dout_data, // dout
|
|
q_data, // q
|
|
k_data, // k
|
|
v_data, // v
|
|
out_data, // out
|
|
softmax_lse_data, // softmax_lse
|
|
dq_data, // dq
|
|
dk_data, // dk
|
|
dv_data, // dv
|
|
lod_seqlen_q, // lod_seqlens_q
|
|
lod_seqlen_k, // lod_seqlens_k
|
|
max_seqlen_q, // max_seqlen_q
|
|
max_seqlen_k, // max_seqlen_k
|
|
num_heads, // head_num
|
|
num_heads_k, // head_num_k
|
|
head_size, // head_dim
|
|
real_scale, // softmax_scale
|
|
dropout, // p_dropout
|
|
static_cast<int32_t>(seed_offset_data[0]), // seed
|
|
causal, // is_causal
|
|
nullptr, // attn_mask
|
|
bias_data, // bias
|
|
nullptr, // q_maxptr
|
|
nullptr, // k_maxptr
|
|
nullptr, // v_maxptr
|
|
nullptr, // o_maxptr
|
|
nullptr, // dq_maxptr
|
|
nullptr, // dk_maxptr
|
|
nullptr, // dv_maxptr
|
|
nullptr, // do_maxptr
|
|
false, // is_qkv_fusion
|
|
false, // is_dqkv_fusion
|
|
fa_layout, // qkv_layout
|
|
nullptr, // alibi_slopes
|
|
{}, // alibi_slopes_shape
|
|
-1, // window_size_left
|
|
-1, // window_size_right
|
|
head_size_v, // v_head_dim
|
|
(const int*)downstart_row_indices_data,
|
|
(const int*)downend_row_indices_data,
|
|
(const int*)upstart_row_indices_data,
|
|
(const int*)upend_row_indices_data,
|
|
is_flashmask ? startend_row_indices->dims()[1] : 0,
|
|
nullptr,
|
|
is_flashmask ? flashmask_stream : nullptr);
|
|
PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");
|
|
if (is_flashmask && flashmask_stream != nullptr) {
|
|
r = xpu_wait(flashmask_stream);
|
|
PADDLE_ENFORCE_XPU_SUCCESS(r);
|
|
xpu_stream_destroy(flashmask_stream);
|
|
}
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnUnpaddedGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& cu_seqlens_q,
|
|
const DenseTensor& cu_seqlens_k,
|
|
const DenseTensor& out,
|
|
const DenseTensor& softmax_lse,
|
|
const DenseTensor& seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const DenseTensor& dout,
|
|
const Scalar& max_seqlen_q,
|
|
const Scalar& max_seqlen_k,
|
|
float scale,
|
|
float dropout,
|
|
bool causal,
|
|
DenseTensor* dq,
|
|
DenseTensor* dk,
|
|
DenseTensor* dv) {
|
|
dev_ctx.template Alloc<T>(dq);
|
|
dev_ctx.template Alloc<T>(dk);
|
|
dev_ctx.template Alloc<T>(dv);
|
|
auto dims = q.dims();
|
|
|
|
const int64_t batch_size = cu_seqlens_q.numel() - 1;
|
|
const int64_t num_heads = dims[1];
|
|
const int64_t head_size = dims[2];
|
|
const int64_t head_size_v = v.dims()[2];
|
|
const int64_t num_heads_k = k.dims()[1];
|
|
|
|
api::VectorParam<int> qlod{cu_seqlens_q.data<int>(),
|
|
static_cast<int64_t>(cu_seqlens_q.numel()),
|
|
nullptr};
|
|
api::VectorParam<int> kvlod{cu_seqlens_k.data<int>(),
|
|
static_cast<int64_t>(cu_seqlens_k.numel()),
|
|
nullptr};
|
|
|
|
FlashAttnGradKernelBase<T>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
qlod,
|
|
kvlod,
|
|
out,
|
|
softmax_lse,
|
|
seed_offset,
|
|
attn_mask,
|
|
paddle::none,
|
|
dout,
|
|
batch_size,
|
|
max_seqlen_q,
|
|
max_seqlen_k,
|
|
num_heads,
|
|
num_heads_k,
|
|
head_size,
|
|
head_size_v,
|
|
scale,
|
|
dropout,
|
|
causal,
|
|
dq,
|
|
dk,
|
|
dv);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashAttnGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& out,
|
|
const DenseTensor& softmax_lse,
|
|
const DenseTensor& seed_offset,
|
|
const optional<DenseTensor>& attn_mask,
|
|
const DenseTensor& dout,
|
|
float dropout,
|
|
bool causal,
|
|
DenseTensor* dq,
|
|
DenseTensor* dk,
|
|
DenseTensor* dv) {
|
|
dev_ctx.template Alloc<T>(dq);
|
|
dev_ctx.template Alloc<T>(dk);
|
|
dev_ctx.template Alloc<T>(dv);
|
|
|
|
// q, k, v [batch_size, seq_len, num_heads, head_dim]
|
|
const auto& dims = q.dims();
|
|
|
|
const int64_t batch_size = dims[0];
|
|
const int64_t seqlen_q = dims[1];
|
|
const int64_t num_heads = dims[2];
|
|
const int64_t head_size_og = dout.dims()[3];
|
|
const int64_t head_size = dims[3];
|
|
const int64_t head_size_v = v.dims()[3];
|
|
const int64_t seqlen_k = k.dims()[1];
|
|
const int64_t num_heads_k = k.dims()[2];
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
head_size_og,
|
|
head_size_v,
|
|
common::errors::InvalidArgument(
|
|
"flash_attn_bwd receive input with head_size_og == head_size_v"));
|
|
|
|
// lod info
|
|
std::vector<int> qlod_vec = {0};
|
|
std::vector<int> kvlod_vec = {0};
|
|
for (int batch_idx = 1; batch_idx <= batch_size; ++batch_idx) {
|
|
qlod_vec.push_back(seqlen_q * batch_idx);
|
|
kvlod_vec.push_back(seqlen_k * batch_idx);
|
|
}
|
|
api::VectorParam<int> qlod{
|
|
qlod_vec.data(), static_cast<int64_t>(qlod_vec.size()), nullptr};
|
|
api::VectorParam<int> kvlod{
|
|
kvlod_vec.data(), static_cast<int64_t>(kvlod_vec.size()), nullptr};
|
|
|
|
FlashAttnGradKernelBase<T>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
qlod,
|
|
kvlod,
|
|
out,
|
|
softmax_lse,
|
|
seed_offset,
|
|
attn_mask,
|
|
paddle::none,
|
|
dout,
|
|
batch_size,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
num_heads,
|
|
num_heads_k,
|
|
head_size,
|
|
head_size_v,
|
|
0.0,
|
|
dropout,
|
|
causal,
|
|
dq,
|
|
dk,
|
|
dv);
|
|
}
|
|
|
|
template <typename T, typename Context>
|
|
void FlashMaskGradKernel(const Context& dev_ctx,
|
|
const DenseTensor& q,
|
|
const DenseTensor& k,
|
|
const DenseTensor& v,
|
|
const DenseTensor& startend_row_indices,
|
|
const DenseTensor& out,
|
|
const DenseTensor& softmax_lse,
|
|
const DenseTensor& seed_offset,
|
|
const DenseTensor& dout,
|
|
float dropout,
|
|
bool causal,
|
|
DenseTensor* dq,
|
|
DenseTensor* dk,
|
|
DenseTensor* dv) {
|
|
dev_ctx.template Alloc<T>(dq);
|
|
dev_ctx.template Alloc<T>(dk);
|
|
dev_ctx.template Alloc<T>(dv);
|
|
|
|
// q, k, v [batch_size, seq_len, num_heads, head_dim]
|
|
const auto& dims = q.dims();
|
|
|
|
const int64_t batch_size = dims[0];
|
|
const int64_t seqlen_q = dims[1];
|
|
const int64_t num_heads = dims[2];
|
|
const int64_t head_size_og = dout.dims()[3];
|
|
const int64_t head_size = dims[3];
|
|
const int64_t head_size_v = v.dims()[3];
|
|
const int64_t seqlen_k = k.dims()[1];
|
|
const int64_t num_heads_k = k.dims()[2];
|
|
|
|
PADDLE_ENFORCE_EQ(
|
|
head_size_og,
|
|
head_size_v,
|
|
common::errors::InvalidArgument(
|
|
"flash_attn_bwd receive input with head_size_og == head_size_v"));
|
|
|
|
// lod info
|
|
std::vector<int> qlod_vec = {0};
|
|
std::vector<int> kvlod_vec = {0};
|
|
for (int batch_idx = 1; batch_idx <= batch_size; ++batch_idx) {
|
|
qlod_vec.push_back(seqlen_q * batch_idx);
|
|
kvlod_vec.push_back(seqlen_k * batch_idx);
|
|
}
|
|
api::VectorParam<int> qlod{
|
|
qlod_vec.data(), static_cast<int64_t>(qlod_vec.size()), nullptr};
|
|
api::VectorParam<int> kvlod{
|
|
kvlod_vec.data(), static_cast<int64_t>(kvlod_vec.size()), nullptr};
|
|
FlashAttnGradKernelBase<T>(dev_ctx,
|
|
q,
|
|
k,
|
|
v,
|
|
qlod,
|
|
kvlod,
|
|
out,
|
|
softmax_lse,
|
|
seed_offset,
|
|
paddle::none,
|
|
startend_row_indices,
|
|
dout,
|
|
batch_size,
|
|
seqlen_q,
|
|
seqlen_k,
|
|
num_heads,
|
|
num_heads_k,
|
|
head_size,
|
|
head_size_v,
|
|
0.0,
|
|
dropout,
|
|
causal,
|
|
dq,
|
|
dk,
|
|
dv);
|
|
}
|
|
} // namespace phi
|
|
|
|
PD_REGISTER_KERNEL(flash_attn_unpadded_grad,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashAttnUnpaddedGradKernel,
|
|
float,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
kernel->InputAt(3).SetBackend(phi::Backend::CPU); // cu_seqlens_q
|
|
kernel->InputAt(4).SetBackend(phi::Backend::CPU); // cu_seqlens_k
|
|
kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND); // seed_offset
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(flash_attn_grad,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashAttnGradKernel,
|
|
phi::bfloat16,
|
|
float,
|
|
phi::float16) {
|
|
kernel->InputAt(5).SetBackend(phi::Backend::ALL_BACKEND); // seed_offset
|
|
}
|
|
|
|
PD_REGISTER_KERNEL(flashmask_attention_grad,
|
|
XPU,
|
|
ALL_LAYOUT,
|
|
phi::FlashMaskGradKernel,
|
|
phi::float16,
|
|
phi::bfloat16) {
|
|
kernel->InputAt(6).SetBackend(
|
|
phi::Backend::ALL_BACKEND); // fixed_seed_offset
|
|
}
|