#include "../../../devices/nvidia/nvidia_common.cuh"
#include "add_rms_norm_nvidia.cuh"

#include "../../../devices/nvidia/nvidia_kernel_common.cuh"
#include <cub/block/block_reduce.cuh>

#include "../../../reduce/cuda/reduce.cuh"

#include "../cuda/kernel.cuh"

// DIM=4096, block=1024, BF16: nv_bfloat162 + float regs; pair idx = tid + i*1024 (same reduction order as scalar fast path).
// (Contiguous longlong2 tiling changed per-thread partial sums order vs CUB reduce and broke bit-level match with reference runs.)
__device__ void add_rmsnormBlock_dim4096_bs1024_bf162_vec(
    __nv_bfloat16 *__restrict__ y,
    __nv_bfloat16 *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const __nv_bfloat16 *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const __nv_bfloat16 *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const __nv_bfloat16 *__restrict__ w,
    size_t nhead,
    float epsilon) {
    constexpr unsigned int BS = 1024;
    constexpr float DIM_F = 4096.0f;

    const size_t batch_idx = blockIdx.x / nhead;
    const size_t head_idx = blockIdx.x % nhead;

    const __nv_bfloat16 *a_base = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead;
    const __nv_bfloat16 *b_base = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead;
    __nv_bfloat16 *res_base = residual_out + batch_idx * stride_residual_out_batch + head_idx * stride_residual_out_nhead;
    __nv_bfloat16 *y_base = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead;

    const __nv_bfloat162 *a_ptr2 = reinterpret_cast<const __nv_bfloat162 *>(a_base);
    const __nv_bfloat162 *b_ptr2 = reinterpret_cast<const __nv_bfloat162 *>(b_base);
    const __nv_bfloat162 *w_ptr2 = reinterpret_cast<const __nv_bfloat162 *>(w);
    __nv_bfloat162 *res_ptr2 = reinterpret_cast<__nv_bfloat162 *>(res_base);
    __nv_bfloat162 *y_ptr2 = reinterpret_cast<__nv_bfloat162 *>(y_base);

    float sum_squared = 0.0f;
    float s1_reg[2];
    float s2_reg[2];

#pragma unroll
    for (int i = 0; i < 2; ++i) {
        const int idx = static_cast<int>(threadIdx.x) + i * static_cast<int>(BS);
        const __nv_bfloat162 val_a = a_ptr2[idx];
        const __nv_bfloat162 val_b = b_ptr2[idx];

        const float f_a1 = __low2float(val_a);
        const float f_a2 = __high2float(val_a);
        const float f_b1 = __low2float(val_b);
        const float f_b2 = __high2float(val_b);

        const float t1 = f_a1 + f_b1;
        const float t2 = f_a2 + f_b2;
        s1_reg[i] = t1;
        s2_reg[i] = t2;

        res_ptr2[idx] = __floats2bfloat162_rn(t1, t2);
        sum_squared += t1 * t1 + t2 * t2;
    }

    using BlockReduce = cub::BlockReduce<float, BS>;
    __shared__ typename BlockReduce::TempStorage temp_storage;
    sum_squared = BlockReduce(temp_storage).Sum(sum_squared);

    __shared__ float rms;
    if (threadIdx.x == 0) {
        rms = rsqrtf(sum_squared / DIM_F + epsilon);
    }
    __syncthreads();

#pragma unroll
    for (int i = 0; i < 2; ++i) {
        const int idx = static_cast<int>(threadIdx.x) + i * static_cast<int>(BS);
        const __nv_bfloat162 val_w = w_ptr2[idx];
        const float f_w1 = __low2float(val_w);
        const float f_w2 = __high2float(val_w);
        const float y1 = s1_reg[i] * f_w1 * rms;
        const float y2 = s2_reg[i] * f_w2 * rms;
        y_ptr2[idx] = __floats2bfloat162_rn(y1, y2);
    }
}

INFINIOP_CUDA_KERNEL add_rmsnormKernel_dim4096_bs1024_bf162_vec(
    __nv_bfloat16 *__restrict__ y,
    __nv_bfloat16 *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const __nv_bfloat16 *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const __nv_bfloat16 *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const __nv_bfloat16 *__restrict__ w,
    size_t nhead,
    float epsilon) {
    add_rmsnormBlock_dim4096_bs1024_bf162_vec(
        y, residual_out,
        stride_y_batch, stride_y_nhead,
        stride_residual_out_batch, stride_residual_out_nhead,
        a, stride_a_batch, stride_a_nhead,
        b, stride_b_batch, stride_b_nhead,
        w, nhead, epsilon);
}

// DIM=8192, block=1024: 4x nv_bfloat162 per thread; pair idx = tid + i*1024 (same as scalar tiling; avoids longlong2 reorder issues).
__device__ void add_rmsnormBlock_dim8192_bs1024_bf162_vec(
    __nv_bfloat16 *__restrict__ y,
    __nv_bfloat16 *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const __nv_bfloat16 *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const __nv_bfloat16 *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const __nv_bfloat16 *__restrict__ w,
    size_t nhead,
    float epsilon) {
    constexpr unsigned int BS = 1024;
    constexpr float DIM_F = 8192.0f;

    const size_t batch_idx = blockIdx.x / nhead;
    const size_t head_idx = blockIdx.x % nhead;

    const __nv_bfloat16 *a_base = a + batch_idx * stride_a_batch + head_idx * stride_a_nhead;
    const __nv_bfloat16 *b_base = b + batch_idx * stride_b_batch + head_idx * stride_b_nhead;
    __nv_bfloat16 *res_base = residual_out + batch_idx * stride_residual_out_batch + head_idx * stride_residual_out_nhead;
    __nv_bfloat16 *y_base = y + batch_idx * stride_y_batch + head_idx * stride_y_nhead;

    const __nv_bfloat162 *a_ptr2 = reinterpret_cast<const __nv_bfloat162 *>(a_base);
    const __nv_bfloat162 *b_ptr2 = reinterpret_cast<const __nv_bfloat162 *>(b_base);
    const __nv_bfloat162 *w_ptr2 = reinterpret_cast<const __nv_bfloat162 *>(w);
    __nv_bfloat162 *res_ptr2 = reinterpret_cast<__nv_bfloat162 *>(res_base);
    __nv_bfloat162 *y_ptr2 = reinterpret_cast<__nv_bfloat162 *>(y_base);

    float sum_squared = 0.0f;
    float s1_reg[4];
    float s2_reg[4];

#pragma unroll
    for (int i = 0; i < 4; ++i) {
        const int idx = static_cast<int>(threadIdx.x) + i * static_cast<int>(BS);
        const __nv_bfloat162 val_a = a_ptr2[idx];
        const __nv_bfloat162 val_b = b_ptr2[idx];

        const float f_a1 = __low2float(val_a);
        const float f_a2 = __high2float(val_a);
        const float f_b1 = __low2float(val_b);
        const float f_b2 = __high2float(val_b);

        const float t1 = f_a1 + f_b1;
        const float t2 = f_a2 + f_b2;
        s1_reg[i] = t1;
        s2_reg[i] = t2;

        res_ptr2[idx] = __floats2bfloat162_rn(t1, t2);
        sum_squared += t1 * t1 + t2 * t2;
    }

    using BlockReduce = cub::BlockReduce<float, BS>;
    __shared__ typename BlockReduce::TempStorage temp_storage;
    sum_squared = BlockReduce(temp_storage).Sum(sum_squared);

    __shared__ float rms;
    if (threadIdx.x == 0) {
        rms = rsqrtf(sum_squared / DIM_F + epsilon);
    }
    __syncthreads();

#pragma unroll
    for (int i = 0; i < 4; ++i) {
        const int idx = static_cast<int>(threadIdx.x) + i * static_cast<int>(BS);
        const __nv_bfloat162 val_w = w_ptr2[idx];
        const float f_w1 = __low2float(val_w);
        const float f_w2 = __high2float(val_w);
        const float y1 = s1_reg[i] * f_w1 * rms;
        const float y2 = s2_reg[i] * f_w2 * rms;
        y_ptr2[idx] = __floats2bfloat162_rn(y1, y2);
    }
}

INFINIOP_CUDA_KERNEL add_rmsnormKernel_dim8192_bs1024_bf162_vec(
    __nv_bfloat16 *__restrict__ y,
    __nv_bfloat16 *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const __nv_bfloat16 *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const __nv_bfloat16 *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const __nv_bfloat16 *__restrict__ w,
    size_t nhead,
    float epsilon) {
    add_rmsnormBlock_dim8192_bs1024_bf162_vec(
        y, residual_out,
        stride_y_batch, stride_y_nhead,
        stride_residual_out_batch, stride_residual_out_nhead,
        a, stride_a_batch, stride_a_nhead,
        b, stride_b_batch, stride_b_nhead,
        w, nhead, epsilon);
}

template <unsigned int BLOCK_SIZE, typename Tcompute, typename Tdata, typename Tweight>
INFINIOP_CUDA_KERNEL add_rmsnormKernel(
    Tdata *__restrict__ y,
    Tdata *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const Tdata *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const Tdata *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const Tweight *__restrict__ w,
    size_t nhead,
    size_t dim,
    float epsilon) {
    add_rmsnormBlock<BLOCK_SIZE, Tcompute>(
        y, residual_out,
        stride_y_batch, stride_y_nhead,
        stride_residual_out_batch, stride_residual_out_nhead,
        a, stride_a_batch, stride_a_nhead,
        b, stride_b_batch, stride_b_nhead,
        w, nhead, dim, epsilon);
}

template <typename Tcompute, typename Tdata, typename Tweight>
INFINIOP_CUDA_KERNEL add_rmsnormKernel_dim4096_bs1024(
    Tdata *__restrict__ y,
    Tdata *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const Tdata *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const Tdata *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const Tweight *__restrict__ w,
    size_t nhead,
    float epsilon) {
    add_rmsnormBlock_dim4096_bs1024<Tcompute, Tdata, Tweight>(
        y, residual_out,
        stride_y_batch, stride_y_nhead,
        stride_residual_out_batch, stride_residual_out_nhead,
        a, stride_a_batch, stride_a_nhead,
        b, stride_b_batch, stride_b_nhead,
        w, nhead, epsilon);
}

template <typename Tcompute, typename Tdata, typename Tweight>
INFINIOP_CUDA_KERNEL add_rmsnormKernel_dim8192_bs1024(
    Tdata *__restrict__ y,
    Tdata *__restrict__ residual_out,
    ptrdiff_t stride_y_batch,
    ptrdiff_t stride_y_nhead,
    ptrdiff_t stride_residual_out_batch,
    ptrdiff_t stride_residual_out_nhead,
    const Tdata *__restrict__ a,
    ptrdiff_t stride_a_batch,
    ptrdiff_t stride_a_nhead,
    const Tdata *__restrict__ b,
    ptrdiff_t stride_b_batch,
    ptrdiff_t stride_b_nhead,
    const Tweight *__restrict__ w,
    size_t nhead,
    float epsilon) {
    add_rmsnormBlock_dim8192_bs1024<Tcompute, Tdata, Tweight>(
        y, residual_out,
        stride_y_batch, stride_y_nhead,
        stride_residual_out_batch, stride_residual_out_nhead,
        a, stride_a_batch, stride_a_nhead,
        b, stride_b_batch, stride_b_nhead,
        w, nhead, epsilon);
}

namespace op::add_rms_norm::nvidia {

struct Descriptor::Opaque {
    std::shared_ptr<device::nvidia::Handle::Internal> internal;
};

Descriptor::~Descriptor() {
    delete _opaque;
}

infiniStatus_t Descriptor::create(
    infiniopHandle_t handle,
    Descriptor **desc_ptr,
    infiniopTensorDescriptor_t y_desc,
    infiniopTensorDescriptor_t residual_out_desc,
    infiniopTensorDescriptor_t a_desc,
    infiniopTensorDescriptor_t b_desc,
    infiniopTensorDescriptor_t weight_desc,
    float epsilon) {
    auto result = AddRMSNormInfo::create(y_desc, residual_out_desc, a_desc, b_desc, weight_desc, epsilon);
    CHECK_RESULT(result);
    auto info = result.take();

    *desc_ptr = new Descriptor(
        new Opaque{reinterpret_cast<device::nvidia::Handle *>(handle)->internal()},
        std::move(info),
        0,
        handle->device, handle->device_id);
    return INFINI_STATUS_SUCCESS;
}

// launch kernel with different data types
template <unsigned int BLOCK_SIZE>
infiniStatus_t launchKernel(
    uint32_t batch_size, size_t nhead, size_t dim,
    void *y, infiniDtype_t atype, ptrdiff_t stride_y_batch, ptrdiff_t stride_y_nhead,
    void *residual_out, ptrdiff_t stride_residual_out_batch, ptrdiff_t stride_residual_out_nhead,
    const void *a, ptrdiff_t stride_a_batch, ptrdiff_t stride_a_nhead,
    const void *b, ptrdiff_t stride_b_batch, ptrdiff_t stride_b_nhead,
    const void *w, infiniDtype_t wtype,
    float epsilon,
    cudaStream_t cuda_stream) {

#define LAUNCH_KERNEL(Tdata, Tweight, Tcompute)                                                                  \
    add_rmsnormKernel<BLOCK_SIZE, Tcompute, Tdata, Tweight><<<batch_size * nhead, BLOCK_SIZE, 0, cuda_stream>>>( \
        reinterpret_cast<Tdata *>(y),                                                                            \
        reinterpret_cast<Tdata *>(residual_out),                                                                 \
        stride_y_batch,                                                                                          \
        stride_y_nhead,                                                                                          \
        stride_residual_out_batch,                                                                               \
        stride_residual_out_nhead,                                                                               \
        reinterpret_cast<const Tdata *>(a),                                                                      \
        stride_a_batch,                                                                                          \
        stride_a_nhead,                                                                                          \
        reinterpret_cast<const Tdata *>(b),                                                                      \
        stride_b_batch,                                                                                          \
        stride_b_nhead,                                                                                          \
        reinterpret_cast<const Tweight *>(w),                                                                    \
        nhead,                                                                                                   \
        dim,                                                                                                     \
        epsilon)

#define LAUNCH_KERNEL_DIM4096_BS1024(Tdata, Tweight, Tcompute)                                                      \
    add_rmsnormKernel_dim4096_bs1024<Tcompute, Tdata, Tweight><<<batch_size * nhead, CUDA_BLOCK_SIZE_1024, 0, cuda_stream>>>( \
        reinterpret_cast<Tdata *>(y),                                                                               \
        reinterpret_cast<Tdata *>(residual_out),                                                                  \
        stride_y_batch,                                                                                           \
        stride_y_nhead,                                                                                           \
        stride_residual_out_batch,                                                                                  \
        stride_residual_out_nhead,                                                                                  \
        reinterpret_cast<const Tdata *>(a),                                                                         \
        stride_a_batch,                                                                                           \
        stride_a_nhead,                                                                                           \
        reinterpret_cast<const Tdata *>(b),                                                                         \
        stride_b_batch,                                                                                           \
        stride_b_nhead,                                                                                           \
        reinterpret_cast<const Tweight *>(w),                                                                       \
        nhead,                                                                                                      \
        epsilon)

#define LAUNCH_KERNEL_DIM4096_BS1024_BF162_VEC                                                                     \
    add_rmsnormKernel_dim4096_bs1024_bf162_vec<<<batch_size * nhead, CUDA_BLOCK_SIZE_1024, 0, cuda_stream>>>(       \
        reinterpret_cast<__nv_bfloat16 *>(y),                                                                       \
        reinterpret_cast<__nv_bfloat16 *>(residual_out),                                                            \
        stride_y_batch,                                                                                             \
        stride_y_nhead,                                                                                             \
        stride_residual_out_batch,                                                                                  \
        stride_residual_out_nhead,                                                                                  \
        reinterpret_cast<const __nv_bfloat16 *>(a),                                                                 \
        stride_a_batch,                                                                                             \
        stride_a_nhead,                                                                                             \
        reinterpret_cast<const __nv_bfloat16 *>(b),                                                                 \
        stride_b_batch,                                                                                             \
        stride_b_nhead,                                                                                             \
        reinterpret_cast<const __nv_bfloat16 *>(w),                                                                 \
        nhead,                                                                                                      \
        epsilon)

#define LAUNCH_KERNEL_DIM8192_BS1024(Tdata, Tweight, Tcompute)                                                      \
    add_rmsnormKernel_dim8192_bs1024<Tcompute, Tdata, Tweight><<<batch_size * nhead, CUDA_BLOCK_SIZE_1024, 0, cuda_stream>>>( \
        reinterpret_cast<Tdata *>(y),                                                                               \
        reinterpret_cast<Tdata *>(residual_out),                                                                  \
        stride_y_batch,                                                                                           \
        stride_y_nhead,                                                                                           \
        stride_residual_out_batch,                                                                                  \
        stride_residual_out_nhead,                                                                                  \
        reinterpret_cast<const Tdata *>(a),                                                                         \
        stride_a_batch,                                                                                           \
        stride_a_nhead,                                                                                           \
        reinterpret_cast<const Tdata *>(b),                                                                         \
        stride_b_batch,                                                                                           \
        stride_b_nhead,                                                                                           \
        reinterpret_cast<const Tweight *>(w),                                                                       \
        nhead,                                                                                                      \
        epsilon)

#define LAUNCH_KERNEL_DIM8192_BS1024_BF162_VEC                                                                     \
    add_rmsnormKernel_dim8192_bs1024_bf162_vec<<<batch_size * nhead, CUDA_BLOCK_SIZE_1024, 0, cuda_stream>>>(       \
        reinterpret_cast<__nv_bfloat16 *>(y),                                                                       \
        reinterpret_cast<__nv_bfloat16 *>(residual_out),                                                            \
        stride_y_batch,                                                                                             \
        stride_y_nhead,                                                                                             \
        stride_residual_out_batch,                                                                                  \
        stride_residual_out_nhead,                                                                                  \
        reinterpret_cast<const __nv_bfloat16 *>(a),                                                                 \
        stride_a_batch,                                                                                             \
        stride_a_nhead,                                                                                             \
        reinterpret_cast<const __nv_bfloat16 *>(b),                                                                 \
        stride_b_batch,                                                                                             \
        stride_b_nhead,                                                                                             \
        reinterpret_cast<const __nv_bfloat16 *>(w),                                                                 \
        nhead,                                                                                                      \
        epsilon)

    if (dim == 4096 && BLOCK_SIZE == CUDA_BLOCK_SIZE_1024) {
        if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) {
            LAUNCH_KERNEL_DIM4096_BS1024(half, half, float);
        } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) {
            LAUNCH_KERNEL_DIM4096_BS1024(half, __nv_bfloat16, float);
        } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) {
            LAUNCH_KERNEL_DIM4096_BS1024(half, float, float);
        } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) {
            LAUNCH_KERNEL_DIM4096_BS1024_BF162_VEC;
        } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) {
            LAUNCH_KERNEL_DIM4096_BS1024(__nv_bfloat16, half, float);
        } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) {
            LAUNCH_KERNEL_DIM4096_BS1024(__nv_bfloat16, float, float);
        } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) {
            LAUNCH_KERNEL_DIM4096_BS1024(float, float, float);
        } else {
            return INFINI_STATUS_BAD_TENSOR_DTYPE;
        }
    } else if (dim == 8192 && BLOCK_SIZE == CUDA_BLOCK_SIZE_1024) {
        if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) {
            LAUNCH_KERNEL_DIM8192_BS1024(half, half, float);
        } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) {
            LAUNCH_KERNEL_DIM8192_BS1024(half, __nv_bfloat16, float);
        } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) {
            LAUNCH_KERNEL_DIM8192_BS1024(half, float, float);
        } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) {
            LAUNCH_KERNEL_DIM8192_BS1024_BF162_VEC;
        } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) {
            LAUNCH_KERNEL_DIM8192_BS1024(__nv_bfloat16, half, float);
        } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) {
            LAUNCH_KERNEL_DIM8192_BS1024(__nv_bfloat16, float, float);
        } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) {
            LAUNCH_KERNEL_DIM8192_BS1024(float, float, float);
        } else {
            return INFINI_STATUS_BAD_TENSOR_DTYPE;
        }
    } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F16) {
        LAUNCH_KERNEL(half, half, float);
    } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_BF16) {
        LAUNCH_KERNEL(half, __nv_bfloat16, float);
    } else if (atype == INFINI_DTYPE_F16 && wtype == INFINI_DTYPE_F32) {
        LAUNCH_KERNEL(half, float, float);
    } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_BF16) {
        LAUNCH_KERNEL(__nv_bfloat16, __nv_bfloat16, float);
    } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F16) {
        LAUNCH_KERNEL(__nv_bfloat16, half, float);
    } else if (atype == INFINI_DTYPE_BF16 && wtype == INFINI_DTYPE_F32) {
        LAUNCH_KERNEL(__nv_bfloat16, float, float);
    } else if (atype == INFINI_DTYPE_F32 && wtype == INFINI_DTYPE_F32) {
        LAUNCH_KERNEL(float, float, float);
    } else {
        return INFINI_STATUS_BAD_TENSOR_DTYPE;
    }

#undef LAUNCH_KERNEL_DIM8192_BS1024_BF162_VEC
#undef LAUNCH_KERNEL_DIM8192_BS1024
#undef LAUNCH_KERNEL_DIM4096_BS1024_BF162_VEC
#undef LAUNCH_KERNEL_DIM4096_BS1024
#undef LAUNCH_KERNEL

    return INFINI_STATUS_SUCCESS;
}

infiniStatus_t Descriptor::calculate(
    void *workspace, size_t workspace_size,
    void *y, void *residual_out, const void *a, const void *b, const void *weight,
    void *stream) const {

    if (workspace_size < _workspace_size) {
        return INFINI_STATUS_INSUFFICIENT_WORKSPACE;
    }

    auto stride_a_batch = _info.a_strides[0];
    auto stride_a_nhead = _info.a_strides[1];
    auto stride_b_batch = _info.b_strides[0];
    auto stride_b_nhead = _info.b_strides[1];
    auto stride_y_batch = _info.y_strides[0];
    auto stride_y_nhead = _info.y_strides[1];
    auto stride_residual_out_batch = _info.residual_out_strides[0];
    auto stride_residual_out_nhead = _info.residual_out_strides[1];
    auto dim = _info.dim();
    uint32_t batch_size = static_cast<uint32_t>(_info.shape[0]);
    size_t nhead = _info.shape.size() > 2 ? _info.shape[1] : 1;
    auto cuda_stream = reinterpret_cast<cudaStream_t>(stream);

    // launch kernel with different block sizes
    if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_512) {
        CHECK_STATUS(launchKernel<CUDA_BLOCK_SIZE_512>(
            batch_size, nhead, dim,
            y, _info.atype, stride_y_batch, stride_y_nhead,
            residual_out, stride_residual_out_batch, stride_residual_out_nhead,
            a, stride_a_batch, stride_a_nhead,
            b, stride_b_batch, stride_b_nhead,
            weight, _info.wtype, _info.epsilon, cuda_stream));
    } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_1024) {
        CHECK_STATUS(launchKernel<CUDA_BLOCK_SIZE_1024>(
            batch_size, nhead, dim,
            y, _info.atype, stride_y_batch, stride_y_nhead,
            residual_out, stride_residual_out_batch, stride_residual_out_nhead,
            a, stride_a_batch, stride_a_nhead,
            b, stride_b_batch, stride_b_nhead,
            weight, _info.wtype, _info.epsilon, cuda_stream));
    } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_2048) {
        CHECK_STATUS(launchKernel<CUDA_BLOCK_SIZE_2048>(
            batch_size, nhead, dim,
            y, _info.atype, stride_y_batch, stride_y_nhead,
            residual_out, stride_residual_out_batch, stride_residual_out_nhead,
            a, stride_a_batch, stride_a_nhead,
            b, stride_b_batch, stride_b_nhead,
            weight, _info.wtype, _info.epsilon, cuda_stream));
    } else if (_opaque->internal->maxThreadsPerBlock() == CUDA_BLOCK_SIZE_4096) {
        CHECK_STATUS(launchKernel<CUDA_BLOCK_SIZE_4096>(
            batch_size, nhead, dim,
            y, _info.atype, stride_y_batch, stride_y_nhead,
            residual_out, stride_residual_out_batch, stride_residual_out_nhead,
            a, stride_a_batch, stride_a_nhead,
            b, stride_b_batch, stride_b_nhead,
            weight, _info.wtype, _info.epsilon, cuda_stream));
    } else {
        return INFINI_STATUS_DEVICE_ARCHITECTURE_NOT_SUPPORTED;
    }
    return INFINI_STATUS_SUCCESS;
}
} // namespace op::add_rms_norm::nvidia
