// Copyright (c) 2024 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

#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
#include "paddle/phi/backends/dynload/flashattn.h"
#endif

#include "paddle/phi/backends/gpu/gpu_info.h"
#include "paddle/phi/common/memory_utils.h"
#include "paddle/phi/kernels/arange_kernel.h"
#include "paddle/phi/kernels/funcs/broadcast_function.h"
#include "paddle/phi/kernels/funcs/elementwise_functor.h"
#include "paddle/phi/kernels/funcs/reduce_function.h"
#include "paddle/phi/kernels/funcs/transpose_function.cuh"
#include "paddle/phi/kernels/fusion/gpu/attn_gemm.h"
#include "paddle/phi/kernels/gpudnn/softmax_gpudnn.h"

namespace phi {
namespace funcs {

template <typename T>
__global__ void SimpleScaleKernel(int64_t numel, float scale, T* inout) {
  CUDA_KERNEL_LOOP_TYPE(i, numel, int64_t) {
    inout[i] = static_cast<T>(scale * static_cast<float>(inout[i]));
  }
}

inline std::string MemoryDebugString(const DenseTensor& t) {
  int device_id = phi::backends::gpu::GetCurrentDeviceId();
  int64_t allocated =
      phi::memory_utils::DeviceMemoryStatCurrentValue("Allocated", device_id);
  int64_t reserved =
      phi::memory_utils::DeviceMemoryStatCurrentValue("Reserved", device_id);

  std::stringstream ss;
  ss << "shape=[" << t.dims()
     << "], size=" << static_cast<float>(t.memory_size()) / (1 << 20)
     << " MB, ptr=" << t.data()
     << "; [MEMORY] allocated=" << static_cast<float>(allocated) / (1 << 20)
     << " MB"
     << ", reserved=" << static_cast<float>(reserved) / (1 << 20) << " MB";
  return ss.str();
}

template <typename T>
void AllocWithDebugInfo(const GPUContext& dev_ctx,
                        const std::string& info,
                        DenseTensor* t) {
  dev_ctx.Alloc<T>(t, t->numel() * sizeof(T));
  if (VLOG_IS_ON(4)) {
    VLOG(4) << info << ": " << MemoryDebugString(*t);
  }
}

inline std::string TensorDebugString(const DenseTensor* t,
                                     const std::string& info) {
  std::stringstream ss;
  ss << info << ": ";
  if (t) {
    if (t->initialized()) {
      ss << "shape=[" << t->dims() << "], ptr=" << t->data();
    } else {
      ss << "not initialized";
    }
  } else {
    ss << "nullptr";
  }
  return ss.str();
}

inline void WaitWithDebugInfo(const GPUContext& dev_ctx) {
  if (VLOG_IS_ON(5)) {
    dev_ctx.Wait();
    VLOG(5) << "[Flash attn Synchronize] ";
  }
}

template <typename T>
inline void TypeDebugInfo() {
  if (VLOG_IS_ON(4)) {
    if (std::is_same<T, phi::float16>::value) {
      VLOG(4) << "[Grad]: T is phi::float16.";
    } else if (std::is_same<T, phi::bfloat16>::value) {
      VLOG(4) << "[Grad]: T is phi::bfloat16.";
    } else if (std::is_same<T, float>::value) {
      VLOG(4) << "[Grad]: T is float.";
    }
  }
}

template <typename T>
struct TernaryAddFunctor {
  inline HOSTDEVICE T operator()(T a, T b, T c) const { return a + b + c; }
};

template <typename T>
struct GateAttentionConfig {
 public:
  const GPUContext& dev_ctx;

  bool merge_qkv;
  bool has_gating;
  bool use_flash_attn;

  int64_t batch_size;
  int64_t seq_len_m;
  int64_t seq_len_r;
  int64_t q_dim;
  int64_t kv_dim;
  int64_t head_dim;
  int64_t m_size;
  int64_t num_heads;

  DDim qkv_out_dims;
  DDim qkv_transpose_out_dims;

  DDim q_out_dims;
  DDim kv_out_dims;
  DDim q_transpose_out_dims;
  DDim kv_transpose_out_dims;

  DDim qk_out_dims;
  DDim softmax_out_dims;
  DDim qktv_out_dims;
  DDim gate_out_dims;

  GateAttentionConfig(const GPUContext& dev_ctx,
                      const DenseTensor* query,
                      const DenseTensor* key,
                      const DenseTensor* query_weight,
                      const DenseTensor* qkv_weight,
                      bool merge_qkv,
                      bool has_gating,
                      bool use_flash_attn)
      : dev_ctx(dev_ctx),
        merge_qkv(merge_qkv),
        has_gating(has_gating),
        use_flash_attn(use_flash_attn) {
    // query: shape=[batch_size, seq_len_m, seq_len_r, q_dim]
    batch_size = query->dims()[0];
    seq_len_m = query->dims()[1];
    seq_len_r = query->dims()[2];
    q_dim = query->dims()[3];

    if (merge_qkv) {
      PADDLE_ENFORCE_NOT_NULL(
          qkv_weight,
          common::errors::NotFound("The input qkv_weight can not be nullptr "
                                   "when merge_qkv is true."));

      // When q_dim == kv_dim, QKV matmul can be computed merged.
      // qkv_weight: shape=[3, num_heads, head_dim, q_dim]
      num_heads = qkv_weight->dims()[1];
      head_dim = qkv_weight->dims()[2];
      m_size = seq_len_r;
      kv_dim = q_dim;

      qkv_out_dims = {batch_size, seq_len_m, seq_len_r, 3, num_heads, head_dim};
      qkv_transpose_out_dims = {
          3, batch_size, seq_len_m, num_heads, seq_len_r, head_dim};
    } else {
      PADDLE_ENFORCE_NOT_NULL(
          key,
          common::errors::NotFound(
              "The input key can not be nullptr when merge_qkv is false."));
      PADDLE_ENFORCE_NOT_NULL(
          query_weight,
          common::errors::NotFound("The input query_weight can not be "
                                   "nullptr when merge_qkv is false."));

      // When q_dim != kv_dim, QKV matmul must be computed saparately.
      // key: shape=[batch_size, seq_len_m, m_size, kv_dim]
      // query_w: shape=[q_dim, num_heads, head_dim]
      num_heads = query_weight->dims()[1];
      head_dim = query_weight->dims()[2];
      m_size = key->dims()[2];
      kv_dim = key->dims()[3];

      q_out_dims = {batch_size, seq_len_m, seq_len_r, num_heads, head_dim};
      kv_out_dims = {batch_size, seq_len_m, m_size, num_heads, head_dim};
      q_transpose_out_dims = {
          batch_size, seq_len_m, num_heads, seq_len_r, head_dim};
      kv_transpose_out_dims = {
          batch_size, seq_len_m, num_heads, m_size, head_dim};
    }

    qk_out_dims = {batch_size, seq_len_m, num_heads, seq_len_r, m_size};
    softmax_out_dims = {batch_size, seq_len_m, num_heads, seq_len_r, m_size};
    qktv_out_dims = {batch_size, seq_len_m, num_heads, seq_len_r, head_dim};
    gate_out_dims = {batch_size, seq_len_m, seq_len_r, num_heads, head_dim};
  }

  bool CanUseFlashAttn() const {
#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
    if (!std::is_same<T, phi::bfloat16>::value &&
        !std::is_same<T, phi::float16>::value) {
      return false;
    }

    if (merge_qkv && batch_size == 1) {
      if (head_dim == 32 || head_dim == 64 || head_dim == 128) {
        return use_flash_attn;
      }
    }
#endif
    return false;
  }

  int64_t GetQuerySize() const {
    return batch_size * seq_len_m * seq_len_r * num_heads * head_dim;
  }

  DenseTensor* GetQKVOut() {
    if (!qkv_out.IsInitialized()) {
      qkv_out.Resize(qkv_out_dims);
      AllocWithDebugInfo<T>(dev_ctx, "qkv_out", &qkv_out);
    }
    return &qkv_out;
  }

  DenseTensor* GetQueryOut() {
    if (!query_out.IsInitialized()) {
      query_out.Resize(q_out_dims);
      AllocWithDebugInfo<T>(dev_ctx, "query_out", &query_out);
    }
    return &query_out;
  }

  DenseTensor* GetKeyOut() {
    if (!key_out.IsInitialized()) {
      key_out.Resize(kv_out_dims);
      AllocWithDebugInfo<T>(dev_ctx, "key_out", &key_out);
    }
    return &key_out;
  }

  DenseTensor* GetValueOut() {
    if (!value_out.IsInitialized()) {
      value_out.Resize(kv_out_dims);
      AllocWithDebugInfo<T>(dev_ctx, "value_out", &value_out);
    }
    return &value_out;
  }

  DenseTensor* GetQKOut(DenseTensor* softmax_out) {
    // softmax_dim = qk_out_dim[-1] = qk_out_dim[rank - 1]
    int softmax_dim = m_size;
    if (!softmax_out || phi::UseCudnnSoftmax<T>(dev_ctx, softmax_dim, true)) {
      // Not sure whether cudnn softmax can execute inplace.
      if (!qkv_out.IsInitialized()) {
        qk_out.Resize(qk_out_dims);
        AllocWithDebugInfo<T>(dev_ctx, "qk_out", &qk_out);
      }
      return &qk_out;
    } else {
      // Enable inplace softmax.
      return softmax_out;
    }
  }

  DenseTensor* GetQKTVOut(DenseTensor* gate_out) {
    if (has_gating && gate_out) {
      // Reuse gate_out.
      gate_out->Resize(qktv_out_dims);
      return gate_out;
    } else {
      if (!qktv_out.IsInitialized()) {
        qktv_out.Resize(qktv_out_dims);
        AllocWithDebugInfo<T>(dev_ctx, "qktv_out", &qktv_out);
      }
      return &qktv_out;
    }
  }

  void ClearQKVOut() {
    if (qkv_out.IsInitialized()) {
      qkv_out.clear();
    }
  }

  void ClearQKOut() {
    if (qk_out.IsInitialized()) {
      qk_out.clear();
    }
  }

  void ClearQKTVOut() {
    if (qktv_out.IsInitialized()) {
      qktv_out.clear();
    }
  }

 protected:
  DenseTensor qkv_out;
  DenseTensor query_out;
  DenseTensor key_out;
  DenseTensor value_out;
  // qk_out = BatchedGEMM(Q, K^T)
  // qk_out: shape=[batch_size, seq_len_m, num_heads, seq_len_r, m_size]
  // softmax_out = softmax(qk_out + nonbatched_bias + src_mask)
  // The shape of qk_out, softmax_out is the same, thus can be called inplace.
  DenseTensor qk_out;
  // qktv_out may reuse gate_out.
  DenseTensor qktv_out;
};

template <typename T>
struct GateAttentionGradConfig : public GateAttentionConfig<T> {
 public:
  GateAttentionGradConfig(const GPUContext& dev_ctx,
                          const DenseTensor* query,
                          const DenseTensor* key,
                          const DenseTensor* query_weight,
                          const DenseTensor* qkv_weight,
                          bool merge_qkv,
                          bool has_gating,
                          bool use_flash_attn)
      : GateAttentionConfig<T>(dev_ctx,
                               query,
                               key,
                               query_weight,
                               qkv_weight,
                               merge_qkv,
                               has_gating,
                               use_flash_attn) {}

  DenseTensor* GetQKVOutGrad() {
    if (!qkv_out_grad.IsInitialized()) {
      qkv_out_grad.Resize(this->qkv_out_dims);
      AllocWithDebugInfo<T>(this->dev_ctx, "qkv_out_grad", &qkv_out_grad);
    }
    return &qkv_out_grad;
  }

  DenseTensor* GetQueryOutGrad() {
    if (!query_out_grad.IsInitialized()) {
      query_out_grad.Resize(this->q_out_dims);
      AllocWithDebugInfo<T>(this->dev_ctx, "query_out_grad", &query_out_grad);
    }
    return &query_out_grad;
  }

  DenseTensor* GetKeyOutGrad() {
    if (!key_out_grad.IsInitialized()) {
      key_out_grad.Resize(this->kv_out_dims);
      AllocWithDebugInfo<T>(this->dev_ctx, "key_out_grad", &key_out_grad);
    }
    return &key_out_grad;
  }

  DenseTensor* GetValueOutGrad() {
    if (!value_out_grad.IsInitialized()) {
      value_out_grad.Resize(this->kv_out_dims);
      AllocWithDebugInfo<T>(this->dev_ctx, "value_out_grad", &value_out_grad);
    }
    return &value_out_grad;
  }

  DenseTensor* GetQKOutGrad(DenseTensor* softmax_out_grad) {
    // softmax_dim = qk_out_dim[-1] = qk_out_dim[rank - 1]
    int softmax_dim = this->m_size;
    if (!softmax_out_grad ||
        phi::UseCudnnSoftmax<T>(this->dev_ctx, softmax_dim, true)) {
      if (!qk_out_grad.IsInitialized()) {
        qk_out_grad.Resize(this->qk_out_dims);
        AllocWithDebugInfo<T>(this->dev_ctx, "qk_out_grad", &qk_out_grad);
      }
      return &qk_out_grad;
    } else {
      return softmax_out_grad;
    }
  }

 protected:
  DenseTensor qkv_out_grad;
  DenseTensor query_out_grad;
  DenseTensor key_out_grad;
  DenseTensor value_out_grad;
  DenseTensor qk_out_grad;
};

template <typename T>
class FMHAGateRef {
 public:
  FMHAGateRef(const GPUContext& dev_ctx, bool merge_qkv)
      : dev_ctx_(dev_ctx), merge_qkv_(merge_qkv) {}

  void ComputeForward(const DenseTensor* nonbatched_bias,
                      const DenseTensor* src_mask,
                      DenseTensor* q_transpose_out,
                      DenseTensor* k_transpose_out,
                      DenseTensor* v_transpose_out,
                      DenseTensor* qkv_transpose_out,
                      DenseTensor* softmax_out,
                      DenseTensor* fmha_out,
                      DenseTensor* gate_out,
                      GateAttentionConfig<T>* config) {
    T* q_ptr = nullptr;
    T* k_ptr = nullptr;
    T* v_ptr = nullptr;

    if (merge_qkv_) {
      // qkv_transpose_out = transpose(qkv_out)
      PADDLE_ENFORCE_NOT_NULL(
          qkv_transpose_out,
          common::errors::NotFound("The input qkv_transpose_out can not be "
                                   "nullptr when merge_qkv is true."));

      DenseTensor* qkv_out = config->GetQKVOut();
      ComputeQKVTransposeForward(*qkv_out, qkv_transpose_out);
      config->ClearQKVOut();

      // q_size == k_size
      int64_t q_size = config->GetQuerySize();
      q_ptr = qkv_transpose_out->data<T>();
      k_ptr = q_ptr + q_size;
      v_ptr = k_ptr + q_size;
    } else {
      PADDLE_ENFORCE_NOT_NULL(
          q_transpose_out,
          common::errors::NotFound("The input q_transpose_out can not be "
                                   "nullptr when merge_qkv is false."));
      PADDLE_ENFORCE_NOT_NULL(
          k_transpose_out,
          common::errors::NotFound("The input k_transpose_out can not be "
                                   "nullptr when merge_qkv is false."));
      PADDLE_ENFORCE_NOT_NULL(
          v_transpose_out,
          common::errors::NotFound("The input v_transpose_out can not be "
                                   "nullptr when merge_qkv is false."));

      DenseTensor* query_out = config->GetQueryOut();
      DenseTensor* key_out = config->GetKeyOut();
      DenseTensor* value_out = config->GetValueOut();
      ComputeQKVTransposeForward(*query_out,
                                 *key_out,
                                 *value_out,
                                 q_transpose_out,
                                 k_transpose_out,
                                 v_transpose_out);

      // q_size != k_size
      q_ptr = q_transpose_out->data<T>();
      k_ptr = k_transpose_out->data<T>();
      v_ptr = v_transpose_out->data<T>();
    }
    // qk_out = BatchedGEMM(Q, K^T)
    // [batch_size, seq_len_m, num_heads, seq_len_r, head_dim] *
    //                [batch_size, seq_len_m, num_heads, m_size, head_dim]
    // -> [batch_size, seq_len_m, num_heads, seq_len_r, m_size]
    DenseTensor* qk_out = config->GetQKOut(softmax_out);
    T* qk_out_ptr = qk_out->data<T>();

    int64_t gemm_batch_size =
        config->batch_size * config->seq_len_m * config->num_heads;
    int64_t gemm_m = config->seq_len_r;
    int64_t gemm_n = config->m_size;
    int64_t gemm_k = config->head_dim;
    T alpha = static_cast<T>(1.0 / sqrt(config->head_dim));
    // attn = matmul(q, k.transpose(-1, -2))
    ComputeBatchedGEMM(q_ptr,
                       k_ptr,
                       qk_out_ptr,
                       false,
                       true,
                       gemm_m,
                       gemm_n,
                       gemm_k,
                       gemm_batch_size,
                       alpha);

    // attn = softmax_dropout(attn, 0, self.training, mask=mask, bias=bias)
    // softmax_out = softmax(qk_out + nonbatched_bias + src_mask)
    ComputeBiasMaskSoftmaxForward(
        nonbatched_bias, src_mask, qk_out, softmax_out);
    config->ClearQKOut();

    // qktv_out = BatchedGEMM(softmax_out, V)
    // [batch_size, seq_len_m, num_heads, seq_len_r, m_size] *
    // [batch_size, seq_len_m, num_heads, m_size,    head_dim]
    // -> [batch_size, seq_len_m, num_heads, seq_len_r, head_dim]
    DenseTensor* qktv_out = config->GetQKTVOut(gate_out);
    T* qktv_out_ptr = qktv_out->data<T>();

    gemm_m = config->seq_len_r;
    gemm_n = config->head_dim;
    gemm_k = config->m_size;

    // o = matmul(attn, v)
    T* softmax_out_ptr = softmax_out->data<T>();
    ComputeBatchedGEMM(softmax_out_ptr,
                       v_ptr,
                       qktv_out_ptr,
                       false,
                       false,
                       gemm_m,
                       gemm_n,
                       gemm_k,
                       gemm_batch_size);

    // fmha_out = transpose(qktv_out)
    // o = o.transpose(-2, -3).contiguous()
    ComputeQKTVTransposeForward(*qktv_out, fmha_out);

    config->ClearQKTVOut();
    if (config->has_gating) {
      gate_out->Resize(config->gate_out_dims);
    }
  }

  void ComputeBackward(const DenseTensor* q_transpose_out,
                       const DenseTensor* k_transpose_out,
                       const DenseTensor* v_transpose_out,
                       const DenseTensor* qkv_transpose_out,
                       const DenseTensor* softmax_out,
                       const DenseTensor* fmha_out_grad,
                       DenseTensor* src_mask_grad,
                       DenseTensor* nonbatched_bias_grad,
                       GateAttentionGradConfig<T>* config) {
    const T* q_ptr = nullptr;
    const T* k_ptr = nullptr;
    const T* v_ptr = nullptr;

    T* q_grad_ptr = nullptr;
    T* k_grad_ptr = nullptr;
    T* v_grad_ptr = nullptr;

    DenseTensor q_transpose_out_grad;
    DenseTensor k_transpose_out_grad;
    DenseTensor v_transpose_out_grad;
    DenseTensor qkv_transpose_out_grad;

    if (merge_qkv_) {
      PADDLE_ENFORCE_NOT_NULL(
          qkv_transpose_out,
          common::errors::NotFound("The input qkv_transpose_out can not be "
                                   "nullptr when merge_qkv is true."));

      int64_t q_size = config->GetQuerySize();
      q_ptr = qkv_transpose_out->data<T>();
      k_ptr = q_ptr + q_size;
      v_ptr = k_ptr + q_size;

      qkv_transpose_out_grad.Resize(config->qkv_transpose_out_dims);
      AllocWithDebugInfo<T>(
          dev_ctx_, "qkv_transpose_out_grad", &qkv_transpose_out_grad);

      q_grad_ptr = qkv_transpose_out_grad.data<T>();
      k_grad_ptr = q_grad_ptr + q_size;
      v_grad_ptr = k_grad_ptr + q_size;
    } else {
      PADDLE_ENFORCE_NOT_NULL(
          q_transpose_out,
          common::errors::NotFound("The input q_transpose_out can not be "
                                   "nullptr when merge_qkv is false."));
      PADDLE_ENFORCE_NOT_NULL(
          k_transpose_out,
          common::errors::NotFound("The input k_transpose_out can not be "
                                   "nullptr when merge_qkv is false."));
      PADDLE_ENFORCE_NOT_NULL(
          v_transpose_out,
          common::errors::NotFound("The input v_transpose_out can not be "
                                   "nullptr when merge_qkv is false."));

      q_ptr = q_transpose_out->data<T>();
      k_ptr = k_transpose_out->data<T>();
      v_ptr = v_transpose_out->data<T>();

      q_transpose_out_grad.Resize(config->q_transpose_out_dims);
      k_transpose_out_grad.Resize(config->kv_transpose_out_dims);
      v_transpose_out_grad.Resize(config->kv_transpose_out_dims);

      q_grad_ptr = dev_ctx_.Alloc<T>(&q_transpose_out_grad,
                                     q_transpose_out_grad.numel() * sizeof(T));
      k_grad_ptr = dev_ctx_.Alloc<T>(&k_transpose_out_grad,
                                     k_transpose_out_grad.numel() * sizeof(T));
      v_grad_ptr = dev_ctx_.Alloc<T>(&v_transpose_out_grad,
                                     v_transpose_out_grad.numel() * sizeof(T));
    }

    DenseTensor softmax_out_grad;
    softmax_out_grad.Resize(config->softmax_out_dims);
    AllocWithDebugInfo<T>(dev_ctx_, "softmax_out_grad", &softmax_out_grad);

    int64_t gemm_batch_size =
        config->batch_size * config->seq_len_m * config->num_heads;
    {
      // Forward: fmha_out = transpose(qktv_out)
      DenseTensor qktv_out_grad;
      qktv_out_grad.Resize(config->qktv_out_dims);
      AllocWithDebugInfo<T>(dev_ctx_, "qktv_out_grad", &qktv_out_grad);
      ComputeQKTVTransposeBackward(*fmha_out_grad, &qktv_out_grad);

      // Forward: qktv_out = BatchedGEMM(softmax_out, V)
      // Backward:
      //  V_grad = BatchedGEMM(softmax_out^T, qktv_out_grad) (dy = x^T * dout)
      int64_t gemm_m = config->m_size;
      int64_t gemm_n = config->head_dim;
      int64_t gemm_k = config->seq_len_r;

      const T* softmax_out_ptr = softmax_out->data<T>();
      const T* qktv_out_grad_ptr = qktv_out_grad.data<T>();
      ComputeBatchedGEMM(softmax_out_ptr,
                         qktv_out_grad_ptr,
                         v_grad_ptr,
                         true,
                         false,
                         gemm_m,
                         gemm_n,
                         gemm_k,
                         gemm_batch_size);

      // Backward: softmax_out_grad = qktv_out_grad * V^T (dx = dout * y^T)
      gemm_m = config->seq_len_r;
      gemm_n = config->m_size;
      gemm_k = config->head_dim;

      T* softmax_out_grad_ptr = softmax_out_grad.data<T>();
      ComputeBatchedGEMM(qktv_out_grad_ptr,
                         v_ptr,
                         softmax_out_grad_ptr,
                         false,
                         true,
                         gemm_m,
                         gemm_n,
                         gemm_k,
                         gemm_batch_size);
    }

    DenseTensor* qk_out_grad = config->GetQKOutGrad(&softmax_out_grad);
    ComputeBiasMaskSoftmaxBackward(&softmax_out_grad,
                                   softmax_out,
                                   src_mask_grad,
                                   qk_out_grad,
                                   nonbatched_bias_grad);

    // Forward: qk_out = BatchedGEMM(Q, K^T)
    // Backward: k_grad = BatchedGEMM(qk_out_grad^T, Q) (dy = dout^t * x)
    int64_t gemm_m = config->m_size;
    int64_t gemm_n = config->head_dim;
    int64_t gemm_k = config->seq_len_r;
    T alpha = static_cast<T>(1.0 / sqrt(config->head_dim));

    T* qk_out_grad_ptr = qk_out_grad->data<T>();
    ComputeBatchedGEMM(qk_out_grad_ptr,
                       q_ptr,
                       k_grad_ptr,
                       true,
                       false,
                       gemm_m,
                       gemm_n,
                       gemm_k,
                       gemm_batch_size,
                       alpha);

    // Backward: q_grad = BatchedGEMM(qk_out_grad, K) (dx = dout * y)
    gemm_m = config->seq_len_r;
    gemm_n = config->head_dim;
    gemm_k = config->m_size;
    ComputeBatchedGEMM(qk_out_grad_ptr,
                       k_ptr,
                       q_grad_ptr,
                       false,
                       false,
                       gemm_m,
                       gemm_n,
                       gemm_k,
                       gemm_batch_size,
                       alpha);

    if (merge_qkv_) {
      DenseTensor* qkv_out_grad = config->GetQKVOutGrad();
      ComputeQKVTransposeBackward(qkv_transpose_out_grad, qkv_out_grad);
    } else {
      DenseTensor* q_out_grad = config->GetQueryOutGrad();
      DenseTensor* k_out_grad = config->GetKeyOutGrad();
      DenseTensor* v_out_grad = config->GetValueOutGrad();
      ComputeQKVTransposeBackward(q_transpose_out_grad,
                                  k_transpose_out_grad,
                                  v_transpose_out_grad,
                                  q_out_grad,
                                  k_out_grad,
                                  v_out_grad);
    }
  }

  void ComputeQKVTransposeForward(const DenseTensor& q_out,
                                  const DenseTensor& k_out,
                                  const DenseTensor& v_out,
                                  DenseTensor* q_transpose_out,
                                  DenseTensor* k_transpose_out,
                                  DenseTensor* v_transpose_out) {
    std::vector<int> perm = {0, 1, 3, 2, 4};
    funcs::TransposeGPUKernelDriver<T>(dev_ctx_, q_out, perm, q_transpose_out);
    funcs::TransposeGPUKernelDriver<T>(dev_ctx_, k_out, perm, k_transpose_out);
    funcs::TransposeGPUKernelDriver<T>(dev_ctx_, v_out, perm, v_transpose_out);
  }

  void ComputeQKVTransposeBackward(const DenseTensor& q_transpose_out_grad,
                                   const DenseTensor& k_transpose_out_grad,
                                   const DenseTensor& v_transpose_out_grad,
                                   DenseTensor* q_out_grad,
                                   DenseTensor* k_out_grad,
                                   DenseTensor* v_out_grad) {
    std::vector<int> perm = {0, 1, 3, 2, 4};
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, q_transpose_out_grad, perm, q_out_grad);
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, k_transpose_out_grad, perm, k_out_grad);
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, v_transpose_out_grad, perm, v_out_grad);
  }

  // [batch_size, seq_len_m, seq_len_r, 3, num_heads, head_dim] ->
  //         [3, batch_size, seq_len_m, num_heads, seq_len_r, head_dim]
  void ComputeQKVTransposeForward(const DenseTensor& qkv_out,
                                  DenseTensor* qkv_transpose_out) {
    std::vector<int> perm = {3, 0, 1, 4, 2, 5};
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, qkv_out, perm, qkv_transpose_out);
  }

  void ComputeQKVTransposeBackward(const DenseTensor& qkv_transpose_out_grad,
                                   DenseTensor* qkv_out_grad) {
    std::vector<int> perm = {1, 2, 4, 0, 3, 5};
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, qkv_transpose_out_grad, perm, qkv_out_grad);
  }

  // [batch_size, seq_len_m, num_head, seq_len_r, c] ->
  //         [batch_size, seq_len_m, seq_len_r, num_head, c]
  void ComputeQKTVTransposeForward(const DenseTensor& qktv_out,
                                   DenseTensor* fmha_out) {
    std::vector<int> perm = {0, 1, 3, 2, 4};
    funcs::TransposeGPUKernelDriver<T>(dev_ctx_, qktv_out, perm, fmha_out);
  }

  void ComputeQKTVTransposeBackward(const DenseTensor& fmha_out_grad,
                                    DenseTensor* qktv_out_grad) {
    std::vector<int> perm = {0, 1, 3, 2, 4};
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, fmha_out_grad, perm, qktv_out_grad);
  }

  // qk_out = qk_out + nonbatched_bias + src_mask
  // softmax_out = softmax(src_mask_out)
  void ComputeBiasMaskSoftmaxForward(const DenseTensor* nonbatched_bias,
                                     const DenseTensor* src_mask,
                                     DenseTensor* qk_out,
                                     DenseTensor* softmax_out) {
    if (nonbatched_bias) {
      std::vector<const DenseTensor*> ins = {qk_out, src_mask, nonbatched_bias};
      std::vector<DenseTensor*> outs = {qk_out};
      funcs::BroadcastKernel<T>(dev_ctx_, ins, &outs, TernaryAddFunctor<T>());
    } else {
      std::vector<const DenseTensor*> ins = {qk_out, src_mask};
      std::vector<DenseTensor*> outs = {qk_out};
      funcs::BroadcastKernel<T>(dev_ctx_, ins, &outs, funcs::AddFunctor<T>());
    }
    phi::SoftmaxForwardCUDAKernelDriver<T>(dev_ctx_, *qk_out, -1, softmax_out);
  }

  // src_mask_out = qk_out + nonbatched_bias + src_mask
  // softmax_out = softmax(src_mask_out)
  void ComputeBiasMaskSoftmaxBackward(const DenseTensor* softmax_out_grad,
                                      const DenseTensor* softmax_out,
                                      DenseTensor* src_mask_grad,
                                      DenseTensor* qk_out_grad,
                                      DenseTensor* nonbatched_bias_grad) {
    PADDLE_ENFORCE_NOT_NULL(
        qk_out_grad,
        common::errors::NotFound("The qk_out_grad can not be nullptr."));

    PADDLE_ENFORCE_EQ(qk_out_grad->dims(),
                      softmax_out->dims(),
                      common::errors::InvalidArgument(
                          "The shape of qk_out_grad and softmax_out is "
                          "expected to be the same. But received qk_out_grad's "
                          "shape = %s, softmax_out's shape = %s.",
                          qk_out_grad->dims(),
                          softmax_out->dims()));

    PADDLE_ENFORCE_EQ(src_mask_grad,
                      nullptr,
                      common::errors::InvalidArgument(
                          "src_mask_grad is expected to be nullptr."));

    phi::SoftmaxBackwardCUDAKernelDriver<T>(
        dev_ctx_, *softmax_out, *softmax_out_grad, -1, qk_out_grad);

    if (nonbatched_bias_grad) {
      // [batch_size, seq_len_m, num_heads, seq_len_r, m_size] ->
      //      [batch_size, 1, num_heads, seq_len_r, m_size]
      funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
          dev_ctx_,
          *qk_out_grad,
          nonbatched_bias_grad,
          kps::IdentityFunctor<T>(),
          {1});
    }
  }

 private:
  void ComputeBatchedGEMM(const T* a_ptr,
                          const T* b_ptr,
                          T* c_ptr,
                          bool trans_a,
                          bool trans_b,
                          int64_t m,
                          int64_t n,
                          int64_t k,
                          int64_t batch_size,
                          T alpha = static_cast<T>(1.0),
                          T beta = static_cast<T>(0.0)) {
    int64_t stride_a = m * k;
    int64_t stride_b = k * n;

    CBLAS_TRANSPOSE cblas_trans_a = trans_a ? CblasTrans : CblasNoTrans;
    CBLAS_TRANSPOSE cblas_trans_b = trans_b ? CblasTrans : CblasNoTrans;
    auto blas = funcs::GetBlas<GPUContext, T>(dev_ctx_);
    blas.BatchedGEMM(cblas_trans_a,
                     cblas_trans_b,
                     m,
                     n,
                     k,
                     alpha,
                     a_ptr,
                     b_ptr,
                     beta,
                     c_ptr,
                     batch_size,
                     stride_a,
                     stride_b);
  }

  const GPUContext& dev_ctx_;
  bool merge_qkv_;
};

template <typename T>
class FlashAttnWithGating {
 public:
  FlashAttnWithGating(const GPUContext& dev_ctx, bool merge_qkv)
      : dev_ctx_(dev_ctx), merge_qkv_(merge_qkv) {}

  void ComputeForward(const DenseTensor* nonbatched_bias,
                      const DenseTensor* src_mask,
                      DenseTensor* qkv_transpose_out,
                      DenseTensor* softmax_lse,
                      DenseTensor* fmha_out,
                      GateAttentionConfig<T>* config) {
#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
    bool is_bf16 =
        qkv_transpose_out->dtype() == DataType::BFLOAT16 ? true : false;
    TypeDebugInfo<T>();

    PADDLE_ENFORCE_NOT_NULL(
        qkv_transpose_out,
        common::errors::NotFound("The input qkv_transpose_out can not be "
                                 "nullptr when merge_qkv is true."));

    // 1. Transpose qkv_out for flash_attn.
    DenseTensor* qkv_out = config->GetQKVOut();
    ComputeQKVTransposeForward(*qkv_out, qkv_transpose_out);
    config->ClearQKVOut();

    // q_size == k_size
    int64_t q_size = config->GetQuerySize();
    T* q_ptr = qkv_transpose_out->data<T>();
    T* k_ptr = q_ptr + q_size;
    T* v_ptr = k_ptr + q_size;

    // 2. Scale Q: q_ptr = alpha * q_ptr
    ComputeScaleQ(q_size, config->head_dim, q_ptr);

    // 3. flash_attn parameter setting.
    DenseTensor cu_seq_q;
    DenseTensor cu_seq_k;
    InitArgumentsAndSeqTensors(config, &cu_seq_q, &cu_seq_k);

    std::vector<int64_t> temp_mask_dim = GetCompressedDim(src_mask);
    std::vector<int64_t> temp_bias_dim = GetCompressedDim(nonbatched_bias);

    softmax_lse->Resize({fa_batch_size_, fa_num_heads_, fa_softmax_lse_dim_});
    AllocWithDebugInfo<float>(dev_ctx_, "softmax_lse", softmax_lse);

    if (VLOG_IS_ON(6)) {
      VLOG(6) << "temp_mask_dim={" << make_ddim(temp_mask_dim) << "}";
      VLOG(6) << "temp_bias_dim={" << make_ddim(temp_bias_dim) << "}";
      VLOG(6) << TensorDebugString(&cu_seq_q, "cu_seq_q");
      VLOG(6) << TensorDebugString(&cu_seq_k, "cu_seq_k");
      VLOG(6) << TensorDebugString(nonbatched_bias, "nonbatched_bias");
      VLOG(6) << TensorDebugString(src_mask, "src_mask");
      VLOG(6) << TensorDebugString(qkv_transpose_out, "qkv_transpose_out");
      VLOG(6) << TensorDebugString(softmax_lse, "softmax_lse");
      VLOG(6) << TensorDebugString(fmha_out, "fmha_out");
    }

    // 4. Get workspace size and run the flash-attention kernel.
    uint64_t workspace_size = 0;
    DenseTensor workspace;
    cudaStream_t stream = dev_ctx_.stream();
    for (bool need_calc : {false, true}) {
      // first calling, need_calc=false, set out_ptr to nullptr to calculate
      // workspace size second calling, need_calc=true, run flash-attention
      // kernel.
      void* out_ptr =
          need_calc ? static_cast<void*>(fmha_out->data()) : nullptr;
      void* workspace_ptr = nullptr;
      if (need_calc) {
        VLOG(6) << "Step 2: Call the flash-attention kernel";
        if (workspace_size > 0) {
          workspace = CreateWorkspace(workspace_size);
          workspace_ptr = static_cast<void*>(workspace.data());
        }
      } else {
        VLOG(6) << "Step 1: Calculate the workspace_size";
      }
      bool succ = phi::dynload::flash_attn_fwd_with_bias_and_mask(
          static_cast<const void*>(q_ptr),
          static_cast<const void*>(k_ptr),
          static_cast<const void*>(v_ptr),
          out_ptr,  // set out to nullptr to calculate workspace size
          cu_seq_q.data<int32_t>(),
          cu_seq_k.data<int32_t>(),
          fa_total_q_,
          fa_total_k_,
          fa_batch_size_,
          fa_num_heads_,
          fa_head_size_,
          fa_max_seqlen_q_,
          fa_max_seqlen_k_,
          fa_dropout_prob_,
          fa_softmax_scale_,
          fa_zero_tensors_,
          is_bf16,
          fa_num_splits_,
          softmax_lse->data(),
          workspace_ptr,
          &workspace_size,
          stream,
          fa_seed_,
          fa_offset_,
          src_mask ? src_mask->data() : nullptr,
          nonbatched_bias ? nonbatched_bias->data() : nullptr,
          src_mask ? temp_mask_dim.data() : nullptr,
          nonbatched_bias ? temp_bias_dim.data() : nullptr);
      PADDLE_ENFORCE_EQ(
          succ,
          true,
          common::errors::External(phi::dynload::flash_attn_error()));
      WaitWithDebugInfo(dev_ctx_);
    }
#else
    PADDLE_THROW(common::errors::Unimplemented(
        "FlashAttention is unsupported, please set use_flash_attn to false."));
#endif
  }

  void ComputeBackward(const DenseTensor* qkv_transpose_out,
                       const DenseTensor* src_mask,
                       const DenseTensor* nonbatched_bias,
                       const DenseTensor* softmax_lse,
                       const DenseTensor* fmha_out,
                       const DenseTensor* fmha_out_grad,
                       DenseTensor* src_mask_grad,
                       DenseTensor* nonbatched_bias_grad,
                       GateAttentionGradConfig<T>* config) {
#if defined(PADDLE_WITH_FLASHATTN) && !defined(PADDLE_WITH_HIP)
    bool is_bf16 =
        qkv_transpose_out->dtype() == DataType::BFLOAT16 ? true : false;
    TypeDebugInfo<T>();

    PADDLE_ENFORCE_NOT_NULL(
        qkv_transpose_out,
        common::errors::NotFound("The input qkv_transpose_out can not be "
                                 "nullptr when merge_qkv is true."));

    int64_t q_size = config->GetQuerySize();
    const T* q_ptr = qkv_transpose_out->data<T>();
    const T* k_ptr = q_ptr + q_size;
    const T* v_ptr = k_ptr + q_size;

    DenseTensor qkv_transpose_out_grad;
    qkv_transpose_out_grad.Resize({3,
                                   config->batch_size,
                                   config->seq_len_m,
                                   config->seq_len_r,
                                   config->num_heads,
                                   config->head_dim});
    AllocWithDebugInfo<T>(
        dev_ctx_, "qkv_transpose_out_grad", &qkv_transpose_out_grad);

    T* q_grad_ptr = qkv_transpose_out_grad.data<T>();
    T* k_grad_ptr = q_grad_ptr + q_size;
    T* v_grad_ptr = k_grad_ptr + q_size;
    WaitWithDebugInfo(dev_ctx_);

    // 1. flash_attn parameter setting.
    DenseTensor cu_seq_q;
    DenseTensor cu_seq_k;
    InitArgumentsAndSeqTensors(config, &cu_seq_q, &cu_seq_k);
    const int32_t* cu_seq_q_ptr = cu_seq_q.data<int32_t>();
    const int32_t* cu_seq_k_ptr = cu_seq_k.data<int32_t>();

    std::vector<int64_t> temp_mask_dim = GetCompressedDim(src_mask);
    std::vector<int64_t> temp_bias_dim = GetCompressedDim(nonbatched_bias);

    DenseTensor softmax_d;
    softmax_d.Resize(softmax_lse->dims());
    AllocWithDebugInfo<float>(dev_ctx_, "d_softmax_lse", &softmax_d);

    DenseTensor bias_d;
    if (nonbatched_bias) {
      bias_d.Resize(
          {fa_batch_size_, fa_num_heads_, fa_max_seqlen_q_, fa_max_seqlen_k_});
      AllocWithDebugInfo<T>(dev_ctx_, "d_bias", &bias_d);
    }

    if (VLOG_IS_ON(6)) {
      VLOG(6) << TensorDebugString(fmha_out, "fmha_out");
      VLOG(6) << TensorDebugString(fmha_out_grad, "fmha_out_grad");
      VLOG(6) << TensorDebugString(softmax_lse, "softmax_lse");
      VLOG(6) << TensorDebugString(&softmax_d, "softmax_d");
      VLOG(6) << TensorDebugString(nonbatched_bias, "nonbatched_bias");
      VLOG(6) << TensorDebugString(&bias_d, "bias_d");
    }

    // 2. Get workspace size and run the flash-attention kernel.
    uint64_t workspace_size = 0;
    DenseTensor workspace;
    cudaStream_t stream = dev_ctx_.stream();
    for (bool need_calc : {false, true}) {
      // first calling, need_calc=false, set out_ptr to nullptr to calculate
      // workspace size second calling, need_calc=true, run flash-attention
      // kernel.
      const void* out_ptr =
          need_calc ? static_cast<const void*>(fmha_out->data()) : nullptr;
      void* workspace_ptr = nullptr;
      if (need_calc) {
        VLOG(6) << "Step 2: Call the flash-attention kernel";
        if (workspace_size > 0) {
          workspace = CreateWorkspace(workspace_size);
          workspace_ptr = static_cast<void*>(workspace.data());
        }
      } else {
        VLOG(6) << "Step 1: Calculate the workspace_size";
      }

      bool succ = phi::dynload::flash_attn_bwd_with_bias_and_mask(
          static_cast<const void*>(q_ptr),
          static_cast<const void*>(k_ptr),
          static_cast<const void*>(v_ptr),
          static_cast<void*>(q_grad_ptr),
          static_cast<void*>(k_grad_ptr),
          static_cast<void*>(v_grad_ptr),
          out_ptr,  // set out to nullptr to calculate workspace size
          static_cast<const void*>(fmha_out_grad->data()),
          cu_seq_q_ptr,
          cu_seq_k_ptr,
          fa_total_q_,
          fa_total_k_,
          fa_batch_size_,
          fa_num_heads_,
          fa_head_size_,
          fa_max_seqlen_q_,
          fa_max_seqlen_k_,
          fa_dropout_prob_,
          fa_softmax_scale_,
          fa_zero_tensors_,
          is_bf16,
          fa_num_splits_,
          softmax_lse->data(),
          softmax_d.data(),
          nonbatched_bias ? bias_d.data() : nullptr,
          workspace_ptr,
          &workspace_size,
          stream,
          fa_seed_,
          fa_offset_,
          src_mask ? src_mask->data() : nullptr,
          nonbatched_bias ? nonbatched_bias->data() : nullptr,
          src_mask ? temp_mask_dim.data() : nullptr,
          nonbatched_bias ? temp_bias_dim.data() : nullptr);
      PADDLE_ENFORCE_EQ(
          succ,
          true,
          common::errors::External(phi::dynload::flash_attn_error()));
      WaitWithDebugInfo(dev_ctx_);
    }

    if (nonbatched_bias) {
      // compare block reduce
      auto dbias_first_dim = bias_d.numel() / nonbatched_bias->numel();
      bias_d.Resize({dbias_first_dim,
                     temp_bias_dim[0],
                     temp_bias_dim[1],
                     temp_bias_dim[2],
                     temp_bias_dim[3]});
      funcs::ReduceKernel<T, T, kps::AddFunctor, kps::IdentityFunctor<T>>(
          dev_ctx_,
          bias_d,
          nonbatched_bias_grad,
          kps::IdentityFunctor<T>(),
          {0});
    }

    // 3. Scale Q's grad: q_grad_ptr = alpha * q_grad_ptr
    ComputeScaleQ(q_size, config->head_dim, q_grad_ptr);

    // 4. Compute the grad of qkv_out.
    DenseTensor* qkv_out_grad = config->GetQKVOutGrad();
    ComputeQKVTransposeBackward(qkv_transpose_out_grad, qkv_out_grad);
#else
    PADDLE_THROW(common::errors::Unimplemented(
        "FlashAttention is unsupported, please set use_flash_attn to false."));
#endif
  }

 private:
  std::vector<int64_t> GetCompressedDim(const DenseTensor* tensor) {
    std::vector<int64_t> compressed_dims;
    if (tensor) {
      int64_t first_dim = 1;
      const auto& origin_dims = tensor->dims();
      auto rank = origin_dims.size();
      for (int i = 0; i < rank - 3; ++i) {
        first_dim *= origin_dims[i];
      }
      compressed_dims = {first_dim,
                         origin_dims[rank - 3],
                         origin_dims[rank - 2],
                         origin_dims[rank - 1]};
    }
    return compressed_dims;
  }

  DenseTensor CreateWorkspace(uint64_t workspace_size) {
    DenseTensor workspace;
    if (workspace_size > 0) {
      workspace = Empty<float, GPUContext>(
          dev_ctx_, {int64_t(workspace_size / sizeof(float))});
    }
    VLOG(5) << "Allocate workspace: workspace_size=" << workspace_size;
    return workspace;
  }

  void GenerateSeedAndOffset(int64_t batch_size, int64_t num_heads) {
    auto gen = dev_ctx_.GetGenerator();
    uint64_t inc = batch_size * num_heads * 32;
    auto seed_offset_pair = gen->IncrementOffset(inc);
    fa_seed_ = seed_offset_pair.first;
    fa_offset_ = seed_offset_pair.second;
  }

  void InitArgumentsAndSeqTensors(GateAttentionConfig<T>* config,
                                  DenseTensor* cu_seq_q,
                                  DenseTensor* cu_seq_k) {
    fa_batch_size_ = static_cast<int>(config->batch_size) *
                     static_cast<int>(config->seq_len_m);
    fa_num_heads_ = static_cast<int>(config->num_heads);  // qkv_dims[2];
    fa_head_size_ = static_cast<int>(config->head_dim);   // qkv_dims[3];
    fa_max_seqlen_q_ = config->seq_len_r;
    fa_max_seqlen_k_ = config->m_size;
    fa_total_q_ = fa_batch_size_ * fa_max_seqlen_q_;
    fa_total_k_ = fa_batch_size_ * fa_max_seqlen_k_;

    // 0 for an internal heuristic, which is optimal
    fa_num_splits_ = 0;
    fa_zero_tensors_ = false;

    fa_softmax_lse_dim_ = ((fa_max_seqlen_q_ + 16 - 1) / 16) * 16;
    fa_softmax_scale_ = 1.0f;
    fa_dropout_prob_ = 0.0f;
    GenerateSeedAndOffset(fa_batch_size_, fa_num_heads_);

    phi::ArangeNullaryKernel<int32_t, GPUContext>(
        dev_ctx_,
        0,
        (fa_batch_size_ + 1) * fa_max_seqlen_q_,
        fa_max_seqlen_q_,
        cu_seq_q);
    phi::ArangeNullaryKernel<int32_t, GPUContext>(
        dev_ctx_,
        0,
        (fa_batch_size_ + 1) * fa_max_seqlen_k_,
        fa_max_seqlen_k_,
        cu_seq_k);

    if (VLOG_IS_ON(6)) {
      VLOG(6) << "fa_batch_size       : " << fa_batch_size_;
      VLOG(6) << "fa_total_q          : " << fa_total_q_;
      VLOG(6) << "fa_total_k          : " << fa_total_k_;
      VLOG(6) << "fa_num_heads        : " << fa_num_heads_;
      VLOG(6) << "fa_head_size        : " << fa_head_size_;
      VLOG(6) << "fa_max_seqlen_q     : " << fa_max_seqlen_q_;
      VLOG(6) << "fa_max_seqlen_k     : " << fa_max_seqlen_k_;
      VLOG(6) << "fa_num_splits       : " << fa_num_splits_;
      VLOG(6) << "fa_softmax_lse_dim  : " << fa_softmax_lse_dim_;
      VLOG(6) << "fa_softmax_scale    : " << fa_softmax_scale_;
      VLOG(6) << "fa_dropout_prob     : " << fa_dropout_prob_;
    }
  }

  // [batch_size, seq_len_m, seq_len_r, 3, num_heads, head_dim] ->
  //         [3, batch_size, seq_len_m, seq_len_r, num_heads, head_dim]
  void ComputeQKVTransposeForward(const DenseTensor& qkv_out,
                                  DenseTensor* qkv_transpose_out) {
    std::vector<int> perm = {3, 0, 1, 2, 4, 5};
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, qkv_out, perm, qkv_transpose_out);
  }

  // [3, batch_size, seq_len_m, seq_len_r, num_heads, head_dim] ->
  //        [batch_size, seq_len_m, seq_len_r, 3, num_heads, head_dim]
  void ComputeQKVTransposeBackward(const DenseTensor& qkv_transpose_out_grad,
                                   DenseTensor* qkv_out_grad) {
    std::vector<int> perm = {1, 2, 3, 0, 4, 5};
    funcs::TransposeGPUKernelDriver<T>(
        dev_ctx_, qkv_transpose_out_grad, perm, qkv_out_grad);
  }

  void ComputeScaleQ(int64_t numel, int64_t head_dim, T* ptr) {
    float scale = static_cast<float>(1.0f / std::sqrt(head_dim));
    VLOG(6) << "[ComputeScaleQ] numel=" << numel << ", scale=" << scale;

    auto gpu_config =
        phi::backends::gpu::GetGpuLaunchConfig1D(dev_ctx_, numel, 1);
    SimpleScaleKernel<T><<<gpu_config.block_per_grid,
                           gpu_config.thread_per_block,
                           0,
                           dev_ctx_.stream()>>>(numel, scale, ptr);
  }

  const GPUContext& dev_ctx_;
  bool merge_qkv_;

  int fa_batch_size_;
  int fa_total_q_;
  int fa_total_k_;
  int fa_num_heads_;
  int fa_head_size_;
  int fa_max_seqlen_q_;
  int fa_max_seqlen_k_;
  int fa_num_splits_;
  int fa_softmax_lse_dim_;
  float fa_softmax_scale_{1.0f};
  float fa_dropout_prob_{0.0f};
  uint64_t fa_seed_{0};
  uint64_t fa_offset_{0};
  bool fa_zero_tensors_{false};
};

}  // namespace funcs
}  // namespace phi