// 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 <cstddef>
#include "glog/logging.h"  // For VLOG()
#include "paddle/common/enforce.h"
#include "paddle/common/flags.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/core/dense_tensor.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/funcs/elementwise_base.h"
#include "paddle/phi/kernels/gpu/flash_attn_utils.h"
#include "paddle/phi/kernels/reduce_sum_kernel.h"
#include "paddle/phi/kernels/slice_kernel.h"
#ifdef PADDLE_WITH_FLASHATTN_V3
#include "paddle/phi/kernels/gpu/flash_attn_v3_grad_kernel.h"
#endif

COMMON_DECLARE_bool(cudnn_deterministic);
COMMON_DECLARE_int32(flash_attn_version);

namespace phi {

int get_num_split() {
  // 0 for an internal heuristic, which is optimal
  return FLAGS_cudnn_deterministic ? 1 : 0;
}

template <typename T, uint64_t HeaddimDiv32>
static __global__ void SumStridedKV(const T* src,
                                    T* dst,
                                    const uint64_t sRowDim1,
                                    const uint64_t sRowDim2,
                                    const uint64_t sRowDim3,
                                    const uint64_t sColDim,
                                    const uint64_t sRowStride1,
                                    const uint64_t sRowStride2,
                                    const uint64_t sColStride,
                                    const uint64_t dRowStride1,
                                    const uint64_t dRowStride2) {
  // SrcShape [seqlen, num_heads_k, num_heads/num_heads_k, headdim]
  // AxisName [row1  , row2       , col                  , row3   ]
  // LoopMap  [blockx, thready    , serialreduce         , threadx]
  // Ensure blockDim.x == 32 && blockDim.z == 1
  // Ensure sRowStride3 == dRowStride3 == 1 (headdim dim is contiguous)
  using IndexType = uint64_t;
  constexpr IndexType BlockDimX = 32;
  const IndexType SRow1Begin = blockIdx.x * sRowStride1;
  const IndexType SRow1End = sRowDim1 * sRowStride1;
  const IndexType SRow1Stride = gridDim.x * sRowStride1;

  const IndexType SRow2Begin = threadIdx.y * sRowStride2;
  const IndexType SRow2End = sRowDim2 * sRowStride2;
  const IndexType SRow2Stride = blockDim.y * sRowStride2;

  // const IndexType SRow3Begin = threadIdx.x * sRowStride3;
  // const IndexType SRow3End = sRowDim3 * sRowStride3;
  // const IndexType SRow3Stride = BlockDimX * sRowStride3;

  constexpr IndexType SColBegin = 0;
  const IndexType SColEnd = sColDim * sColStride;
  const IndexType SColStride = sColStride;

  const IndexType DRow1Begin = blockIdx.x * dRowStride1;
  const IndexType DRow1Stride = gridDim.x * dRowStride1;

  const IndexType DRow2Begin = threadIdx.y * dRowStride2;
  const IndexType DRow2Stride = dRowStride2;

  // const IndexType DRow3Begin = threadIdx.x * dRowStride3;
  // const IndexType DRow3Stride = blockDim.x * dRowStride3;

  for (auto row1 = SRow1Begin, drow1 = DRow1Begin; row1 < SRow1End;
       row1 += SRow1Stride, drow1 += DRow1Stride) {
    for (auto row2 = SRow2Begin, drow2 = DRow2Begin; row2 < SRow2End;
         row2 += SRow2Stride, drow2 += DRow2Stride) {
      const auto i1 = row1 + row2 + threadIdx.x;
      const auto di1 = drow1 + drow2 + threadIdx.x;
      T v[HeaddimDiv32];
#pragma unroll
      for (auto i = IndexType(0); i < HeaddimDiv32; i++) {
        v[i] = T{0};
      }
      for (auto col = SColBegin; col < SColEnd; col += SColStride) {
        const auto i2 = i1 + col;
#pragma unroll
        for (auto i = IndexType(0); i < HeaddimDiv32; i++) {
          v[i] += src[i2 + i * BlockDimX];
        }
      }
#pragma unroll
      for (auto i = IndexType(0); i < HeaddimDiv32; i++) {
        dst[di1 + i * BlockDimX] = v[i];
      }
    }
  }
}

template <typename T>
static auto selectSumkernel(int64_t headdim) {
  PADDLE_ENFORCE_LE(headdim,
                    256,
                    common::errors::InvalidArgument(
                        "FlashAttention only support headdim <= 256"));
  PADDLE_ENFORCE_EQ(headdim % 32,
                    0,
                    common::errors::InvalidArgument(
                        "FlashAttention only support headdim %% 32 == 0"));
  PADDLE_ENFORCE_NE(
      headdim, 0, common::errors::InvalidArgument("Headdim can't be zero"));
#define CASEN(n) \
  case n:        \
    return SumStridedKV<T, n>;
  switch (headdim / 32) {
    CASEN(1);
    CASEN(2);
    CASEN(3);
    CASEN(4);
    CASEN(5);
    CASEN(6);
    CASEN(7);
    CASEN(8);
  }
  PADDLE_FATAL("Unreachable in selectSumKernel");
#undef CASEN
}

template <typename T, typename Context>
static void kvReduceForGQA(const Context& dev_ctx,
                           const DenseTensor& dk_tmp,
                           DenseTensor* dk) {
  PADDLE_ENFORCE_EQ(
      dk->strides()[2],
      1,
      common::errors::InvalidArgument("headdim dimension must be contiguous"));
  PADDLE_ENFORCE_EQ(
      dk_tmp.strides()[3],
      1,
      common::errors::InvalidArgument("headdim dimension must be contiguous"));
  const int64_t reduceDimSize = dk_tmp.dims()[2];
  const size_t blockNum =
      std::min((static_cast<int64_t>(dk_tmp.dims()[0] + 31) / 32),
               static_cast<int64_t>(1024l));
  const dim3 threadNum{32, 4, 1};
  auto sumkernel = selectSumkernel<T>(dk_tmp.dims()[3]);
  sumkernel<<<blockNum, threadNum, 0, dev_ctx.stream()>>>(
      reinterpret_cast<const T*>(dk_tmp.data()),
      reinterpret_cast<T*>(dk->data()),
      dk_tmp.dims()[0],
      dk_tmp.dims()[1],
      dk_tmp.dims()[3],
      dk_tmp.dims()[2],
      dk_tmp.strides()[0],
      dk_tmp.strides()[1],
      // dk_tmp.strides()[3],
      dk_tmp.strides()[2],
      dk->strides()[0],
      dk->strides()[1]
      // dk->strides()[2]
  );
}
template <typename T, typename Context>
static void kvReduceBatchedForGQA(const Context& dev_ctx,
                                  const DenseTensor& dk_tmp,
                                  DenseTensor* dk) {
  PADDLE_ENFORCE_EQ(
      dk->strides()[3],
      1,
      common::errors::InvalidArgument("headdim dimension must be contiguous"));
  PADDLE_ENFORCE_EQ(
      dk_tmp.strides()[4],
      1,
      common::errors::InvalidArgument("headdim dimension must be contiguous"));
  PADDLE_ENFORCE_EQ(dk->strides()[0],
                    dk->strides()[1] * dk->dims()[1],
                    common::errors::InvalidArgument(
                        "batchsize dimension must be contiguous"));
  PADDLE_ENFORCE_EQ(dk_tmp.strides()[0],
                    dk_tmp.strides()[1] * dk_tmp.dims()[1],
                    common::errors::InvalidArgument(
                        "batchsize dimension must be contiguous"));
  const int64_t reduceDimSize = dk_tmp.dims()[3];
  const size_t blockNum = std::min(
      (static_cast<int64_t>(dk_tmp.dims()[0] * dk_tmp.dims()[1] + 31) / 32),
      static_cast<int64_t>(1024l));
  const dim3 threadNum{32, 4, 1};
  auto sumkernel = selectSumkernel<T>(dk_tmp.dims()[4]);
  // here implicitly flat [batch,seqlen], and require batch dim to be contiguous
  sumkernel<<<blockNum, threadNum, 0, dev_ctx.stream()>>>(
      reinterpret_cast<const T*>(dk_tmp.data()),
      reinterpret_cast<T*>(dk->data()),
      dk_tmp.dims()[0] * dk_tmp.dims()[1],
      dk_tmp.dims()[2],
      dk_tmp.dims()[4],
      dk_tmp.dims()[3],
      dk_tmp.strides()[1],
      dk_tmp.strides()[2],
      // dk_tmp.strides()[4],
      dk_tmp.strides()[3],
      dk->strides()[1],
      dk->strides()[2]
      // dk->strides()[3]
  );
}

template <typename T, typename Context>
void FlashAttnUnpaddedGradBaseKernel(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,
                                     bool varlen_padded) {
#ifdef PADDLE_WITH_FLASHATTN
  // q,k,v [total_*, num_heads, head_dim]
  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_v = dout.dims()[2];
  const int64_t head_size = dims[2];
  const int64_t total_k = k.dims()[0];
  const int64_t num_heads_k = k.dims()[1];
  const int64_t total_q = dims[0];

  bool is_mha = (num_heads == num_heads_k);

  DenseTensor* kdq = dq;
  DenseTensor dq_tmp;
  if (!dq) {
    dq_tmp.Resize(dims);
    dev_ctx.template Alloc<T>(&dq_tmp);
    kdq = &dq_tmp;
  }

  std::initializer_list<int64_t> dk_dv_shape = {
      total_k, num_heads_k, num_heads / num_heads_k, head_size};

  DenseTensor *kdk = dk, *kdv = dv;
  DenseTensor dk_tmp;
  if (!dk || !is_mha) {
    dk_tmp.Resize(dk_dv_shape);
    dev_ctx.template Alloc<T>(&dk_tmp);
    kdk = &dk_tmp;
  }

  DenseTensor dv_tmp;
  if (!dv || !is_mha) {
    dv_tmp.Resize(dk_dv_shape);
    dev_ctx.template Alloc<T>(&dv_tmp);
    kdv = &dv_tmp;
  }

#ifdef PADDLE_WITH_HIP
  const hipStream_t stream = dev_ctx.stream();
#else
  const cudaStream_t stream = dev_ctx.stream();
#endif

  int num_splits = get_num_split();

  // TODO(umiswing): add shape check
  PADDLE_ENFORCE_EQ(head_size,
                    head_size_v,
                    common::errors::InvalidArgument(
                        "This kernel does not support headdim != headdim_v, "
                        "but got headdim = %d and headdim_v = %d",
                        head_size,
                        head_size_v));

  int64_t max_seqlen_q = max_seqlen_q_.to<int64_t>();
  int64_t max_seqlen_k = max_seqlen_k_.to<int64_t>();

  FlashAttnBwdParamsV2 params =
      FlashAttnBwdParamsV2(dev_ctx,
                           /*version=*/2,
                           batch_size,
                           max_seqlen_q,
                           max_seqlen_k,
                           num_heads,
                           num_heads_k,
                           head_size,
                           dropout,
                           scale,
                           causal,
                           q.dtype(),
                           attn_mask,
                           nullptr,  // startend_row_indices,
                           seed_offset.data<int64_t>(),
                           /*unpadded_lse*/ true,
                           total_q);

  VLOG(10) << "FlashAttn bwd seed: " << params.seed
           << ", offset: " << params.offset;
  bool succ = dynload::flash_attn_varlen_bwd(
      dout.data(),
      q.data(),
      k.data(),
      v.data(),
      out.data(),
      params.softmax_d.data(),
      softmax_lse.data(),
      cu_seqlens_q.data<int32_t>(),
      cu_seqlens_k.data<int32_t>(),
      params.rng_state.data(),
      kdq->data(),
      kdk->data(),
      kdv->data(),
      params.dq_accum.data(),
      params.batch_size,
      params.max_seqlen_q,
      params.max_seqlen_k,
      params.seqlen_q_rounded,
      params.seqlen_k_rounded,
      params.num_heads,
      params.num_heads_k,
      params.head_size,
      params.head_size_rounded,
      params.dropout,
      params.softmax_scale,
      1.0f / params.softmax_scale,
      params.causal,
      params.is_bf16,
      num_splits,
      stream,
      params.seed,
      params.offset,
      params.attn_mask_tensor ? params.attn_mask_tensor->data() : nullptr,
      params.attn_mask_tensor ? params.mask_dims.data() : nullptr,
      q.strides()[0],
      k.strides()[0],
      v.strides()[0],
      q.strides()[1],
      k.strides()[1],
      v.strides()[1],
      out.strides()[0],
      out.strides()[1],
      max_seqlen_q * q.strides()[0],
      max_seqlen_k * k.strides()[0],
      max_seqlen_k * v.strides()[0],
      max_seqlen_q * out.strides()[0],
      kdq->strides()[0],
      kdk->strides()[0],
      kdv->strides()[0],
      kdq->strides()[1],
      kdk->strides()[kdk->strides().size() - 2],
      kdv->strides()[kdv->strides().size() - 2],
      dout.strides()[0],
      dout.strides()[1],
      max_seqlen_q * kdq->strides()[0],
      max_seqlen_k * kdk->strides()[0],
      max_seqlen_k * kdv->strides()[0],
      max_seqlen_q * dout.strides()[0],
#ifdef PADDLE_WITH_CUDA
      varlen_padded,
      params.total_q
#else
      varlen_padded
#endif
  );
  CheckFlashAttnStatus(succ);
  if (!is_mha) {
    if (dk) {
      if (dk->meta().is_contiguous())
        SumKernel<T, Context>(dev_ctx, dk_tmp, {2}, dk->type(), false, dk);
      else
        kvReduceForGQA<T, Context>(dev_ctx, dk_tmp, dk);
    }
    if (dv) {
      if (dv->meta().is_contiguous())
        SumKernel<T, Context>(dev_ctx, dv_tmp, {2}, dv->type(), false, dv);
      else
        kvReduceForGQA<T, Context>(dev_ctx, dv_tmp, dv);
    }
  }
#else
  RaiseNotSupportedError();
#endif
}

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) {
#ifdef PADDLE_WITH_FLASHATTN
  if (dq) {
    dev_ctx.template Alloc<T>(dq);
  }
  if (dk) {
    dev_ctx.template Alloc<T>(dk);
  }
  if (dv) {
    dev_ctx.template Alloc<T>(dv);
  }
  FlashAttnUnpaddedGradBaseKernel<T>(dev_ctx,
                                     q,
                                     k,
                                     v,
                                     cu_seqlens_q,
                                     cu_seqlens_k,
                                     out,
                                     softmax_lse,
                                     seed_offset,
                                     attn_mask,
                                     dout,
                                     max_seqlen_q,
                                     max_seqlen_k,
                                     scale,
                                     dropout,
                                     causal,
                                     dq,
                                     dk,
                                     dv,
                                     false /*varlen_padded*/);
#else
  RaiseNotSupportedError();
#endif
}

static void sliceFlattenView(const DenseTensor& in,
                             DenseTensor* out,
                             int axis,
                             int64_t offset,
                             int64_t sliceLength) {
  PADDLE_ENFORCE_LT(
      axis,
      in.dims().size(),
      common::errors::InvalidArgument("sliceView receive axis out of bound"));
  std::array<int64_t, DDim::kMaxRank> dimArr;
  std::array<int64_t, DDim::kMaxRank> strideArr;
  auto id = dimArr.begin(), is = strideArr.begin();
  for (int i = 0; i < in.dims().size(); i++) {
    if (i == axis) continue;
    if (i == axis + 1)
      *id = in.dims()[i] * sliceLength;
    else
      *id = in.dims()[i];
    *is = in.strides()[i];
    id++;
    is++;
  }
  *out = DenseTensor{
      in.Holder(),
      DenseTensorMeta{in.dtype(),
                      DDim{dimArr.data(), in.dims().size() - 1},
                      DDim(strideArr.data(), in.dims().size() - 1)}};
  out->set_offset(in.offset() +
                  offset * in.strides()[axis] * SizeOf(out->dtype()));
}
template <typename OutT>
struct ZeroFunctor {
  __device__ __forceinline__ OutT operator()() const {
    return static_cast<OutT>(0);
  }
};
template <typename T, typename Context>
void FlashAttnVarlenQKVPackedGradKernel(const Context& dev_ctx,
                                        const DenseTensor& qkv,
                                        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,
                                        bool varlen_padded,
                                        DenseTensor* dqkv) {
#ifdef PADDLE_WITH_FLASHATTN
  // q,k,v [total_*, num_heads, head_dim]
  const auto head_groupnum = qkv.dims()[1];  // nheads/nheads_k + 1 + 1
  DenseTensor q, k, v;
  sliceFlattenView(qkv, &q, 1, 0, head_groupnum - 2);
  sliceFlattenView(qkv, &k, 1, head_groupnum - 2, 1);
  sliceFlattenView(qkv, &v, 1, head_groupnum - 1, 1);
  // DenseTensor dqkv_tmp;
  if (!dqkv) {
    return;
    // dqkv is the only output. No need to compute if no dqkv
    // dqkv_tmp.Resize(qkv.dims());
    // dqkv = &dqkv_tmp;
  }
  dev_ctx.template Alloc<T>(dqkv);
  {
    std::vector<const DenseTensor*> inputs{};
    std::vector<DenseTensor*> outputs{dqkv};
    funcs::ElementwiseKernel<T>(dev_ctx, inputs, &outputs, ZeroFunctor<T>());
  }
  DenseTensor dq, dk, dv;
  sliceFlattenView(*dqkv, &dq, 1, 0, head_groupnum - 2);
  sliceFlattenView(*dqkv, &dk, 1, head_groupnum - 2, 1);
  sliceFlattenView(*dqkv, &dv, 1, head_groupnum - 1, 1);
  FlashAttnUnpaddedGradBaseKernel<T>(dev_ctx,
                                     q,
                                     k,
                                     v,
                                     cu_seqlens_q,
                                     cu_seqlens_k,
                                     out,
                                     softmax_lse,
                                     seed_offset,
                                     attn_mask,
                                     dout,
                                     max_seqlen_q,
                                     max_seqlen_k,
                                     scale,
                                     dropout,
                                     causal,
                                     &dq,
                                     &dk,
                                     &dv,
                                     varlen_padded);
#else
  RaiseNotSupportedError();
#endif
}
template <typename T, typename Context>
void FlashAttnGradBaseKernel(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 optional<DenseTensor>& startend_row_indices,
                             const DenseTensor& dout,
                             float dropout,
                             bool causal,
                             DenseTensor* dq,
                             DenseTensor* dk,
                             DenseTensor* dv) {
#ifdef PADDLE_WITH_FLASHATTN
  // 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_v = v.dims()[3];
  const int64_t head_size = dims[3];
  const int64_t seqlen_k = k.dims()[1];
  const int64_t num_heads_k = k.dims()[2];

  bool is_mha = (num_heads == num_heads_k);

  std::initializer_list<int64_t> dk_dv_shape = {
      batch_size, seqlen_k, num_heads_k, num_heads / num_heads_k, head_size};

  DenseTensor* kdq = dq;
  DenseTensor dq_tmp;
  if (!dq) {
    dq_tmp.Resize(dims);
    dev_ctx.template Alloc<T>(&dq_tmp);
    kdq = &dq_tmp;
  }

  DenseTensor *kdk = dk, *kdv = dv;
  DenseTensor dk_tmp;
  if (!dk || !is_mha) {
    dk_tmp.Resize(dk_dv_shape);
    dev_ctx.template Alloc<T>(&dk_tmp);
    kdk = &dk_tmp;
  }

  DenseTensor dv_tmp;
  if (!dv || !is_mha) {
    dv_tmp.Resize(dk_dv_shape);
    dev_ctx.template Alloc<T>(&dv_tmp);
    kdv = &dv_tmp;
  }

#ifdef PADDLE_WITH_HIP
  const hipStream_t stream = dev_ctx.stream();
#else
  const cudaStream_t stream = dev_ctx.stream();
#endif

  // TODO(umiswing): add shape check
  PADDLE_ENFORCE_EQ(head_size,
                    head_size_v,
                    common::errors::InvalidArgument(
                        "This kernel does not support headdim != headdim_v, "
                        "but got headdim = %d and headdim_v = %d",
                        head_size,
                        head_size_v));

  const float softmax_scale = 1.0f / std::sqrt(head_size);
  const float softmax_unscale = std::sqrt(head_size);

  int version = FLAGS_flash_attn_version == 3 && FLAGS_cudnn_deterministic &&
                        head_size > 128
                    ? 2
                    : FLAGS_flash_attn_version;
  FlashAttnBwdParamsV2 params =
      FlashAttnBwdParamsV2(dev_ctx,
                           version,
                           batch_size,
                           seqlen_q,
                           seqlen_k,
                           num_heads,
                           num_heads_k,
                           head_size,
                           dropout,
                           softmax_scale,
                           causal,
                           q.dtype(),
                           attn_mask,
                           startend_row_indices,
                           seed_offset.data<int64_t>(),
                           /*unpadded_lse*/ false,
                           /*total_q*/ 0);

  VLOG(10) << "[FlashAttn Backward" << version << "] q.shape=[" << q.dims()
           << "], k.shape=[" << k.dims() << "], v.shape=[" << v.dims() << "]";
  VLOG(10) << "[FlashAttn Backward" << version << "] dropout=" << dropout
           << ", seed=" << params.seed << ", offset=" << params.offset;
  VLOG(10) << "[FlashAttn Backward" << version
           << "] softmax_scale=" << softmax_scale
           << ", softmax_unscale=" << softmax_unscale;
  if (attn_mask.get_ptr()) {
    VLOG(10) << "[FlashAttn Backward" << version << "] attn_mask.shape=["
             << (attn_mask.get_ptr())->dims() << "]";
  }

  int num_splits = get_num_split();

  DenseTensor flashmask_maxmin, 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 = params.startend_row_indices != nullptr;
  if (is_flashmask) {
    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]"));
    auto flashmask_maxmin_shape = params.startend_row_indices->dims();
    flashmask_maxmin_shape[2] = (flashmask_maxmin_shape[2] + 31) / 32 * 8;
    flashmask_maxmin.set_type(DataType::INT32);
    flashmask_maxmin.Resize(flashmask_maxmin_shape);
    dev_ctx.template Alloc<T>(&flashmask_maxmin);

    downstart_row_indices =
        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 =
            Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
        upend_row_indices_data = upend_row_indices.data();
      } else {
        downend_row_indices =
            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 =
          Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {3}, {4});
      upend_row_indices_data = upend_row_indices.data();
      downend_row_indices =
          Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
      downend_row_indices_data = downend_row_indices.data();
      upstart_row_indices =
          Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {2}, {3});
      upstart_row_indices_data = upstart_row_indices.data();
    }
  }

#ifdef PADDLE_WITH_HIP
  bool succ = dynload::flash_attn_bwd(
      dout.data(),
      q.data(),
      k.data(),
      v.data(),
      out.data(),
      params.softmax_d.data(),
      softmax_lse.data(),
      params.rng_state.data(),
      kdq->data(),
      kdk->data(),
      kdv->data(),
      params.dq_accum.data(),
      params.batch_size,
      params.max_seqlen_q,
      params.max_seqlen_k,
      params.seqlen_q_rounded,
      params.seqlen_k_rounded,
      params.num_heads,
      params.num_heads_k,
      params.head_size,
      params.head_size_rounded,
      params.dropout,
      params.softmax_scale,
      softmax_unscale,
      params.causal,
      params.is_bf16,
      num_splits,
      stream,
      params.seed,
      params.offset,
      params.attn_mask_tensor ? params.attn_mask_tensor->data() : nullptr,
      params.attn_mask_tensor ? params.mask_dims.data() : nullptr,
      is_flashmask ? downstart_row_indices_data : nullptr,
      is_flashmask ? params.startend_row_indices_dims.data() : nullptr,
      is_flashmask ? upend_row_indices_data : nullptr,
      is_flashmask ? downend_row_indices_data : nullptr,
      is_flashmask ? upstart_row_indices_data : nullptr,
      is_flashmask ? flashmask_maxmin.data() : nullptr,
      q.strides()[1],
      k.strides()[1],
      v.strides()[1],
      q.strides()[2],
      k.strides()[2],
      v.strides()[2],
      out.strides()[1],
      out.strides()[2],
      q.strides()[0],
      k.strides()[0],
      v.strides()[0],
      out.strides()[0],
      kdq->strides()[1],
      kdk->strides()[1],
      kdv->strides()[1],
      kdq->strides()[2],
      kdk->strides()[kdk->strides().size() - 2],
      kdv->strides()[kdv->strides().size() - 2],
      dout.strides()[1],
      dout.strides()[2],
      kdq->strides()[0],
      kdk->strides()[0],
      kdv->strides()[0],
      dout.strides()[0]);
#else
  bool succ;
  int arch =
      backends::gpu::GetGPUComputeCapability(dev_ctx.GetPlace().GetDeviceId());

  if (arch == 80 && version == 3) {
    RaiseNotSupportedError(3);
  }

  if (arch == 90 && version == 3) {
#ifdef PADDLE_WITH_FLASHATTN_V3
    if (is_flashmask || params.attn_mask_tensor) {
      PADDLE_THROW(common::errors::Unimplemented(
          "FlashMask or Dense Mask is unsupported in FlashAttention V3"));
    }

    FlashAttnV3GradKernel<T, Context>(dev_ctx,
                                      q,
                                      k,
                                      v,
                                      out,
                                      softmax_lse,
                                      dout,
                                      params.softmax_scale,
                                      causal,
                                      -1,   // window_size_left
                                      -1,   // window_size_right
                                      0.f,  // softcap
                                      0,    // sm_margin
                                      dq,
                                      dk,
                                      dv);
#else
    RaiseNotSupportedError(3);
#endif
  } else {
    succ = dynload::flash_attn_bwd(
        dout.data(),
        q.data(),
        k.data(),
        v.data(),
        out.data(),
        params.softmax_d.data(),
        softmax_lse.data(),
        params.rng_state.data(),
        kdq->data(),
        kdk->data(),
        kdv->data(),
        params.dq_accum.data(),
        params.batch_size,
        params.max_seqlen_q,
        params.max_seqlen_k,
        params.seqlen_q_rounded,
        params.seqlen_k_rounded,
        params.num_heads,
        params.num_heads_k,
        params.head_size,
        params.head_size_rounded,
        params.dropout,
        params.softmax_scale,
        softmax_unscale,
        params.causal,
        params.is_bf16,
        num_splits,
        stream,
        params.seed,
        params.offset,
        params.attn_mask_tensor ? params.attn_mask_tensor->data() : nullptr,
        params.attn_mask_tensor ? params.mask_dims.data() : nullptr,
        is_flashmask ? downstart_row_indices_data : nullptr,
        is_flashmask ? params.startend_row_indices_dims.data() : nullptr,
        is_flashmask ? upend_row_indices_data : nullptr,
        is_flashmask ? downend_row_indices_data : nullptr,
        is_flashmask ? upstart_row_indices_data : nullptr,
        is_flashmask ? flashmask_maxmin.data() : nullptr,
        q.strides()[1],
        k.strides()[1],
        v.strides()[1],
        q.strides()[2],
        k.strides()[2],
        v.strides()[2],
        out.strides()[1],
        out.strides()[2],
        q.strides()[0],
        k.strides()[0],
        v.strides()[0],
        out.strides()[0],
        kdq->strides()[1],
        kdk->strides()[1],
        kdv->strides()[1],
        kdq->strides()[2],
        kdk->strides()[kdk->strides().size() - 2],
        kdv->strides()[kdv->strides().size() - 2],
        dout.strides()[1],
        dout.strides()[2],
        kdq->strides()[0],
        kdk->strides()[0],
        kdv->strides()[0],
        dout.strides()[0]);
  }
#endif
  if (version != 3) {
    CheckFlashAttnStatus(succ);  // umiswing: no return status in fa3
    if (!is_mha) {
      if (dk) {
        if (dk->meta().is_contiguous())
          SumKernel<T, Context>(dev_ctx, dk_tmp, {3}, dk->type(), false, dk);
        else
          kvReduceBatchedForGQA<T, Context>(dev_ctx, dk_tmp, dk);
      }

      if (dv) {
        if (dv->meta().is_contiguous())
          SumKernel<T, Context>(dev_ctx, dv_tmp, {3}, dv->type(), false, dv);
        else
          kvReduceBatchedForGQA<T, Context>(dev_ctx, dv_tmp, dv);
      }
    }
  }
#else
  RaiseNotSupportedError();
#endif
}

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) {
  if (dq) {
    dev_ctx.template Alloc<T>(dq);
  }
  if (dk) {
    dev_ctx.template Alloc<T>(dk);
  }
  if (dv) {
    dev_ctx.template Alloc<T>(dv);
  }
  if (dout.numel() == 0) {
    if (dq) Full<T, Context>(dev_ctx, dq->dims(), 0, dq);
    if (dk) Full<T, Context>(dev_ctx, dk->dims(), 0, dk);
    if (dv) Full<T, Context>(dev_ctx, dv->dims(), 0, dv);
    return;
  }
  FlashAttnGradBaseKernel<T, Context>(dev_ctx,
                                      q,
                                      k,
                                      v,
                                      out,
                                      softmax_lse,
                                      seed_offset,
                                      attn_mask,
                                      paddle::none,
                                      dout,
                                      dropout,
                                      causal,
                                      dq,
                                      dk,
                                      dv);
}

template <typename T, typename Context>
void FlashAttnQKVPackedGradKernel(const Context& dev_ctx,
                                  const DenseTensor& qkv,
                                  const DenseTensor& out,
                                  const DenseTensor& softmax_lse,
                                  const DenseTensor& seed_offset,
                                  const optional<DenseTensor>& attn_mask,
                                  const DenseTensor& dout,
                                  float dropout,
                                  bool causal,
                                  DenseTensor* dqkv) {
#ifdef PADDLE_WITH_FLASHATTN
  // qkv [batchsize, seqlen, nheads/nheads_k+2, nheads_k, head_dim]
  const auto head_groupnum = qkv.dims()[2];  // nheads/nheads_k + 1 + 1
  DenseTensor q, k, v;
  sliceFlattenView(qkv, &q, 2, 0, head_groupnum - 2);
  sliceFlattenView(qkv, &k, 2, head_groupnum - 2, 1);
  sliceFlattenView(qkv, &v, 2, head_groupnum - 1, 1);
  // DenseTensor dqkv_tmp;
  if (!dqkv) {
    return;
    // dqkv is the only output. No need to compute if no dqkv
    // dqkv_tmp.Resize(qkv.dims());
    // dqkv = &dqkv_tmp;
  }
  dev_ctx.template Alloc<T>(dqkv);
  DenseTensor dq, dk, dv;
  sliceFlattenView(*dqkv, &dq, 2, 0, head_groupnum - 2);
  sliceFlattenView(*dqkv, &dk, 2, head_groupnum - 2, 1);
  sliceFlattenView(*dqkv, &dv, 2, head_groupnum - 1, 1);
  FlashAttnGradBaseKernel<T, Context>(dev_ctx,
                                      q,
                                      k,
                                      v,
                                      out,
                                      softmax_lse,
                                      seed_offset,
                                      attn_mask,
                                      paddle::none,
                                      dout,
                                      dropout,
                                      causal,
                                      &dq,
                                      &dk,
                                      &dv);
#else
  RaiseNotSupportedError();
#endif
}

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) {
  if (dq) {
    dev_ctx.template Alloc<T>(dq);
  }
  if (dk) {
    dev_ctx.template Alloc<T>(dk);
  }
  if (dv) {
    dev_ctx.template Alloc<T>(dv);
  }
  FlashAttnGradBaseKernel<T, Context>(dev_ctx,
                                      q,
                                      k,
                                      v,
                                      out,
                                      softmax_lse,
                                      seed_offset,
                                      paddle::none,
                                      startend_row_indices,
                                      dout,
                                      dropout,
                                      causal,
                                      dq,
                                      dk,
                                      dv);
}
}  // namespace phi

PD_REGISTER_KERNEL(flash_attn_unpadded_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashAttnUnpaddedGradKernel,
                   phi::float16,
                   phi::bfloat16) {
  kernel->InputAt(7).SetBackend(phi::Backend::CPU);  // seed_offset
}

PD_REGISTER_KERNEL(flash_attn_varlen_qkvpacked_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashAttnVarlenQKVPackedGradKernel,
                   phi::float16,
                   phi::bfloat16) {
  kernel->InputAt(5).SetBackend(phi::Backend::CPU);  // seed_offset
}

PD_REGISTER_KERNEL(flash_attn_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashAttnGradKernel,
                   phi::float16,
                   phi::bfloat16) {
  kernel->InputAt(5).SetBackend(phi::Backend::CPU);  // seed_offset
}

PD_REGISTER_KERNEL(flash_attn_qkvpacked_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashAttnQKVPackedGradKernel,
                   phi::float16,
                   phi::bfloat16) {
  kernel->InputAt(3).SetBackend(phi::Backend::CPU);  // seed_offset
}

PD_REGISTER_KERNEL(flashmask_attention_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashMaskGradKernel,
                   phi::float16,
                   phi::bfloat16) {
  kernel->InputAt(6).SetBackend(phi::Backend::CPU);  // seed_offset
}