// Copyright (c) 2025 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/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/platform/device_context.h"
#include "paddle/phi/core/tensor_utils.h"
#include "paddle/phi/kernels/empty_kernel.h"
#include "paddle/phi/kernels/funcs/math_function.h"
#include "paddle/phi/kernels/gpu/flash_attn_utils.h"
#include "paddle/phi/kernels/gpu/flash_attn_v3_utils.h"

#include "paddle/phi/kernels/concat_kernel.h"
#include "paddle/phi/kernels/gpu/flash_attn_v3_grad_kernel.h"
#include "paddle/phi/kernels/slice_kernel.h"

COMMON_DECLARE_bool(cudnn_deterministic);

namespace phi {

// b: batch_size
// s_q: seqlen_q
// s_k: seqlen_k
// h: num_heads
// h_k: num_heads_k
// d: head_size
template <typename T, typename Context>
void FlashAttnV3GradBaseKernel(
    const Context &dev_ctx,
    const DenseTensor
        &dout,  // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
    const DenseTensor
        &q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
    const DenseTensor
        &k,  // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
    const DenseTensor
        &v,  // (b, s_k, h_k, dv) or (total_k, h_k, dv) if there is cu_seqlens_k
    const DenseTensor
        &out,  // (b, s_q, h, dv) or (total_q, h, dv) if there is cu_seqlens_q
    const DenseTensor
        &softmax_lse,  // (b, h, s_q) or (h, total_q) if there is cu_seqlens_q
    const optional<DenseTensor>
        &dq_,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
    const optional<DenseTensor>
        &dk_,  // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
    const optional<DenseTensor> &dv_,  // (b, s_k, h_k, dv) or (total_k, h_k,
                                       // dv) if there is cu_seqlens_k
    const optional<DenseTensor> &cu_seqlens_q_,  // b+1
    const optional<DenseTensor> &cu_seqlens_k_,  // b+1
    const optional<DenseTensor>
        &seqused_q_,  // b. If given, only this many elements of each batch
                      // element's queries and outputs are used.
    const optional<DenseTensor>
        &seqused_k_,  // b. If given, only this many elements of each batch
                      // element's keys are used.
    int max_seqlen_q_,
    int max_seqlen_k_,
    float const softmax_scale,
    bool is_causal,
    int window_size_left,
    int window_size_right,
    float const softcap,
    bool const deterministic,
    int const sm_margin,
    DenseTensor *dq,
    DenseTensor *dk,
    DenseTensor *dv,
    DenseTensor *softmax_d,
    DenseTensor *softmax_lse_log2,
    DenseTensor *dq_accum,
    DenseTensor *dk_accum,
    DenseTensor *dv_accum) {
#ifdef PADDLE_WITH_FLASHATTN_V3

  // TODO(umiswing): support ampere
  int device_id = dev_ctx.GetPlace().GetDeviceId();
  auto dprops = paddle::platform::GetDeviceProperties(device_id);
  const bool is_sm90 = dprops.major == 9 && dprops.minor == 0;
  PADDLE_ENFORCE_EQ(is_sm90,
                    true,
                    common::errors::Unavailable(
                        "FlashAttention-3 only supports Hopper GPUs."));

  auto q_type = q.dtype();
  PADDLE_ENFORCE_EQ(
      (q_type == DataType::FLOAT16 || q_type == DataType::BFLOAT16),
      true,
      common::errors::InvalidArgument(
          "FlashAttention-3 bwd only support fp16 and bf16 data type"));
  PADDLE_ENFORCE_EQ(k.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and key must have the same dtype"));
  PADDLE_ENFORCE_EQ(v.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and value must have the same dtype"));
  PADDLE_ENFORCE_EQ(out.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and out must have the same dtype"));
  PADDLE_ENFORCE_EQ(dout.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and dout must have the same dtype"));

  CHECK_DEVICE(q);
  CHECK_DEVICE(k);
  CHECK_DEVICE(v);
  CHECK_DEVICE(out);
  CHECK_DEVICE(dout);
  CHECK_DEVICE(softmax_lse);

  PADDLE_ENFORCE_EQ(q.strides()[q.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "Input tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(k.strides()[k.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "Input tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(v.strides()[v.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "Input tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(out.strides()[out.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "out tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(dout.strides()[dout.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "dout tensor must have contiguous last dimension"));

  DenseTensor cu_seqlens_q;
  bool const is_varlen_q = cu_seqlens_q_.is_initialized();
  if (is_varlen_q) {
    cu_seqlens_q = cu_seqlens_q_.get();
    CHECK_DEVICE(cu_seqlens_q);
    CHECK_CONTIGUOUS(cu_seqlens_q);
    PADDLE_ENFORCE_EQ(cu_seqlens_q.dtype(),
                      DataType::INT32,
                      common::errors::InvalidArgument(
                          "cu_seqlens_q must have dtype paddle.int32"));
    PADDLE_ENFORCE_GT(
        max_seqlen_q_,
        0,
        common::errors::InvalidArgument(
            "max_seqlen_q must be provided if cu_seqlens_q is provided"));
  }
  DenseTensor cu_seqlens_k;
  bool const is_varlen_k = cu_seqlens_k_.is_initialized();
  if (is_varlen_k) {
    cu_seqlens_k = cu_seqlens_k_.get();
    CHECK_DEVICE(cu_seqlens_k);
    CHECK_CONTIGUOUS(cu_seqlens_k);
    PADDLE_ENFORCE_EQ(cu_seqlens_k.dtype(),
                      DataType::INT32,
                      common::errors::InvalidArgument(
                          "cu_seqlens_k must have dtype paddle.int32"));
    PADDLE_ENFORCE_GT(
        max_seqlen_k_,
        0,
        common::errors::InvalidArgument(
            "max_seqlen_k must be provided if cu_seqlens_k is provided"));
  }
  // This is what we will template on
  bool const is_varlen = is_varlen_q || is_varlen_k ||
                         seqused_q_.is_initialized() ||
                         seqused_k_.is_initialized();
#ifdef FLASHATTENTION_DISABLE_VARLEN
  PADDLE_ENFORCE_EQ(!is_varlen,
                    true,
                    common::errors::Unavailable(
                        "This flash attention build does not support varlen."));
#endif

  auto const sizes = q.dims();
  int const batch_size = !is_varlen_q ? sizes[0] : cu_seqlens_q.dims()[0] - 1;
  int const seqlen_q = !is_varlen_q ? sizes[1] : max_seqlen_q_;
  int const total_q = !is_varlen_q ? batch_size * sizes[1] : sizes[0];
  int const num_heads = q.dims()[q.dims().size() - 2];
  int const head_size = q.dims()[q.dims().size() - 1];
  int const head_size_v = v.dims()[v.dims().size() - 1];
  int const seqlen_k = !is_varlen_k ? k.dims()[1] : max_seqlen_k_;
  int const total_k = !is_varlen_k ? batch_size * k.dims()[1] : k.dims()[0];
  int const num_heads_k = k.dims()[k.dims().size() - 2];
  PADDLE_ENFORCE_EQ(
      head_size % 8,
      0,
      common::errors::InvalidArgument("head_size should be a multiple of 8"));
  int const max_headdim = get_max_headdim();
  PADDLE_ENFORCE_EQ(
      head_size_v % 8,
      0,
      common::errors::InvalidArgument("head_size_v should be a multiple of 8"));
  PADDLE_ENFORCE_LE(
      std::max(head_size, head_size_v),
      max_headdim,
      common::errors::InvalidArgument(
          "FlashAttention forward only supports head dimension at most %d",
          max_headdim));
  PADDLE_ENFORCE_EQ(
      num_heads % num_heads_k,
      0,
      common::errors::InvalidArgument(
          "Number of heads in key/value must divide number of heads in query"));

  // This needs to go before kBlockM & kBlockN since we rely on the correct
  // window_size and is_causal to set kBlockM
  if (window_size_left >= seqlen_k - 1) {
    window_size_left = -1;
  }
  if (window_size_right >= seqlen_q - 1) {
    window_size_right = -1;
  }
  if (is_causal) {
    window_size_right = 0;
  }
  // There's a case where is_causal=false, window_size=(-1, 0). Then
  // set_params_bprop will set params.is_causal=true. If we don't have is_causal
  // here matching params.is_causal, we might get the wrong kBlockM (and cause
  // IMA).
  is_causal = window_size_left < 0 && window_size_right == 0;

  int const arch = dprops.major * 10 + dprops.minor;
  int const head_size_rounded =
      round_up_headdim(std::max(head_size, head_size_v));
  int const head_size_v_rounded = head_size_rounded;
  // Very important that these match the kernel configs
  bool const is_local =
      (window_size_left >= 0 || window_size_right >= 0) && !is_causal;
  int const kBlockM_sm90 =
      head_size_rounded <= 64
          ? (is_causal && softcap > 0.0 ? 96 : 128)
          : (head_size_rounded <= 96
                 ? 64
                 : (head_size_rounded <= 128
                        ? (is_causal || is_local || softcap > 0.0 ? 64 : 80)
                        : 64));
  int const kBlockM_sm80 = head_size_rounded <= 64 ? 128 : 64;
  int const kBlockM_sm86 = head_size_rounded <= 192 ? 64 : 32;
  int const kBlockM =
      arch >= 90 ? kBlockM_sm90
                 : (arch == 86 || arch == 89 ? kBlockM_sm86 : kBlockM_sm80);
  int const kBlockN_sm90 =
      head_size_rounded <= 128 ? 128 : (head_size_rounded <= 192 ? 96 : 80);
  int const kBlockN_sm80 =
      head_size_rounded <= 128 ? 128 : (head_size_rounded <= 192 ? 80 : 64);
  int const kBlockN_sm86 =
      head_size_rounded <= 64
          ? 128
          : (head_size_rounded <= 96
                 ? 128
                 : (head_size_rounded <= 128
                        ? 96
                        : (head_size_rounded <= 192 ? 64 : 64)));
  int const kBlockN =
      arch >= 90 ? kBlockN_sm90
                 : (arch == 86 || arch == 89 ? kBlockN_sm86 : kBlockN_sm80);
  auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
  int const seqlen_q_rounded = round_multiple(seqlen_q, kBlockM);
  int const seqlen_k_rounded = round_multiple(seqlen_k, kBlockN);
  int const total_q_padded_rounded =
      round_multiple(total_q + batch_size * kBlockM, kBlockM);
  int const total_k_padded_rounded =
      round_multiple(total_k + batch_size * kBlockN, kBlockN);

  if (!is_varlen_q) {
    CHECK_SHAPE(q, batch_size, seqlen_q, num_heads, head_size);
    CHECK_SHAPE(out, batch_size, seqlen_q, num_heads, head_size_v);
    CHECK_SHAPE(dout, batch_size, seqlen_q, num_heads, head_size_v);
  } else {
    CHECK_SHAPE(q, total_q, num_heads, head_size);
    CHECK_SHAPE(out, total_q, num_heads, head_size_v);
    CHECK_SHAPE(dout, total_q, num_heads, head_size_v);
    CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
  }
  if (!is_varlen_k) {
    CHECK_SHAPE(k, batch_size, seqlen_k, num_heads_k, head_size);
    CHECK_SHAPE(v, batch_size, seqlen_k, num_heads_k, head_size_v);
  } else {
    CHECK_SHAPE(k, total_k, num_heads_k, head_size);
    CHECK_SHAPE(v, total_k, num_heads_k, head_size_v);
    CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
  }

  if (seqused_q_.is_initialized()) {
    auto seqused_q = seqused_q_.get();
    PADDLE_ENFORCE_EQ(
        seqused_q.dtype(),
        DataType::INT32,
        common::errors::InvalidArgument("seqused_q must have dtype int32"));
    CHECK_DEVICE(seqused_q);
    CHECK_CONTIGUOUS(seqused_q);
    CHECK_SHAPE(seqused_q, batch_size);
  }
  if (seqused_k_.is_initialized()) {
    auto seqused_k = seqused_k_.get();
    PADDLE_ENFORCE_EQ(
        seqused_k.dtype(),
        DataType::INT32,
        common::errors::InvalidArgument("seqused_k must have dtype int32"));
    CHECK_DEVICE(seqused_k);
    CHECK_CONTIGUOUS(seqused_k);
    CHECK_SHAPE(seqused_k, batch_size);
  }

  if (dq_.is_initialized()) {
    *dq = dq_.get();
    PADDLE_ENFORCE_EQ(
        dq->dtype(),
        q_type,
        common::errors::InvalidArgument("dq must have the same dtype as q"));
    CHECK_DEVICE((*dq));
    PADDLE_ENFORCE_EQ(dq->strides()[dq->strides().size() - 1],
                      1,
                      common::errors::InvalidArgument(
                          "dq must have contiguous last dimension"));
    if (!is_varlen_q) {
      CHECK_SHAPE((*dq), batch_size, seqlen_q, num_heads, head_size);
    } else {
      CHECK_SHAPE((*dq), total_q, num_heads, head_size);
    }
  } else {
    *dq = EmptyLike<T, Context>(dev_ctx, q);
  }
  if (dk_.is_initialized()) {
    *dk = dk_.get();
    PADDLE_ENFORCE_EQ(
        dk->dtype(),
        q_type,
        common::errors::InvalidArgument("dk must have the same dtype as q"));
    CHECK_DEVICE((*dk));
    PADDLE_ENFORCE_EQ(dk->strides()[dk->strides().size() - 1],
                      1,
                      common::errors::InvalidArgument(
                          "dk must have contiguous last dimension"));
    if (!is_varlen_k) {
      CHECK_SHAPE((*dk), batch_size, seqlen_k, num_heads_k, head_size);
    } else {
      CHECK_SHAPE((*dk), total_k, num_heads_k, head_size);
    }
  } else {
    *dk = EmptyLike<T, Context>(dev_ctx, k);
  }
  if (dv_.is_initialized()) {
    *dv = dv_.get();
    PADDLE_ENFORCE_EQ(
        dv->dtype(),
        q_type,
        common::errors::InvalidArgument("dv must have the same dtype as q"));
    CHECK_DEVICE((*dv));
    PADDLE_ENFORCE_EQ(dv->strides()[dv->strides().size() - 1],
                      1,
                      common::errors::InvalidArgument(
                          "dv must have contiguous last dimension"));
    if (!is_varlen_k) {
      CHECK_SHAPE((*dv), batch_size, seqlen_k, num_heads_k, head_size_v);
    } else {
      CHECK_SHAPE((*dv), total_k, num_heads_k, head_size_v);
    }
  } else {
    *dv = EmptyLike<T, Context>(dev_ctx, v);
  }

  // Otherwise the kernel will be launched from cuda:0 device
  // Cast to char to avoid compiler warning about narrowing

  // Need softmax_d to have total_q_padded_rounded since we want its address to
  // be aligned by 16/8 bytes for TMA / LDG.64
  if (!is_varlen) {
    if (softmax_d) {
      // Need softmax_d to have seqlen_q_rounded since we want its address to be
      // aligned by 16/8 bytes for TMA / LDG.64
      softmax_d->Resize({batch_size, num_heads, seqlen_q_rounded});
    }
    if (softmax_lse_log2) {
      softmax_lse_log2->Resize(
          make_ddim({batch_size, num_heads, seqlen_q_rounded}));
    }
  } else {
    if (softmax_d) {
      softmax_d->Resize({num_heads, total_q_padded_rounded});
    }
    if (softmax_lse_log2) {
      softmax_lse_log2->Resize({num_heads, total_q_padded_rounded});
    }
  }
  if (softmax_d) {
    dev_ctx.template Alloc<float>(softmax_d);
  }
  if (softmax_lse_log2) {
    dev_ctx.template Alloc<float>(softmax_lse_log2);
  }
  if (dq_accum) {
    if (!is_varlen) {
      dq_accum->Resize(make_ddim(
          {batch_size, num_heads, seqlen_q_rounded * head_size_rounded}));
    } else {
      dq_accum->Resize(
          make_ddim({num_heads, total_q_padded_rounded * head_size_rounded}));
    }
    dev_ctx.template Alloc<float>(dq_accum);
  }
  if (num_heads_k != num_heads) {  // MQA / GQA
    if (!is_varlen) {
      if (dk_accum) {
        dk_accum->Resize(make_ddim(
            {batch_size, num_heads_k, seqlen_k_rounded * head_size_rounded}));
      }
      if (dv_accum) {
        dv_accum->Resize(make_ddim(
            {batch_size, num_heads_k, seqlen_k_rounded * head_size_v_rounded}));
      }
    } else {
      if (dk_accum) {
        dk_accum->Resize(make_ddim(
            {num_heads_k, total_k_padded_rounded, head_size_rounded}));
      }
      if (dv_accum) {
        dv_accum->Resize(make_ddim(
            {num_heads_k, total_k_padded_rounded, head_size_v_rounded}));
      }
    }
    if (dk_accum) {
      dev_ctx.template Alloc<float>(dk_accum);
    }
    if (dv_accum) {
      dev_ctx.template Alloc<float>(dv_accum);
    }
    funcs::SetConstant<Context, float> set_zero;

    if (dk_accum) {
      set_zero(dev_ctx, dk_accum, float{0});
    }
    if (dv_accum) {
      set_zero(dev_ctx, dv_accum, float{0});
    }
  }

  Flash_bwd_params *params_handle = get_flash_bwd_params_handle();
  dynload::fa3_clear_bwd_params_handle(params_handle);
  set_params_dgrad(
      params_handle,
      batch_size,
      seqlen_q,
      seqlen_k,
      seqlen_q_rounded,
      seqlen_k_rounded,
      num_heads,
      num_heads_k,
      head_size,
      head_size_rounded,
      q,
      k,
      v,
      out,
      dout,
      dq,
      dk,
      dv,
      !is_varlen_q ? nullptr : cu_seqlens_q.data(),
      !is_varlen_k ? nullptr : cu_seqlens_k.data(),
      seqused_q_.is_initialized() ? const_cast<void *>(seqused_q_.get().data())
                                  : nullptr,
      seqused_k_.is_initialized() ? const_cast<void *>(seqused_k_.get().data())
                                  : nullptr,
      dq_accum ? dq_accum->data() : nullptr,
      num_heads_k != num_heads && dk_accum ? dk_accum->data() : nullptr,
      num_heads_k != num_heads && dv_accum ? dv_accum->data() : nullptr,
      const_cast<void *>(softmax_lse.data()),
      softmax_d ? const_cast<void *>(softmax_d->data()) : nullptr,
      /*p_dropout=*/0.f,
      softmax_scale,
      window_size_left,
      window_size_right,
      dprops,
      softcap,
      deterministic,
      sm_margin);
  dynload::fa3_bwd_params_set_total_q(params_handle, total_q);
  dynload::fa3_bwd_params_set_total_k(params_handle, total_k);
  dynload::fa3_bwd_params_set_softmax_lse_log2_ptr(
      params_handle, softmax_lse_log2 ? softmax_lse_log2->data() : nullptr);
  dynload::fa3_bwd_params_set_dv(params_handle, head_size_v);
  dynload::fa3_bwd_params_set_dv_rounded(params_handle, head_size_v_rounded);
  // auto tile_count_semaphore = (params.is_causal || params.is_local) ?
  // paddle::zeros({1}, opts.dtype(torch::kInt32)) : torch::empty({1},
  // opts.dtype(torch::kInt32)); params.tile_count_semaphore =
  // tile_count_semaphore.data_ptr<int>(); Will be zero'ed out in the backward
  // preprocess kernel
  DenseTensor dq_semaphore = Empty<int32_t>(
      dev_ctx, {(seqlen_q + kBlockM - 1) / kBlockM, batch_size, num_heads});
  dynload::fa3_bwd_params_set_dq_semaphore(params_handle,
                                           dq_semaphore.data<int>());
  DenseTensor dk_semaphore = Empty<int32_t>(
      dev_ctx, {(seqlen_k + kBlockN - 1) / kBlockN, batch_size, num_heads_k});
  DenseTensor dv_semaphore = Empty<int32_t>(
      dev_ctx, {(seqlen_k + kBlockN - 1) / kBlockN, batch_size, num_heads_k});
  if (num_heads_k != num_heads &&
      dynload::fa3_bwd_params_get_deterministic(params_handle)) {
    funcs::SetConstant<Context, int32_t> set_zero_dk;
    set_zero_dk(dev_ctx, &dk_semaphore, static_cast<int32_t>(0));
    funcs::SetConstant<Context, int32_t> set_zero_dv;
    set_zero_dv(dev_ctx, &dv_semaphore, static_cast<int32_t>(0));
    dynload::fa3_bwd_params_set_dk_semaphore(params_handle,
                                             dk_semaphore.data<int>());
    dynload::fa3_bwd_params_set_dv_semaphore(params_handle,
                                             dv_semaphore.data<int>());
  }

#ifdef FLASHATTENTION_DISABLE_LOCAL
  PADDLE_ENABLE_EQ(
      !dynload::fa3_bwd_params_get_is_local(params_handle),
      true,
      "This flash attention build does not support local attention.");
#endif
#ifdef FLASHATTENTION_DISABLE_SOFTCAP
  PADDLE_ENABLE_EQ(
      dynload::fa3_bwd_params_get_softcap(params_handle),
      0.0,
      "This flash attention build does not support tanh softcapping.");
#endif

  if (total_q > 0 && total_k > 0 && num_heads_k > 0) {
    dynload::fa3_run_mha_bwd(params_handle, dev_ctx.stream());
  } else if (total_k > 0 && num_heads_k > 0) {
    // If seqlen_q == 0, then we have an empty tensor. We need to set the output
    // to 0.
    funcs::SetConstant<Context, T> set_zero;
    set_zero(dev_ctx, dk, T{0});
    set_zero(dev_ctx, dv, T{0});
    if (softmax_d) {
      funcs::SetConstant<Context, float> set_zero_fp32;
      set_zero_fp32(dev_ctx, softmax_d, float{0});
    }
  } else if (total_q > 0 && num_heads_k > 0) {
    funcs::SetConstant<Context, T> set_zero;
    set_zero(dev_ctx, dq, T{0});
    if (softmax_d) {
      funcs::SetConstant<Context, float> set_zero_fp32;
      set_zero_fp32(dev_ctx, softmax_d, float{0});
    }
  }
#else
  RaiseNotSupportedError();
#endif
}

template <typename T, typename Context>
void FlashAttnV3GradKernel(const Context &dev_ctx,
                           const DenseTensor &q,
                           const DenseTensor &k,
                           const DenseTensor &v,
                           const DenseTensor &out,
                           const DenseTensor &softmax_lse,
                           const DenseTensor &out_grad,
                           float const softmax_scale,
                           bool is_causal,
                           int window_size_left,
                           int window_size_right,
                           float const softcap,
                           int const sm_margin,
                           DenseTensor *dq,
                           DenseTensor *dk,
                           DenseTensor *dv) {
#ifdef PADDLE_WITH_FLASHATTN_V3
  PADDLE_ENFORCE_EQ(softcap,
                    0,
                    common::errors::InvalidArgument(
                        "softcap is not supported, please set softcap to 0"));
  PADDLE_ENFORCE_EQ(
      sm_margin,
      0,
      common::errors::InvalidArgument(
          "sm_margin is not supported, please set sm_margin to 0"));
  // umiswing: fake grad tensor for FlashAttnV3GradBaseKernel
  DenseTensor softmax_d;
  DenseTensor softmax_lse_log2;
  DenseTensor dq_accum;
  DenseTensor dk_accum;
  DenseTensor dv_accum;
  const int64_t b = q.dims()[0];
  const int64_t s_q = q.dims()[1];
  const int64_t s_k = k.dims()[1];
  const int64_t h_q = q.dims()[2];
  const int64_t h_k = k.dims()[2];
  const int64_t d_q = q.dims()[3];
  const int64_t d_v = v.dims()[3];
  const bool is_mla =
      q.dims()[q.dims().size() - 1] > v.dims()[v.dims().size() - 1];
  if (is_mla) {
    PADDLE_ENFORCE_EQ(v.dims()[v.dims().size() - 1],
                      out.dims()[out.dims().size() - 1],
                      common::errors::InvalidArgument(
                          "head_dim_v and head_dim_o must be equal"));
    PADDLE_ENFORCE_EQ(v.dims()[v.dims().size() - 2],
                      out.dims()[out.dims().size() - 2],
                      common::errors::InvalidArgument(
                          "num_heads_v and num_heads_o must be equal"));
    PADDLE_ENFORCE_EQ(
        v.dims()[v.dims().size() - 3],
        out.dims()[out.dims().size() - 3],
        common::errors::InvalidArgument("seqlen_v and seqlen_o must be equal"));
  }
  FlashAttnV3GradBaseKernel<T, Context>(dev_ctx,
                                        out_grad,
                                        q,
                                        k,
                                        v,
                                        out,
                                        softmax_lse,
                                        paddle::none,
                                        paddle::none,
                                        paddle::none,
                                        paddle::none,
                                        paddle::none,
                                        paddle::none,
                                        paddle::none,
                                        0,
                                        0,
                                        softmax_scale,
                                        is_causal,
                                        window_size_left,
                                        window_size_right,
                                        softcap,
                                        FLAGS_cudnn_deterministic,
                                        sm_margin,
                                        dq,
                                        dk,
                                        dv,
                                        &softmax_d,
                                        &softmax_lse_log2,
                                        &dq_accum,
                                        &dk_accum,
                                        &dv_accum);
  // umiswing: some branch in upstream fa3 could have padded the head
  // dimension
  PADDLE_ENFORCE_EQ(
      dq->dims()[dq->dims().size() - 1],
      q.dims()[q.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dq != head dimension of q (%d != %d)",
          dq->dims()[dq->dims().size() - 1],
          q.dims()[q.dims().size() - 1]));

  PADDLE_ENFORCE_EQ(
      dk->dims()[dk->dims().size() - 1],
      k.dims()[k.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dk != head dimension of k (%d != %d)",
          dk->dims()[dk->dims().size() - 1],
          k.dims()[k.dims().size() - 1]));

  PADDLE_ENFORCE_EQ(
      dv->dims()[dv->dims().size() - 1],
      v.dims()[v.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dv != head dimension of v (%d != %d)",
          dv->dims()[dv->dims().size() - 1],
          v.dims()[v.dims().size() - 1]));
#else
  RaiseNotSupportedError();
#endif
}

template <typename T, typename Context>
void FlashAttnV3VarlenGradKernel(const Context &dev_ctx,
                                 const DenseTensor &q,
                                 const DenseTensor &k,
                                 const DenseTensor &v,
                                 const DenseTensor &out,
                                 const DenseTensor &softmax_lse,
                                 const DenseTensor &cu_seqlens_q,
                                 const DenseTensor &cu_seqlens_k,
                                 const optional<DenseTensor> &seqused_q,
                                 const optional<DenseTensor> &seqused_k,
                                 const DenseTensor &out_grad,
                                 float const softmax_scale,
                                 const Scalar &max_seqlen_q,
                                 const Scalar &max_seqlen_k,
                                 bool const causal,
                                 int const window_size_left,
                                 int const window_size_right,
                                 float const softcap,
                                 int const sm_margin,
                                 DenseTensor *dq,
                                 DenseTensor *dk,
                                 DenseTensor *dv) {
#ifdef PADDLE_WITH_FLASHATTN_V3
  PADDLE_ENFORCE_EQ(softcap,
                    0,
                    common::errors::InvalidArgument(
                        "softcap is not supported, please set softcap to 0"));
  PADDLE_ENFORCE_EQ(
      sm_margin,
      0,
      common::errors::InvalidArgument(
          "sm_margin is not supported, please set sm_margin to 0"));

  const int64_t head_size = q.dims()[q.dims().size() - 1];
  const int64_t head_size_v = v.dims()[v.dims().size() - 1];
  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));

  // umiswing: fake grad tensor for FlashAttnV3GradBaseKernel
  DenseTensor softmax_d;
  DenseTensor softmax_lse_log2;
  DenseTensor dq_accum;
  DenseTensor dk_accum;
  DenseTensor dv_accum;
  const int64_t max_seqlen_q_ = max_seqlen_q.to<int64_t>();
  const int64_t max_seqlen_k_ = max_seqlen_k.to<int64_t>();
  FlashAttnV3GradBaseKernel<T, Context>(dev_ctx,
                                        out_grad,
                                        q,
                                        k,
                                        v,
                                        out,
                                        softmax_lse,
                                        paddle::none,  // dq_
                                        paddle::none,  // dk_
                                        paddle::none,  // dv_
                                        cu_seqlens_q,
                                        cu_seqlens_k,
                                        seqused_q,
                                        seqused_k,
                                        max_seqlen_q_,
                                        max_seqlen_k_,
                                        softmax_scale,
                                        causal,
                                        window_size_left,
                                        window_size_right,
                                        softcap,
                                        FLAGS_cudnn_deterministic,
                                        sm_margin,
                                        dq,
                                        dk,
                                        dv,
                                        &softmax_d,
                                        &softmax_lse_log2,
                                        &dq_accum,
                                        &dk_accum,
                                        &dv_accum);

  // umiswing: some branch in upstream fa3 could have padded the head dimension
  PADDLE_ENFORCE_EQ(
      dq->dims()[dq->dims().size() - 1],
      q.dims()[q.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dq != head dimension of q (%d != %d)",
          dq->dims()[dq->dims().size() - 1],
          q.dims()[q.dims().size() - 1]));

  PADDLE_ENFORCE_EQ(
      dk->dims()[dk->dims().size() - 1],
      k.dims()[k.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dk != head dimension of k (%d != %d)",
          dk->dims()[dk->dims().size() - 1],
          k.dims()[k.dims().size() - 1]));

  PADDLE_ENFORCE_EQ(
      dv->dims()[dv->dims().size() - 1],
      v.dims()[v.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dv != head dimension of v (%d != %d)",
          dv->dims()[dv->dims().size() - 1],
          v.dims()[v.dims().size() - 1]));

#else
  RaiseNotSupportedError();
#endif
}

template <typename T, typename Context>
void FlashMaskV2GradBaseKernel(
    const Context &dev_ctx,
    const DenseTensor
        &dout,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
    const DenseTensor
        &q,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
    const DenseTensor
        &k,  // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
    const DenseTensor
        &v,  // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
    const DenseTensor
        &out,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
    const DenseTensor
        &softmax_lse,  // (b, h, s_q) or (h, total_q) if there is cu_seqlens_q
    const optional<DenseTensor>
        &dq_,  // (b, s_q, h, d) or (total_q, h, d) if there is cu_seqlens_q
    const optional<DenseTensor>
        &dk_,  // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
    const optional<DenseTensor>
        &dv_,  // (b, s_k, h_k, d) or (total_k, h_k, d) if there is cu_seqlens_k
    const optional<DenseTensor> &cu_seqlens_q_,  // b+1
    const optional<DenseTensor> &cu_seqlens_k_,  // b+1
    const optional<DenseTensor>
        &seqused_q_,  // b. If given, only this many elements of each batch
                      // element's queries and outputs are used.
    const optional<DenseTensor>
        &seqused_k_,  // b. If given, only this many elements of each batch
                      // element's keys are used.
    const optional<DenseTensor> &startend_row_indices_,
    const optional<DenseTensor> &block_mask_,  // （(b,h,s//128,s//128)
    int max_seqlen_q_,
    int max_seqlen_k_,
    float const softmax_scale,
    bool is_causal,
    int window_size_left,
    int window_size_right,
    float const softcap,
    bool const deterministic,
    int const sm_margin,
    int const rank,
    int const nranks,
    DenseTensor *dq,
    DenseTensor *dk,
    DenseTensor *dv,
    DenseTensor *softmax_d,
    DenseTensor *softmax_lse_log2,
    DenseTensor *dq_accum,
    DenseTensor *dk_accum,
    DenseTensor *dv_accum) {
#ifdef PADDLE_WITH_FLASHATTN_V3

  // TODO(umiswing): support ampere
  int device_id = dev_ctx.GetPlace().GetDeviceId();
  auto dprops = paddle::platform::GetDeviceProperties(device_id);
  const bool is_sm90 = dprops.major == 9 && dprops.minor == 0;
  PADDLE_ENFORCE_EQ(is_sm90,
                    true,
                    common::errors::Unavailable(
                        "FlashAttention-3 only supports Hopper GPUs."));

  auto q_type = q.dtype();
  PADDLE_ENFORCE_EQ(
      (q_type == DataType::FLOAT16 || q_type == DataType::BFLOAT16),
      true,
      common::errors::InvalidArgument(
          "FlashAttention-3 bwd only support fp16 and bf16 data type"));
  PADDLE_ENFORCE_EQ(k.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and key must have the same dtype"));
  PADDLE_ENFORCE_EQ(v.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and value must have the same dtype"));
  PADDLE_ENFORCE_EQ(out.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and out must have the same dtype"));
  PADDLE_ENFORCE_EQ(dout.dtype(),
                    q_type,
                    common::errors::InvalidArgument(
                        "query and dout must have the same dtype"));

  CHECK_DEVICE(q);
  CHECK_DEVICE(k);
  CHECK_DEVICE(v);
  CHECK_DEVICE(out);
  CHECK_DEVICE(dout);
  CHECK_DEVICE(softmax_lse);

  PADDLE_ENFORCE_EQ(q.strides()[q.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "Input tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(k.strides()[k.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "Input tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(v.strides()[v.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "Input tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(out.strides()[out.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "out tensor must have contiguous last dimension"));
  PADDLE_ENFORCE_EQ(dout.strides()[dout.strides().size() - 1],
                    1,
                    common::errors::InvalidArgument(
                        "dout tensor must have contiguous last dimension"));

  DenseTensor cu_seqlens_q;
  bool const is_varlen_q = cu_seqlens_q_.is_initialized();
  if (is_varlen_q) {
    cu_seqlens_q = cu_seqlens_q_.get();
    CHECK_DEVICE(cu_seqlens_q);
    CHECK_CONTIGUOUS(cu_seqlens_q);
    PADDLE_ENFORCE_EQ(cu_seqlens_q.dtype(),
                      DataType::INT32,
                      common::errors::InvalidArgument(
                          "cu_seqlens_q must have dtype paddle.int32"));
    PADDLE_ENFORCE_GT(
        max_seqlen_q_,
        0,
        common::errors::InvalidArgument(
            "max_seqlen_q must be provided if cu_seqlens_q is provided"));
  }
  DenseTensor cu_seqlens_k;
  bool const is_varlen_k = cu_seqlens_k_.is_initialized();
  if (is_varlen_k) {
    cu_seqlens_k = cu_seqlens_k_.get();
    CHECK_DEVICE(cu_seqlens_k);
    CHECK_CONTIGUOUS(cu_seqlens_k);
    PADDLE_ENFORCE_EQ(cu_seqlens_k.dtype(),
                      DataType::INT32,
                      common::errors::InvalidArgument(
                          "cu_seqlens_k must have dtype paddle.int32"));
    PADDLE_ENFORCE_GT(
        max_seqlen_k_,
        0,
        common::errors::InvalidArgument(
            "max_seqlen_k must be provided if cu_seqlens_k is provided"));
  }
  // This is what we will template on
  bool const is_varlen = is_varlen_q || is_varlen_k ||
                         seqused_q_.is_initialized() ||
                         seqused_k_.is_initialized();
#ifdef FLASHATTENTION_DISABLE_VARLEN
  PADDLE_ENFORCE_EQ(!is_varlen,
                    true,
                    common::errors::Unavailable(
                        "This flash attention build does not support varlen."));
#endif

  auto const sizes = q.dims();
  int const batch_size = !is_varlen_q ? sizes[0] : cu_seqlens_q.dims()[0] - 1;
  int const seqlen_q = !is_varlen_q ? sizes[1] : max_seqlen_q_;
  int const total_q = !is_varlen_q ? batch_size * sizes[1] : sizes[0];
  int const num_heads = q.dims()[q.dims().size() - 2];
  int const head_size = q.dims()[q.dims().size() - 1];
  int const seqlen_k = !is_varlen_k ? k.dims()[1] : max_seqlen_k_;
  int const total_k = !is_varlen_k ? batch_size * k.dims()[1] : k.dims()[0];
  int const num_heads_k = k.dims()[k.dims().size() - 2];
  PADDLE_ENFORCE_EQ(
      head_size % 8,
      0,
      common::errors::InvalidArgument("head_size should be a multiple of 8"));
  int const max_headdim = flashmaskv2_get_max_headdim();
  PADDLE_ENFORCE_LE(
      head_size,
      max_headdim,
      common::errors::InvalidArgument(
          "FlashAttention forward only supports head dimension at most %d",
          max_headdim));
  PADDLE_ENFORCE_EQ(
      num_heads % num_heads_k,
      0,
      common::errors::InvalidArgument(
          "Number of heads in key/value must divide number of heads in query"));

  // This needs to go before kBlockM & kBlockN since we rely on the correct
  // window_size and is_causal to set kBlockM
  if (window_size_left >= seqlen_k - 1) {
    window_size_left = -1;
  }
  if (window_size_right >= seqlen_q - 1) {
    window_size_right = -1;
  }
  if (is_causal) {
    window_size_right = 0;
  }
  // There's a case where is_causal=false, window_size=(-1, 0). Then
  // set_params_bprop will set params.is_causal=true. If we don't have is_causal
  // here matching params.is_causal, we might get the wrong kBlockM (and cause
  // IMA).
  is_causal = window_size_left < 0 && window_size_right == 0;

  int const arch = dprops.major * 10 + dprops.minor;
  int const head_size_rounded = flashmaskv2_round_up_headdim(head_size);
  // Very important that these match the kernel configs
  bool const is_local =
      (window_size_left >= 0 || window_size_right >= 0) && !is_causal;
  bool const is_flashmask = startend_row_indices_.is_initialized();
  DenseTensor startend_row_indices;
  if (is_flashmask) startend_row_indices = startend_row_indices_.get();
  bool const has_softcap = softcap > 0.0;

  // flashmask
  DenseTensor flashmask_maxmin, lt_start_row_indices, lt_end_row_indices,
      ut_start_row_indices, ut_end_row_indices;
  if (is_flashmask) {
    PADDLE_ENFORCE_EQ(
        startend_row_indices.dtype(),
        DataType::INT32,
        common::errors::InvalidArgument(
            "flashmask_attention startend_row_indices must be INT32 type"));
    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 = startend_row_indices.dims();
    // TODO(umiswing): refine this block constraint (kBlockN % 32), since some
    // of kBlockN is not divisible by 32 flashmask_maxmin_shape[2] =
    // (flashmask_maxmin_shape[2] + 31) / 32 * 8;
    flashmask_maxmin_shape[2] =
        ((flashmask_maxmin_shape[2] + 31) / 32 + 3) / 4 * 4;
    flashmask_maxmin_shape[3] = 8;

    flashmask_maxmin.set_type(DataType::INT32);
    flashmask_maxmin.Resize(flashmask_maxmin_shape);
    dev_ctx.template Alloc<int32_t>(&flashmask_maxmin);

    lt_start_row_indices =
        Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {0}, {1});
    if (startend_row_indices.dims()[3] == 2) {
      if (!is_causal) {
        ut_end_row_indices =
            Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {1}, {2});
      } else {
        lt_end_row_indices =
            Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {1}, {2});
      }
    } else if (startend_row_indices.dims()[3] == 4) {
      ut_end_row_indices =
          Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {3}, {4});
      lt_end_row_indices =
          Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {1}, {2});
      ut_start_row_indices =
          Slice<int32_t>(dev_ctx, startend_row_indices, {3}, {2}, {3});
    }
  }

  bool const is_blockmask = block_mask_.is_initialized();
  DenseTensor block_mask;
  if (is_blockmask) block_mask = block_mask_.get();

  if (is_blockmask) {
    PADDLE_ENFORCE_EQ(
        is_flashmask,
        true,
        common::errors::InvalidArgument(
            "blockmask should be used with flashmask at the same time "));

    PADDLE_ENFORCE_EQ(block_mask.dims().size(),
                      4,
                      common::errors::InvalidArgument(
                          "blockmask receive blockmask_indices with dim "
                          "[batch_size, num_heads, blocklen_q, blocklen_k]"));

    PADDLE_ENFORCE_EQ(block_mask.dims()[2],
                      (seqlen_q + 127) / 128,
                      common::errors::InvalidArgument(
                          "blockmask only supports blockdim_q = 128 now"));

    PADDLE_ENFORCE_EQ(block_mask.dims()[3],
                      (seqlen_k + 127) / 128,
                      common::errors::InvalidArgument(
                          "blockmask only supports blockdim_k = 128 now"));

    PADDLE_ENFORCE_EQ(
        block_mask.dims()[1],
        startend_row_indices.dims()[1],
        common::errors::InvalidArgument(
            "blockmask only supports same dim num_heads with flashmask now"));

    PADDLE_ENFORCE_LE(seqlen_k,
                      1024 * 128,
                      common::errors::InvalidArgument(
                          "blockmask only supports seqlen <= 128k in bwd now"));

    PADDLE_ENFORCE_LE(seqlen_q,
                      1024 * 128,
                      common::errors::InvalidArgument(
                          "blockmask only supports seqlen <= 128k in bwd now"));
  }

  const bool has_lt_start = lt_start_row_indices.initialized();
  const bool has_lt_end = lt_end_row_indices.initialized();
  const bool has_ut_start = ut_start_row_indices.initialized();
  const bool has_ut_end = ut_end_row_indices.initialized();

  // umiswing: The tile dispatch for flashmask is now different from fa3.
  // Replacing the original ternary operator with lambda makes the code
  // easier to reason about and less error-prone.
  const auto [kBlockM_sm90, kBlockN_sm90] = [&]() -> std::pair<int, int> {
    if (head_size_rounded <= 64) {
      if (is_flashmask && !is_causal) {
        return {64, 96};
      } else if (is_causal && has_softcap || is_flashmask) {
        return {96, 128};
      } else {
        return {128, 128};
      }
    } else if (head_size_rounded <= 128) {
      // umiswing: by now, we reuse template instantiation of head dim 128 for
      // head dim in range (64, 128], and therefore no separate dispatch for
      // head dim in range (64, 96]
      if (is_causal || is_local || has_softcap) {
        return {64, 128};
      } else {
        if ((seqlen_q >= 1024 || seqlen_k >= 1024) &&
            !(has_lt_end && has_ut_start)) {
          return {64, 128};
        } else {
          return {64, 64};
        }
      }
    } else if (head_size_rounded <= 256) {
      // umiswing: by now, we reuse template instantiation of head dim 256 for
      // head dim in range (128, 256], and therefore no separate dispatch for
      // head dim in range (128, 192]
      if (has_lt_end && has_ut_start) {
        return {64, 32};
      } else {
        return {64, 64};
      }
    } else {
      PADDLE_THROW(
          common::errors::Unimplemented("head dim is rounded to %d, which is "
                                        "not supported in FlashMask V3 now.",
                                        head_size_rounded));
      return {0, 0};
    }
  }();

  int const kBlockM_sm80 = head_size_rounded <= 64 ? 128 : 64;
  int const kBlockM_sm86 = head_size_rounded <= 192 ? 64 : 32;
  int const kBlockM =
      arch >= 90 ? kBlockM_sm90
                 : (arch == 86 || arch == 89 ? kBlockM_sm86 : kBlockM_sm80);
  int const kBlockN_sm80 =
      head_size_rounded <= 128 ? 128 : (head_size_rounded <= 192 ? 80 : 64);
  int const kBlockN_sm86 =
      head_size_rounded <= 64
          ? 128
          : (head_size_rounded <= 96
                 ? 128
                 : (head_size_rounded <= 128
                        ? 96
                        : (head_size_rounded <= 192 ? 64 : 64)));
  int const kBlockN =
      arch >= 90 ? kBlockN_sm90
                 : (arch == 86 || arch == 89 ? kBlockN_sm86 : kBlockN_sm80);
  auto round_multiple = [](int x, int m) { return (x + m - 1) / m * m; };
  int const seqlen_q_rounded = round_multiple(seqlen_q, kBlockM);

  // if KV head >= 4, we will consider using RS overlap
  PADDLE_ENFORCE_LE(
      nranks,
      64,
      common::errors::InvalidArgument(
          "nranks for FlashMask overlap should <= 64, got: %d", nranks));
  PADDLE_ENFORCE_EQ(
      (nranks == 1) || (seqlen_k >= 4096 && seqlen_k <= 131072 &&
                        (seqlen_k & (seqlen_k - 1)) == 0),
      true,
      common::errors::InvalidArgument(
          "If nranks > 1 (using overlap), currently only [4, 8, 16, 32, 64, "
          "128]K seqlen_k is supported, got nranks = %d, seqlen_k = %d",
          nranks,
          seqlen_k));
  const int chunks_per_seg = dynload::flashmaskv2_get_num_chunks_per_stage(
      seqlen_k, nranks, num_heads_k);
  PADDLE_ENFORCE_GT(
      chunks_per_seg,
      0,
      common::errors::InvalidArgument(
          "chunks_per_seg should be at least 1, but got: %d. Check whether "
          "WITH_NVSHMEM is on for this Paddle compile.",
          chunks_per_seg));
  bool const use_rs_overlap = nranks > 1;
  VLOG(6) << "FlashMask RS overlap: use rs: " << use_rs_overlap
          << ", num chunk: " << chunks_per_seg;
  int const dkv_accum_s_scaler =
      use_rs_overlap ? chunks_per_seg : nranks;  // * cp_size
  int const dkv_s_scaler =
      use_rs_overlap ? 1 : nranks;  // dk, dv remains local seqlen
  int const seqlen_k_rounded_cp =
      round_multiple(dkv_accum_s_scaler * seqlen_k, kBlockN);
  int const total_q_padded_rounded =
      round_multiple(total_q + batch_size * kBlockM, kBlockM);
  int const total_k_padded_rounded =
      round_multiple(total_k + batch_size * kBlockN, kBlockN);

  if (!is_varlen_q) {
    CHECK_SHAPE(q, batch_size, seqlen_q, num_heads, head_size);
    CHECK_SHAPE(out, batch_size, seqlen_q, num_heads, head_size);
    CHECK_SHAPE(dout, batch_size, seqlen_q, num_heads, head_size);
  } else {
    CHECK_SHAPE(q, total_q, num_heads, head_size);
    CHECK_SHAPE(out, total_q, num_heads, head_size);
    CHECK_SHAPE(dout, total_q, num_heads, head_size);
    CHECK_SHAPE(cu_seqlens_q, batch_size + 1);
  }
  if (!is_varlen_k) {
    CHECK_SHAPE(k, batch_size, seqlen_k, num_heads_k, head_size);
    CHECK_SHAPE(v, batch_size, seqlen_k, num_heads_k, head_size);
  } else {
    CHECK_SHAPE(k, total_k, num_heads_k, head_size);
    CHECK_SHAPE(v, total_k, num_heads_k, head_size);
    CHECK_SHAPE(cu_seqlens_k, batch_size + 1);
  }

  if (seqused_q_.is_initialized()) {
    auto seqused_q = seqused_q_.get();
    PADDLE_ENFORCE_EQ(
        seqused_q.dtype(),
        DataType::INT32,
        common::errors::InvalidArgument("seqused_q must have dtype int32"));
    CHECK_DEVICE(seqused_q);
    CHECK_CONTIGUOUS(seqused_q);
    CHECK_SHAPE(seqused_q, batch_size);
  }
  if (seqused_k_.is_initialized()) {
    auto seqused_k = seqused_k_.get();
    PADDLE_ENFORCE_EQ(
        seqused_k.dtype(),
        DataType::INT32,
        common::errors::InvalidArgument("seqused_k must have dtype int32"));
    CHECK_DEVICE(seqused_k);
    CHECK_CONTIGUOUS(seqused_k);
    CHECK_SHAPE(seqused_k, batch_size);
  }

  if (dq_.is_initialized()) {
    *dq = dq_.get();
    PADDLE_ENFORCE_EQ(
        dq->dtype(),
        q_type,
        common::errors::InvalidArgument("dq must have the same dtype as q"));
    CHECK_DEVICE((*dq));
    PADDLE_ENFORCE_EQ(dq->strides()[dq->strides().size() - 1],
                      1,
                      common::errors::InvalidArgument(
                          "dq must have contiguous last dimension"));
    if (!is_varlen_q) {
      CHECK_SHAPE((*dq), batch_size, seqlen_q, num_heads, head_size);
    } else {
      CHECK_SHAPE((*dq), total_q, num_heads, head_size);
    }
  } else {
    *dq = EmptyLike<T, Context>(dev_ctx, q);
  }

  PADDLE_ENFORCE_GT(nranks,
                    0,
                    common::errors::InvalidArgument(
                        "nranks should be at least 1, but got: %d", nranks));

  if (nranks > 1) {
    PADDLE_ENFORCE_EQ(
        is_varlen_k,
        false,
        common::errors::InvalidArgument(
            "when nranks > 1, FlashMask does not support varlen k."));
    PADDLE_ENFORCE_LT(rank,
                      nranks,
                      common::errors::InvalidArgument(
                          "FlashMask distributed overlap requires "
                          "rank < nranks, but got rank = %d >= nranks %d.",
                          rank,
                          nranks));
  }

  auto GradTensorCheckSetter = [&](const DenseTensor &t,
                                   const paddle::optional<DenseTensor> &dt_,
                                   DenseTensor *dt,
                                   const char *name) {
    if (dt_.is_initialized()) {
      *dt = dt_.get();
      PADDLE_ENFORCE_EQ(dt->dtype(),
                        q_type,
                        common::errors::InvalidArgument(
                            "%s must have the same dtype as q", name));
      CHECK_DEVICE((*dt));
      PADDLE_ENFORCE_EQ(dt->strides()[dt->strides().size() - 1],
                        1,
                        common::errors::InvalidArgument(
                            "%s must have contiguous last dimension", name));
      if (!is_varlen_k) {
        CHECK_SHAPE(
            (*dt), batch_size, seqlen_k * dkv_s_scaler, num_heads_k, head_size);
      } else {
        CHECK_SHAPE((*dt), total_k, num_heads_k, head_size);
      }
    } else {
      // nranks > 1: using distributed overlap will actually compute with
      // complete size. If nrank == 1, dkv_s_scaler will be 1
      *dt = phi::Empty<T, Context>(
          dev_ctx,
          {batch_size, seqlen_k * dkv_s_scaler, num_heads_k, head_size});
    }
  };

  GradTensorCheckSetter(k, dk_, dk, "dk");
  GradTensorCheckSetter(v, dv_, dv, "dv");

  // Otherwise the kernel will be launched from cuda:0 device
  // Cast to char to avoid compiler warning about narrowing

  // Need softmax_d to have total_q_padded_rounded since we want its address to
  // be aligned by 16/8 bytes for TMA / LDG.64
  if (!is_varlen) {
    if (softmax_d) {
      // Need softmax_d to have seqlen_q_rounded since we want its address to be
      // aligned by 16/8 bytes for TMA / LDG.64
      softmax_d->Resize({batch_size, num_heads, seqlen_q_rounded});
    }
    if (softmax_lse_log2) {
      softmax_lse_log2->Resize(
          make_ddim({batch_size, num_heads, seqlen_q_rounded}));
    }
  } else {
    if (softmax_d) {
      softmax_d->Resize({num_heads, total_q_padded_rounded});
    }
    if (softmax_lse_log2) {
      softmax_lse_log2->Resize({num_heads, total_q_padded_rounded});
    }
  }
  if (softmax_d) {
    dev_ctx.template Alloc<float>(softmax_d);
  }
  if (softmax_lse_log2) {
    dev_ctx.template Alloc<float>(softmax_lse_log2);
  }
  if (dq_accum) {
    if (!is_varlen) {
      dq_accum->Resize(make_ddim(
          {batch_size, num_heads, seqlen_q_rounded * head_size_rounded}));
    } else {
      dq_accum->Resize(
          make_ddim({num_heads, total_q_padded_rounded * head_size_rounded}));
    }
    dev_ctx.template Alloc<float>(dq_accum);
  }
  if (num_heads_k != num_heads) {  // MQA / GQA
    if (!is_varlen) {
      // dk and dv accum should directly account for CP overlap
      if (dk_accum) {
        dk_accum->Resize(make_ddim({batch_size,
                                    num_heads_k,
                                    seqlen_k_rounded_cp * head_size_rounded}));
      }
      if (dv_accum) {
        dv_accum->Resize(make_ddim({batch_size,
                                    num_heads_k,
                                    seqlen_k_rounded_cp * head_size_rounded}));
      }
    } else {
      if (dk_accum) {
        dk_accum->Resize(make_ddim(
            {num_heads_k, total_k_padded_rounded, head_size_rounded}));
      }
      if (dv_accum) {
        dv_accum->Resize(make_ddim(
            {num_heads_k, total_k_padded_rounded, head_size_rounded}));
      }
    }
    if (dk_accum) {
      dev_ctx.template Alloc<float>(dk_accum);
    }
    if (dv_accum) {
      dev_ctx.template Alloc<float>(dv_accum);
    }
    funcs::SetConstant<Context, float> set_zero;

    if (dk_accum) {
      set_zero(dev_ctx, dk_accum, float{0});
    }
    if (dv_accum) {
      set_zero(dev_ctx, dv_accum, float{0});
    }
  }

  FlashMask_bwd_params *params_handle = get_flashmask_bwd_params_handle();
  dynload::flashmaskv2_clear_bwd_params_handle(params_handle);
  set_flashmaskv2_params_dgrad(
      params_handle,
      batch_size,
      seqlen_q,
      seqlen_k,
      seqlen_q_rounded,
      seqlen_k_rounded_cp,  // length of grad accum will be scaled for CP
                            // overlap
      num_heads,
      num_heads_k,
      head_size,
      head_size_rounded,
      q,
      k,
      v,
      out,
      dout,
      dq,
      dk,
      dv,
      !is_varlen_q ? nullptr : cu_seqlens_q.data(),
      !is_varlen_k ? nullptr : cu_seqlens_k.data(),
      seqused_q_.is_initialized() ? const_cast<void *>(seqused_q_.get().data())
                                  : nullptr,
      seqused_k_.is_initialized() ? const_cast<void *>(seqused_k_.get().data())
                                  : nullptr,
      dq_accum ? dq_accum->data() : nullptr,
      num_heads_k != num_heads && dk_accum ? dk_accum->data() : nullptr,
      num_heads_k != num_heads && dv_accum ? dv_accum->data() : nullptr,
      const_cast<void *>(softmax_lse.data()),
      softmax_d ? (softmax_d->data()) : nullptr,
      /*p_dropout=*/0.f,
      softmax_scale,
      window_size_left,
      window_size_right,
      dprops,
      softcap,
      deterministic,
      sm_margin);
  dynload::flashmaskv2_bwd_params_set_total_q(params_handle, total_q);
  dynload::flashmaskv2_bwd_params_set_total_k(params_handle, total_k);
  dynload::flashmaskv2_bwd_params_set_softmax_lse_log2_ptr(
      params_handle, softmax_lse_log2 ? softmax_lse_log2->data() : nullptr);
  dynload::flashmaskv2_bwd_params_set_dv(
      params_handle,
      head_size);  // We don't support hdim_v being
                   // different from hdim_qk for now
  DenseTensor tile_count_semaphore;
  if (arch >= 90) {
    tile_count_semaphore = Full<int32_t, Context>(dev_ctx, {1}, 0);
    dynload::flashmaskv2_bwd_params_set_tile_count_semaphore(
        params_handle, tile_count_semaphore.data<int>());
  } else {
    dynload::flashmaskv2_bwd_params_set_tile_count_semaphore(params_handle,
                                                             nullptr);
  }

  DenseTensor dq_semaphore = Empty<int32_t>(
      dev_ctx, {(seqlen_q + kBlockM - 1) / kBlockM, batch_size, num_heads});
  dynload::flashmaskv2_bwd_params_set_dq_semaphore(params_handle,
                                                   dq_semaphore.data<int>());
  // dk_semaphore should have the same seqlen with dk_accum
  DenseTensor dk_semaphore =
      Empty<int32_t>(dev_ctx,
                     {(seqlen_k * dkv_accum_s_scaler + kBlockN - 1) / kBlockN,
                      batch_size,
                      num_heads_k});
  DenseTensor dv_semaphore =
      Empty<int32_t>(dev_ctx,
                     {(seqlen_k * dkv_accum_s_scaler + kBlockN - 1) / kBlockN,
                      batch_size,
                      num_heads_k});
  if (num_heads_k != num_heads &&
      dynload::flashmaskv2_bwd_params_get_deterministic(params_handle)) {
    // xiangrui: we need to zero them out
    funcs::SetConstant<Context, int32_t> set_zero_dk;
    set_zero_dk(dev_ctx, &dk_semaphore, static_cast<int32_t>(0));
    funcs::SetConstant<Context, int32_t> set_zero_dv;
    set_zero_dv(dev_ctx, &dv_semaphore, static_cast<int32_t>(0));
    dynload::flashmaskv2_bwd_params_set_dk_semaphore(params_handle,
                                                     dk_semaphore.data<int>());
    dynload::flashmaskv2_bwd_params_set_dv_semaphore(params_handle,
                                                     dv_semaphore.data<int>());
  }

  if (is_flashmask) {
    if (lt_start_row_indices.initialized())
      dynload::flashmaskv2_bwd_params_set_lt_start_ptr(
          params_handle, (lt_start_row_indices.data<int32_t>()));
    else
      dynload::flashmaskv2_bwd_params_set_lt_start_ptr(params_handle, nullptr);

    if (lt_end_row_indices.initialized())
      dynload::flashmaskv2_bwd_params_set_lt_end_ptr(
          params_handle, (lt_end_row_indices.data<int32_t>()));
    else
      dynload::flashmaskv2_bwd_params_set_lt_end_ptr(params_handle, nullptr);

    if (ut_start_row_indices.initialized())
      dynload::flashmaskv2_bwd_params_set_ut_start_ptr(
          params_handle, (ut_start_row_indices.data<int32_t>()));
    else
      dynload::flashmaskv2_bwd_params_set_ut_start_ptr(params_handle, nullptr);

    if (ut_end_row_indices.initialized())
      dynload::flashmaskv2_bwd_params_set_ut_end_ptr(
          params_handle, (ut_end_row_indices.data<int32_t>()));
    else
      dynload::flashmaskv2_bwd_params_set_ut_end_ptr(params_handle, nullptr);

    if (flashmask_maxmin.initialized())
      dynload::flashmaskv2_bwd_params_set_flashmask_maxmin_ptr(
          params_handle, (flashmask_maxmin.data<int32_t>()));
    else
      dynload::flashmaskv2_bwd_params_set_flashmask_maxmin_ptr(params_handle,
                                                               nullptr);

    dynload::flashmaskv2_bwd_params_set_h_flashmask(
        params_handle, startend_row_indices.dims()[1]);
    dynload::flashmaskv2_bwd_params_set_h_h_flashmask_ratio(
        params_handle, num_heads / startend_row_indices.dims()[1]);

#ifdef PADDLE_WITH_NVSHMEM
    // only when NVSHMEM is compiled in paddle, can we use the following
    dynload::flashmaskv2_bwd_params_set_rank(params_handle, rank);
    dynload::flashmaskv2_bwd_params_set_nranks(params_handle, nranks);
#endif  // PADDLE_WITH_NVSHMEM
  } else {
    dynload::flashmaskv2_bwd_params_set_lt_start_ptr(params_handle, nullptr);
    dynload::flashmaskv2_bwd_params_set_lt_end_ptr(params_handle, nullptr);
    dynload::flashmaskv2_bwd_params_set_ut_start_ptr(params_handle, nullptr);
    dynload::flashmaskv2_bwd_params_set_ut_end_ptr(params_handle, nullptr);
    dynload::flashmaskv2_bwd_params_set_flashmask_maxmin_ptr(params_handle,
                                                             nullptr);
    dynload::flashmaskv2_bwd_params_set_h_flashmask(params_handle, 0);
    dynload::flashmaskv2_bwd_params_set_h_h_flashmask_ratio(params_handle, 0);
  }

  if (is_blockmask) {
    // xhy: blockmask is now only support blockdim_q k = 128
    dynload::flashmaskv2_bwd_params_set_m_block_dim(params_handle, 128);
    dynload::flashmaskv2_bwd_params_set_n_block_dim(params_handle, 128);
    dynload::flashmaskv2_bwd_params_set_block_mask_ptr(
        params_handle, (block_mask.data<int32_t>()));
  }
#ifdef FLASHATTENTION_DISABLE_LOCAL
  PADDLE_ENABLE_EQ(
      !dynload::flashmaskv2_bwd_params_get_is_local(params_handle),
      true,
      "This flash attention build does not support local attention.");
#endif
#ifdef FLASHATTENTION_DISABLE_SOFTCAP
  PADDLE_ENABLE_EQ(
      dynload::flashmaskv2_bwd_params_get_softcap(params_handle),
      0.0,
      "This flash attention build does not support tanh softcapping.");
#endif

  if (total_q > 0 && total_k > 0 && num_heads_k > 0) {
    dynload::flashmaskv2_run_mha_bwd(params_handle, dev_ctx.stream());
  } else if (total_k > 0 && num_heads_k > 0) {
    // If seqlen_q == 0, then we have an empty tensor. We need to set the output
    // to 0.
    funcs::SetConstant<Context, T> set_zero;
    set_zero(dev_ctx, dk, T{0});
    set_zero(dev_ctx, dv, T{0});
    if (softmax_d) {
      funcs::SetConstant<Context, float> set_zero_fp32;
      set_zero_fp32(dev_ctx, softmax_d, float{0});
    }
  } else if (total_q > 0 && num_heads_k > 0) {
    funcs::SetConstant<Context, T> set_zero;
    set_zero(dev_ctx, dq, T{0});
    if (softmax_d) {
      funcs::SetConstant<Context, float> set_zero_fp32;
      set_zero_fp32(dev_ctx, softmax_d, float{0});
    }
  }
#else
  RaiseNotSupportedError();
#endif
}

template <typename T, typename Context>
void FlashMaskV2GradKernel(
    const Context &dev_ctx,
    const DenseTensor &q,
    const DenseTensor &k,
    const DenseTensor &v,
    const DenseTensor &out,
    const DenseTensor &softmax_lse,
    const DenseTensor &startend_row_indices,  // TODO(xiehaoyang): remove this
    const optional<DenseTensor> &block_mask,
    const DenseTensor &out_grad,
    float const softmax_scale,
    bool is_causal,
    int rank,
    int nranks,
    DenseTensor *dq,
    DenseTensor *dk,
    DenseTensor *dv) {
#ifdef PADDLE_WITH_FLASHATTN_V3

  PADDLE_ENFORCE_EQ(
      q.dims()[q.dims().size() - 1],
      v.dims()[v.dims().size() - 1],
      common::errors::InvalidArgument("head_dim_q != head_dim_v (%d != %d)",
                                      q.dims()[q.dims().size() - 1],
                                      v.dims()[v.dims().size() - 1]));

  // umiswing: fake grad tensor for FlashAttnV3GradBaseKernel
  DenseTensor softmax_d;
  DenseTensor softmax_lse_log2;
  DenseTensor dq_accum;
  DenseTensor dk_accum;
  DenseTensor dv_accum;
  FlashMaskV2GradBaseKernel<T, Context>(
      dev_ctx,
      out_grad,
      q,
      k,
      v,
      out,
      softmax_lse,
      paddle::none,  // dq_
      paddle::none,  // dk_
      paddle::none,  // dv_
      paddle::none,
      paddle::none,
      paddle::none,
      paddle::none,
      startend_row_indices,
      block_mask,
      0,  // max_seqlen_q,
      0,  // max_seqlen_k,
      softmax_scale,
      is_causal,
      -1,                         // window_size_left,
      -1,                         // window_size_right,
      0,                          // softcap,
      FLAGS_cudnn_deterministic,  // deterministic,
      0,                          // sm_margin,
      rank,
      nranks,
      dq,
      dk,
      dv,
      &softmax_d,
      &softmax_lse_log2,
      &dq_accum,
      &dk_accum,
      &dv_accum);

  // umiswing: some branch in upstream fa3 could have padded the head dimension
  PADDLE_ENFORCE_EQ(
      dq->dims()[dq->dims().size() - 1],
      out_grad.dims()[out_grad.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dq != head dimension of out_grad (%d != %d)",
          dq->dims()[dq->dims().size() - 1],
          out_grad.dims()[out_grad.dims().size() - 1]));

  PADDLE_ENFORCE_EQ(
      dk->dims()[dk->dims().size() - 1],
      out_grad.dims()[out_grad.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dk != head dimension of out_grad (%d != %d)",
          dk->dims()[dk->dims().size() - 1],
          out_grad.dims()[out_grad.dims().size() - 1]));

  PADDLE_ENFORCE_EQ(
      dv->dims()[dv->dims().size() - 1],
      out_grad.dims()[out_grad.dims().size() - 1],
      common::errors::InvalidArgument(
          "head dimension of dv != head dimension of out_grad (%d != %d)",
          dv->dims()[dv->dims().size() - 1],
          out_grad.dims()[out_grad.dims().size() - 1]));

#else
  RaiseNotSupportedError();
#endif
}

}  // namespace phi

PD_REGISTER_KERNEL(flash_attn_v3_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashAttnV3GradKernel,
                   phi::float16,
                   phi::bfloat16) {}

PD_REGISTER_KERNEL(flash_attn_v3_varlen_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashAttnV3VarlenGradKernel,
                   phi::float16,
                   phi::bfloat16) {}

PD_REGISTER_KERNEL(flashmask_attention_v2_grad,
                   GPU,
                   ALL_LAYOUT,
                   phi::FlashMaskV2GradKernel,
                   phi::float16,
                   phi::bfloat16) {}