// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

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

#include "paddle/phi/backends/xpu/enforce_xpu.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/slice_kernel.h"
#include "paddle/phi/kernels/xpu/flash_attn_utils.h"
#include "xfa/flash_api.h"
namespace phi {
#define MHA_VARLEN_BWD_TYPES_AND_ARGS(T, TACCUM, TGEMM, TID)                  \
  xpu::Context *ctx, const T *dout, const T *q, const T *k, const T *v,       \
      const T *out, const TACCUM *softmax_lse, T *dq, T *dk, T *dv,           \
      const xpu::VectorParam<TID>&lod_seqlens_q,                              \
      const xpu::VectorParam<TID>&lod_seqlens_k, int64_t max_seqlen_q,        \
      int64_t max_seqlen_k, int64_t head_num, int64_t head_num_k,             \
      int64_t head_dim, const float softmax_scale, const float p_dropout,     \
      int seed, const bool is_causal, const TACCUM *attn_mask,                \
      const TACCUM *bias, const float *q_maxptr, const float *k_maxptr,       \
      const float *v_maxptr, const float *o_maxptr, float *dq_maxptr,         \
      float *dk_maxptr, float *dv_maxptr, const float *do_maxptr,             \
      const bool is_qkv_fusion, const bool is_dqkv_fusion,                    \
      const int64_t qkv_layout, const float *alibi_slopes,                    \
      const std::vector<int64_t>&alibi_slopes_shape, int window_size_left,    \
      int window_size_right, int64_t v_head_dim,                              \
      const int *downstart_row_indices_data,                                  \
      const int *downend_row_indices_data,                                    \
      const int *upstart_row_indices_data, const int *upend_row_indices_data, \
      const int flash_mask_head_num, int *flashmask_maxmin,                   \
      XPUStream side_stream

#define MHA_VARLEN_BWD_ARGS                                                    \
  ctx, dout, q, k, v, out, softmax_lse, dq, dk, dv, lod_seqlens_q,             \
      lod_seqlens_k, max_seqlen_q, max_seqlen_k, head_num, head_num_k,         \
      head_dim, softmax_scale, p_dropout, seed, is_causal, attn_mask, bias,    \
      q_maxptr, k_maxptr, v_maxptr, o_maxptr, dq_maxptr, dk_maxptr, dv_maxptr, \
      do_maxptr, is_qkv_fusion, is_dqkv_fusion, qkv_layout, alibi_slopes,      \
      alibi_slopes_shape, window_size_left, window_size_right, v_head_dim,     \
      downstart_row_indices_data, downend_row_indices_data,                    \
      upstart_row_indices_data, upend_row_indices_data, flash_mask_head_num,   \
      flashmask_maxmin, side_stream

template <typename T, typename TACCUM, typename TGEMM, typename TID>
int mha_varlen_bwd_wrapper(
    MHA_VARLEN_BWD_TYPES_AND_ARGS(T, TACCUM, TGEMM, TID)) {
  PADDLE_THROW(
      "Unsupported template params combination for mha_varlen_bwd, should not "
      "reach here.");
}

#define DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(T, TACCUM, TGEMM, TID) \
  template <>                                                         \
  int mha_varlen_bwd_wrapper<T, TACCUM, TGEMM, TID>(                  \
      MHA_VARLEN_BWD_TYPES_AND_ARGS(T, TACCUM, TGEMM, TID)) {         \
    return baidu::xpu::xfa::mha_varlen_bwd<T, TACCUM, TGEMM, TID>(    \
        MHA_VARLEN_BWD_ARGS);                                         \
  }

DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(float, float, tfloat32, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(float, float, float, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, float, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, XPUTypeFP16, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, tfloat32, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16,
                                       XPUTypeFP16,
                                       XPUTypeFP16,
                                       int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16, float, tfloat32, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16, float, float, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16, float, XPUTypeFP16, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeFP16,
                                       XPUTypeFP16,
                                       XPUTypeFP16,
                                       int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16, float, XPUTypeBF16, int);
DECLARE_SUPPORTED_MHA_VARLEN_BWD_TYPES(XPUTypeBF16,
                                       XPUTypeFP16,
                                       XPUTypeBF16,
                                       int);

template <typename T, typename Context>
void FlashAttnGradKernelBase(const Context& dev_ctx,
                             const DenseTensor& q,
                             const DenseTensor& k,
                             const DenseTensor& v,
                             const api::VectorParam<int>& lod_seqlen_q,
                             const api::VectorParam<int>& lod_seqlen_k,
                             const DenseTensor& out,
                             const DenseTensor& softmax_lse,
                             const DenseTensor& seed_offset,
                             const optional<DenseTensor>& attn_mask,
                             const optional<DenseTensor>& startend_row_indices,
                             const DenseTensor& dout,
                             const int64_t batch_size,
                             const Scalar& max_seqlen_q_,
                             const Scalar& max_seqlen_k_,
                             const int64_t num_heads,
                             const int64_t num_heads_k,
                             const int64_t head_size,
                             const int64_t head_size_v,
                             float scale,
                             float dropout,
                             bool causal,
                             DenseTensor* dq,
                             DenseTensor* dk,
                             DenseTensor* dv) {
  xpu::ctx_guard RAII_GUARD(dev_ctx.x_context());

  using XPUType = typename XPUTypeTrait<T>::Type;
  using XPUTypeFP16 = typename XPUTypeTrait<phi::float16>::Type;
  const XPUType* q_data = reinterpret_cast<const XPUType*>(q.data<T>());
  const XPUType* k_data = reinterpret_cast<const XPUType*>(k.data<T>());
  const XPUType* v_data = reinterpret_cast<const XPUType*>(v.data<T>());
  const XPUType* out_data = reinterpret_cast<const XPUType*>(out.data<T>());
  const float* softmax_lse_data = softmax_lse.data<float>();
  const XPUType* dout_data = reinterpret_cast<const XPUType*>(dout.data<T>());

  float real_scale = scale == 0.0f ? 1.0f / std::sqrt(head_size) : scale;

  const float* bias_data = nullptr;
  int64_t fa_layout = AttnQKVLayout_t::ATTN_BLHD;
  DenseTensor downstart_row_indices, upend_row_indices, downend_row_indices,
      upstart_row_indices;
  void *downstart_row_indices_data = nullptr, *upend_row_indices_data = nullptr,
       *downend_row_indices_data = nullptr, *upstart_row_indices_data = nullptr;
  bool is_flashmask = startend_row_indices.get_ptr() != nullptr;
  XPUStream flashmask_stream;
  if (is_flashmask) {
    xpu_stream_create(&flashmask_stream);
    PADDLE_ENFORCE_EQ(
        startend_row_indices->dims().size(),
        4,
        common::errors::InvalidArgument(
            "flashmask_attention receive startend_row_indices with dim "
            "[batch_size, num_heads,seq_len, mask_bounds]"));
    PADDLE_ENFORCE_EQ(startend_row_indices->dims()[3] == 1 ||
                          startend_row_indices->dims()[3] == 2 ||
                          startend_row_indices->dims()[3] == 4,
                      true,
                      common::errors::InvalidArgument(
                          "flashmask_attention startend_row_indices "
                          "mask_bounds must in [1,2,4]"));
    downstart_row_indices =
        phi::Slice<int32_t>(dev_ctx, startend_row_indices.get(), {3}, {0}, {1});
    downstart_row_indices_data = downstart_row_indices.data();
    if (startend_row_indices->dims()[3] == 2) {
      if (!causal) {
        upend_row_indices = phi::Slice<int32_t>(
            dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
        upend_row_indices_data = upend_row_indices.data();
      } else {
        downend_row_indices = phi::Slice<int32_t>(
            dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
        downend_row_indices_data = downend_row_indices.data();
      }
    } else if (startend_row_indices->dims()[3] == 4) {
      upend_row_indices = phi::Slice<int32_t>(
          dev_ctx, startend_row_indices.get(), {3}, {3}, {4});
      upend_row_indices_data = upend_row_indices.data();
      downend_row_indices = phi::Slice<int32_t>(
          dev_ctx, startend_row_indices.get(), {3}, {1}, {2});
      downend_row_indices_data = downend_row_indices.data();
      upstart_row_indices = phi::Slice<int32_t>(
          dev_ctx, startend_row_indices.get(), {3}, {2}, {3});
      upstart_row_indices_data = upstart_row_indices.data();
    }
  } else if (attn_mask.get_ptr() != nullptr) {
    const auto& mask_dims = attn_mask->dims();
    if (mask_dims.size() == 3 || (mask_dims[1] == 1 && mask_dims.size() == 4)) {
      fa_layout |= AttnQKVLayout_t::BIAS_BLL;
    } else {
      PADDLE_ENFORCE_EQ(mask_dims.size(),
                        4,
                        common::errors::InvalidArgument(
                            "flash_attn_bwd requires mask's shape "
                            "like [b,l,l] or [b, h, l, l]"));
    }
    if (attn_mask->dtype() == DataType::FLOAT32) {
      bias_data = attn_mask->data<float>();
    } else if (attn_mask->dtype() == DataType::FLOAT16 ||
               attn_mask->dtype() == DataType::BFLOAT16) {
      float* bias_tmp = RAII_GUARD.alloc_l3_or_gm<float>(attn_mask->numel());
      int r = xpu::cast<XPUType, float>(
          dev_ctx.x_context(),
          reinterpret_cast<const XPUType*>(attn_mask->data<T>()),
          bias_tmp,
          attn_mask->numel());
      PADDLE_ENFORCE_XDNN_SUCCESS(r, "cast");
      bias_data = bias_tmp;
    } else {
      errors::Unimplemented(
          "Unsupported dtype for attention_mask in xpu flash attention, only "
          "float32, float16 and "
          "bfloat16 are supported.");
    }
  }
  // output
  XPUType* dq_data = reinterpret_cast<XPUType*>(dq->data<T>());
  XPUType* dk_data = reinterpret_cast<XPUType*>(dk->data<T>());
  XPUType* dv_data = reinterpret_cast<XPUType*>(dv->data<T>());

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

  // get seed offset
  const int64_t* seed_offset_data = seed_offset.data<int64_t>();
  int fa_tgemm = get_flash_attn_tgemm<XPUType>();
  auto flash_attention_grad_kernel =
      mha_varlen_bwd_wrapper<XPUType, float, tfloat32, int>;
  if (fa_tgemm == XPU_FA_DTYPE::FA_FLOAT) {
    flash_attention_grad_kernel =
        mha_varlen_bwd_wrapper<XPUType, float, float, int>;
  } else if (fa_tgemm == XPU_FA_DTYPE::FA_FLOAT16) {
    flash_attention_grad_kernel =
        mha_varlen_bwd_wrapper<XPUType, float, XPUTypeFP16, int>;
  }
  // template<typename T, typename TACCUM, typename TGEMM, typename TID = int>
  // int mha_varlen_bwd(xdnn::Context* xpu_ctx, const T* dout, const T* q, const
  // T* k, const T* v, const T* out, const TACCUM* softmax_lse, T* dq, T* dk, T*
  // dv, const xdnn::VectorParam<TID>& lod_seqlens_q, const
  // xdnn::VectorParam<TID>& lod_seqlens_k, int64_t max_seqlen_q, int64_t
  // max_seqlen_k, int64_t head_num, int64_t head_num_k, int64_t head_dim, const
  // float softmax_scale = 0.0f, const float p_dropout = 0.0f, int seed =
  // 0x45678901, const bool is_causal = true, const TACCUM* attn_mask = nullptr,
  // const TACCUM* bias = nullptr, const float* q_maxptr = nullptr, const float*
  // k_maxptr = nullptr, const float* v_maxptr = nullptr, const float* o_maxptr
  // = nullptr, float* dq_maxptr = nullptr, float* dk_maxptr = nullptr, float*
  // dv_maxptr = nullptr, const float* do_maxptr = nullptr, const bool
  // is_qkv_fusion = false, const bool is_dqkv_fusion = false, const int64_t
  // qkv_layout = AttnQKVLayout_t::ATTN_BLHD, const float* alibi_slopes =
  // nullptr, const std::vector<int64_t>& alibi_slopes_shape = {}, int
  // window_size_left = -1, int window_size_right = -1, int64_t v_head_dim =
  // -1, const int* downstart_row_indices_data = nullptr,
  // const int* downend_row_indices_data = nullptr,
  // const int* upstart_row_indices_data = nullptr,
  // const int* upend_row_indices_data = nullptr,
  // const int flash_mask_head_num = 0,
  // int* flashmask_maxmin = nullptr,
  // XPUStream side_stream = nullptr);
  int r = flash_attention_grad_kernel(
      dev_ctx.x_context(),
      dout_data,                                  // dout
      q_data,                                     // q
      k_data,                                     // k
      v_data,                                     // v
      out_data,                                   // out
      softmax_lse_data,                           // softmax_lse
      dq_data,                                    // dq
      dk_data,                                    // dk
      dv_data,                                    // dv
      lod_seqlen_q,                               // lod_seqlens_q
      lod_seqlen_k,                               // lod_seqlens_k
      max_seqlen_q,                               // max_seqlen_q
      max_seqlen_k,                               // max_seqlen_k
      num_heads,                                  // head_num
      num_heads_k,                                // head_num_k
      head_size,                                  // head_dim
      real_scale,                                 // softmax_scale
      dropout,                                    // p_dropout
      static_cast<int32_t>(seed_offset_data[0]),  // seed
      causal,                                     // is_causal
      nullptr,                                    // attn_mask
      bias_data,                                  // bias
      nullptr,                                    // q_maxptr
      nullptr,                                    // k_maxptr
      nullptr,                                    // v_maxptr
      nullptr,                                    // o_maxptr
      nullptr,                                    // dq_maxptr
      nullptr,                                    // dk_maxptr
      nullptr,                                    // dv_maxptr
      nullptr,                                    // do_maxptr
      false,                                      // is_qkv_fusion
      false,                                      // is_dqkv_fusion
      fa_layout,                                  // qkv_layout
      nullptr,                                    // alibi_slopes
      {},                                         // alibi_slopes_shape
      -1,                                         // window_size_left
      -1,                                         // window_size_right
      head_size_v,                                // v_head_dim
      (const int*)downstart_row_indices_data,
      (const int*)downend_row_indices_data,
      (const int*)upstart_row_indices_data,
      (const int*)upend_row_indices_data,
      is_flashmask ? startend_row_indices->dims()[1] : 0,
      nullptr,
      is_flashmask ? flashmask_stream : nullptr);
  PADDLE_ENFORCE_XDNN_SUCCESS(r, "mha_varlen_bwd");
  if (is_flashmask && flashmask_stream != nullptr) {
    r = xpu_wait(flashmask_stream);
    PADDLE_ENFORCE_XPU_SUCCESS(r);
    xpu_stream_destroy(flashmask_stream);
  }
}

template <typename T, typename Context>
void FlashAttnUnpaddedGradKernel(const Context& dev_ctx,
                                 const DenseTensor& q,
                                 const DenseTensor& k,
                                 const DenseTensor& v,
                                 const DenseTensor& cu_seqlens_q,
                                 const DenseTensor& cu_seqlens_k,
                                 const DenseTensor& out,
                                 const DenseTensor& softmax_lse,
                                 const DenseTensor& seed_offset,
                                 const optional<DenseTensor>& attn_mask,
                                 const DenseTensor& dout,
                                 const Scalar& max_seqlen_q,
                                 const Scalar& max_seqlen_k,
                                 float scale,
                                 float dropout,
                                 bool causal,
                                 DenseTensor* dq,
                                 DenseTensor* dk,
                                 DenseTensor* dv) {
  dev_ctx.template Alloc<T>(dq);
  dev_ctx.template Alloc<T>(dk);
  dev_ctx.template Alloc<T>(dv);
  auto dims = q.dims();

  const int64_t batch_size = cu_seqlens_q.numel() - 1;
  const int64_t num_heads = dims[1];
  const int64_t head_size = dims[2];
  const int64_t head_size_v = v.dims()[2];
  const int64_t num_heads_k = k.dims()[1];

  api::VectorParam<int> qlod{cu_seqlens_q.data<int>(),
                             static_cast<int64_t>(cu_seqlens_q.numel()),
                             nullptr};
  api::VectorParam<int> kvlod{cu_seqlens_k.data<int>(),
                              static_cast<int64_t>(cu_seqlens_k.numel()),
                              nullptr};

  FlashAttnGradKernelBase<T>(dev_ctx,
                             q,
                             k,
                             v,
                             qlod,
                             kvlod,
                             out,
                             softmax_lse,
                             seed_offset,
                             attn_mask,
                             paddle::none,
                             dout,
                             batch_size,
                             max_seqlen_q,
                             max_seqlen_k,
                             num_heads,
                             num_heads_k,
                             head_size,
                             head_size_v,
                             scale,
                             dropout,
                             causal,
                             dq,
                             dk,
                             dv);
}

template <typename T, typename Context>
void FlashAttnGradKernel(const Context& dev_ctx,
                         const DenseTensor& q,
                         const DenseTensor& k,
                         const DenseTensor& v,
                         const DenseTensor& out,
                         const DenseTensor& softmax_lse,
                         const DenseTensor& seed_offset,
                         const optional<DenseTensor>& attn_mask,
                         const DenseTensor& dout,
                         float dropout,
                         bool causal,
                         DenseTensor* dq,
                         DenseTensor* dk,
                         DenseTensor* dv) {
  dev_ctx.template Alloc<T>(dq);
  dev_ctx.template Alloc<T>(dk);
  dev_ctx.template Alloc<T>(dv);

  // q, k, v [batch_size, seq_len, num_heads, head_dim]
  const auto& dims = q.dims();

  const int64_t batch_size = dims[0];
  const int64_t seqlen_q = dims[1];
  const int64_t num_heads = dims[2];
  const int64_t head_size_og = dout.dims()[3];
  const int64_t head_size = dims[3];
  const int64_t head_size_v = v.dims()[3];
  const int64_t seqlen_k = k.dims()[1];
  const int64_t num_heads_k = k.dims()[2];

  PADDLE_ENFORCE_EQ(
      head_size_og,
      head_size_v,
      common::errors::InvalidArgument(
          "flash_attn_bwd receive input with head_size_og == head_size_v"));

  // lod info
  std::vector<int> qlod_vec = {0};
  std::vector<int> kvlod_vec = {0};
  for (int batch_idx = 1; batch_idx <= batch_size; ++batch_idx) {
    qlod_vec.push_back(seqlen_q * batch_idx);
    kvlod_vec.push_back(seqlen_k * batch_idx);
  }
  api::VectorParam<int> qlod{
      qlod_vec.data(), static_cast<int64_t>(qlod_vec.size()), nullptr};
  api::VectorParam<int> kvlod{
      kvlod_vec.data(), static_cast<int64_t>(kvlod_vec.size()), nullptr};

  FlashAttnGradKernelBase<T>(dev_ctx,
                             q,
                             k,
                             v,
                             qlod,
                             kvlod,
                             out,
                             softmax_lse,
                             seed_offset,
                             attn_mask,
                             paddle::none,
                             dout,
                             batch_size,
                             seqlen_q,
                             seqlen_k,
                             num_heads,
                             num_heads_k,
                             head_size,
                             head_size_v,
                             0.0,
                             dropout,
                             causal,
                             dq,
                             dk,
                             dv);
}

template <typename T, typename Context>
void FlashMaskGradKernel(const Context& dev_ctx,
                         const DenseTensor& q,
                         const DenseTensor& k,
                         const DenseTensor& v,
                         const DenseTensor& startend_row_indices,
                         const DenseTensor& out,
                         const DenseTensor& softmax_lse,
                         const DenseTensor& seed_offset,
                         const DenseTensor& dout,
                         float dropout,
                         bool causal,
                         DenseTensor* dq,
                         DenseTensor* dk,
                         DenseTensor* dv) {
  dev_ctx.template Alloc<T>(dq);
  dev_ctx.template Alloc<T>(dk);
  dev_ctx.template Alloc<T>(dv);

  // q, k, v [batch_size, seq_len, num_heads, head_dim]
  const auto& dims = q.dims();

  const int64_t batch_size = dims[0];
  const int64_t seqlen_q = dims[1];
  const int64_t num_heads = dims[2];
  const int64_t head_size_og = dout.dims()[3];
  const int64_t head_size = dims[3];
  const int64_t head_size_v = v.dims()[3];
  const int64_t seqlen_k = k.dims()[1];
  const int64_t num_heads_k = k.dims()[2];

  PADDLE_ENFORCE_EQ(
      head_size_og,
      head_size_v,
      common::errors::InvalidArgument(
          "flash_attn_bwd receive input with head_size_og == head_size_v"));

  // lod info
  std::vector<int> qlod_vec = {0};
  std::vector<int> kvlod_vec = {0};
  for (int batch_idx = 1; batch_idx <= batch_size; ++batch_idx) {
    qlod_vec.push_back(seqlen_q * batch_idx);
    kvlod_vec.push_back(seqlen_k * batch_idx);
  }
  api::VectorParam<int> qlod{
      qlod_vec.data(), static_cast<int64_t>(qlod_vec.size()), nullptr};
  api::VectorParam<int> kvlod{
      kvlod_vec.data(), static_cast<int64_t>(kvlod_vec.size()), nullptr};
  FlashAttnGradKernelBase<T>(dev_ctx,
                             q,
                             k,
                             v,
                             qlod,
                             kvlod,
                             out,
                             softmax_lse,
                             seed_offset,
                             paddle::none,
                             startend_row_indices,
                             dout,
                             batch_size,
                             seqlen_q,
                             seqlen_k,
                             num_heads,
                             num_heads_k,
                             head_size,
                             head_size_v,
                             0.0,
                             dropout,
                             causal,
                             dq,
                             dk,
                             dv);
}
}  // namespace phi

PD_REGISTER_KERNEL(flash_attn_unpadded_grad,
                   XPU,
                   ALL_LAYOUT,
                   phi::FlashAttnUnpaddedGradKernel,
                   float,
                   phi::float16,
                   phi::bfloat16) {
  kernel->InputAt(3).SetBackend(phi::Backend::CPU);          // cu_seqlens_q
  kernel->InputAt(4).SetBackend(phi::Backend::CPU);          // cu_seqlens_k
  kernel->InputAt(7).SetBackend(phi::Backend::ALL_BACKEND);  // seed_offset
}

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

PD_REGISTER_KERNEL(flashmask_attention_grad,
                   XPU,
                   ALL_LAYOUT,
                   phi::FlashMaskGradKernel,
                   phi::float16,
                   phi::bfloat16) {
  kernel->InputAt(6).SetBackend(
      phi::Backend::ALL_BACKEND);  // fixed_seed_offset
}