{"flashmask_attention",

XPUKernelSet({phi::DataType::BFLOAT16, phi::DataType::FLOAT16})},

{"flashmask_attention_grad",

XPUKernelSet({phi::DataType::BFLOAT16, phi::DataType::FLOAT16})},

#include "paddle/phi/kernels/slice_kernel.h"

void FlashAttnGradKernelBase(

const Context& 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 paddle::optional<DenseTensor>& attn_mask,

const paddle::optional<DenseTensor>& startend_row_indices,

const DenseTensor& dout,

const int batch_size,

const Scalar& max_seqlen_q_,

const Scalar& max_seqlen_k_,

const int num_heads,

const int num_heads_k,

const int head_size,

const int head_size_v,

float scale,

float dropout,

bool causal,

DenseTensor* dq,

DenseTensor* dk,

DenseTensor* dv) {

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>(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>(ctx, startend_row_indices.get(), {3}, {1}, {2});

upend_row_indices_data = upend_row_indices.data();

} else {

downend_row_indices =

phi::Slice<int32_t>(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>(ctx, startend_row_indices.get(), {3}, {3}, {4});

upend_row_indices_data = upend_row_indices.data();

downend_row_indices =

phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {1}, {2});

downend_row_indices_data = downend_row_indices.data();

upstart_row_indices =

phi::Slice<int32_t>(ctx, startend_row_indices.get(), {3}, {2}, {3});

upstart_row_indices_data = upstart_row_indices.data();

}

} else if (attn_mask.get_ptr() != nullptr) {

(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);

if (is_flashmask && flashmask_stream != nullptr) {

xpu_stream_destroy(flashmask_stream);

}

paddle::none,

paddle::none,

template <typename T, typename Context>

void FlashMaskGradKernel(const Context& 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) {

#ifdef PADDLE_WITH_XPU_XRE5

ctx.template Alloc<T>(dq);

ctx.template Alloc<T>(dk);

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>(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);

#else

PADDLE_THROW(common::errors::Unimplemented(

"re-compile using -DWITH_XPU_XRE5=ON to use FlashMaskGradKernel"));

#endif

}

PD_REGISTER_KERNEL(flashmask_attention_grad,

XPU,

ALL_LAYOUT,

phi::FlashMaskGradKernel,

phi::dtype::float16,

phi::dtype::bfloat16) {

kernel->InputAt(6).SetBackend(

phi::Backend::ALL_BACKEND); // fixed_seed_offset

}