/*
 * Copyright (c) 2020-2023, NVIDIA CORPORATION.  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 "src/fastertransformer/kernels/xlnet_attention_kernels.h"
#include "src/fastertransformer/utils/cuda_type_utils.cuh"
namespace fastertransformer {
template<typename T>
int numPerThread()
{
    return sizeof(float) / sizeof(T);
}

int next_pow2(int a)
{
    int rval = 32;
    if (a > 32) {
        while (rval < a) {
            rval <<= 1;
        }
    }
    return rval;
}

__device__ half2 cH2(const half* ptr, int offset)
{
    return __ldg((half2*)(ptr + offset));
}

#ifdef ENABLE_BF16
__device__ __nv_bfloat162 cH2(const __nv_bfloat16* ptr, int offset)
{
    return ldg((__nv_bfloat162*)(ptr + offset));
}
#endif

__forceinline__ __device__ unsigned lane_id()
{
    // unsigned ret;
    // asm volatile("mov.u32 %0, %laneid;" : "=r"(ret));
    unsigned ret = __lane_id();
    return ret;
}

__forceinline__ __device__ unsigned warp_id()
{
    // this is not equal to threadIdx.x / 32
    // unsigned ret;
    // asm volatile("mov.u32 %0, %warpid;" : "=r"(ret));
    int32_t threadid = blockIdx.x * blockDim.x + threadIdx.x;
    int32_t warp_id = threadid / warpSize;
    // asm volatile ("v_mov_b32 %0, %1": "=v"(ret) : "v"(warp_id));
    return warp_id;
}

template<typename T>
__inline__ __device__ T warpReduceSum(T val)
{
    for (int mask = 16; mask > 0; mask >>= 1) {
        val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
    }
    return val;
}

template<typename T>
__inline__ __device__ T warpReduceMax(T val)
{
    for (int mask = 16; mask > 0; mask >>= 1) {
        val = max(val, __shfl_xor_sync(FINAL_MASK, val, mask, 32));
    }
    return val;
}

template<typename T>
__inline__ __device__ T blockReduceSum(T val)
{
    static __shared__ T shared[32];
    int                 lane = threadIdx.x & 0x1f;
    int                 wid  = threadIdx.x >> 5;

    val = warpReduceSum<T>(val);

    if (lane == 0) {
        shared[wid] = val;
    }

    __syncthreads();

    val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : (T)(0.0f);
    val = warpReduceSum<T>(val);

    return val;
}

template<typename T>
__inline__ __device__ T blockReduceMax(T val)
{
    static __shared__ T shared[32];
    int                 lane = threadIdx.x & 0x1f;  // in-warp idx
    int                 wid  = threadIdx.x >> 5;    // warp idx

    val = warpReduceMax(val);  // get maxx in each warp

    if (lane == 0) {  // record in-warp maxx by warp Idx
        shared[wid] = val;
    }

    __syncthreads();

    val = (threadIdx.x < (blockDim.x >> 5)) ? shared[lane] : -1e30f;
    val = warpReduceMax(val);

    return val;
}

/********************** Kernels ************************/

// Applied to half and bfloat16
template<typename T>
void __global__ prepareMatrixes(T*        q_buf,
                                T*        q_buf_bd,
                                T*        q_buf_ef,
                                T*        k_buf,
                                T*        k_buf_bd,
                                T*        k_buf_ef,
                                const T*  query_buf,
                                const T*  key_buf,
                                const T*  k_head_r,
                                const T*  attr_seg_embed,
                                const T*  attr_bias_Q_w,
                                const T*  attr_bias_Q_r,
                                const T*  attr_bias_Q_s,
                                const int off0,
                                const int i_off1,
                                const int o_off1,
                                int       off2)
{
    using T2 = typename TypeConverter<T>::Type;  // half2 or bfloat162

    int batch    = blockIdx.y;
    int seq      = blockIdx.x;
    int head_loc = threadIdx.x * 2;

    T2 tmp;
    if (head_loc < i_off1) {
        int head     = head_loc / off2;
        int loc      = head_loc % off2;
        int h2_index = (batch * off0 + seq * i_off1 + head_loc) >> 1;

        tmp              = ((T2*)query_buf)[h2_index];
        int h2_index_out = (batch * off0 + head * o_off1 + seq * off2 + loc) >> 1;
        // left matrix
        ((T2*)q_buf)[h2_index_out]    = hadd2(tmp, cH2(attr_bias_Q_w, head_loc));
        ((T2*)q_buf_bd)[h2_index_out] = hadd2(tmp, cH2(attr_bias_Q_r, head_loc));
        ((T2*)q_buf_ef)[h2_index_out] = hadd2(tmp, cH2(attr_bias_Q_s, head_loc));

        // right matrix
        ((T2*)k_buf)[h2_index_out] = ((T2*)key_buf)[h2_index];  // ac

        // bd
        h2_index     = (seq * i_off1 + head_loc) >> 1;  //(seq, head_loc)
        tmp          = ((T2*)k_head_r)[h2_index];
        h2_index_out = (batch * off0 * 2 + head * o_off1 * 2 + seq * off2 + loc) >> 1;  //(batch, head,seq,loc)
        ((T2*)k_buf_bd)[h2_index_out] = tmp;

        h2_index = (seq * i_off1 + head_loc + off0) >> 1;  //(seq+seq_len, head_loc)
        tmp      = ((T2*)k_head_r)[h2_index];
        h2_index_out =
            (batch * off0 * 2 + (head * 2 + 1) * o_off1 + seq * off2 + loc) >> 1;  //(batch, head,seq+seq_len,loc)
        ((T2*)k_buf_bd)[h2_index_out] = tmp;

        // ef
        if (seq <= 1) {
            h2_index                      = (seq * i_off1 + head_loc) >> 1;  //(seq, head, loc)
            tmp                           = ((T2*)attr_seg_embed)[h2_index];
            h2_index_out                  = (batch * 2 * i_off1 + (head * 2 + seq) * off2 + loc) >> 1;  //(head,seq,loc)
            ((T2*)k_buf_ef)[h2_index_out] = tmp;
        }
    }
}

template<>
void __global__ prepareMatrixes(float*       q_buf,
                                float*       q_buf_bd,
                                float*       q_buf_ef,
                                float*       k_buf,
                                float*       k_buf_bd,
                                float*       k_buf_ef,
                                const float* query_buf,
                                const float* key_buf,
                                const float* k_head_r,
                                const float* attr_seg_embed,
                                const float* attr_bias_Q_w,
                                const float* attr_bias_Q_r,
                                const float* attr_bias_Q_s,
                                const int    off0,
                                const int    i_off1,
                                const int    o_off1,
                                int          off2)
{
    int batch    = blockIdx.y;
    int seq      = blockIdx.x;
    int head_loc = threadIdx.x;

    float tmp;
    if (head_loc < i_off1) {
        int head = head_loc / off2;
        int loc  = head_loc % off2;

        int index     = batch * off0 + seq * i_off1 + head_loc;
        tmp           = query_buf[index];
        int index_out = batch * off0 + head * o_off1 + seq * off2 + loc;
        // left matrix
        q_buf[index_out]    = tmp + __ldg(attr_bias_Q_w + head_loc);  // tex2D(t_attr_bias_Q_w, loc, head);
        q_buf_bd[index_out] = tmp + __ldg(attr_bias_Q_r + head_loc);  // tex2D(t_attr_bias_Q_r, loc, head);
        q_buf_ef[index_out] = tmp + __ldg(attr_bias_Q_s + head_loc);  // tex2D(t_attr_bias_Q_s, loc, head);

        // right matrix
        k_buf[index_out] = key_buf[index];  // ac

        // bd
        index               = seq * i_off1 + head_loc;  //(seq, head_loc)
        tmp                 = k_head_r[index];
        index_out           = index_out + batch * off0 + head * o_off1;  //(batch, head,seq,loc)
        k_buf_bd[index_out] = tmp;

        index               = index + off0;  //(seq+seq_len, head_loc)
        tmp                 = k_head_r[index];
        index_out           = index_out + o_off1;  //(batch, head,seq+seq_len,loc)
        k_buf_bd[index_out] = tmp;

        // ef
        if (seq <= 1) {
            index               = seq * i_off1 + head_loc;  //(seq, head, loc)
            tmp                 = attr_seg_embed[index];
            index_out           = batch * 2 * i_off1 + (head * 2 + seq) * off2 + loc;  //(head,seq,loc)
            k_buf_ef[index_out] = tmp;
        }
    }
}

// Applied to half and bfloat16
template<typename T>
void __global__ transpose102(T* dst, const T* src, const int off0, const int i_off1, const int o_off1, const int off2)
{
    using T2 = typename TypeConverter<T>::Type;  // half2 or bfloat162
    int x[4] = {0};
    x[0]     = blockIdx.x;   //[0,7]
    x[1]     = blockIdx.y;   //[0,11]
    x[2]     = threadIdx.x;  //[0,127]

    int input_index = (x[0] * off0 + x[1] * i_off1 + x[2] * off2) >> 1;  // [batch, 0, 1, 2]=[d0,d1,d2,d3]

    int out_index = (x[0] * off0 + x[2] * o_off1 + x[1] * off2) >> 1;  // [batch, 1, 0, 2]=[d0,d2,d1,d3]

    ((T2*)dst)[out_index] = ((T2*)src)[input_index];
}

template<>
void __global__
transpose102(float* dst, const float* src, const int off0, const int i_off1, const int o_off1, const int off2)
{
    int x[4] = {0};
    x[0]     = blockIdx.x;   //[0,7]
    x[1]     = blockIdx.y;   //[0,11]
    x[2]     = threadIdx.x;  //[0,127]
    x[3]     = threadIdx.y;  //[0,1]

    int input_index = x[0] * off0 + x[1] * i_off1 + x[2] * off2 + x[3];  // [batch, 0, 1, 2]=[d0,d1,d2,d3]

    int out_index = x[0] * off0 + x[2] * o_off1 + x[1] * off2 + x[3];  // [batch, 1, 0, 2]=[d0,d2,d1,d3]

    dst[out_index] = src[input_index];
}

template<typename T>
void __global__ transpose201(
    T* dst, const T* src, const int off0, const int i_off1, const int head_num, const int o_off1, const int seq_len)
{
    int batch = blockIdx.x;
    int d0    = blockIdx.y;
    int d1    = threadIdx.x;

    extern __shared__ float sdata[];
    int                     i = 0;

    // Read data into shared memory
    int index  = batch * off0 + d0 * i_off1;  // d1*i_off2+d2
    int offset = d1;
    src        = src + index;
    int row    = offset / head_num;
    int col    = offset % head_num;
    for (i = 0; i < head_num; i++) {
        sdata[row * (head_num + 1) + col] = cuda_cast<float>(src[offset]);
        offset += seq_len;
        row = offset / head_num;
        col = offset % head_num;
    }

    __syncthreads();

    index  = batch * off0 + d0 * seq_len + d1;
    offset = 0;
    dst    = dst + index;
    for (i = 0; i < head_num; i++) {
        dst[offset] = cuda_cast<T>(sdata[d1 * (head_num + 1) + i]);
        offset += o_off1;
    }
}

/*int threads=512;
  seq_dim1=threads/seq_len
  seq_dim2=seq_len/dimx
  dim3 grid_rel(batch_size, head_num, seq_dim2);
  dim3 block_rel(seq_dim1, seq_len);*/
template<typename T>
void __global__ relShiftBd(T* outMatrix, const T* inputMatrix, const int off0, const int off1, const int seq_len)
{
    using T2  = typename TypeConverter<T>::Type;        // half2 or bfloat162
    int batch = blockIdx.x;                             //[0,7]
    int head  = blockIdx.y;                             //[0,11]
    int row   = blockIdx.z * blockDim.x + threadIdx.x;  //[0,127]
    int col   = threadIdx.y * 2;                        //[0,255]

    int input_index = (batch * off0 + head * off1 + row * seq_len) * 2 + col;
    int out_index;
    int out_row;
    int out_col;
    int tmp_index;
    T2  idata;
    if (col >= seq_len || row != 0) {
        idata = ((T2*)inputMatrix)[input_index >> 1];
        // int tmp_index=row*(2*seq_len-1)+row+col-seq_len;
        tmp_index = row * 2 * seq_len + col - seq_len;
        out_row   = tmp_index / (seq_len * 2 - 1);
        out_col   = tmp_index % (seq_len * 2 - 1);
        if (out_col < seq_len) {
            out_index            = (batch * off0 + head * off1 + out_row * seq_len + out_col);
            outMatrix[out_index] = idata.x;
        }
        tmp_index += 1;
        out_row = tmp_index / (seq_len * 2 - 1);
        out_col = tmp_index % (seq_len * 2 - 1);
        if (out_col < seq_len) {
            out_index            = (batch * off0 + head * off1 + out_row * seq_len + out_col);
            outMatrix[out_index] = idata.y;
        }
    }
}

/*dim3 grid_shift(batch_size, head_num, seq_len);
  dim3 block_shift(seq_len*2);
  int off0=head_num*seq_len*seq_len;
  int off1=seq_len*seq_len; */
template<>
void __global__
relShiftBd(float* outMatrix, const float* inputMatrix, const int off0, const int off1, const int seq_len)
{
    int batch = blockIdx.x;   //[0,7]
    int head  = blockIdx.y;   //[0,11]
    int row   = blockIdx.z;   //[0,127]
    int col   = threadIdx.x;  //[0,255]

    int input_index = (batch * off0 + head * off1 + row * seq_len) * 2 + col;
    if (col >= seq_len || row != 0) {
        float idata = inputMatrix[input_index];
        // int tmp_index=row*(2*seq_len-1)+row+col-seq_len;
        int tmp_index = row * 2 * seq_len + col - seq_len;
        int out_row   = tmp_index / (seq_len * 2 - 1);
        int out_col   = tmp_index % (seq_len * 2 - 1);
        if (out_col < seq_len) {
            int out_index        = batch * off0 + head * off1 + out_row * seq_len + out_col;
            outMatrix[out_index] = idata;
        }
    }
}

/*dim3 grid_score(batch_size,head_num,seq_len);
  dim3 block_score(next_pow2(seq_len));
  int off0=head_num*seq_len*seq_len;
  int off1=seq_len*seq_len;
  float p=(1/(pow(size_per_head,0.5)));

  int voff0=head_num*seq_len*size_per_head;
  int v_o_off1=seq_len*size_per_head;
  int voff2=size_per_head;
  int v_i_off1=head_num*size_per_head;*/
template<typename T>
__global__ void calAttnScore_valueBuf(T*          attn_score,
                                      const T*    ac,
                                      const T*    bd,
                                      const T*    ef,
                                      const T*    attn_mask,
                                      const int   off0,
                                      const int   off1,
                                      const int   seq_len,
                                      const float p,
                                      T*          value_buf_trans,
                                      const T*    value_buf,
                                      const int   voff0,
                                      const int   v_i_off1,
                                      const int   v_o_off1,
                                      const int   voff2)
{
    int batch = blockIdx.x;
    int head  = blockIdx.y;
    int seq1  = blockIdx.z;
    int seq2  = threadIdx.x;

    int offset = batch * off0 + head * off1 + seq1 * seq_len;
    int index  = offset + seq2;
    int out_index;
    T   score;
    T   mask;
    T   large_value = -1e4;
    if (sizeof(T) == 4) {
        large_value = -1e30f;
    }
    if (seq2 < seq_len) {
        score = ac[index] + bd[index] + ef[index];
        score = score * p;

        out_index = batch * off1 + seq1 * seq_len + seq2;
        mask      = attn_mask[out_index] * (large_value);
        score     = score + mask;
    }
    // softmax(attn_score+offset,seq_len, seq2);
    __shared__ float s_sum, s_max;
    float            tmp     = seq2 < seq_len ? score : large_value;
    float            max_val = blockReduceMax<float>(tmp);
    if (seq2 == 0) {
        s_max = max_val;
    }
    __syncthreads();
    float qk_tmp  = seq2 < seq_len ? __expf((float)(tmp - s_max)) : 0.0f;
    float sum_val = blockReduceSum<float>(qk_tmp);
    __syncthreads();
    if (seq2 == 0) {
        s_sum = sum_val;
    }
    __syncthreads();
    if (seq2 < seq_len) {
        attn_score[index] = (T)(qk_tmp / s_sum);
    }
    // end softmax

    offset = seq2;
    while (offset < voff2) {
        out_index                  = batch * voff0 + head * v_o_off1 + seq1 * voff2 + offset;
        index                      = batch * voff0 + seq1 * v_i_off1 + head * voff2 + offset;
        value_buf_trans[out_index] = value_buf[index];
        offset += seq_len;
    }
}

// Applied to half and bfloat16
template<typename T>
void __global__ calAttnScore_valueBuf_small(T*          attn_score,
                                            const T*    ac,
                                            const T*    bd,
                                            const T*    ef,
                                            const T*    attn_mask,
                                            const int   off0,
                                            const int   off1,
                                            const int   seq_len,
                                            int         n_seq1,
                                            const float p,
                                            T*          value_buf_trans,
                                            const T*    value_buf,
                                            const int   voff0,
                                            const int   v_i_off1,
                                            const int   v_o_off1,
                                            const int   voff2)
{
    using T2 = typename TypeConverter<T>::Type;  // half2 or bfloat162
    int lid  = lane_id();
    int tid  = threadIdx.x;
    int wid  = tid / 32;
    int seq2 = lid << 1;

    int batch = blockIdx.x;
    int head  = blockIdx.y;
    int seq1  = blockIdx.z * n_seq1 + wid;

    int    offset = batch * off0 + head * off1 + seq1 * seq_len;
    int    index  = (offset + seq2) >> 1;
    int    out_index;
    float2 tmp1, tmp2;

    // Data prepare section
    if (seq2 < seq_len) {
        tmp1 = cuda_cast<float2>(((T2*)ac)[index]);
        tmp2 = cuda_cast<float2>(((T2*)bd)[index]);
        tmp1.x += tmp2.x;
        tmp1.y += tmp2.y;
        // tmp1=__hadd2(tmp1,tmp2);
        tmp2 = cuda_cast<float2>(((T2*)ef)[index]);
        tmp1.x += tmp2.x;
        tmp1.y += tmp2.y;

        // half2 score=__hadd2(tmp1, tmp2);
        tmp1.x = tmp1.x * p;
        tmp1.y = tmp1.y * p;

        out_index = (batch * off1 + seq1 * seq_len + seq2) >> 1;
        tmp2      = cuda_cast<float2>(((T2*)attn_mask)[out_index]);

        tmp1.x = tmp1.x + -1e4 * tmp2.x;
        tmp1.y = tmp1.y + -1e4 * tmp2.y;
    }
    else {
        tmp1.x = tmp1.y = -1e4f;
    }

    // Softmax section
    float tmp = tmp1.x > tmp1.y ? tmp1.x : tmp1.y;
    for (int mask = 16; mask > 0; mask >>= 1) {
        tmp = max(tmp, __shfl_xor_sync(FINAL_MASK, tmp, mask, 32));
    }
    tmp = __shfl_sync(FINAL_MASK, tmp, 0);

    /// normalize the input
    tmp1.x = seq2 < seq_len ? __expf((float)(tmp1.x - tmp)) : 0.0f;
    tmp1.y = seq2 < seq_len ? __expf((float)(tmp1.y - tmp)) : 0.0f;
    tmp    = tmp1.x + tmp1.y;
    /// get sum of the normalized value
    for (int mask = 16; mask > 0; mask >>= 1) {
        tmp = tmp + __shfl_xor_sync(FINAL_MASK, tmp, mask, 32);
    }
    if (seq2 == 0) {
        tmp = tmp;
    }
    tmp = __shfl_sync(FINAL_MASK, tmp, 0);

    /// set the value
    if (seq2 < seq_len) {
        tmp1.x                   = tmp1.x / tmp;
        tmp1.y                   = tmp1.y / tmp;
        ((T2*)attn_score)[index] = cuda_cast<T2>(tmp1);
    }

    // value_buf section
    offset = seq2;
    while (offset < voff2) {
        index = (batch * voff0 + seq1 * v_i_off1 + head * voff2 + offset) >> 1;
        T2 v  = ((T2*)value_buf)[index];

        out_index                         = (batch * voff0 + head * v_o_off1 + seq1 * voff2 + offset) >> 1;
        ((T2*)value_buf_trans)[out_index] = v;
        offset += seq_len;
    }
}

// Applied to half and bfloat16
template<typename T>
void __global__ calAttnScore_valueBuf_large(T*          attn_score,
                                            const T*    ac,
                                            const T*    bd,
                                            const T*    ef,
                                            const T*    attn_mask,
                                            const int   off0,
                                            const int   off1,
                                            const int   seq_len,
                                            const float p,
                                            T*          value_buf_trans,
                                            const T*    value_buf,
                                            const int   voff0,
                                            const int   v_i_off1,
                                            const int   v_o_off1,
                                            const int   voff2)
{
    using T2  = typename TypeConverter<T>::Type;  // half2 or bfloat162
    int batch = blockIdx.x;
    int head  = blockIdx.y;
    int seq1  = blockIdx.z;

    int lid  = lane_id();
    int tid  = threadIdx.x;
    int wid  = tid / 32;
    int seq2 = tid << 1;

    int              offset = batch * off0 + head * off1 + seq1 * seq_len;
    int              index  = (offset + seq2) >> 1;
    int              out_index;
    float2           tmp1, tmp2;
    __shared__ float sdata[32];
    __shared__ float s_max;
    __shared__ float s_sum;
    // Data prepare section
    if (seq2 < seq_len) {
        tmp1 = cuda_cast<float2>(((T2*)ac)[index]);
        tmp2 = cuda_cast<float2>(((T2*)bd)[index]);
        tmp1.x += tmp2.x;
        tmp1.y += tmp2.y;
        // tmp1=__hadd2(tmp1,tmp2);
        tmp2 = cuda_cast<float2>(((T2*)ef)[index]);
        tmp1.x += tmp2.x;
        tmp1.y += tmp2.y;

        // half2 score=__hadd2(tmp1, tmp2);
        tmp1.x = tmp1.x * p;
        tmp1.y = tmp1.y * p;

        out_index = (batch * off1 + seq1 * seq_len + seq2) >> 1;
        tmp2      = cuda_cast<float2>(((T2*)attn_mask)[out_index]);

        tmp1.x = tmp1.x + -1e4 * tmp2.x;
        tmp1.y = tmp1.y + -1e4 * tmp2.y;
    }
    else {
        tmp1.x = tmp1.y = -1e4f;
    }
    // Softmax section
    float tmp = tmp1.x > tmp1.y ? tmp1.x : tmp1.y;
    for (int mask = 16; mask > 0; mask /= 2) {
        tmp = max(tmp, __shfl_xor_sync(FINAL_MASK, tmp, mask, 32));
    }
    if (wid == 0) {
        sdata[lid] = -1e4f;
    }
    __syncthreads();
    if (lid == 0) {
        sdata[wid] = tmp;
    }
    __syncthreads();

    if (wid == 0) {
        tmp = sdata[lid];
        for (int mask = 16; mask > 0; mask /= 2) {
            tmp = max(tmp, __shfl_xor_sync(FINAL_MASK, tmp, mask, 32));
        }
    }
    if (tid == 0) {
        s_max = tmp;
    }
    __syncthreads();

    /// normalize the input

    tmp1.x = seq2 < seq_len ? __expf((float)(tmp1.x - s_max)) : 0.0f;
    tmp1.y = seq2 < seq_len ? __expf((float)(tmp1.y - s_max)) : 0.0f;
    tmp    = tmp1.x + tmp1.y;

    /// get sum of the normalized value
    for (int mask = 16; mask > 0; mask /= 2) {
        tmp = tmp + __shfl_xor_sync(FINAL_MASK, tmp, mask, 32);
    }
    if (wid == 0) {
        sdata[lid] = 0;
    }
    __syncthreads();
    if (lid == 0) {
        sdata[wid] = tmp;
    }
    __syncthreads();

    if (wid == 0) {
        tmp = sdata[tid];
        for (int mask = 16; mask > 0; mask /= 2) {
            tmp = tmp + __shfl_xor_sync(FINAL_MASK, tmp, mask, 32);
        }
    }
    if (tid == 0) {
        s_sum = tmp;
    }
    __syncthreads();

    /// set the value
    if (seq2 < seq_len) {
        tmp1.x                   = tmp1.x / s_sum;
        tmp1.y                   = tmp1.y / s_sum;
        ((T2*)attn_score)[index] = cuda_cast<T2>(tmp1);
    }
    // value_buf section
    offset = seq2;
    while (offset < voff2) {
        index = (batch * voff0 + seq1 * v_i_off1 + head * voff2 + offset) >> 1;
        T2 v  = ((T2*)value_buf)[index];

        out_index                         = (batch * voff0 + head * v_o_off1 + seq1 * voff2 + offset) >> 1;
        ((T2*)value_buf_trans)[out_index] = v;
        offset += seq_len;
    }
}

// Applied to half or bfloat16
template<typename T>
__global__ void
transpose102_v2(T* dst, const T* src, const int off0, const int i_off1, const int o_off1, const int off2)
{
    using T2 = typename TypeConverter<T>::Type;  // half2 or bfloat162
    int x[4] = {0};
    x[0]     = blockIdx.x;              //[0,7] batch_size
    x[1]     = threadIdx.x * 2 / off2;  // head_num
    x[2]     = blockIdx.y;              // seq_len
    x[3]     = threadIdx.x * 2 % off2;  //[0,63] size_per_head

    if (x[3] < off2) {
        T2  tmp;
        int in_index  = (x[0] * off0 + x[1] * i_off1 + x[2] * off2 + x[3]) >> 1;  // [batch, 0, 1, 2]=[d0,d1,d2,d3]
        tmp           = ((T2*)src)[in_index];
        int out_index = (x[0] * off0 + x[2] * o_off1 + x[1] * off2 + x[3]) >> 1;  // [batch, 1, 0, 2]=[d0,d2,d1,d3]
        ((T2*)dst)[out_index] = tmp;
    }
}

// dim3 grid_trans_v(batch_size,seq_len, head_num);
// dim3 block_trans_v(size_per_head);
template<>
__global__ void
transpose102_v2(float* dst, const float* src, const int off0, const int i_off1, const int o_off1, const int off2)
{
    int x[4] = {0};
    x[0]     = blockIdx.x;
    x[1]     = threadIdx.x / off2;
    x[2]     = blockIdx.y;          //[0,128] seq_len
    x[3]     = threadIdx.x % off2;  //[0,31] size_per_head

    float tmp;
    if (x[3] < off2) {
        int input_index = x[0] * off0 + x[1] * i_off1 + x[2] * off2 + x[3];  // [batch, 0, 1, 2]=[d0,d1,d2,d3]
        tmp             = src[input_index];
        int out_index   = x[0] * off0 + x[2] * o_off1 + x[1] * off2 + x[3];  // [batch, 1, 0, 2]=[d0,d2,d1,d3]

        dst[out_index] = tmp;
    }
}

// Applied to half and bfloat16
template<typename T>
__global__ void
addBias_layerNorm(T* out, const T* input, const T* bias, const T* gamma, const T* beta, int m, int n, float epsilon)
{
    using T2             = typename TypeConverter<T>::Type;  // half2 or bfloat162
    int              tid = threadIdx.x;
    __shared__ float s_mean;
    __shared__ float s_variance;
    float            mean     = 0.0f;
    float            variance = 0.0f;
    float2           local_out_fp2;

    T2*       out_ptr   = (T2*)out;
    const T2* input_ptr = (const T2*)input;
    const T2* bias_ptr  = (const T2*)bias;  // blockIdx.x * n + i
    const T2* gamma_ptr = (const T2*)gamma;
    const T2* beta_ptr  = (const T2*)beta;

    float local_out = 0.0f;
    int   id        = (blockIdx.x * n + tid * 2) >> 1;
    local_out_fp2   = cuda_cast<float2>(hadd2(input_ptr[id], ldg(&bias_ptr[id])));

    local_out += local_out_fp2.x;
    local_out += local_out_fp2.y;

    mean = blockReduceSum<float>(local_out);
    if (threadIdx.x == 0) {
        s_mean = mean / n;
    }
    __syncthreads();

    variance = (local_out_fp2.x - s_mean) * (local_out_fp2.x - s_mean);
    variance += (local_out_fp2.y - s_mean) * (local_out_fp2.y - s_mean);
    variance = blockReduceSum<float>(variance);
    if (threadIdx.x == 0)
        s_variance = rsqrtf(variance / n + epsilon);
    __syncthreads();

    float2 gamma_val = cuda_cast<float2>(ldg(&gamma_ptr[tid]));
    float2 beta_val  = cuda_cast<float2>(ldg(&beta_ptr[tid]));
    local_out_fp2.x  = (local_out_fp2.x - s_mean) * s_variance * gamma_val.x + beta_val.x;
    local_out_fp2.y  = (local_out_fp2.y - s_mean) * s_variance * gamma_val.y + beta_val.y;
    out_ptr[id]      = cuda_cast<T2>(local_out_fp2);
}

template<>
__global__ void addBias_layerNorm(float*       out,
                                  const float* input,
                                  const float* bias,
                                  const float* gamma,
                                  const float* beta,
                                  int          m,
                                  int          n,
                                  float        epsilon)
{
    int tid = threadIdx.x;

    __shared__ float s_mean;
    __shared__ float s_variance;
    float            mean     = 0.0f;
    float            variance = 0.0f;

    float local_out = 0.0f;
    int   id        = blockIdx.x * n + tid;
    local_out += (float)(input[id]);
    local_out += (float)(__ldg(&bias[id]));

    mean = blockReduceSum<float>(local_out);
    if (threadIdx.x == 0) {
        s_mean = mean / n;
    }
    __syncthreads();

    float diff = local_out - s_mean;
    variance   = blockReduceSum<float>(diff * diff);
    if (threadIdx.x == 0) {
        s_variance = rsqrtf(variance / n + epsilon);
    }
    __syncthreads();

    out[id] = (float)(((local_out - s_mean) * s_variance) * (float)(__ldg(&gamma[tid])) + (float)(__ldg(&beta[tid])));
}

/*width=hidden_dim_ff;
  height=seq_len;
  dim3 block(1024);
  dim3 grid(batch_size, seq_len);
  gelu_bias_loop<<<grid, block, 0, stream>>>(output_fc1, output_fc1, attr_fc1_bias, hidden_dim_ff,seq_len); */
template<typename T>
__global__ void gelu_bias_loop(T* src, const T* bias, int width, int height)
{
    int batch = blockIdx.x;
    int x     = blockIdx.y;
    int y     = threadIdx.x;

    if (x < height) {
        int   index = batch * width * height + x * width;
        float v_src;
        float v_bias;
        float v;
        for (; y < width; y = y + blockDim.x) {
            v_bias = bias[y];
            v_src  = src[index + y];
            v      = v_src + v_bias;

            src[index + y] = (T)(0.5f * v * (1.0f + tanhf(0.79788456f * (v + 0.044715f * v * v * v))));
        }
    }
}
template<>
__global__ void gelu_bias_loop(__half* src, const __half* bias, int width, int height)
{
    int batch = blockIdx.x;
    int x     = blockIdx.y;
    int y     = threadIdx.x * 2;

    if (x < height) {
        int    index = batch * width * height + x * width;
        half2  v_src;
        half2  v_bias;
        half2  v;
        float2 t;
        for (; y < width; y = y + blockDim.x * 2) {
            v_bias = ((half2*)bias)[y >> 1];
            v_src  = ((half2*)src)[(index + y) >> 1];
            v      = __hadd2(v_src, v_bias);
            t      = __half22float2(v);
            t.x    = (0.5f * t.x * (1.0f + tanhf(0.79788456f * (t.x + 0.044715f * t.x * t.x * t.x))));
            t.y    = (0.5f * t.y * (1.0f + tanhf(0.79788456f * (t.y + 0.044715f * t.y * t.y * t.y))));

            ((half2*)src)[(index + y) >> 1] = __float22half2_rn(t);
        }
    }
}

/*********************Invoke Functions***********************/

template<typename T>
void invokePrepareMatrixes(int          batch_size,
                           int          seq_len,
                           int          hidden_dim,
                           int          size_per_head,
                           T*           q_buf,
                           T*           q_buf_bd,
                           T*           q_buf_ef,
                           T*           k_buf,
                           T*           k_buf_bd,
                           T*           k_buf_ef,
                           T*           query_buf,
                           T*           key_buf,
                           T*           k_head_r,
                           T*           attr_seg_embed,
                           T*           attr_bias_Q_w,
                           T*           attr_bias_Q_r,
                           T*           attr_bias_Q_s,
                           cudaStream_t stream)
{
    int off0   = seq_len * hidden_dim;  // seq_len*head_num*size_per_head
    int i_off1 = hidden_dim;            // head_num*size_per_head
    int o_off1 = seq_len * size_per_head;
    int off2   = size_per_head;

    dim3 grid(seq_len, batch_size);
    dim3 block(next_pow2(hidden_dim) / numPerThread<T>());
    prepareMatrixes<<<grid, block, 0, stream>>>(q_buf,
                                                q_buf_bd,
                                                q_buf_ef,
                                                k_buf,
                                                k_buf_bd,
                                                k_buf_ef,
                                                query_buf,
                                                key_buf,
                                                k_head_r,
                                                attr_seg_embed,
                                                attr_bias_Q_w,
                                                attr_bias_Q_r,
                                                attr_bias_Q_s,
                                                off0,
                                                i_off1,
                                                o_off1,
                                                off2);
}

template<typename T>
void invokeTranspose102(
    int batch_size, int seq_len, int head_num, T* qk_buf_ef_trans, T* qk_buf_ef, cudaStream_t stream)
{
    dim3 grid_trans_v(batch_size, head_num);

    // dim3 block_trans_v(seq_len,2);//float
    // dim3 block_trans_v(seq_len);//__half
    dim3 block_trans_v(seq_len, 2 / (numPerThread<T>()));

    int toff0   = head_num * seq_len * 2;
    int ti_off1 = seq_len * 2;
    int to_off1 = head_num * 2;
    int toff2   = 2;

    transpose102<<<grid_trans_v, block_trans_v, 0, stream>>>(
        qk_buf_ef_trans, qk_buf_ef, toff0, ti_off1, to_off1, toff2);
}

template<typename T>
void invokeTranspose201(
    int batch_size, int seq_len, int head_num, T* qk_buf_ef_seg_trans, T* qk_buf_ef_seg, cudaStream_t stream)
{
    dim3 grid_trans2(batch_size, seq_len);
    dim3 block_trans2(seq_len);
    int  t2_off0   = seq_len * seq_len * head_num;
    int  t2_i_off1 = seq_len * head_num;
    int  t2_o_off1 = seq_len * seq_len;
    int  t2_i_off2 = head_num;
    int  t2_o_off2 = seq_len;

    transpose201<<<grid_trans2, block_trans2, seq_len*(head_num + 1) * sizeof(float), stream>>>(
        qk_buf_ef_seg_trans, qk_buf_ef_seg, t2_off0, t2_i_off1, t2_i_off2, t2_o_off1, t2_o_off2);
}

// applied to half and bfloat16
template<typename T>
void blockRelShiftBd(dim3& grid, dim3& block, int batch_size, int head_num, int seq_len)
{
    int threads  = 512;
    int seq_dim1 = threads / seq_len;
    int seq_dim2 = seq_len / seq_dim1;

    grid.x = batch_size;
    grid.y = head_num;
    grid.z = seq_dim2;

    block.x = seq_dim1;
    block.y = seq_len;
}

template<>
void blockRelShiftBd<float>(dim3& grid, dim3& block, int batch_size, int head_num, int seq_len)
{
    grid.x = batch_size;
    grid.y = head_num;
    grid.z = seq_len;

    block.x = seq_len * 2;
}

template<typename T>
void invokeRelShiftBd(int batch_size, int head_num, int seq_len, T* qk_buf_bd_shift, T* qk_buf_bd, cudaStream_t stream)
{
    dim3 grid_shift;
    dim3 block_shift;
    blockRelShiftBd<T>(grid_shift, block_shift, batch_size, head_num, seq_len);

    int off0 = head_num * seq_len * seq_len;
    int off1 = seq_len * seq_len;

    relShiftBd<<<grid_shift, block_shift, 0, stream>>>(qk_buf_bd_shift, qk_buf_bd, off0, off1, seq_len);
}

template<typename T>
void invokeCalAttnScore(int          batch_size,
                        int          head_num,
                        int          seq_len,
                        int          size_per_head,
                        float        q_scaling,
                        T*           attn_score,
                        T*           qk_buf,
                        T*           qk_buf_bd_shift,
                        T*           qk_buf_ef_seg_trans,
                        T*           attn_mask,
                        T*           value_buf_trans,
                        T*           value_buf,
                        cudaStream_t stream)
{
    int   off0 = head_num * seq_len * seq_len;
    int   off1 = seq_len * seq_len;
    float p    = 1 / ((pow(size_per_head, 0.5)) * q_scaling);

    int voff0    = head_num * seq_len * size_per_head;
    int v_o_off1 = seq_len * size_per_head;
    int voff2    = size_per_head;
    int v_i_off1 = head_num * size_per_head;
    if (seq_len <= 32) {
        dim3 grid_score(batch_size, head_num, 2);
        dim3 block_score(seq_len / 2 * next_pow2(seq_len / 2));

        calAttnScore_valueBuf_small<<<grid_score, block_score, 0, stream>>>(attn_score,
                                                                            qk_buf,
                                                                            qk_buf_bd_shift,
                                                                            qk_buf_ef_seg_trans,
                                                                            attn_mask,
                                                                            off0,
                                                                            off1,
                                                                            seq_len,
                                                                            seq_len / 2,
                                                                            p,
                                                                            value_buf_trans,
                                                                            value_buf,
                                                                            voff0,
                                                                            v_i_off1,
                                                                            v_o_off1,
                                                                            voff2);
    }
    else if (seq_len <= 64) {
        dim3 grid_score(batch_size, head_num, seq_len / 2);
        dim3 block_score(2 * next_pow2(seq_len / 2));

        calAttnScore_valueBuf_small<<<grid_score, block_score, 0, stream>>>(attn_score,
                                                                            qk_buf,
                                                                            qk_buf_bd_shift,
                                                                            qk_buf_ef_seg_trans,
                                                                            attn_mask,
                                                                            off0,
                                                                            off1,
                                                                            seq_len,
                                                                            2,
                                                                            p,
                                                                            value_buf_trans,
                                                                            value_buf,
                                                                            voff0,
                                                                            v_i_off1,
                                                                            v_o_off1,
                                                                            voff2);
    }
    else {
        dim3 grid_score(batch_size, head_num, seq_len);
        dim3 block_score(next_pow2(seq_len / 2));

        calAttnScore_valueBuf_large<<<grid_score, block_score, 0, stream>>>(attn_score,
                                                                            qk_buf,
                                                                            qk_buf_bd_shift,
                                                                            qk_buf_ef_seg_trans,
                                                                            attn_mask,
                                                                            off0,
                                                                            off1,
                                                                            seq_len,
                                                                            p,
                                                                            value_buf_trans,
                                                                            value_buf,
                                                                            voff0,
                                                                            v_i_off1,
                                                                            v_o_off1,
                                                                            voff2);
    }
}

template<>
void invokeCalAttnScore(int          batch_size,
                        int          head_num,
                        int          seq_len,
                        int          size_per_head,
                        float        q_scaling,
                        float*       attn_score,
                        float*       qk_buf,
                        float*       qk_buf_bd_shift,
                        float*       qk_buf_ef_seg_trans,
                        float*       attn_mask,
                        float*       value_buf_trans,
                        float*       value_buf,
                        cudaStream_t stream)
{
    int   off0 = head_num * seq_len * seq_len;
    int   off1 = seq_len * seq_len;
    float p    = 1 / ((pow(size_per_head, 0.5)) * q_scaling);

    int voff0    = head_num * seq_len * size_per_head;
    int v_o_off1 = seq_len * size_per_head;
    int voff2    = size_per_head;
    int v_i_off1 = head_num * size_per_head;

    dim3 grid_score(batch_size, head_num, seq_len);
    dim3 block_score(next_pow2(seq_len));
    calAttnScore_valueBuf<<<grid_score, block_score, 0, stream>>>(attn_score,
                                                                  qk_buf,
                                                                  qk_buf_bd_shift,
                                                                  qk_buf_ef_seg_trans,
                                                                  attn_mask,
                                                                  off0,
                                                                  off1,
                                                                  seq_len,
                                                                  p,
                                                                  value_buf_trans,
                                                                  value_buf,
                                                                  voff0,
                                                                  v_i_off1,
                                                                  v_o_off1,
                                                                  voff2);
}

template<typename T>
void invokeTranspose102v2(
    int batch_size, int seq_len, int head_num, int size_per_head, T* attn_vec_trans, T* attn_vec, cudaStream_t stream)
{
    dim3 grid_trans_v(batch_size, seq_len);
    dim3 block_trans_v(head_num * size_per_head / numPerThread<T>());

    int off0   = head_num * seq_len * size_per_head;
    int i_off1 = seq_len * size_per_head;
    int o_off1 = head_num * size_per_head;
    int off2   = size_per_head;

    transpose102_v2<<<grid_trans_v, block_trans_v, 0, stream>>>(attn_vec_trans, attn_vec, off0, i_off1, o_off1, off2);
}

template<typename T>
void invokeAddResidualLayerNorm(int          batch_size,
                                int          seq_len,
                                int          hidden_dim,
                                T*           attn_layernorm,
                                T*           attn_out,
                                const T*     in_tensor,
                                const T*     attr_layernorm_gamma,
                                const T*     attr_layernorm_beta,
                                cudaStream_t stream)
{
    dim3 grid(batch_size * seq_len);
    dim3 block(hidden_dim / numPerThread<T>());
    addBias_layerNorm<T><<<grid, block, 0, stream>>>(attn_layernorm,
                                                     attn_out,
                                                     in_tensor,
                                                     attr_layernorm_gamma,
                                                     attr_layernorm_beta,
                                                     batch_size * seq_len,
                                                     hidden_dim,
                                                     epsilon);
}

template<typename T>
void invokeGelu(
    int batch_size, int seq_len, int hidden_dim_ff, T* output_fc1, const T* attr_fc1_bias, cudaStream_t stream)
{
    dim3 block(1024 / numPerThread<T>());
    dim3 grid(batch_size, seq_len);
    gelu_bias_loop<<<grid, block, 0, stream>>>(output_fc1, attr_fc1_bias, hidden_dim_ff, seq_len);
}

/*********************The explicit instantiation part***********************/
template void invokePrepareMatrixes<float>(int          batch_size,
                                           int          seq_len,
                                           int          hidden_dim,
                                           int          size_per_head,
                                           float*       q_buf,
                                           float*       q_buf_bd,
                                           float*       q_buf_ef,
                                           float*       k_buf,
                                           float*       k_buf_bd,
                                           float*       k_buf_ef,
                                           float*       query_buf,
                                           float*       key_buf,
                                           float*       k_head_r,
                                           float*       attr_seg_embed,
                                           float*       attr_bias_Q_w,
                                           float*       attr_bias_Q_r,
                                           float*       attr_bias_Q_s,
                                           cudaStream_t stream);

template void invokePrepareMatrixes<half>(int          batch_size,
                                          int          seq_len,
                                          int          hidden_dim,
                                          int          size_per_head,
                                          half*        q_buf,
                                          half*        q_buf_bd,
                                          half*        q_buf_ef,
                                          half*        k_buf,
                                          half*        k_buf_bd,
                                          half*        k_buf_ef,
                                          half*        query_buf,
                                          half*        key_buf,
                                          half*        k_head_r,
                                          half*        attr_seg_embed,
                                          half*        attr_bias_Q_w,
                                          half*        attr_bias_Q_r,
                                          half*        attr_bias_Q_s,
                                          cudaStream_t stream);

#ifdef ENABLE_BF16
template void invokePrepareMatrixes<__nv_bfloat16>(int            batch_size,
                                                   int            seq_len,
                                                   int            hidden_dim,
                                                   int            size_per_head,
                                                   __nv_bfloat16* q_buf,
                                                   __nv_bfloat16* q_buf_bd,
                                                   __nv_bfloat16* q_buf_ef,
                                                   __nv_bfloat16* k_buf,
                                                   __nv_bfloat16* k_buf_bd,
                                                   __nv_bfloat16* k_buf_ef,
                                                   __nv_bfloat16* query_buf,
                                                   __nv_bfloat16* key_buf,
                                                   __nv_bfloat16* k_head_r,
                                                   __nv_bfloat16* attr_seg_embed,
                                                   __nv_bfloat16* attr_bias_Q_w,
                                                   __nv_bfloat16* attr_bias_Q_r,
                                                   __nv_bfloat16* attr_bias_Q_s,
                                                   cudaStream_t   stream);
#endif

template void invokeTranspose102<float>(
    int batch_size, int seq_len, int head_num, float* qk_buf_ef_trans, float* qk_buf_ef, cudaStream_t stream);

template void invokeTranspose102<half>(
    int batch_size, int seq_len, int head_num, half* qk_buf_ef_trans, half* qk_buf_ef, cudaStream_t stream);

#ifdef ENABLE_BF16
template void invokeTranspose102<__nv_bfloat16>(int            batch_size,
                                                int            seq_len,
                                                int            head_num,
                                                __nv_bfloat16* qk_buf_ef_trans,
                                                __nv_bfloat16* qk_buf_ef,
                                                cudaStream_t   stream);
#endif

template void invokeTranspose201<float>(
    int batch_size, int seq_len, int head_num, float* qk_buf_ef_seg_trans, float* qk_buf_ef_seg, cudaStream_t stream);

template void invokeTranspose201<half>(
    int batch_size, int seq_len, int head_num, half* qk_buf_ef_seg_trans, half* qk_buf_ef_seg, cudaStream_t stream);

#ifdef ENABLE_BF16
template void invokeTranspose201<__nv_bfloat16>(int            batch_size,
                                                int            seq_len,
                                                int            head_num,
                                                __nv_bfloat16* qk_buf_ef_seg_trans,
                                                __nv_bfloat16* qk_buf_ef_seg,
                                                cudaStream_t   stream);
#endif

template void invokeRelShiftBd<float>(
    int batch_size, int head_num, int seq_len, float* qk_buf_bd_shift, float* qk_buf_bd, cudaStream_t stream);

template void invokeRelShiftBd<half>(
    int batch_size, int head_num, int seq_len, half* qk_buf_bd_shift, half* qk_buf_bd, cudaStream_t stream);

#ifdef ENABLE_BF16
template void invokeRelShiftBd<__nv_bfloat16>(int            batch_size,
                                              int            head_num,
                                              int            seq_len,
                                              __nv_bfloat16* qk_buf_bd_shift,
                                              __nv_bfloat16* qk_buf_bd,
                                              cudaStream_t   stream);
#endif

template void invokeCalAttnScore<float>(int          batch_size,
                                        int          head_num,
                                        int          seq_len,
                                        int          size_per_head,
                                        float        q_scaling,
                                        float*       attn_score,
                                        float*       qk_buf,
                                        float*       qk_buf_bd_shift,
                                        float*       qk_buf_ef_seg_trans,
                                        float*       attn_mask,
                                        float*       value_buf_trans,
                                        float*       value_buf,
                                        cudaStream_t stream);

template void invokeCalAttnScore<half>(int          batch_size,
                                       int          head_num,
                                       int          seq_len,
                                       int          size_per_head,
                                       float        q_scaling,
                                       half*        attn_score,
                                       half*        qk_buf,
                                       half*        qk_buf_bd_shift,
                                       half*        qk_buf_ef_seg_trans,
                                       half*        attn_mask,
                                       half*        value_buf_trans,
                                       half*        value_buf,
                                       cudaStream_t stream);

#ifdef ENABLE_BF16
template void invokeCalAttnScore<__nv_bfloat16>(int            batch_size,
                                                int            head_num,
                                                int            seq_len,
                                                int            size_per_head,
                                                float          q_scaling,
                                                __nv_bfloat16* attn_score,
                                                __nv_bfloat16* qk_buf,
                                                __nv_bfloat16* qk_buf_bd_shift,
                                                __nv_bfloat16* qk_buf_ef_seg_trans,
                                                __nv_bfloat16* attn_mask,
                                                __nv_bfloat16* value_buf_trans,
                                                __nv_bfloat16* value_buf,
                                                cudaStream_t   stream);
#endif

template void invokeTranspose102v2<float>(int          batch_size,
                                          int          seq_len,
                                          int          head_num,
                                          int          size_per_head,
                                          float*       attn_vec_trans,
                                          float*       attn_vec,
                                          cudaStream_t stream);
template void invokeTranspose102v2<half>(int          batch_size,
                                         int          seq_len,
                                         int          head_num,
                                         int          size_per_head,
                                         half*        attn_vec_trans,
                                         half*        attn_vec,
                                         cudaStream_t stream);
#ifdef ENABLE_BF16
template void invokeTranspose102v2<__nv_bfloat16>(int            batch_size,
                                                  int            seq_len,
                                                  int            head_num,
                                                  int            size_per_head,
                                                  __nv_bfloat16* attn_vec_trans,
                                                  __nv_bfloat16* attn_vec,
                                                  cudaStream_t   stream);
#endif

template void invokeAddResidualLayerNorm<float>(int          batch_size,
                                                int          seq_len,
                                                int          hidden_dim,
                                                float*       attn_layernorm,
                                                float*       attn_out,
                                                const float* in_tensor,
                                                const float* attr_layernorm_gamma,
                                                const float* attr_layernorm_beta,
                                                cudaStream_t stream);

template void invokeAddResidualLayerNorm<half>(int          batch_size,
                                               int          seq_len,
                                               int          hidden_dim,
                                               half*        attn_layernorm,
                                               half*        attn_out,
                                               const half*  in_tensor,
                                               const half*  attr_layernorm_gamma,
                                               const half*  attr_layernorm_beta,
                                               cudaStream_t stream);
#ifdef ENABLE_BF16
template void invokeAddResidualLayerNorm<__nv_bfloat16>(int                  batch_size,
                                                        int                  seq_len,
                                                        int                  hidden_dim,
                                                        __nv_bfloat16*       attn_layernorm,
                                                        __nv_bfloat16*       attn_out,
                                                        const __nv_bfloat16* in_tensor,
                                                        const __nv_bfloat16* attr_layernorm_gamma,
                                                        const __nv_bfloat16* attr_layernorm_beta,
                                                        cudaStream_t         stream);
#endif

template void invokeGelu<float>(
    int batch_size, int seq_len, int hidden_dim_ff, float* output_fc1, const float* attr_fc1_bias, cudaStream_t stream);

template void invokeGelu<half>(
    int batch_size, int seq_len, int hidden_dim_ff, half* output_fc1, const half* attr_fc1_bias, cudaStream_t stream);

}  // namespace fastertransformer
