custom_ops.cu

/*************************************************************************
 * Copyright (c) 2022-2024, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 *
 * See LICENSE for license information.
 ************************************************************************/

#include <cub/cub.cuh>
#include <map>
#include <vector>

#include "common.h"
#include "common/common.h"

namespace transformer_engine {
namespace paddle_ext {

// convert bias type to enum
NVTE_Bias_Type get_nvte_bias_type(const std::string bias_type) {
  if (bias_type == "no_bias") {
    return NVTE_Bias_Type::NVTE_NO_BIAS;
  } else if (bias_type == "pre_scale_bias") {
    return NVTE_Bias_Type::NVTE_PRE_SCALE_BIAS;
  } else if (bias_type == "post_scale_bias") {
    return NVTE_Bias_Type::NVTE_POST_SCALE_BIAS;
  } else {
    NVTE_ERROR("Invalid bias type. \n");
  }
}

// convert attn mask type to enum
NVTE_Mask_Type get_nvte_mask_type(const std::string mask_type) {
  if (mask_type == "padding") {
    return NVTE_Mask_Type::NVTE_PADDING_MASK;
  } else if (mask_type == "causal") {
    return NVTE_Mask_Type::NVTE_CAUSAL_MASK;
  } else if (mask_type == "no_mask") {
    return NVTE_Mask_Type::NVTE_NO_MASK;
  } else {
    NVTE_ERROR("Invalid attention mask type. \n");
  }
}

void cast_to_fp8(const paddle::Tensor &input, const paddle::Tensor &scale,
                 paddle::Tensor &output,     // NOLINT
                 paddle::Tensor &amax,       // NOLINT
                 paddle::Tensor &scale_inv,  // NOLINT
                 int64_t index, int64_t otype) {
  auto shape = GetShapeArray(input);

  auto input_cu = MakeNvteTensor(input);
  auto output_cu = MakeNvteTensor(
      output.data(), shape, Int2NvteDType(otype), GetDataPtr<float>(amax, index),
      const_cast<void *>(GetDataPtr<float>(scale, index)), GetDataPtr<float>(scale_inv, index));

  nvte_fp8_quantize(input_cu.data(), output_cu.data(), input.stream());
}

std::vector<paddle::Tensor> cast_from_fp8(const paddle::Tensor &input,
                                          const paddle::Tensor &scale_inv, int64_t index,
                                          int64_t itype, int64_t otype) {
  auto shape = GetShapeArray(input);

  auto output = paddle::empty_like(input, Nvte2PaddleDType(Int2NvteDType(otype)));
  auto input_cu =
      MakeNvteTensor(const_cast<void *>(input.data()), shape, Int2NvteDType(itype), nullptr,
                     nullptr, const_cast<void *>(GetDataPtr<float>(scale_inv, index)));
  auto output_cu = MakeNvteTensor(output);

  nvte_fp8_dequantize(input_cu.data(), output_cu.data(), input.stream());

  return {output};
}

std::vector<paddle::Tensor> te_transpose(const paddle::Tensor &input, int64_t otype) {
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the input to have 2 dimensions.");
  size_t M = shape[0];
  size_t N = shape[1];

  auto output = paddle::empty({input.shape()[1], input.shape()[0]}, input.dtype(), input.place());

  auto input_cu = MakeNvteTensor(const_cast<void *>(input.data()), {M, N}, Int2NvteDType(otype));
  auto output_cu = MakeNvteTensor(output.data(), {N, M}, Int2NvteDType(otype));

  nvte_transpose(input_cu.data(), output_cu.data(), input.stream());

  return {output};
}

void te_cast_transpose(const paddle::Tensor &input, const paddle::Tensor &scale,
                       paddle::Tensor &output_cast,       // NOLINT
                       paddle::Tensor &output_transpose,  // NOLINT
                       paddle::Tensor &amax,              // NOLINT
                       paddle::Tensor &scale_inv,         // NOLINT
                       int64_t index, int64_t otype) {
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the input to have 2 dimensions.");

  size_t M = shape[0];
  size_t N = shape[1];

  auto input_cu = MakeNvteTensor(input);
  void *amax_data = GetDataPtr<float>(amax, index);
  void *scale_data = const_cast<void *>(GetDataPtr<float>(scale, index));
  void *scale_inv_data = GetDataPtr<float>(scale_inv, index);
  auto output_cast_cu = MakeNvteTensor(output_cast.data(), {M, N}, Int2NvteDType(otype), amax_data,
                                       scale_data, scale_inv_data);
  auto output_transpose_cu = MakeNvteTensor(output_transpose.data(), {N, M}, Int2NvteDType(otype),
                                            amax_data, scale_data, scale_inv_data);

  nvte_cast_transpose(input_cu.data(), output_cast_cu.data(), output_transpose_cu.data(),
                      input.stream());
}

std::vector<paddle::Tensor> te_cast_transpose_bgrad(const paddle::Tensor &grad_output,
                                                    const paddle::Tensor &scale,
                                                    paddle::Tensor &amax,       // NOLINT
                                                    paddle::Tensor &scale_inv,  // NOLINT
                                                    int64_t index, int64_t otype) {
  auto shape = GetShapeArray(grad_output);
  NVTE_CHECK(shape.size() == 2, "Expect the input to have 2 dimensions.");

  size_t M = shape[0];
  size_t N = shape[1];

  auto grad_bias =
      paddle::empty({grad_output.shape()[1]}, grad_output.dtype(), grad_output.place());
  auto grad_output_cast =
      paddle::empty_like(grad_output, Nvte2PaddleDType(Int2NvteDType(otype)), grad_output.place());
  auto grad_output_transpose =
      paddle::empty({grad_output.shape()[1], grad_output.shape()[0]},
                    Nvte2PaddleDType(Int2NvteDType(otype)), grad_output.place());

  auto input_cu = MakeNvteTensor(grad_output);
  void *amax_data = GetDataPtr<float>(amax, index);
  void *scale_data = const_cast<void *>(GetDataPtr<float>(scale, index));
  void *scale_inv_data = GetDataPtr<float>(scale_inv, index);
  auto output_cast_cu = MakeNvteTensor(grad_output_cast.data(), {M, N}, Int2NvteDType(otype),
                                       amax_data, scale_data, scale_inv_data);
  auto output_transpose_cu =
      MakeNvteTensor(grad_output_transpose.data(), {N, M}, Int2NvteDType(otype), amax_data,
                     scale_data, scale_inv_data);
  auto dbias_cu = MakeNvteTensor(grad_bias);
  transformer_engine::TensorWrapper workspace;

  nvte_cast_transpose_dbias(input_cu.data(), output_cast_cu.data(), output_transpose_cu.data(),
                            dbias_cu.data(), workspace.data(), grad_output.stream());

  // Fill workspace
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), grad_output.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  nvte_cast_transpose_dbias(input_cu.data(), output_cast_cu.data(), output_transpose_cu.data(),
                            dbias_cu.data(), workspace.data(), grad_output.stream());

  return {grad_bias, grad_output_cast, grad_output_transpose};
}

void te_gemm(const paddle::Tensor &A, const paddle::optional<paddle::Tensor> &A_scale_inverse,
             const paddle::Tensor &B, const paddle::optional<paddle::Tensor> &B_scale_inverse,
             const paddle::optional<paddle::Tensor> &bias, paddle::Tensor &D,            // NOLINT
             paddle::optional<paddle::Tensor> &D_scale,                                  // NOLINT
             paddle::optional<paddle::Tensor> &D_amax,                                   // NOLINT
             paddle::optional<paddle::Tensor> &pre_gelu_out, paddle::Tensor &workspace,  // NOLINT
             int64_t A_index, int64_t B_index, int64_t D_index, int64_t A_type, int64_t B_type,
             int64_t D_type, int64_t bias_type, bool transa, bool transb, bool grad,
             int64_t workspace_size, bool accumulate, bool use_split_accumulator,
             int64_t math_sm_count) {
  auto te_A = MakeNvteTensor(
      const_cast<void *>(A.data()), GetShapeArray(A), Int2NvteDType(A_type), nullptr, nullptr,
      const_cast<void *>(GetOptionalDataPtr<float>(A_scale_inverse, A_index)));
  auto te_B = MakeNvteTensor(
      const_cast<void *>(B.data()), GetShapeArray(B), Int2NvteDType(B_type), nullptr, nullptr,
      const_cast<void *>(GetOptionalDataPtr<float>(B_scale_inverse, B_index)));
  auto te_D = MakeNvteTensor(D.data(), GetShapeArray(D), Int2NvteDType(D_type),
                             GetOptionalDataPtr<float>(D_amax, D_index),
                             GetOptionalDataPtr<float>(D_scale, D_index), nullptr);

  auto te_bias = MakeNvteTensor(const_cast<void *>(GetOptionalDataPtr(bias)), GetShapeArray(bias),
                                Int2NvteDType(bias_type));

  DType gelu_dtype = pre_gelu_out ? Paddle2NvteDType(pre_gelu_out->dtype()) : Int2NvteDType(D_type);
  auto te_pre_gelu_out =
      MakeNvteTensor(GetOptionalDataPtr(pre_gelu_out), GetShapeArray(pre_gelu_out), gelu_dtype);
  auto te_workspace =
      MakeNvteTensor(workspace.data(), {static_cast<size_t>(workspace_size)}, DType::kByte);

  nvte_cublas_gemm(te_A.data(), te_B.data(), te_D.data(), te_bias.data(), te_pre_gelu_out.data(),
                   transa, transb, grad, te_workspace.data(), accumulate, use_split_accumulator,
                   math_sm_count, A.stream());
}

std::vector<paddle::Tensor> te_gelu_fp8(const paddle::Tensor &input, const paddle::Tensor &scale,
                                        paddle::Tensor &amax,       // NOLINT
                                        paddle::Tensor &scale_inv,  // NOLINT
                                        int64_t index, int64_t otype) {
  auto output = paddle::empty_like(input, Nvte2PaddleDType(DType::kByte), input.place());

  auto input_cu = MakeNvteTensor(input);
  auto output_cu = MakeNvteTensor(
      output.data(), GetShapeArray(input), Int2NvteDType(otype), GetDataPtr<float>(amax, index),
      const_cast<void *>(GetDataPtr<float>(scale, index)), GetDataPtr<float>(scale_inv, index));

  nvte_gelu(input_cu.data(), output_cu.data(), input.stream());

  return {output};
}

std::vector<paddle::Tensor> te_gelu(const paddle::Tensor &input, int64_t otype) {
  auto output = paddle::empty_like(input, Nvte2PaddleDType(Int2NvteDType(otype)), input.place());

  auto input_cu = MakeNvteTensor(input);
  auto output_cu = MakeNvteTensor(output.data(), GetShapeArray(input), Int2NvteDType(otype));

  nvte_gelu(input_cu.data(), output_cu.data(), input.stream());

  return {output};
}

std::vector<paddle::Tensor> te_swiglu(const paddle::Tensor &input, int64_t otype) {
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the input to have 2 dimensions.");

  size_t M = shape[0];
  size_t N = shape[1];

  auto output = paddle::empty({input.shape()[0], input.shape()[1] / 2},
                              Nvte2PaddleDType(Int2NvteDType(otype)), input.place());

  auto input_cu = MakeNvteTensor(input);
  auto output_cu = MakeNvteTensor(output.data(), GetShapeArray(output), Int2NvteDType(otype));

  nvte_swiglu(input_cu.data(), output_cu.data(), input.stream());

  return {output};
}

std::vector<paddle::Tensor> te_swiglu_fp8(const paddle::Tensor &input, const paddle::Tensor &scale,
                                          paddle::Tensor &amax,       // NOLINT
                                          paddle::Tensor &scale_inv,  // NOLINT
                                          int64_t index, int64_t otype) {
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the input to have 2 dimensions.");

  size_t M = shape[0];
  size_t N = shape[1];

  auto output = paddle::empty({input.shape()[0], input.shape()[1] / 2},
                              Nvte2PaddleDType(Int2NvteDType(otype)), input.place());

  auto input_cu = MakeNvteTensor(input);
  auto output_cu = MakeNvteTensor(
      output.data(), GetShapeArray(output), Int2NvteDType(otype), GetDataPtr<float>(amax, index),
      const_cast<void *>(GetDataPtr<float>(scale, index)), GetDataPtr<float>(scale_inv, index));

  nvte_swiglu(input_cu.data(), output_cu.data(), input.stream());

  return {output};
}

std::vector<paddle::Tensor> te_dswiglu(const paddle::Tensor &grad, const paddle::Tensor &input,
                                       int64_t otype) {
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the input to have 2 dimensions.");

  size_t M = shape[0];
  size_t N = shape[1];

  auto output = paddle::empty_like(input, Nvte2PaddleDType(Int2NvteDType(otype)), input.place());

  auto input_cu = MakeNvteTensor(input.data(), {M, N}, Paddle2NvteDType(input.dtype()));
  auto grad_cu = MakeNvteTensor(grad.data(), {M, N / 2}, Paddle2NvteDType(grad.dtype()));
  auto output_cu = MakeNvteTensor(output.data(), {M, N}, Paddle2NvteDType(output.dtype()));

  nvte_dswiglu(grad_cu.data(), input_cu.data(), output_cu.data(), input.stream());

  return {output};
}

std::vector<paddle::Tensor> te_cast_transpose_bgrad_dgelu(const paddle::Tensor &grad_output,
                                                          const paddle::Tensor &gelu_input,
                                                          const paddle::Tensor &scale,
                                                          paddle::Tensor &amax,       // NOLINT
                                                          paddle::Tensor &scale_inv,  // NOLINT
                                                          int64_t index, int64_t otype) {
  auto shape = GetShapeArray(grad_output);
  NVTE_CHECK(shape.size() == 2, "Expect the grad_output to have 2 dimensions.");

  size_t M = shape[0];
  size_t N = shape[1];

  // DType grad_output_type = GetTransformerEngineDType(grad_output.scalar_type());
  auto grad_bias =
      paddle::empty({grad_output.shape()[1]}, grad_output.dtype(), grad_output.place());

  auto dgelu = paddle::empty_like(grad_output, Nvte2PaddleDType(DType::kByte), grad_output.place());

  auto dgelu_transpose = paddle::empty({grad_output.shape()[1], grad_output.shape()[0]},
                                       Nvte2PaddleDType(DType::kByte), grad_output.place());

  void *amax_data = GetDataPtr<float>(amax, index);
  void *scale_data = const_cast<void *>(GetDataPtr<float>(scale, index));
  void *scale_inv_data = GetDataPtr<float>(scale_inv, index);

  TensorWrapper workspace;

  auto gelu_input_cu = MakeNvteTensor(gelu_input);
  auto input_cu = MakeNvteTensor(grad_output);
  auto cast_output_cu = MakeNvteTensor(dgelu.data(), {M, N}, Int2NvteDType(otype), amax_data,
                                       scale_data, scale_inv_data);
  auto transposed_output_cu = MakeNvteTensor(dgelu_transpose.data(), {N, M}, Int2NvteDType(otype),
                                             amax_data, scale_data, scale_inv_data);
  auto dbias_cu = MakeNvteTensor(grad_bias);

  nvte_cast_transpose_dbias_dgelu(input_cu.data(), gelu_input_cu.data(), cast_output_cu.data(),
                                  transposed_output_cu.data(), dbias_cu.data(), workspace.data(),
                                  grad_output.stream());

  // Fill workspace
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), grad_output.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  nvte_cast_transpose_dbias_dgelu(input_cu.data(), gelu_input_cu.data(), cast_output_cu.data(),
                                  transposed_output_cu.data(), dbias_cu.data(), workspace.data(),
                                  grad_output.stream());

  return {dgelu, dgelu_transpose, grad_bias};
}

std::vector<paddle::Tensor> te_layernorm_fwd_fp8(const paddle::Tensor &input,
                                                 const paddle::Tensor &weight,
                                                 const paddle::Tensor &bias,
                                                 const paddle::Tensor &scale,
                                                 paddle::Tensor &amax,       // NOLINT
                                                 paddle::Tensor &scale_inv,  // NOLINT
                                                 float eps, int64_t index, int64_t otype,
                                                 int64_t sm_margin, bool zero_centered_gamma) {
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the grad_output to have 2 dimensions.");

  size_t N = shape[0];
  size_t H = shape[1];

  auto ln_out = paddle::empty_like(input, Nvte2PaddleDType(Int2NvteDType(otype)), input.place());
  auto mu = paddle::empty({static_cast<int64_t>(N)}, paddle::DataType::FLOAT32, input.place());
  auto rsigma = paddle::empty({static_cast<int64_t>(N)}, paddle::DataType::FLOAT32, input.place());
  auto input_cu = MakeNvteTensor(input);
  auto gamma_cu = MakeNvteTensor(weight);
  auto beta_cu = MakeNvteTensor(bias);
  auto z_cu = MakeNvteTensor(
      ln_out.data(), {N, H}, Int2NvteDType(otype), GetDataPtr<float>(amax, index),
      const_cast<void *>(GetDataPtr<float>(scale, index)), GetDataPtr<float>(scale_inv, index));
  auto mu_cu = MakeNvteTensor(mu);
  auto rsigma_cu = MakeNvteTensor(rsigma);
  TensorWrapper workspace, barrier;

  auto num_sm = cudaDevicePropertiesManager::Instance().GetMultiProcessorCount();

  // This call populates workspace and barrier tensors with the required config
  const auto func = zero_centered_gamma ? nvte_layernorm1p_fwd : nvte_layernorm_fwd;
  func(input_cu.data(), gamma_cu.data(), beta_cu.data(), eps, z_cu.data(), mu_cu.data(),
       rsigma_cu.data(), input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  // Fill workspace and barrier
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), input.place());
  auto barrier_data = AllocateSpace(barrier.shape(), barrier.dtype(), input.place(), true);
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());
  barrier = MakeNvteTensor(barrier_data.data(), barrier.shape(), barrier.dtype());

  // Actual call to fwd kernel
  func(input_cu.data(), gamma_cu.data(), beta_cu.data(), eps, z_cu.data(), mu_cu.data(),
       rsigma_cu.data(), input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  return {ln_out, mu, rsigma};
}

std::vector<paddle::Tensor> te_layernorm_fwd(const paddle::Tensor &input,
                                             const paddle::Tensor &weight,
                                             const paddle::Tensor &bias, float eps, int64_t otype,
                                             int64_t sm_margin, bool zero_centered_gamma) {
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the grad_output to have 2 dimensions.");

  size_t N = shape[0];
  size_t H = shape[1];

  auto ln_out = paddle::empty_like(input, input.dtype(), input.place());
  auto mu = paddle::empty({static_cast<int64_t>(N)}, paddle::DataType::FLOAT32, input.place());
  auto rsigma = paddle::empty({static_cast<int64_t>(N)}, paddle::DataType::FLOAT32, input.place());
  auto input_cu = MakeNvteTensor(input);
  auto gamma_cu = MakeNvteTensor(weight);
  auto beta_cu = MakeNvteTensor(bias);
  auto z_cu = MakeNvteTensor(ln_out.data(), {N, H}, Int2NvteDType(otype));
  auto mu_cu = MakeNvteTensor(mu);
  auto rsigma_cu = MakeNvteTensor(rsigma);
  TensorWrapper workspace, barrier;

  auto num_sm = cudaDevicePropertiesManager::Instance().GetMultiProcessorCount();

  // This call populates workspace and barrier tensors with the required config
  const auto func = zero_centered_gamma ? nvte_layernorm1p_fwd : nvte_layernorm_fwd;
  func(input_cu.data(), gamma_cu.data(), beta_cu.data(), eps, z_cu.data(), mu_cu.data(),
       rsigma_cu.data(), input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  // Fill workspace and barrier
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), input.place());
  auto barrier_data = AllocateSpace(barrier.shape(), barrier.dtype(), input.place(), true);
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());
  barrier = MakeNvteTensor(barrier_data.data(), barrier.shape(), barrier.dtype());

  // Actual call to fwd kernel
  func(input_cu.data(), gamma_cu.data(), beta_cu.data(), eps, z_cu.data(), mu_cu.data(),
       rsigma_cu.data(), input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  return {ln_out, mu, rsigma};
}

std::vector<paddle::Tensor> te_layernorm_bwd(const paddle::Tensor &dz, const paddle::Tensor &x,
                                             const paddle::Tensor &mu, const paddle::Tensor &rsigma,
                                             const paddle::Tensor &gamma, int64_t sm_margin,
                                             bool zero_centered_gamma) {
  auto dx = paddle::empty_like(x, x.dtype(), x.place());
  auto dgamma = paddle::empty_like(gamma, gamma.dtype(), gamma.place());
  auto dbeta = paddle::empty_like(gamma, gamma.dtype(), gamma.place());

  TensorWrapper workspace, barrier, dgamma_part, dbeta_part;

  auto dz_cu = MakeNvteTensor(dz);
  auto x_cu = MakeNvteTensor(x);
  auto mu_cu = MakeNvteTensor(mu);
  auto rsigma_cu = MakeNvteTensor(rsigma);
  auto gamma_cu = MakeNvteTensor(gamma);
  auto dx_cu = MakeNvteTensor(dx);
  auto dgamma_cu = MakeNvteTensor(dgamma);
  auto dbeta_cu = MakeNvteTensor(dbeta);

  auto num_sm = cudaDevicePropertiesManager::Instance().GetMultiProcessorCount();

  // This call populates tensors with the required config.
  const auto bwd_fun = zero_centered_gamma ? nvte_layernorm1p_bwd : nvte_layernorm_bwd;
  bwd_fun(dz_cu.data(), x_cu.data(), mu_cu.data(), rsigma_cu.data(), gamma_cu.data(), dx_cu.data(),
          dgamma_cu.data(), dbeta_cu.data(), dgamma_part.data(), dbeta_part.data(), dz.stream(),
          num_sm - sm_margin, workspace.data(), barrier.data());

  // Alloc space for Tensors.
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), x.place());
  auto barrier_data = AllocateSpace(barrier.shape(), barrier.dtype(), x.place(), true);
  auto dgamma_part_data = AllocateSpace(dgamma_part.shape(), dgamma_part.dtype(), x.place());
  auto dbeta_part_data = AllocateSpace(dbeta_part.shape(), dbeta_part.dtype(), x.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());
  barrier = MakeNvteTensor(barrier_data.data(), barrier.shape(), barrier.dtype());
  dgamma_part = MakeNvteTensor(dgamma_part_data.data(), dgamma_part.shape(), dgamma_part.dtype());
  dbeta_part = MakeNvteTensor(dbeta_part_data.data(), dbeta_part.shape(), dbeta_part.dtype());

  // Actual call to bwd kernel.
  bwd_fun(dz_cu.data(), x_cu.data(), mu_cu.data(), rsigma_cu.data(), gamma_cu.data(), dx_cu.data(),
          dgamma_cu.data(), dbeta_cu.data(), dgamma_part.data(), dbeta_part.data(), dz.stream(),
          num_sm - sm_margin, workspace.data(), barrier.data());

  return {dx, dgamma, dbeta};
}

std::vector<paddle::Tensor> te_rmsnorm_fwd(const paddle::Tensor &input,
                                           const paddle::Tensor &weight, float eps, int64_t otype,
                                           int64_t sm_margin, bool zero_centered_gamma) {
  NVTE_CHECK(zero_centered_gamma == false, "zero_centered_gamma is not supported yet for RMSNorm.");
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the grad_output to have 2 dimensions.");

  size_t N = shape[0];
  size_t H = shape[1];

  auto ln_out = paddle::empty_like(input, input.dtype(), input.place());
  auto rsigma = paddle::empty({static_cast<int64_t>(N)}, paddle::DataType::FLOAT32, input.place());
  auto input_cu = MakeNvteTensor(input);
  auto gamma_cu = MakeNvteTensor(weight);
  auto z_cu = MakeNvteTensor(ln_out.data(), {N, H}, Int2NvteDType(otype));
  auto rsigma_cu = MakeNvteTensor(rsigma);
  TensorWrapper workspace, barrier;

  auto num_sm = cudaDevicePropertiesManager::Instance().GetMultiProcessorCount();

  // This call populates workspace and barrier tensors with the required config

  nvte_rmsnorm_fwd(input_cu.data(), gamma_cu.data(), eps, z_cu.data(), rsigma_cu.data(),
                   input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  // Fill workspace and barrier
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), input.place());
  auto barrier_data = AllocateSpace(barrier.shape(), barrier.dtype(), input.place(), true);
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());
  barrier = MakeNvteTensor(barrier_data.data(), barrier.shape(), barrier.dtype());

  // Actual call to fwd kernel
  nvte_rmsnorm_fwd(input_cu.data(), gamma_cu.data(), eps, z_cu.data(), rsigma_cu.data(),
                   input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  return {ln_out, rsigma};
}

std::vector<paddle::Tensor> te_rmsnorm_fwd_fp8(const paddle::Tensor &input,
                                               const paddle::Tensor &weight,
                                               const paddle::Tensor &scale,
                                               paddle::Tensor &amax,       // NOLINT
                                               paddle::Tensor &scale_inv,  // NOLINT
                                               float eps, int64_t index, int64_t otype,
                                               int64_t sm_margin, bool zero_centered_gamma) {
  NVTE_CHECK(zero_centered_gamma == false, "zero_centered_gamma is not supported yet for RMSNorm.");
  auto shape = GetShapeArray(input);
  NVTE_CHECK(shape.size() == 2, "Expect the grad_output to have 2 dimensions.");

  size_t N = shape[0];
  size_t H = shape[1];

  auto ln_out = paddle::empty_like(input, Nvte2PaddleDType(Int2NvteDType(otype)), input.place());
  auto rsigma = paddle::empty({static_cast<int64_t>(N)}, paddle::DataType::FLOAT32, input.place());
  auto input_cu = MakeNvteTensor(input);
  auto gamma_cu = MakeNvteTensor(weight);
  auto z_cu = MakeNvteTensor(
      ln_out.data(), {N, H}, Int2NvteDType(otype), GetDataPtr<float>(amax, index),
      const_cast<void *>(GetDataPtr<float>(scale, index)), GetDataPtr<float>(scale_inv, index));
  auto rsigma_cu = MakeNvteTensor(rsigma);
  TensorWrapper workspace, barrier;

  auto num_sm = cudaDevicePropertiesManager::Instance().GetMultiProcessorCount();

  // This call populates workspace and barrier tensors with the required config
  nvte_rmsnorm_fwd(input_cu.data(), gamma_cu.data(), eps, z_cu.data(), rsigma_cu.data(),
                   input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  // Fill workspace and barrier
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), input.place());
  auto barrier_data = AllocateSpace(barrier.shape(), barrier.dtype(), input.place(), true);
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());
  barrier = MakeNvteTensor(barrier_data.data(), barrier.shape(), barrier.dtype());

  // Actual call to fwd kernel
  nvte_rmsnorm_fwd(input_cu.data(), gamma_cu.data(), eps, z_cu.data(), rsigma_cu.data(),
                   input.stream(), num_sm - sm_margin, workspace.data(), barrier.data());

  return {ln_out, rsigma};
}

std::vector<paddle::Tensor> te_rmsnorm_bwd(const paddle::Tensor &dz, const paddle::Tensor &x,
                                           const paddle::Tensor &rsigma,
                                           const paddle::Tensor &gamma, int64_t sm_margin,
                                           bool zero_centered_gamma) {
  NVTE_CHECK(zero_centered_gamma == false, "zero_centered_gamma is not supported yet for RMSNorm.");
  auto dx = paddle::empty_like(x, x.dtype(), x.place());
  auto dgamma = paddle::empty_like(gamma, gamma.dtype(), gamma.place());

  TensorWrapper workspace, barrier, dgamma_part;

  auto dz_cu = MakeNvteTensor(dz);
  auto x_cu = MakeNvteTensor(x);
  auto rsigma_cu = MakeNvteTensor(rsigma);
  auto gamma_cu = MakeNvteTensor(gamma);
  auto dx_cu = MakeNvteTensor(dx);
  auto dgamma_cu = MakeNvteTensor(dgamma);

  auto num_sm = cudaDevicePropertiesManager::Instance().GetMultiProcessorCount();

  // This call populates tensors with the required config.
  nvte_rmsnorm_bwd(dz_cu.data(), x_cu.data(), rsigma_cu.data(), gamma_cu.data(), dx_cu.data(),
                   dgamma_cu.data(), dgamma_part.data(), dz.stream(), num_sm - sm_margin,
                   workspace.data(), barrier.data());

  // Alloc space for Tensors.
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), x.place());
  auto barrier_data = AllocateSpace(barrier.shape(), barrier.dtype(), x.place(), true);
  auto dgamma_part_data = AllocateSpace(dgamma_part.shape(), dgamma_part.dtype(), x.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());
  barrier = MakeNvteTensor(barrier_data.data(), barrier.shape(), barrier.dtype());
  dgamma_part = MakeNvteTensor(dgamma_part_data.data(), dgamma_part.shape(), dgamma_part.dtype());

  // Actual call to bwd kernel.
  nvte_rmsnorm_bwd(dz_cu.data(), x_cu.data(), rsigma_cu.data(), gamma_cu.data(), dx_cu.data(),
                   dgamma_cu.data(), dgamma_part.data(), dz.stream(), num_sm - sm_margin,
                   workspace.data(), barrier.data());

  return {dx, dgamma};
}

__global__ void set_rng_state(std::pair<uint64_t, uint64_t> seed_offset, int64_t *rng_state_ptr) {
  rng_state_ptr[0] = static_cast<int64_t>(seed_offset.first);
  rng_state_ptr[1] = static_cast<int64_t>(seed_offset.second);
}

void te_fused_attn_fwd_qkvpacked(const paddle::Tensor &QKV, const paddle::Tensor &cu_seqlens,
                                 const paddle::optional<paddle::Tensor> &Bias,
                                 paddle::Tensor &O,                              // NOLINT
                                 paddle::optional<paddle::Tensor> &softmax_aux,  // NOLINT
                                 paddle::Tensor &rng_state,                      // NOLINT
                                 int64_t b, int64_t h, int64_t d, int64_t total_seqs,
                                 int64_t max_seqlen, bool is_training, float attn_scale,
                                 float p_dropout, const std::string &qkv_layout,
                                 const std::string &bias_type, const std::string &attn_mask_type,
                                 const int64_t qkv_type, int64_t rng_elts_per_thread) {
  if (is_training && !softmax_aux) {
    NVTE_ERROR("softmax_aux must be provided when training. \n");
  }

  auto qkv_dtype = Int2NvteDType(qkv_type);
  // construct NVTE tensors
  TensorWrapper te_QKV, te_S, te_O, te_Bias, te_cu_seqlens;
  if (qkv_dtype == DType::kBFloat16 || qkv_dtype == DType::kFloat16) {
    // BF16 or FP16
    te_QKV = MakeNvteTensor(QKV);
    te_S = MakeNvteTensor(nullptr, std::vector<size_t>{0}, DType::kFloat32);
    te_O = MakeNvteTensor(O);
  } else {  // TODO: support fp8
    NVTE_ERROR("Fused attention only supports BF16/FP16 data types. \n");
  }
  if ((bias_type != "no_bias") && Bias) {
    auto bias_shape = Bias->shape();
    std::vector<size_t> shape{bias_shape.begin(), bias_shape.end()};
    te_Bias = MakeNvteTensor(GetOptionalDataPtr(Bias), shape, DType::kFloat32);
  }
  te_cu_seqlens = MakeNvteTensor(cu_seqlens.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);

  // convert strings to enums
  NVTE_QKV_Layout qkv_layout_enum = get_nvte_qkv_layout(qkv_layout);
  NVTE_Bias_Type bias_type_enum = get_nvte_bias_type(bias_type);
  NVTE_Mask_Type attn_mask_type_enum = get_nvte_mask_type(attn_mask_type);

  // extract random number generator seed and offset
  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(QKV.place());
  auto gen_cuda = dev_ctx->GetGenerator();
  auto seed_offset = gen_cuda->IncrementOffset(rng_elts_per_thread);
  set_rng_state<<<1, 1, 0, QKV.stream()>>>(seed_offset, static_cast<int64_t *>(rng_state.data()));

  auto te_rng_state = MakeNvteTensor(rng_state);

  // create auxiliary output tensors
  NVTETensorPack nvte_aux_tensor_pack;
  nvte_tensor_pack_create(&nvte_aux_tensor_pack);

  // create workspace
  TensorWrapper workspace;

  auto dummy_seq_offsets = TensorWrapper(nullptr, {static_cast<size_t>(b + 1)}, DType::kInt32);
  // populate tensors with appropriate shapes and dtypes
  nvte_fused_attn_fwd_qkvpacked(
      te_QKV.data(), te_Bias.data(), te_S.data(), te_O.data(), &nvte_aux_tensor_pack,
      te_cu_seqlens.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
      dummy_seq_offsets.data(), dummy_seq_offsets.data(), te_rng_state.data(), max_seqlen,
      is_training, attn_scale, p_dropout, qkv_layout_enum, bias_type_enum, attn_mask_type_enum,
      workspace.data(), QKV.stream());

  // allocate memory for workspace and auxiliary output tensors
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), QKV.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  auto *output_s = reinterpret_cast<transformer_engine::Tensor *>(nvte_aux_tensor_pack.tensors[0]);
  output_s->data.dptr = GetOptionalDataPtr(softmax_aux);

  // execute the kernel
  nvte_fused_attn_fwd_qkvpacked(
      te_QKV.data(), te_Bias.data(), te_S.data(), te_O.data(), &nvte_aux_tensor_pack,
      te_cu_seqlens.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
      dummy_seq_offsets.data(), dummy_seq_offsets.data(), te_rng_state.data(), max_seqlen,
      is_training, attn_scale, p_dropout, qkv_layout_enum, bias_type_enum, attn_mask_type_enum,
      workspace.data(), QKV.stream());

  // destroy tensor wrappers, but not allocated memory
  nvte_tensor_pack_destroy(&nvte_aux_tensor_pack);
}

// fused attention BWD with packed QKV
void te_fused_attn_bwd_qkvpacked(const paddle::Tensor &QKV, const paddle::Tensor &cu_seqlens,
                                 const paddle::Tensor &O, const paddle::Tensor &dO,
                                 const paddle::Tensor &softmax_aux,
                                 paddle::Tensor &dQKV,                     // NOLINT
                                 paddle::optional<paddle::Tensor> &dBias,  // NOLINT
                                 paddle::Tensor &rng_state,                // NOLINT
                                 int64_t b, int64_t h, int64_t d, int64_t total_seqs,
                                 int64_t max_seqlen, float attn_scale, float p_dropout,
                                 const std::string &qkv_layout, const std::string &bias_type,
                                 const std::string &attn_mask_type, int64_t qkv_type) {
  TensorWrapper te_dBias;
  if (bias_type != "no_bias" && dBias) {
    auto bias_shape = dBias->shape();
    std::vector<size_t> shape{bias_shape.begin(), bias_shape.end()};
    te_dBias = MakeNvteTensor(GetOptionalDataPtr(dBias), shape, DType::kFloat32);
  }

  auto qkv_dtype = Int2NvteDType(qkv_type);
  // construct NVTE tensors
  TensorWrapper te_QKV, te_O, te_dO, te_S, te_dP, te_dQKV;
  if (qkv_dtype == DType::kBFloat16 || qkv_dtype == DType::kFloat16) {
    // BF16 or FP16
    te_QKV = MakeNvteTensor(QKV);
    te_O = MakeNvteTensor(O);
    te_dO = MakeNvteTensor(dO);
    te_S = MakeNvteTensor(nullptr, std::vector<size_t>(0), DType::kFloat32);
    te_dP = MakeNvteTensor(nullptr, std::vector<size_t>(0), DType::kFloat32);
    te_dQKV = MakeNvteTensor(dQKV);
  } else {
    NVTE_ERROR("Fused attention only supports BF16/FP16 data types. \n");
  }

  // convert strings to enums
  NVTE_QKV_Layout qkv_layout_enum = get_nvte_qkv_layout(qkv_layout);
  NVTE_Bias_Type bias_type_enum = get_nvte_bias_type(bias_type);
  NVTE_Mask_Type attn_mask_type_enum = get_nvte_mask_type(attn_mask_type);

  // convert auxiliary tensors from forward into NVTETensors
  NVTETensorPack nvte_aux_tensor_pack;
  nvte_tensor_pack_create(&nvte_aux_tensor_pack);

  nvte_aux_tensor_pack.size = 2;  // 1. softmax_aux  2. rng_state
  auto *output_s = reinterpret_cast<Tensor *>(nvte_aux_tensor_pack.tensors[0]);
  auto *fwd_rng_state = reinterpret_cast<Tensor *>(nvte_aux_tensor_pack.tensors[1]);
  output_s->data.shape =
      std::vector<size_t>({static_cast<size_t>(b), static_cast<size_t>(h),
                           static_cast<size_t>(max_seqlen), static_cast<size_t>(max_seqlen)});
  output_s->data.dptr = const_cast<void *>(softmax_aux.data());
  fwd_rng_state->data.shape = std::vector<size_t>({2});
  fwd_rng_state->data.dptr = const_cast<void *>(rng_state.data());

  // create cu_seqlens tensorwrappers
  TensorWrapper te_cu_seqlens;
  te_cu_seqlens = MakeNvteTensor(cu_seqlens.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);

  // create workspace
  TensorWrapper workspace;

  auto dummy_seq_offsets = TensorWrapper(nullptr, {static_cast<size_t>(b + 1)}, DType::kInt32);
  // populate tensors with appropriate shapes and dtypes
  nvte_fused_attn_bwd_qkvpacked(
      te_QKV.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), &nvte_aux_tensor_pack,
      te_dQKV.data(), te_dBias.data(), te_cu_seqlens.data(), dummy_seq_offsets.data(),
      dummy_seq_offsets.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(), max_seqlen,
      attn_scale, p_dropout, qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(),
      QKV.stream());

  // allocate memory for workspace
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), QKV.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  // execute kernel
  nvte_fused_attn_bwd_qkvpacked(
      te_QKV.data(), te_O.data(), te_dO.data(), te_S.data(), te_dP.data(), &nvte_aux_tensor_pack,
      te_dQKV.data(), te_dBias.data(), te_cu_seqlens.data(), dummy_seq_offsets.data(),
      dummy_seq_offsets.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(), max_seqlen,
      attn_scale, p_dropout, qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(),
      QKV.stream());

  // destroy tensor wrappers
  nvte_tensor_pack_destroy(&nvte_aux_tensor_pack);
}

void te_fused_attn_fwd_kvpacked(
    const paddle::Tensor &Q, const paddle::Tensor &KV, const paddle::Tensor &cu_seqlens_q,
    const paddle::Tensor &cu_seqlens_kv, const paddle::optional<paddle::Tensor> &Bias,
    paddle::Tensor &O,                              // NOLINT
    paddle::optional<paddle::Tensor> &softmax_aux,  // NOLINT
    paddle::Tensor &rng_state,                      // NOLINT
    int64_t b, int64_t h, int64_t d, int64_t total_seqs_q, int64_t total_seqs_kv,
    int64_t max_seqlen_q, int64_t max_seqlen_kv, bool is_training, float attn_scale,
    float p_dropout, const std::string &qkv_layout, const std::string &bias_type,
    const std::string &attn_mask_type, const int64_t qkv_type, int64_t rng_elts_per_thread) {
  if (is_training && !softmax_aux) {
    NVTE_ERROR("softmax_aux must be provided when training. \n");
  }

  auto qkv_dtype = Int2NvteDType(qkv_type);

  // construct NVTE tensors
  TensorWrapper te_Q, te_KV, te_S, te_O, te_Bias, te_cu_seqlens_q, te_cu_seqlens_kv;
  if (qkv_dtype == DType::kBFloat16 || qkv_dtype == DType::kFloat16) {
    // BF16 or FP16
    te_Q = MakeNvteTensor(
        Q.data(),
        {static_cast<size_t>(total_seqs_q), static_cast<size_t>(h), static_cast<size_t>(d)},
        qkv_dtype);
    te_KV = MakeNvteTensor(
        KV.data(),
        {static_cast<size_t>(total_seqs_kv), 2, static_cast<size_t>(h), static_cast<size_t>(d)},
        qkv_dtype);
    te_S = MakeNvteTensor(nullptr, std::vector<size_t>{0}, DType::kFloat32);
    te_O = MakeNvteTensor(
        O.data(),
        {static_cast<size_t>(total_seqs_q), static_cast<size_t>(h), static_cast<size_t>(d)},
        qkv_dtype);
  } else {
    NVTE_ERROR("Fused attention only supports BF16/FP16 data types. \n");
  }

  if ((bias_type != "no_bias") && Bias) {
    auto bias_shape = Bias->shape();
    std::vector<size_t> shape{bias_shape.begin(), bias_shape.end()};
    te_Bias = MakeNvteTensor(GetOptionalDataPtr(Bias), shape, DType::kFloat32);
  }

  te_cu_seqlens_q =
      MakeNvteTensor(cu_seqlens_q.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);
  te_cu_seqlens_kv =
      MakeNvteTensor(cu_seqlens_kv.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);

  // convert strings to enums
  NVTE_QKV_Layout qkv_layout_enum = get_nvte_qkv_layout(qkv_layout);
  NVTE_Bias_Type bias_type_enum = get_nvte_bias_type(bias_type);
  NVTE_Mask_Type attn_mask_type_enum = get_nvte_mask_type(attn_mask_type);

  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(Q.place());
  auto gen_cuda = dev_ctx->GetGenerator();
  auto seed_offset = gen_cuda->IncrementOffset(rng_elts_per_thread);
  set_rng_state<<<1, 1, 0, Q.stream()>>>(seed_offset, static_cast<int64_t *>(rng_state.data()));
  auto te_rng_state = MakeNvteTensor(rng_state);

  // create auxiliary output tensors
  NVTETensorPack nvte_aux_tensor_pack;
  nvte_tensor_pack_create(&nvte_aux_tensor_pack);

  // create workspace
  TensorWrapper workspace;

  auto dummy_seq_offsets = TensorWrapper(nullptr, {static_cast<size_t>(b + 1)}, DType::kInt32);
  // populate tensors with appropriate shapes and dtypes
  nvte_fused_attn_fwd_kvpacked(
      te_Q.data(), te_KV.data(), te_Bias.data(), te_S.data(), te_O.data(), &nvte_aux_tensor_pack,
      te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), dummy_seq_offsets.data(),
      dummy_seq_offsets.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
      te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, attn_scale, p_dropout,
      qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(), Q.stream());

  // allocate memory for workspace and auxiliary output tensors
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), Q.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  auto *output_s = reinterpret_cast<transformer_engine::Tensor *>(nvte_aux_tensor_pack.tensors[0]);
  output_s->data.dptr = GetOptionalDataPtr(softmax_aux);

  // execute the kernel
  nvte_fused_attn_fwd_kvpacked(
      te_Q.data(), te_KV.data(), te_Bias.data(), te_S.data(), te_O.data(), &nvte_aux_tensor_pack,
      te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(), dummy_seq_offsets.data(),
      dummy_seq_offsets.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
      te_rng_state.data(), max_seqlen_q, max_seqlen_kv, is_training, attn_scale, p_dropout,
      qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(), Q.stream());

  // destroy tensor wrappers, but not allocated memory
  nvte_tensor_pack_destroy(&nvte_aux_tensor_pack);
}

// fused attention BWD with packed KV
void te_fused_attn_bwd_kvpacked(const paddle::Tensor &Q, const paddle::Tensor &KV,
                                const paddle::Tensor &cu_seqlens_q,
                                const paddle::Tensor &cu_seqlens_kv, const paddle::Tensor &O,
                                const paddle::Tensor &dO, const paddle::Tensor &softmax_aux,
                                paddle::Tensor &dQ,                       // NOLINT
                                paddle::Tensor &dKV,                      // NOLINT
                                paddle::optional<paddle::Tensor> &dBias,  // NOLINT
                                paddle::Tensor &rng_state,                // NOLINT
                                int64_t b, int64_t h, int64_t d, int64_t total_seqs_q,
                                int64_t total_seqs_kv, int64_t max_seqlen_q, int64_t max_seqlen_kv,
                                float attn_scale, float p_dropout, const std::string &qkv_layout,
                                const std::string &bias_type, const std::string &attn_mask_type,
                                int64_t qkv_type) {
  TensorWrapper te_dBias;
  if (bias_type != "no_bias" && dBias) {
    auto bias_shape = dBias->shape();
    std::vector<size_t> shape{bias_shape.begin(), bias_shape.end()};
    te_dBias = MakeNvteTensor(GetOptionalDataPtr(dBias), shape, DType::kFloat32);
  }

  auto qkv_dtype = Int2NvteDType(qkv_type);
  // construct NVTE tensors
  TensorWrapper te_Q, te_KV, te_O, te_dO, te_S, te_dP, te_dQ, te_dKV;
  if (qkv_dtype == DType::kBFloat16 || qkv_dtype == DType::kFloat16) {
    // BF16 or FP16
    te_Q = MakeNvteTensor(Q);
    te_KV = MakeNvteTensor(KV);
    te_O = MakeNvteTensor(O);
    te_dO = MakeNvteTensor(dO);
    te_S = MakeNvteTensor(nullptr, std::vector<size_t>(0), DType::kFloat32);
    te_dP = MakeNvteTensor(nullptr, std::vector<size_t>(0), DType::kFloat32);
    te_dQ = MakeNvteTensor(dQ);
    te_dKV = MakeNvteTensor(dKV);
  } else {
    NVTE_ERROR("Fused attention only supports BF16/FP16 data types. \n");
  }

  // convert strings to enums
  NVTE_QKV_Layout qkv_layout_enum = get_nvte_qkv_layout(qkv_layout);
  NVTE_Bias_Type bias_type_enum = get_nvte_bias_type(bias_type);
  NVTE_Mask_Type attn_mask_type_enum = get_nvte_mask_type(attn_mask_type);

  // convert auxiliary tensors from forward into NVTETensors
  NVTETensorPack nvte_aux_tensor_pack;
  nvte_tensor_pack_create(&nvte_aux_tensor_pack);

  nvte_aux_tensor_pack.size = 2;
  auto *output_s = reinterpret_cast<Tensor *>(nvte_aux_tensor_pack.tensors[0]);
  auto *fwd_rng_state = reinterpret_cast<Tensor *>(nvte_aux_tensor_pack.tensors[1]);
  output_s->data.shape =
      std::vector<size_t>({static_cast<size_t>(b), static_cast<size_t>(h),
                           static_cast<size_t>(max_seqlen_q), static_cast<size_t>(max_seqlen_kv)});
  output_s->data.dptr = const_cast<void *>(softmax_aux.data());
  fwd_rng_state->data.shape = std::vector<size_t>({2});
  fwd_rng_state->data.dptr = const_cast<void *>(rng_state.data());

  // create cu_seqlens tensorwrappers
  TensorWrapper te_cu_seqlens_q, te_cu_seqlens_kv;
  te_cu_seqlens_q =
      MakeNvteTensor(cu_seqlens_q.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);
  te_cu_seqlens_kv =
      MakeNvteTensor(cu_seqlens_kv.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);

  // create workspace
  TensorWrapper workspace;

  auto dummy_seq_offsets = TensorWrapper(nullptr, {static_cast<size_t>(b + 1)}, DType::kInt32);
  // populate tensors with appropriate shapes and dtypes
  nvte_fused_attn_bwd_kvpacked(te_Q.data(), te_KV.data(), te_O.data(), te_dO.data(), te_S.data(),
                               te_dP.data(), &nvte_aux_tensor_pack, te_dQ.data(), te_dKV.data(),
                               te_dBias.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(),
                               dummy_seq_offsets.data(), dummy_seq_offsets.data(),
                               dummy_seq_offsets.data(), dummy_seq_offsets.data(), max_seqlen_q,
                               max_seqlen_kv, attn_scale, p_dropout, qkv_layout_enum,
                               bias_type_enum, attn_mask_type_enum, workspace.data(), Q.stream());

  // allocate memory for workspace
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), Q.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  // execute kernel
  nvte_fused_attn_bwd_kvpacked(te_Q.data(), te_KV.data(), te_O.data(), te_dO.data(), te_S.data(),
                               te_dP.data(), &nvte_aux_tensor_pack, te_dQ.data(), te_dKV.data(),
                               te_dBias.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(),
                               dummy_seq_offsets.data(), dummy_seq_offsets.data(),
                               dummy_seq_offsets.data(), dummy_seq_offsets.data(), max_seqlen_q,
                               max_seqlen_kv, attn_scale, p_dropout, qkv_layout_enum,
                               bias_type_enum, attn_mask_type_enum, workspace.data(), Q.stream());

  // destroy tensor wrappers
  nvte_tensor_pack_destroy(&nvte_aux_tensor_pack);
}

void te_fused_attn_fwd(const paddle::Tensor &Q, const paddle::Tensor &K, const paddle::Tensor &V,
                       const paddle::Tensor &cu_seqlens_q, const paddle::Tensor &cu_seqlens_kv,
                       const paddle::optional<paddle::Tensor> &Bias,
                       paddle::Tensor &O,                              // NOLINT
                       paddle::optional<paddle::Tensor> &softmax_aux,  // NOLINT
                       paddle::Tensor &rng_state,                      // NOLINT
                       int64_t b, int64_t h, int64_t d, int64_t max_seqlen_q, int64_t max_seqlen_kv,
                       bool is_training, float attn_scale, float p_dropout,
                       const std::string &qkv_layout, const std::string &bias_type,
                       const std::string &attn_mask_type, const int64_t qkv_type,
                       int64_t rng_elts_per_thread) {
  if (is_training && !softmax_aux) {
    NVTE_ERROR("softmax_aux must be provided when training. \n");
  }

  auto qkv_dtype = Int2NvteDType(qkv_type);
  // construct NVTE tensors
  TensorWrapper te_Q, te_K, te_V, te_S, te_O, te_Bias, te_cu_seqlens_q, te_cu_seqlens_kv;
  if (qkv_dtype == DType::kBFloat16 || qkv_dtype == DType::kFloat16) {
    // BF16 or FP16
    te_Q = MakeNvteTensor(Q);
    te_K = MakeNvteTensor(K);
    te_V = MakeNvteTensor(V);
    te_S = MakeNvteTensor(nullptr, std::vector<size_t>{0}, DType::kFloat32);
    te_O = MakeNvteTensor(O);
  } else {  // TODO: support fp8
    NVTE_ERROR("Fused attention only supports BF16/FP16 data types. \n");
  }
  if ((bias_type != "no_bias") && Bias) {
    auto bias_shape = Bias->shape();
    std::vector<size_t> shape{bias_shape.begin(), bias_shape.end()};
    te_Bias = MakeNvteTensor(GetOptionalDataPtr(Bias), shape, DType::kFloat32);
  }
  te_cu_seqlens_q =
      MakeNvteTensor(cu_seqlens_q.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);
  te_cu_seqlens_kv =
      MakeNvteTensor(cu_seqlens_kv.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);

  // convert strings to enums
  NVTE_QKV_Layout qkv_layout_enum = get_nvte_qkv_layout(qkv_layout);
  NVTE_Bias_Type bias_type_enum = get_nvte_bias_type(bias_type);
  NVTE_Mask_Type attn_mask_type_enum = get_nvte_mask_type(attn_mask_type);

  // extract random number generator seed and offset
  auto dev_ctx = paddle::experimental::DeviceContextPool::Instance().Get(Q.place());
  auto gen_cuda = dev_ctx->GetGenerator();
  auto seed_offset = gen_cuda->IncrementOffset(rng_elts_per_thread);
  set_rng_state<<<1, 1, 0, Q.stream()>>>(seed_offset, static_cast<int64_t *>(rng_state.data()));

  auto te_rng_state = MakeNvteTensor(rng_state);

  // create auxiliary output tensors
  NVTETensorPack nvte_aux_tensor_pack;
  nvte_tensor_pack_create(&nvte_aux_tensor_pack);

  // create workspace
  TensorWrapper workspace;

  auto dummy_seq_offsets = TensorWrapper(nullptr, {static_cast<size_t>(b + 1)}, DType::kInt32);
  // populate tensors with appropriate shapes and dtypes
  nvte_fused_attn_fwd(te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_S.data(),
                      te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(),
                      te_cu_seqlens_kv.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
                      dummy_seq_offsets.data(), dummy_seq_offsets.data(), te_rng_state.data(),
                      max_seqlen_q, max_seqlen_kv, is_training, attn_scale, p_dropout,
                      qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(),
                      Q.stream());

  // allocate memory for workspace and auxiliary output tensors
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), Q.place());

  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  auto *output_s = reinterpret_cast<transformer_engine::Tensor *>(nvte_aux_tensor_pack.tensors[0]);
  output_s->data.dptr = GetOptionalDataPtr(softmax_aux);

  // execute the kernel
  nvte_fused_attn_fwd(te_Q.data(), te_K.data(), te_V.data(), te_Bias.data(), te_S.data(),
                      te_O.data(), &nvte_aux_tensor_pack, te_cu_seqlens_q.data(),
                      te_cu_seqlens_kv.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
                      dummy_seq_offsets.data(), dummy_seq_offsets.data(), te_rng_state.data(),
                      max_seqlen_q, max_seqlen_kv, is_training, attn_scale, p_dropout,
                      qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(),
                      Q.stream());

  // destroy tensor wrappers, but not allocated memory
  nvte_tensor_pack_destroy(&nvte_aux_tensor_pack);
}

void te_fused_attn_bwd(const paddle::Tensor &Q, const paddle::Tensor &K, const paddle::Tensor &V,
                       const paddle::Tensor &cu_seqlens_q, const paddle::Tensor &cu_seqlens_kv,
                       const paddle::Tensor &O, const paddle::Tensor &dO,
                       const paddle::Tensor &softmax_aux,
                       paddle::Tensor &dQ,                       // NOLINT
                       paddle::Tensor &dK,                       // NOLINT
                       paddle::Tensor &dV,                       // NOLINT
                       paddle::optional<paddle::Tensor> &dBias,  // NOLINT
                       paddle::Tensor &rng_state,                // NOLINT
                       int64_t b, int64_t h, int64_t d, int64_t max_seqlen_q, int64_t max_seqlen_kv,
                       float attn_scale, float p_dropout, const std::string &qkv_layout,
                       const std::string &bias_type, const std::string &attn_mask_type,
                       int64_t qkv_type) {
  TensorWrapper te_dBias;
  if (bias_type != "no_bias" && dBias) {
    auto bias_shape = dBias->shape();
    std::vector<size_t> shape{bias_shape.begin(), bias_shape.end()};
    te_dBias = MakeNvteTensor(GetOptionalDataPtr(dBias), shape, DType::kFloat32);
  }

  auto qkv_dtype = Int2NvteDType(qkv_type);
  // construct NVTE tensors
  TensorWrapper te_Q, te_K, te_V, te_O, te_dO, te_S, te_dP, te_dQ, te_dK, te_dV;
  if (qkv_dtype == DType::kBFloat16 || qkv_dtype == DType::kFloat16) {
    // BF16 or FP16
    te_Q = MakeNvteTensor(Q);
    te_K = MakeNvteTensor(K);
    te_V = MakeNvteTensor(V);
    te_O = MakeNvteTensor(O);
    te_dO = MakeNvteTensor(dO);
    te_S = MakeNvteTensor(nullptr, std::vector<size_t>(0), DType::kFloat32);
    te_dP = MakeNvteTensor(nullptr, std::vector<size_t>(0), DType::kFloat32);
    te_dQ = MakeNvteTensor(dQ);
    te_dK = MakeNvteTensor(dK);
    te_dV = MakeNvteTensor(dV);
  } else {
    NVTE_ERROR("Fused attention only supports BF16/FP16 data types. \n");
  }

  // convert strings to enums
  NVTE_QKV_Layout qkv_layout_enum = get_nvte_qkv_layout(qkv_layout);
  NVTE_Bias_Type bias_type_enum = get_nvte_bias_type(bias_type);
  NVTE_Mask_Type attn_mask_type_enum = get_nvte_mask_type(attn_mask_type);

  // convert auxiliary tensors from forward into NVTETensors
  NVTETensorPack nvte_aux_tensor_pack;
  nvte_tensor_pack_create(&nvte_aux_tensor_pack);

  nvte_aux_tensor_pack.size = 2;
  auto *output_s = reinterpret_cast<Tensor *>(nvte_aux_tensor_pack.tensors[0]);
  auto *fwd_rng_state = reinterpret_cast<Tensor *>(nvte_aux_tensor_pack.tensors[1]);
  output_s->data.shape =
      std::vector<size_t>({static_cast<size_t>(b), static_cast<size_t>(h),
                           static_cast<size_t>(max_seqlen_q), static_cast<size_t>(max_seqlen_kv)});
  output_s->data.dptr = const_cast<void *>(softmax_aux.data());
  fwd_rng_state->data.shape = std::vector<size_t>({2});
  fwd_rng_state->data.dptr = const_cast<void *>(rng_state.data());

  // create cu_seqlens tensorwrappers
  TensorWrapper te_cu_seqlens_q, te_cu_seqlens_kv;
  te_cu_seqlens_q =
      MakeNvteTensor(cu_seqlens_q.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);
  te_cu_seqlens_kv =
      MakeNvteTensor(cu_seqlens_kv.data(), {static_cast<size_t>(b + 1)}, DType::kInt32);

  // create workspace
  TensorWrapper workspace;

  auto dummy_seq_offsets = TensorWrapper(nullptr, {static_cast<size_t>(b + 1)}, DType::kInt32);
  // populate tensors with appropriate shapes and dtypes
  nvte_fused_attn_bwd(te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(),
                      te_dP.data(), &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(),
                      te_dBias.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(),
                      dummy_seq_offsets.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
                      dummy_seq_offsets.data(), max_seqlen_q, max_seqlen_kv, attn_scale, p_dropout,
                      qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(),
                      Q.stream());

  // allocate memory for workspace
  auto workspace_data = AllocateSpace(workspace.shape(), workspace.dtype(), Q.place());
  workspace = MakeNvteTensor(workspace_data.data(), workspace.shape(), workspace.dtype());

  // execute kernel
  nvte_fused_attn_bwd(te_Q.data(), te_K.data(), te_V.data(), te_O.data(), te_dO.data(), te_S.data(),
                      te_dP.data(), &nvte_aux_tensor_pack, te_dQ.data(), te_dK.data(), te_dV.data(),
                      te_dBias.data(), te_cu_seqlens_q.data(), te_cu_seqlens_kv.data(),
                      dummy_seq_offsets.data(), dummy_seq_offsets.data(), dummy_seq_offsets.data(),
                      dummy_seq_offsets.data(), max_seqlen_q, max_seqlen_kv, attn_scale, p_dropout,
                      qkv_layout_enum, bias_type_enum, attn_mask_type_enum, workspace.data(),
                      Q.stream());

  // destroy tensor wrappers
  nvte_tensor_pack_destroy(&nvte_aux_tensor_pack);
}

std::vector<paddle::Tensor> te_scaled_softmax_forward(const paddle::Tensor &input,
                                                      float scale_factor) {
  NVTE_CHECK(input.shape().size() == 4, "expected 4D tensor");
  NVTE_CHECK(
      (input.dtype() == paddle::DataType::FLOAT16) || (input.dtype() == paddle::DataType::BFLOAT16),
      "Only fp16 and bf16 are supported");

  const int batches = input.shape()[0];
  const int attn_heads = input.shape()[1];
  const int query_seq_len = input.shape()[2];
  const int key_seq_len = input.shape()[3];

  NVTE_CHECK(key_seq_len <= 4096);
  NVTE_CHECK(query_seq_len > 1);

  // Output
  auto softmax_results = paddle::empty_like(input, input.dtype(), input.place());

  auto input_cu = MakeNvteTensor(input);
  auto softmax_results_cu = MakeNvteTensor(softmax_results);

  nvte_scaled_softmax_forward(input_cu.data(), softmax_results_cu.data(), scale_factor,
                              input.stream());

  return {softmax_results};
}

void te_scaled_softmax_backward(paddle::Tensor &output_grads,  // NOLINT
                                const paddle::Tensor &softmax_results, float scale_factor) {
  NVTE_CHECK(output_grads.shape().size() == 4, "expected 4D tensor");
  NVTE_CHECK(softmax_results.shape().size() == 4, "expected 4D tensor");

  NVTE_CHECK((output_grads.dtype() == paddle::DataType::FLOAT16) ||
                 (output_grads.dtype() == paddle::DataType::BFLOAT16),
             "Only fp16 and bf16 are supported");
  NVTE_CHECK((softmax_results.dtype() == paddle::DataType::FLOAT16) ||
                 (softmax_results.dtype() == paddle::DataType::BFLOAT16),
             "Only fp16 and bf16 are supported");

  auto output_grads_cu = MakeNvteTensor(output_grads);
  auto softmax_results_cu = MakeNvteTensor(softmax_results);

  // Produce gradients in place.
  nvte_scaled_softmax_backward(output_grads_cu.data(), softmax_results_cu.data(),
                               output_grads_cu.data(), scale_factor, softmax_results.stream());
}

std::vector<paddle::Tensor> te_scaled_masked_softmax_forward(const paddle::Tensor &input,
                                                             const paddle::Tensor &mask,
                                                             float scale_factor) {
  NVTE_CHECK(input.shape().size() == 4, "expected 4D tensor");
  NVTE_CHECK(mask.shape().size() == 4, "expected 4D tensor");
  NVTE_CHECK(
      (input.dtype() == paddle::DataType::FLOAT16) || (input.dtype() == paddle::DataType::BFLOAT16),
      "Only fp16 and bf16 are supported");

  const int batches = input.shape()[0];
  const int pad_batches = mask.shape()[0];
  const int attn_heads = input.shape()[1];
  const int query_seq_len = input.shape()[2];
  const int key_seq_len = input.shape()[3];

  NVTE_CHECK(key_seq_len <= 4096);
  NVTE_CHECK(query_seq_len > 1);
  NVTE_CHECK(pad_batches == 1 || pad_batches == batches);
  NVTE_CHECK(mask.shape()[1] == 1);
  NVTE_CHECK(mask.shape()[2] == query_seq_len);
  NVTE_CHECK(mask.shape()[3] == key_seq_len);

  // Output
  auto softmax_results = paddle::empty_like(input, input.dtype(), input.place());

  auto input_cu = MakeNvteTensor(input);
  auto mask_cu = MakeNvteTensor(mask);
  auto softmax_results_cu = MakeNvteTensor(softmax_results);

  nvte_scaled_masked_softmax_forward(input_cu.data(), mask_cu.data(), softmax_results_cu.data(),
                                     scale_factor, input.stream());

  return {softmax_results};
}

void te_scaled_masked_softmax_backward(paddle::Tensor &output_grads,  // NOLINT
                                       const paddle::Tensor &softmax_results, float scale_factor) {
  NVTE_CHECK(output_grads.shape().size() == 4, "expected 4D tensor");
  NVTE_CHECK(softmax_results.shape().size() == 4, "expected 4D tensor");

  NVTE_CHECK((output_grads.dtype() == paddle::DataType::FLOAT16) ||
                 (output_grads.dtype() == paddle::DataType::BFLOAT16),
             "Only fp16 and bf16 are supported");
  NVTE_CHECK((softmax_results.dtype() == paddle::DataType::FLOAT16) ||
                 (softmax_results.dtype() == paddle::DataType::BFLOAT16),
             "Only fp16 and bf16 are supported");

  auto output_grads_cu = MakeNvteTensor(output_grads);
  auto softmax_results_cu = MakeNvteTensor(softmax_results);

  // Produce gradients in place.
  nvte_scaled_softmax_backward(output_grads_cu.data(), softmax_results_cu.data(),
                               output_grads_cu.data(), scale_factor, softmax_results.stream());
}

std::vector<paddle::Tensor> te_scaled_upper_triang_masked_softmax_forward(
    const paddle::Tensor &input, float scale_factor) {
  NVTE_CHECK(input.shape().size() == 3, "expected 3D tensor");
  NVTE_CHECK(
      (input.dtype() == paddle::DataType::FLOAT16) || (input.dtype() == paddle::DataType::BFLOAT16),
      "Only fp16 and bf16 are supported");

  const int attn_batches = input.shape()[0];
  const int seq_len = input.shape()[1];
  NVTE_CHECK(seq_len <= 2048);

  // Output
  auto softmax_results = paddle::empty_like(input, input.dtype(), input.place());

  auto input_cu = MakeNvteTensor(input);
  auto softmax_results_cu = MakeNvteTensor(softmax_results);

  nvte_scaled_upper_triang_masked_softmax_forward(input_cu.data(), softmax_results_cu.data(),
                                                  scale_factor, input.stream());

  return {softmax_results};
}

void te_scaled_upper_triang_masked_softmax_backward(paddle::Tensor &output_grads,  // NOLINT
                                                    const paddle::Tensor &softmax_results,
                                                    float scale_factor) {
  NVTE_CHECK(output_grads.shape().size() == 3, "expected 3D tensor");
  NVTE_CHECK(softmax_results.shape().size() == 3, "expected 3D tensor");

  NVTE_CHECK((output_grads.dtype() == paddle::DataType::FLOAT16) ||
                 (output_grads.dtype() == paddle::DataType::BFLOAT16),
             "Only fp16 and bf16 are supported");
  NVTE_CHECK((softmax_results.dtype() == paddle::DataType::FLOAT16) ||
                 (softmax_results.dtype() == paddle::DataType::BFLOAT16),
             "Only fp16 and bf16 are supported");
  NVTE_CHECK(output_grads.shape()[1] == output_grads.shape()[2]);

  auto output_grads_cu = MakeNvteTensor(output_grads);
  auto softmax_results_cu = MakeNvteTensor(softmax_results);

  // Produce gradients in place.
  nvte_scaled_upper_triang_masked_softmax_backward(
      output_grads_cu.data(), softmax_results_cu.data(), output_grads_cu.data(), scale_factor,
      softmax_results.stream());
}

constexpr int BLOCK_SIZE = 512;

void amax_and_scale_update_inplace(paddle::Tensor &amax_history,  // NOLINT
                                   paddle::Tensor &scale,         // NOLINT
                                   paddle::Tensor &scale_inv,     // NOLINT
                                   const paddle::Tensor &non_weight_mask, int64_t fp8_dtype,
                                   float margin, const std::string &amax_compute) {
  auto amax_history_ = MakeNvteTensor(amax_history);
  auto scale_ = MakeNvteTensor(scale);
  auto scale_inv_ = MakeNvteTensor(scale_inv);
  const auto non_weight_mask_ = MakeNvteTensor(non_weight_mask);
  nvte_delayed_scaling_recipe_amax_and_scale_update(
      amax_history_.data(), scale_.data(), scale_inv_.data(), non_weight_mask_.data(),
      amax_history_.data(), scale_.data(), scale_inv_.data(), amax_compute.c_str(),
      static_cast<NVTEDType>(fp8_dtype), margin, amax_history.stream());
}

void update_latest_amax_history_inplace(paddle::Tensor &history,  // NOLINT
                                        const paddle::Tensor &amax) {
  // Copy amax to history[0]
  NVTE_CHECK_CUDA(cudaMemcpyAsync(history.data(), amax.data(), amax.numel() * SizeOf(amax.dtype()),
                                  cudaMemcpyDeviceToDevice, amax.stream()));
}

__global__ __launch_bounds__(BLOCK_SIZE) void mask_to_actual_seqlens_kernel(
    const bool *mask, int32_t *q_actual_seqlen, int32_t *kv_actual_seqlen, int q_seqlen,
    int kv_seqlen, bool need_kv) {
  typedef cub::BlockReduce<int, BLOCK_SIZE> BlockReduce;
  __shared__ typename BlockReduce::TempStorage q_smem;
  __shared__ typename BlockReduce::TempStorage kv_smem;
  unsigned int tid = threadIdx.x;
  unsigned int batch_offset = blockIdx.x * q_seqlen * kv_seqlen;

  // load mask, convert to 1/0, do accumulation
  int q = 0, kv = 0;
  for (unsigned int q_idx = tid * kv_seqlen; q_idx < q_seqlen * kv_seqlen;
       q_idx += BLOCK_SIZE * kv_seqlen) {
    q += (mask[q_idx + batch_offset] ? 0 : 1);
  }

  if (need_kv) {
    for (unsigned int kv_idx = tid; kv_idx < kv_seqlen; kv_idx += BLOCK_SIZE) {
      kv += (mask[kv_idx + batch_offset] ? 0 : 1);
    }
  }
  __syncthreads();

  // compute cub::BlockReduce
  int q_sum, kv_sum;
  q_sum = BlockReduce(q_smem).Sum(q);
  if (need_kv) kv_sum = BlockReduce(kv_smem).Sum(kv);

  // write result for this block to global mem
  if (tid == 0) {
    q_actual_seqlen[blockIdx.x + 1] = q_sum;
    if (need_kv) {
      kv_actual_seqlen[blockIdx.x + 1] = kv_sum;
    }
  }
}

__global__ __launch_bounds__(BLOCK_SIZE) void block_prefix_sum_inplace(int32_t *x, int n) {
  typedef cub::BlockScan<int32_t, BLOCK_SIZE> BlockScan;
  __shared__ typename BlockScan::TempStorage smem;
  // +1 to ignore the first element
  int i = blockIdx.x * blockDim.x + threadIdx.x + 1;

  // load data
  int32_t thread_data[1];
  thread_data[0] = i < n ? x[i] : 0;
  __syncthreads();

  // CUB block prefix sum
  BlockScan(smem).InclusiveSum(thread_data, thread_data);
  __syncthreads();

  // write result
  if (i < n) {
    x[i] = thread_data[0];
  }
}

void mask_to_cu_seqlens(const paddle::Tensor &mask,
                        paddle::Tensor &q_cu_seqlen,                     // NOLINT
                        paddle::optional<paddle::Tensor> &kv_cu_seqlen,  // NOLINT
                        int q_seqlen, int kv_seqlen, bool need_kv) {
  if (need_kv) {
    NVTE_CHECK(GetOptionalDataPtr(kv_cu_seqlen) != nullptr,
               "kv_cu_seqlen must be provided when need_kv is true");
  }
  mask_to_actual_seqlens_kernel<<<mask.shape()[0], BLOCK_SIZE, 0, mask.stream()>>>(
      mask.data<bool>(), q_cu_seqlen.data<int32_t>(),
      reinterpret_cast<int32_t *>(GetOptionalDataPtr(kv_cu_seqlen)), q_seqlen, kv_seqlen, need_kv);
  // q_cu_seqlen shape: [bs+1], assume bs is not too large (<=512), so we can use a single block
  // to do prefix sum
  NVTE_CHECK(q_cu_seqlen.numel() - 1 <= BLOCK_SIZE, "batch size too large, kernel may fail");
  block_prefix_sum_inplace<<<1, BLOCK_SIZE, 0, mask.stream()>>>(q_cu_seqlen.data<int32_t>(),
                                                                q_cu_seqlen.numel());
  if (need_kv) {
    block_prefix_sum_inplace<<<1, BLOCK_SIZE, 0, mask.stream()>>>(
        reinterpret_cast<int32_t *>(GetOptionalDataPtr(kv_cu_seqlen)), kv_cu_seqlen->numel());
  }
}

}  // namespace paddle_ext
}  // namespace transformer_engine

PD_BUILD_OP(te_gemm)
    .Inputs({"A", paddle::Optional("A_scale_inverse"), "B", paddle::Optional("B_scale_inverse"),
             paddle::Optional("bias"), "_D", paddle::Optional("_D_scale"),
             paddle::Optional("_D_amax"), paddle::Optional("_pre_gelu_out"), "_workspace"})
    .Outputs({"D", paddle::Optional("D_scale"), paddle::Optional("D_amax"),
              paddle::Optional("pre_gelu_out"), "workspace"})
    .Attrs({"A_index: int64_t", "B_index: int64_t", "D_index: int64_t", "A_type: int64_t",
            "B_type: int64_t", "D_type: int64_t", "bias_type: int64_t", "transa: bool",
            "transb: bool", "grad: bool", "workspace_size: int64_t", "accumulate: bool",
            "use_split_accumulator: bool", "math_sm_count: int64_t"})
    .SetInplaceMap({{"_D", "D"},
                    {paddle::Optional("_D_scale"), paddle::Optional("D_scale")},
                    {paddle::Optional("_D_amax"), paddle::Optional("D_amax")},
                    {paddle::Optional("_pre_gelu_out"), paddle::Optional("pre_gelu_out")},
                    {"_workspace", "workspace"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_gemm));

PD_BUILD_OP(cast_to_fp8)
    .Inputs({"Input", "Scale", "_Output", "_Amax", "_ScaleInv"})
    .Outputs({"Output", "Amax", "ScaleInv"})
    .Attrs({"index: int64_t", "otype: int64_t"})
    .SetInplaceMap({{"_Output", "Output"}, {"_Amax", "Amax"}, {"_ScaleInv", "ScaleInv"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::cast_to_fp8));

PD_BUILD_OP(cast_from_fp8)
    .Inputs({"Input", "ScaleInv"})
    .Outputs({"Output"})
    .Attrs({"index: int64_t", "itype: int64_t", "otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::cast_from_fp8));

PD_BUILD_OP(te_transpose)
    .Inputs({"Input"})
    .Outputs({"Output"})
    .Attrs({"otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_transpose));

PD_BUILD_OP(te_cast_transpose)
    .Inputs({"Input", "Scale", "_CastedOutput", "_TransposedOutput", "_Amax", "_ScaleInv"})
    .Outputs({"CastedOutput", "TransposedOutput", "Amax", "ScaleInv"})
    .SetInplaceMap({{"_CastedOutput", "CastedOutput"},
                    {"_TransposedOutput", "TransposedOutput"},
                    {"_Amax", "Amax"},
                    {"_ScaleInv", "ScaleInv"}})
    .Attrs({"index: int64_t", "otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_cast_transpose));

PD_BUILD_OP(te_cast_transpose_bgrad)
    .Inputs({"GradOutput", "Scale", "_Amax", "_ScaleInv"})
    .Outputs({"dBias", "CastedOutput", "TransposedOutput", "Amax", "ScaleInv"})
    .SetInplaceMap({{"_Amax", "Amax"}, {"_ScaleInv", "ScaleInv"}})
    .Attrs({"index: int64_t", "otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_cast_transpose_bgrad));

PD_BUILD_OP(te_gelu_fp8)
    .Inputs({"Input", "Scale", "_Amax", "_ScaleInv"})
    .Outputs({"Output", "Amax", "ScaleInv"})
    .SetInplaceMap({{"_Amax", "Amax"}, {"_ScaleInv", "ScaleInv"}})
    .Attrs({"index: int64_t", "otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_gelu_fp8));

PD_BUILD_OP(te_gelu)
    .Inputs({"Input"})
    .Outputs({"Output"})
    .Attrs({"otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_gelu));

PD_BUILD_OP(te_swiglu)
    .Inputs({"Input"})
    .Outputs({"Output"})
    .Attrs({"otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_swiglu));

PD_BUILD_OP(te_swiglu_fp8)
    .Inputs({"Input", "Scale", "_Amax", "_ScaleInv"})
    .Outputs({"Output", "Amax", "ScaleInv"})
    .SetInplaceMap({{"_Amax", "Amax"}, {"_ScaleInv", "ScaleInv"}})
    .Attrs({"index: int64_t", "otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_swiglu_fp8));

PD_BUILD_OP(te_dswiglu)
    .Inputs({"Grad", "Input"})
    .Outputs({"Output"})
    .Attrs({"otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_dswiglu));

PD_BUILD_OP(te_cast_transpose_bgrad_dgelu)
    .Inputs({"GradOutput", "GeluInput", "Scale", "_Amax", "_ScaleInv"})
    .Outputs({"CastedDgelu", "TransposedDgelu", "Dbias", "Amax", "ScaleInv"})
    .SetInplaceMap({{"_Amax", "Amax"}, {"_ScaleInv", "ScaleInv"}})
    .Attrs({"index: int64_t", "otype: int64_t"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_cast_transpose_bgrad_dgelu));

PD_BUILD_OP(te_layernorm_fwd_fp8)
    .Inputs({"Input", "Weight", "Bias", "Scale", "_Amax", "_ScaleInv"})
    .Outputs({"Output", "Mu", "Rsigma", "Amax", "ScaleInv"})
    .SetInplaceMap({{"_Amax", "Amax"}, {"_ScaleInv", "ScaleInv"}})
    .Attrs({"eps: float", "index: int64_t", "otype: int64_t", "sm_margin: int64_t",
            "zero_centered_gamma: bool"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_layernorm_fwd_fp8));

PD_BUILD_OP(te_layernorm_fwd)
    .Inputs({"Input", "Weight", "Bias"})
    .Outputs({"Output", "Mu", "Rsigma"})
    .Attrs({"eps: float", "otype: int64_t", "sm_margin: int64_t", "zero_centered_gamma: bool"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_layernorm_fwd));

PD_BUILD_OP(te_layernorm_bwd)
    .Inputs({"Dz", "X", "Mu", "Rsigma", "Gamma"})
    .Outputs({"Dx", "Dgamma", "Dbeta"})
    .Attrs({"sm_margin: int64_t", "zero_centered_gamma: bool"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_layernorm_bwd));

PD_BUILD_OP(te_rmsnorm_fwd)
    .Inputs({"Input", "Weight"})
    .Outputs({"Output", "InvVariance"})
    .Attrs({"eps: float", "otype: int64_t", "sm_margin: int64_t", "zero_centered_gamma: bool"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_rmsnorm_fwd));

PD_BUILD_OP(te_rmsnorm_fwd_fp8)
    .Inputs({"Input", "Weight", "Scale", "_Amax", "_ScaleInv"})
    .Outputs({"Output", "InvVariance", "Amax", "ScaleInv"})
    .SetInplaceMap({{"_Amax", "Amax"}, {"_ScaleInv", "ScaleInv"}})
    .Attrs({"eps: float", "index: int64_t", "otype: int64_t", "sm_margin: int64_t",
            "zero_centered_gamma: bool"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_rmsnorm_fwd_fp8));

PD_BUILD_OP(te_rmsnorm_bwd)
    .Inputs({"Dz", "X", "Rsigma", "Gamma"})
    .Outputs({"Dx", "Dgamma"})
    .Attrs({"sm_margin: int64_t", "zero_centered_gamma: bool"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_rmsnorm_bwd));

PD_BUILD_OP(te_fused_attn_fwd_qkvpacked)
    .Inputs({"QKV", "cu_seqlens", paddle::Optional("Bias"), "_O", paddle::Optional("_softmax_aux"),
             "_rng_state"})
    .Outputs({"O", paddle::Optional("softmax_aux"), "rng_state"})
    .Attrs({"b: int64_t", "h: int64_t", "d: int64_t", "total_seqs: int64_t", "max_seqlen: int64_t",
            "is_training: bool", "attn_scale: float", "p_dropout: float", "qkv_layout: std::string",
            "bias_type: std::string", "attn_mask_type: std::string", "qkv_type: int64_t",
            "rng_elts_per_thread: int64_t"})
    .SetInplaceMap({{"_O", "O"},
                    {paddle::Optional("_softmax_aux"), paddle::Optional("softmax_aux")},
                    {"_rng_state", "rng_state"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_fused_attn_fwd_qkvpacked));

PD_BUILD_OP(te_fused_attn_bwd_qkvpacked)
    .Inputs({"QKV", "cu_seqlens", "O", "dO", "softmax_aux", "_dQKV", paddle::Optional("_dBias"),
             "rng_state"})
    .Outputs({"dQKV", paddle::Optional("dBias")})
    .Attrs({"b: int64_t", "h: int64_t", "d: int64_t", "total_seqs: int64_t", "max_seqlen: int64_t",
            "attn_scale: float", "p_dropout: float", "qkv_layout: std::string",
            "bias_type: std::string", "attn_mask_type: std::string", "qkv_type: int64_t"})
    .SetInplaceMap({{"_dQKV", "dQKV"}, {paddle::Optional("_dBias"), paddle::Optional("dBias")}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_fused_attn_bwd_qkvpacked));

PD_BUILD_OP(te_fused_attn_fwd_kvpacked)
    .Inputs({"Q", "KV", "cu_seqlens_q", "cu_seqlens_kv", paddle::Optional("Bias"), "_O",
             paddle::Optional("_softmax_aux"), "_rng_state"})
    .Outputs({"O", paddle::Optional("softmax_aux"), "rng_state"})
    .Attrs({"b: int64_t", "h: int64_t", "d: int64_t", "total_seqs_q: int64_t",
            "total_seqs_kv: int64_t", "max_seqlen_q: int64_t", "max_seqlen_kv: int64_t",
            "is_training: bool", "attn_scale: float", "p_dropout: float", "qkv_layout: std::string",
            "bias_type: std::string", "attn_mask_type: std::string", "qkv_type: int64_t",
            "rng_elts_per_thread: int64_t"})
    .SetInplaceMap({{"_O", "O"},
                    {paddle::Optional("_softmax_aux"), paddle::Optional("softmax_aux")},
                    {"_rng_state", "rng_state"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_fused_attn_fwd_kvpacked));

PD_BUILD_OP(te_fused_attn_bwd_kvpacked)
    .Inputs({"Q", "KV", "cu_seqlens_q", "cu_seqlens_kv", "O", "dO", "softmax_aux", "_dQ", "_dKV",
             paddle::Optional("_dBias"), "rng_state"})
    .Outputs({"dQ", "dKV", paddle::Optional("dBias")})
    .Attrs({"b: int64_t", "h: int64_t", "d: int64_t", "total_seqs_q: int64_t",
            "total_seqs_kv: int64_t", "max_seqlen_q: int64_t", "max_seqlen_kv: int64_t",
            "attn_scale: float", "p_dropout: float", "qkv_layout: std::string",
            "bias_type: std::string", "attn_mask_type: std::string", "qkv_type: int64_t"})
    .SetInplaceMap({{"_dQ", "dQ"},
                    {"_dKV", "dKV"},
                    {paddle::Optional("_dBias"), paddle::Optional("dBias")}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_fused_attn_bwd_kvpacked));

PD_BUILD_OP(te_fused_attn_fwd)
    .Inputs({"Q", "K", "V", "cu_seqlens_q", "cu_seqlens_kv", paddle::Optional("Bias"), "_O",
             paddle::Optional("_softmax_aux"), "_rng_state"})
    .Outputs({"O", paddle::Optional("softmax_aux"), "rng_state"})
    .Attrs({"b: int64_t", "h: int64_t", "d: int64_t", "max_seqlen_q: int64_t",
            "max_seqlen_kv: int64_t", "is_training: bool", "attn_scale: float", "p_dropout: float",
            "qkv_layout: std::string", "bias_type: std::string", "attn_mask_type: std::string",
            "qkv_type: int64_t", "rng_elts_per_thread: int64_t"})
    .SetInplaceMap({{"_O", "O"},
                    {paddle::Optional("_softmax_aux"), paddle::Optional("softmax_aux")},
                    {"_rng_state", "rng_state"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_fused_attn_fwd));

PD_BUILD_OP(te_fused_attn_bwd)
    .Inputs({"Q", "K", "V", "cu_seqlens_q", "cu_seqlens_kv", "O", "dO", "softmax_aux", "_dQ", "_dK",
             "_dV", paddle::Optional("_dBias"), "rng_state"})
    .Outputs({"dQ", "dK", "dV", paddle::Optional("dBias")})
    .Attrs({"b: int64_t", "h: int64_t", "d: int64_t", "max_seqlen_q: int64_t",
            "max_seqlen_kv: int64_t", "attn_scale: float", "p_dropout: float",
            "qkv_layout: std::string", "bias_type: std::string", "attn_mask_type: std::string",
            "qkv_type: int64_t"})
    .SetInplaceMap({{"_dQ", "dQ"},
                    {"_dK", "dK"},
                    {"_dV", "dV"},
                    {paddle::Optional("_dBias"), paddle::Optional("dBias")}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_fused_attn_bwd));

PD_BUILD_OP(te_scaled_softmax_forward)
    .Inputs({"input"})
    .Outputs({"softmax_results"})
    .Attrs({"scale_factor: float"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_scaled_softmax_forward));

PD_BUILD_OP(te_scaled_softmax_backward)
    .Inputs({"out_grad_", "softmax_results"})
    .Outputs({"out_grad"})
    .Attrs({"scale_factor: float"})
    .SetInplaceMap({{"out_grad_", "out_grad"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_scaled_softmax_backward));

PD_BUILD_OP(te_scaled_masked_softmax_forward)
    .Inputs({"input", "mask"})
    .Outputs({"softmax_results"})
    .Attrs({"scale_factor: float"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_scaled_masked_softmax_forward));

PD_BUILD_OP(te_scaled_masked_softmax_backward)
    .Inputs({"out_grad_", "softmax_results"})
    .Outputs({"out_grad"})
    .Attrs({"scale_factor: float"})
    .SetInplaceMap({{"out_grad_", "out_grad"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::te_scaled_masked_softmax_backward));

PD_BUILD_OP(te_scaled_upper_triang_masked_softmax_forward)
    .Inputs({"input"})
    .Outputs({"softmax_results"})
    .Attrs({"scale_factor: float"})
    .SetKernelFn(
        PD_KERNEL(transformer_engine::paddle_ext::te_scaled_upper_triang_masked_softmax_forward));

PD_BUILD_OP(te_scaled_upper_triang_masked_softmax_backward)
    .Inputs({"out_grad_", "softmax_results"})
    .Outputs({"out_grad"})
    .Attrs({"scale_factor: float"})
    .SetInplaceMap({{"out_grad_", "out_grad"}})
    .SetKernelFn(
        PD_KERNEL(transformer_engine::paddle_ext::te_scaled_upper_triang_masked_softmax_backward));

PD_BUILD_OP(amax_and_scale_update_inplace)
    .Inputs({"_amax_history", "_scale", "_scale_inv", "non_weight_mask"})
    .Outputs({"amax_history", "scale", "scale_inv"})
    .SetInplaceMap({{"_amax_history", "amax_history"},
                    {"_scale", "scale"},
                    {"_scale_inv", "scale_inv"}})
    .Attrs({"fp8_dtype: int64_t", "margin: float", "amax_compute: std::string"})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::amax_and_scale_update_inplace));

PD_BUILD_OP(update_latest_amax_history_inplace)
    .Inputs({"_history", "amax"})
    .Outputs({"history"})
    .SetInplaceMap({{"_history", "history"}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::update_latest_amax_history_inplace));

PD_BUILD_OP(mask_to_cu_seqlens)
    .Inputs({"mask", "_q_cu_seqlen", paddle::Optional("_kv_cu_seqlen")})
    .Outputs({"q_cu_seqlen", paddle::Optional("kv_cu_seqlen")})
    .Attrs({"q_seqlen: int", "kv_seqlen: int", "need_kv: bool"})
    .SetInplaceMap({{"_q_cu_seqlen", "q_cu_seqlen"},
                    {paddle::Optional("_kv_cu_seqlen"), paddle::Optional("kv_cu_seqlen")}})
    .SetKernelFn(PD_KERNEL(transformer_engine::paddle_ext::mask_to_cu_seqlens));