/*
 * SPDX-FileCopyrightText: Copyright (c) 2011-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: NVIDIA TensorRT Source Code License Agreement
 *
 * NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
 * property and proprietary rights in and to this material, related
 * documentation and any modifications thereto. Any use, reproduction,
 * disclosure or distribution of this material and related documentation
 * without an express license agreement from NVIDIA CORPORATION or
 * its affiliates is strictly prohibited.
 */

#include <cfloat>
#include <cstdio>
#include <fmha/numeric_types.h>
#include <fmha/utils.h>

////////////////////////////////////////////////////////////////////////////////////////////////////

// The number of threads per warp.
enum
{
    THREADS_PER_WARP = 32
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename Dst_type, typename Src_type>
struct Softmax_params
{
    // Output pointer.
    Dst_type* dst;
    // Source pointer.
    Src_type const* src;
    // Masks.
    int8_t const* mask;
    // Attention sinks (per head).
    float const* attention_sinks;
    // Softmax sum pointer.
    float* softmax_sum;
    // ALiBi
    bool has_alibi;
    // Dimensions of the problem.
    size_t b, h;
    // Precomputed constants.
    size_t bhs, hs, bs;
    // The scaling factors to apply when we convert to/from float.
    float scale_bmm1, softcapping_scale_bmm1, scale_softmax;
    // The number of reduction warps used by the fused kernel.
    int warps_n;
    int* cu_q_seqlens;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ float to_float(uint16_t const& src, float)
{
    return fmha::half_to_float(src);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

// Disable warning #177-D because this function has not been used elsewhere
#pragma nv_diag_suppress 177

static inline __device__ float to_float(fmha::bf16_t const& src, float)
{
    return __bfloat162float(src);
}

#pragma nv_diag_default 177

////////////////////////////////////////////////////////////////////////////////////////////////////

// Disable warning #177-D because this function has not been used elsewhere
#pragma nv_diag_suppress 177

static inline __device__ float to_float(fmha::e4m3_t const& src, float)
{
    return float(src);
}

#pragma nv_diag_default 177

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ float to_float(float const& src, float)
{
    return src;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ float to_float(int const& src, float scale)
{
    float dst;

    // Convert from int to float.
    dst = static_cast<float>(src);

    // Scale.
    dst *= scale;

    return dst;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ void from_float(uint16_t& dst, float const& src, float)
{
    dst = fmha::float_to_half(src);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ void from_float(fmha::bf16_t& dst, float const& src, float)
{
    dst = fmha::float_to_bf16(src);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ int8_t float_to_int8_rn(float x)
{
    uint32_t dst;
    asm volatile("cvt.rni.sat.s8.f32 %0, %1;" : "=r"(dst) : "f"(x));
    return reinterpret_cast<int8_t const&>(dst);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ void from_float(int8_t& dst, float const& src, float scale)
{
    dst = float_to_int8_rn(src * scale);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ void from_float(fmha::e4m3_t& dst, float const& src, float scale)
{
    dst = fmha::e4m3_t(src * scale);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

static inline __device__ float apply_exp_(float x, float max)
{
    return isinf(x) ? 0.f : __expf(x - max);
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <int N>
static inline __device__ void reduce(float (&data_fp32)[N][1], int8_t const (&mask)[N][1], int warps_n, float& sum_fp32,
    float& max_fp32, float const attention_sink)
{

// Apply the masks.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] = mask[ii][0] ? data_fp32[ii][0] : -HUGE_VALF;
    }

    // Compute the max inside the thread.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        max_fp32 = fmaxf(max_fp32, data_fp32[ii][0]);
    }

// Compute inside the warp.
#pragma unroll
    for (int xor_mask = THREADS_PER_WARP / 2; xor_mask > 0; xor_mask /= 2)
    {
        max_fp32 = fmaxf(max_fp32, __shfl_xor_sync(uint32_t(-1), max_fp32, xor_mask));
    }

// Transform the elements.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] = apply_exp_(data_fp32[ii][0], max_fp32);
    }

    // Compute the max inside the thread.
#if defined(USE_SAME_SUM_ORDER_IN_SOFTMAX_AS_REF_CODE)

#pragma unroll
    for (int ii = 0; ii < N; ii++)
    {
        sum_fp32 += data_fp32[ii][0]; //+0    +64    +128
    }

    // Emulate tmp[0] + tmp[1]
    sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 4);
    __syncwarp();

    // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
    sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 1);
    __syncwarp();
    // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
    sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 2);
    __syncwarp();

    // Emulate final reduction
    sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 8);
    __syncwarp();
    sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 16);
    __syncwarp();

#else
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        sum_fp32 += data_fp32[ii][0];
    }

// Compute inside the warp.
#pragma unroll
    for (int xor_mask = THREADS_PER_WARP / 2; xor_mask > 0; xor_mask /= 2)
    {
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, xor_mask);
    }
#endif

    // // DEBUG.
    // if( blockIdx.z == 1 && threadIdx.y == 0 && threadIdx.x == 5 ) {
    //   printf("elt=%12.8f sum_fp32=%12.8f\n", data_fp32[0].x, sum_fp32);
    // }

    // Fix the sum if needed.
    if (sum_fp32 == 0.f || sum_fp32 != sum_fp32)
    {
        sum_fp32 = 1.f;
    }

    // Normalize.
    float inv_sum_fp32 = 1.f / (sum_fp32 + expf(attention_sink - max_fp32));
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] *= inv_sum_fp32;
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <int N>
static inline __device__ void reduce(float (&data_fp32)[N][2], int8_t const (&mask)[N][2], int warps_n, float& sum_fp32,
    float& max_fp32, float const attention_sink)
{
// Apply the masks.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] = mask[ii][0] ? data_fp32[ii][0] : -HUGE_VALF;
        data_fp32[ii][1] = mask[ii][1] ? data_fp32[ii][1] : -HUGE_VALF;
    }

    // Compute the max inside the thread.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        max_fp32 = fmaxf(max_fp32, data_fp32[ii][0]);
        max_fp32 = fmaxf(max_fp32, data_fp32[ii][1]);
    }

// Compute inside the warp.
#pragma unroll
    for (int xor_mask = THREADS_PER_WARP / 2; xor_mask > 0; xor_mask /= 2)
    {
        max_fp32 = fmaxf(max_fp32, __shfl_xor_sync(uint32_t(-1), max_fp32, xor_mask));
    }

// // DEBUG.
// if( blockIdx.z == 1 && threadIdx.y == 0 && threadIdx.x == 5 ) {
//   printf("elt=%12.8f max_fp32=%12.8f\n", data_fp32[0][0], max_fp32);
// }
// // END OF DEBUG.

// Transform the elements.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] = apply_exp_(data_fp32[ii][0], max_fp32);
        data_fp32[ii][1] = apply_exp_(data_fp32[ii][1], max_fp32);
    }

    // Compute the max inside the thread.
    // float sum_fp32 = 0.f;
#if defined(USE_SAME_SUM_ORDER_IN_SOFTMAX_AS_REF_CODE)
    if (warps_n == 1)
    {
        // TODO not sure if we can improve this on the gmma side without using additional regs.

        // this is intentionally o(n) instead of o(log n)
        // lanes 0 and 1 here represent the first quad.

        // need to account for offset of l0 when addressing absolute lanes.
        int const ti = threadIdx.x % 4;
        float tmp = 0.f;

        for (int ni = 0; ni < N; ni++)
        {
            float x = data_fp32[ni][0] + data_fp32[ni][1];
            tmp += x;

            for (int it = 1; it < 8; it++)
            {
                tmp += __shfl_sync(uint32_t(-1), x, 4 * it + ti);
                __syncwarp();
            }
        }

        // emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
        tmp += __shfl_xor_sync(uint32_t(-1), tmp, 1);
        __syncwarp();
        // emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
        tmp += __shfl_xor_sync(uint32_t(-1), tmp, 2);
        __syncwarp();
        sum_fp32 = __shfl_sync(uint32_t(-1), tmp, 0);
    }
    else if (warps_n == 8)
    {
        // Accumulate warp 0 and warp 4
        float tmp[2] = {0.f, 0.f};
#pragma unroll
        for (int ii = 0; ii < N; ii += 2)
        {
            tmp[0] += data_fp32[ii + 0][0];
            tmp[0] += data_fp32[ii + 0][1];
            tmp[1] += data_fp32[ii + 1][0];
            tmp[1] += data_fp32[ii + 1][1];
        }

        // Emulate tmp[0] + tmp[1]
        tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
        tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 4);

        // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
        tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
        tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 1);
        // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
        tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
        tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 2);
        // Emulate final reduction
        tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
        tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 8);

        tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
        tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 16);
        sum_fp32 = tmp[0] + tmp[1];

        sum_fp32 = __shfl_sync(uint32_t(-1), sum_fp32, 0);
    }
    else
    {

#pragma unroll
        for (int ii = 0; ii < N; ii++)
        {
            sum_fp32 += data_fp32[ii][0] + data_fp32[ii][1];
        }

        // Emulate tmp[0] + tmp[1]
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 4);
        // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 1);
        // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 2);
        // Emulate final reduction
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 8);
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, 16);

        sum_fp32 = __shfl_sync(uint32_t(-1), sum_fp32, 0);
    }

#else
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        sum_fp32 += data_fp32[ii][0];
        sum_fp32 += data_fp32[ii][1];
    }

// Compute inside the warp.
#pragma unroll
    for (int xor_mask = THREADS_PER_WARP / 2; xor_mask > 0; xor_mask /= 2)
    {
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, xor_mask);
    }
#endif

    // // DEBUG.
    // if( blockIdx.z == 1 && threadIdx.y == 0 && threadIdx.x == 5 ) {
    //   printf("elt=%12.8f sum_fp32=%12.8f\n", data_fp32[0][0], sum_fp32);
    // }

    // Fix the sum if needed.
    if (sum_fp32 == 0.f || sum_fp32 != sum_fp32)
    {
        sum_fp32 = 1.f;
    }

    // Normalize.
    float inv_sum_fp32 = 1.f / (sum_fp32 + expf(attention_sink - max_fp32));
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] *= inv_sum_fp32;
        data_fp32[ii][1] *= inv_sum_fp32;
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <int N>
static inline __device__ void reduce(float (&data_fp32)[N][4], int8_t const (&mask)[N][4], int warps_n, float& sum_fp32,
    float& max_fp32, float const attention_sink)
{

// Apply the masks.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] = mask[ii][0] ? data_fp32[ii][0] : -HUGE_VALF;
        data_fp32[ii][1] = mask[ii][1] ? data_fp32[ii][1] : -HUGE_VALF;
        data_fp32[ii][2] = mask[ii][2] ? data_fp32[ii][2] : -HUGE_VALF;
        data_fp32[ii][3] = mask[ii][3] ? data_fp32[ii][3] : -HUGE_VALF;
    }

    // Compute the max inside the thread.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        max_fp32 = fmaxf(max_fp32, data_fp32[ii][0]);
        max_fp32 = fmaxf(max_fp32, data_fp32[ii][1]);
        max_fp32 = fmaxf(max_fp32, data_fp32[ii][2]);
        max_fp32 = fmaxf(max_fp32, data_fp32[ii][3]);
    }

// Compute inside the warp.
#pragma unroll
    for (int xor_mask = THREADS_PER_WARP / 2; xor_mask > 0; xor_mask /= 2)
    {
        max_fp32 = fmaxf(max_fp32, __shfl_xor_sync(uint32_t(-1), max_fp32, xor_mask));
    }

// // DEBUG.
// if( blockIdx.z == 1 && threadIdx.y == 0 && threadIdx.x == 5 ) {
//   printf("elt=%12.8f max_fp32=%12.8f\n", data_fp32[0][0], max_fp32);
// }

// Transform the elements.
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] = apply_exp_(data_fp32[ii][0], max_fp32);
        data_fp32[ii][1] = apply_exp_(data_fp32[ii][1], max_fp32);
        data_fp32[ii][2] = apply_exp_(data_fp32[ii][2], max_fp32);
        data_fp32[ii][3] = apply_exp_(data_fp32[ii][3], max_fp32);
    }

    // Compute the max inside the thread.
    // float sum_fp32 = 0.f;

    // TODO needs refactoring...

#if defined(USE_SAME_SUM_ORDER_IN_SOFTMAX_AS_REF_CODE)
    // Within a thread it should correspond to the operation done in the tmp[0]/[1] loop.

    if (warps_n == 1)
    { // E.g. 4x1: 4 threads iterate over all cores.
        // TODO not sure if we can improve this on the gmma side without using additional regs.

        // this is intentionally o(n) instead of o(log n)
        // lanes 0 and 1 here represent the first quad.

        // need to account for offset of l0 when addressing absolute lanes.
        int const ti = threadIdx.x % 2;
        float tmp[2] = {0.f, 0.f};

        for (int ni = 0; ni < N; ni++)
        {
            // +1
            float x = data_fp32[ni][0] + data_fp32[ni][1];
            float y = data_fp32[ni][2] + data_fp32[ni][3];
            tmp[0] += x;
            tmp[1] += y;

            for (int it = 1; it < 16; it++)
            {
                tmp[0] += __shfl_sync(uint32_t(-1), x, 2 * it + ti);
                __syncwarp();
                tmp[1] += __shfl_sync(uint32_t(-1), y, 2 * it + ti);
                __syncwarp();
            }
        }

        // emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
        tmp[0] += tmp[1];
        // emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
        tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
        __syncwarp();
        sum_fp32 = __shfl_sync(uint32_t(-1), tmp[0], 0);
    }
    else
    {

        // SEQLEN == 128.
        if (N == 1)
        {

            float tmp[2] = {0.f, 0.f};

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 700 // GV100
            // The thread local reduction.
            tmp[0] += data_fp32[0][0];
            tmp[0] += data_fp32[0][1];
            tmp[0] += data_fp32[0][2];
            tmp[0] += data_fp32[0][3];

            // Add threads 0 and 2. Inside a thread in the impl.
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
            __syncwarp();
            // Add threads 0 and 8. Inside the thread.
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
            __syncwarp();
            // Add threads 0 and 16. Inside the thread.
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
            __syncwarp();

            // Add threads 0 and 1.
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
            __syncwarp();

            // Add threads 0 and 4. Inter-warp in the code.
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
            __syncwarp();
#else
            if (warps_n == 2)
            { // 2x2
                tmp[0] += data_fp32[0][0] + data_fp32[0][1];
                tmp[1] += data_fp32[0][2] + data_fp32[0][3];

                // Emulate a_01 += a_23...
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
                __syncwarp();
                tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 2);
                __syncwarp();

                // Emulate a_01 += a_45...
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
                __syncwarp();
                tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 8);
                __syncwarp();

                // Emulate a_01 += a_89...
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
                __syncwarp();
                tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 16);
                __syncwarp();

                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
                tmp[0] += tmp[1];

                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
                __syncwarp();

                // Emulate the final reduction in smem.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
                __syncwarp();
            }
            else
            { // 1x4
                tmp[0] += data_fp32[0][0] + data_fp32[0][1];
                tmp[1] += data_fp32[0][2] + data_fp32[0][3];

                // Add +64.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
                __syncwarp();
                tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 16);
                __syncwarp();

                // T0: Emulate dst[mi] = tmp[mi][0] + tmp[mi][1];
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
                __syncwarp();
                // T1: Emulate dst[mi] = tmp[mi][0] + tmp[mi][1];
                tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 2);
                __syncwarp();

                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
                tmp[0] += tmp[1];

                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
                __syncwarp();
                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 4); __syncwarp();
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
                __syncwarp();
                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 8); __syncwarp();
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
                __syncwarp();
            }

#endif // ! GV100

            // Don't forget to put the value in sum_fp32 :)
            // sum_fp32 = tmp[0];
            sum_fp32 = __shfl_sync(uint32_t(-1), tmp[0], 0);

            // SEQLEN == 256 - compare with 1x4.
        }
        else if (N == 2 || N == 8)
        {

#pragma unroll
            for (int step = 0; step < N; step += 2)
            {

                float tmp[2] = {0.f, 0.f};

#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 700 // GV100

                // The thread local reduction.
                tmp[0] += data_fp32[step + 0][0];
                tmp[0] += data_fp32[step + 0][1];
                tmp[0] += data_fp32[step + 0][2];
                tmp[0] += data_fp32[step + 0][3];

                tmp[1] += data_fp32[step + 1][0];
                tmp[1] += data_fp32[step + 1][1];
                tmp[1] += data_fp32[step + 1][2];
                tmp[1] += data_fp32[step + 1][3];

                // Sum offset 0 and 128 (and so on).
                tmp[0] += tmp[1];

                // Add threads 0 and 2. Inside a thread in the impl.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
                __syncwarp();
                // Add threads 0 and 16. Inside the thread.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
                __syncwarp();

                // Add threads 0 and 1.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
                __syncwarp();

                // Add threads 0 and 4. Inter-warp in the code.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
                __syncwarp();
                // Add threads 0 and 8. Inter-warp in the code.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
                __syncwarp();
#else
                // 0.
                tmp[0] += data_fp32[step + 0][0] + data_fp32[step + 0][1];
                tmp[1] += data_fp32[step + 0][2] + data_fp32[step + 0][3];

                // Add +64.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
                __syncwarp();
                tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 16);
                __syncwarp();

                // Add +128 but use temp storage due to the next round of shfl.
                float xy = data_fp32[step + 1][0] + data_fp32[step + 1][1];
                float zw = data_fp32[step + 1][2] + data_fp32[step + 1][3];

                // Add +128.
                tmp[0] += xy;
                tmp[1] += zw;

                // Add +192.
                tmp[0] += __shfl_xor_sync(uint32_t(-1), xy, 16);
                __syncwarp();
                tmp[1] += __shfl_xor_sync(uint32_t(-1), zw, 16);
                __syncwarp();

                // T0: Emulate dst[mi] = tmp[mi][0] + tmp[mi][1];
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
                __syncwarp();
                // T1: Emulate dst[mi] = tmp[mi][0] + tmp[mi][1];
                tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 2);
                __syncwarp();

                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
                tmp[0] += tmp[1];

                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
                __syncwarp();
                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 4); __syncwarp();
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
                __syncwarp();
                // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 8); __syncwarp();
                tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
                __syncwarp();
#endif // ! GV100

                // Don't forget to put the value in sum_fp32 :)
                sum_fp32 += tmp[0];
            }
            // Emulate taking warp results from position 0, 16, 32, 48, etc.
            sum_fp32 = __shfl_sync(uint32_t(-1), sum_fp32, 0);

            // SEQLEN == 384.
        }
        else if (N == 3)
        {

            float tmp[2] = {0.f, 0.f};

// The reduction inside the thread.
#pragma unroll
            for (int ii = 0; ii < N; ++ii)
            {
                tmp[0] += data_fp32[ii][0];
                tmp[0] += data_fp32[ii][1];
                tmp[1] += data_fp32[ii][2];
                tmp[1] += data_fp32[ii][3];
            }

            // Emulate dst[mi] = tmp[mi][0] + tmp[mi][1];
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
            __syncwarp();
            tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 2);
            __syncwarp();

            // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
            tmp[0] += tmp[1];

            // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
            __syncwarp();

            // Emulate the final summation.
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
            __syncwarp();
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
            __syncwarp();
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
            __syncwarp();

            // Don't forget to put the value in sum_fp32 :)
            sum_fp32 += tmp[0];
            // SEQLEN == 512 - compare with 1x8.
        }
        else if (N >= 4)
        {
            // Emulate thread local
            float tmp[2] = {0.f, 0.f}; // T0, T1
#pragma unroll
            for (int step = 0; step < N; step++)
            {
                tmp[0] += data_fp32[step][0]; // + 0
                tmp[0] += data_fp32[step][1]; // + 1
                tmp[1] += data_fp32[step][2]; // + 2
                tmp[1] += data_fp32[step][3]; // + 3
            }

            // T0: Emulate dst[mi] = tmp[mi][0] + tmp[mi][1];
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 2);
            __syncwarp();
            // T1: Emulate dst[mi] = tmp[mi][0] + tmp[mi][1];
            tmp[1] += __shfl_xor_sync(uint32_t(-1), tmp[1], 2);
            __syncwarp();

            // Emulate intra-thread
            // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 1); __syncwarp();
            tmp[0] += tmp[1];
            // Emulate dst[mi] += __shfl_xor_sync(uint32_t(-1), dst[mi], 2); __syncwarp();
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 1);
            __syncwarp();

            // Emulate inter-thread
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 4);
            __syncwarp();
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 8);
            __syncwarp();
            tmp[0] += __shfl_xor_sync(uint32_t(-1), tmp[0], 16);
            __syncwarp();

            // Don't forget to put the value in sum_fp32 :)
            // sum_fp32 = tmp[0];

            // Emulate taking warp results from position 0, 16, 32, 48, etc.
            sum_fp32 = __shfl_sync(uint32_t(-1), tmp[0], 0);
            // Not supported.
        }
        else
        {
            assert(false);
        }
    } // warps_n ==  1
#else
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        sum_fp32 += data_fp32[ii][0];
        sum_fp32 += data_fp32[ii][1];
        sum_fp32 += data_fp32[ii][2];
        sum_fp32 += data_fp32[ii][3];
    }

// Compute inside the warp.
#pragma unroll
    for (int xor_mask = THREADS_PER_WARP / 2; xor_mask > 0; xor_mask /= 2)
    {
        sum_fp32 += __shfl_xor_sync(uint32_t(-1), sum_fp32, xor_mask);
    }
#endif

    // // DEBUG.
    // if( blockIdx.x == 0 && threadIdx.y == 0 && threadIdx.x == 0 ) {
    //     printf("elt=%12.8f sum_fp32=%12.8f\n", data_fp32[0][0], sum_fp32);
    // }
    // // END OF DEBUG.

    // Fix the sum if needed.
    if (sum_fp32 == 0.f || sum_fp32 != sum_fp32)
    {
        sum_fp32 = 1.f;
    }

    // Normalize.
    float inv_sum_fp32 = 1.f / (sum_fp32 + expf(attention_sink - max_fp32));
#pragma unroll
    for (int ii = 0; ii < N; ++ii)
    {
        data_fp32[ii][0] *= inv_sum_fp32;
        data_fp32[ii][1] *= inv_sum_fp32;
        data_fp32[ii][2] *= inv_sum_fp32;
        data_fp32[ii][3] *= inv_sum_fp32;
    }
}

template <typename Data_type, int X>
struct VecX
{

    using Type = typename fmha::Uint_from_size_in_bytes<X * sizeof(Data_type)>::Type;
    static_assert(sizeof(Type) == X * sizeof(Data_type));

    union Alias
    {
        Type raw;
        Data_type elt[X];
    };

    static __device__ inline void to_floatX(
        float (&dst)[X], Type const& src, float const scale, float const attn_logit_softcapping_scale)
    {
        Alias tmp;
        tmp.raw = src;
#pragma unroll
        for (int it = 0; it < X; it++)
        {
            dst[it] = to_float(tmp.elt[it], scale);
            if (attn_logit_softcapping_scale != 0.f)
            {
                dst[it] = attn_logit_softcapping_scale * fmha::__tanhf(dst[it] / attn_logit_softcapping_scale);
            }
        }
    }

    static __device__ inline void from_floatX(Type& dst, float const (&src)[X], float const scale)
    {
        Alias tmp;
#pragma unroll
        for (int it = 0; it < X; it++)
        {
            from_float(tmp.elt[it], src[it], scale);
        }
        dst = tmp.raw;
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

inline __device__ float get_alibi_head_scaling_factor(int const head_id, int const num_heads)
{
    // Round down to power of 2
    int const num_heads_pow2 = (1u << (31 - __clz(num_heads)));
    if (head_id < num_heads_pow2)
    {
        return exp2f((head_id + 1) * -8.0f / num_heads_pow2);
    }
    else
    {
        float const adjusted_head_id = 2 * (head_id - num_heads_pow2) + 1;
        return exp2f(adjusted_head_id * -4.0f / num_heads_pow2);
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename Dst_type, typename Src_type, int SEQLEN, int WARPS_PER_CTA, int X = 4>
static __global__ void softmax_kernel(Softmax_params<Dst_type, Src_type> params)
{

    // By default, use LDG.64 for the loads and STG.64 for the stores.
    enum
    {
        ELEMENTS_PER_LDG = X,
        ELEMENTS_PER_STG = X
    };

    // The number of Vec_type per thread.
    enum
    {
        VECs_PER_THREAD = SEQLEN / THREADS_PER_WARP / ELEMENTS_PER_LDG
    };

    // DEBUG.
    static_assert(VECs_PER_THREAD * THREADS_PER_WARP * ELEMENTS_PER_LDG == SEQLEN, "");
    // END OF DEBUG.

    using VecO = VecX<Dst_type, X>;
    using VecI = VecX<Src_type, X>;
    using VecM = VecX<int8_t, X>;
    // The vector types.
    using DstX_type = typename VecO::Type;
    using SrcX_type = typename VecI::Type;

    // Make sure the sizes match our expectations.
    static_assert(sizeof(DstX_type) == X * sizeof(Dst_type));
    static_assert(sizeof(SrcX_type) == X * sizeof(Src_type));

    // The type of the mask.
    using MaskX_type = typename VecM::Type;

    // One warp per sequence.
    size_t hi = blockIdx.y * WARPS_PER_CTA + threadIdx.y;
    size_t bi = blockIdx.z;
    size_t si = blockIdx.x;

    // The data offset. Layout is S * B * H * S.
    size_t src_offset = si * params.bhs + bi * params.hs + hi * SEQLEN + threadIdx.x * ELEMENTS_PER_LDG;

    // Load the input elements.
    SrcX_type const* src_ptr = reinterpret_cast<SrcX_type const*>(&params.src[src_offset]);
    SrcX_type data_src[VECs_PER_THREAD];
#pragma unroll
    for (int ii = 0; ii < VECs_PER_THREAD; ++ii)
    {
        if (hi < params.h)
        {
            data_src[ii] = src_ptr[ii * THREADS_PER_WARP];
        }
    }

    // The mask offset. Layout is S * B * S.
    size_t mask_offset = si * params.bs + bi * SEQLEN + threadIdx.x * ELEMENTS_PER_LDG;

    // Load the masks.
    MaskX_type const* mask_ptr = reinterpret_cast<MaskX_type const*>(&params.mask[mask_offset]);
    MaskX_type mask[VECs_PER_THREAD];
#pragma unroll
    for (int ii = 0; ii < VECs_PER_THREAD; ++ii)
    {
        mask[ii] = mask_ptr[ii * THREADS_PER_WARP];
    }

    // Convert the data to float.
    float data_fp32[VECs_PER_THREAD][X];
    int8_t mask_[VECs_PER_THREAD][X];
#pragma unroll
    for (int ii = 0; ii < VECs_PER_THREAD; ++ii)
    {
        VecI::to_floatX(data_fp32[ii], data_src[ii], params.scale_bmm1, params.softcapping_scale_bmm1);

        typename VecM::Alias tmp;
        tmp.raw = mask[ii];
#pragma unroll
        for (int it = 0; it < X; it++)
        {
            mask_[ii][it] = tmp.elt[it];
        }
    }

    if (params.has_alibi)
    {
        float const alibi_factor = get_alibi_head_scaling_factor(hi, params.h);
#pragma unroll
        for (int ii = 0; ii < VECs_PER_THREAD; ii++)
        {
#pragma unroll
            for (int jj = 0; jj < X; jj++)
            {
                int col = ii * THREADS_PER_WARP * X + threadIdx.x * X + jj;
                data_fp32[ii][jj] += alibi_factor * col;
            }
        }
    }

    // The attention sink value.
    float attention_sink = -FLT_MAX;
    if (params.attention_sinks != nullptr)
    {
        attention_sink = params.attention_sinks[hi];
    }

    // Do the reduction.
    float sum_fp32 = 0.f;
    float max_fp32 = -HUGE_VALF;
    reduce(data_fp32, mask_, params.warps_n, sum_fp32, max_fp32, attention_sink);
    if (threadIdx.x == 0)
    {
        int sum_s = params.cu_q_seqlens[bi];
        // [B, S, H, 2] {max, sum} float
        if (hi < params.h)
        {
            params.softmax_sum[(sum_s + si) * params.h * 2 + hi * 2] = max_fp32;
            params.softmax_sum[(sum_s + si) * params.h * 2 + hi * 2 + 1] = sum_fp32;
        }
    }
    // Reconvert to half.
    DstX_type data_dst[VECs_PER_THREAD];
#pragma unroll
    for (int ii = 0; ii < VECs_PER_THREAD; ++ii)
    {
        VecO::from_floatX(data_dst[ii], data_fp32[ii], params.scale_softmax);
    }

    // Store the output elements.
    DstX_type* dst_ptr = reinterpret_cast<DstX_type*>(&params.dst[src_offset]);
#pragma unroll
    for (int ii = 0; ii < VECs_PER_THREAD; ++ii)
    {
        if (hi < params.h)
        {
            dst_ptr[ii * THREADS_PER_WARP] = data_dst[ii];
        }
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename Dst_type, typename Src_type>
void run_softmax(void* dst, void const* src, void const* mask, void const* attention_sinks, void* softmax_sum,
    void* cu_q_seqlens, int s_inner, int s_outer, int b, int h, float scale_bmm1, float scale_softmax,
    float softcapping_scale_bmm1, int warps_n, bool has_alibi)
{

    Softmax_params<Dst_type, Src_type> params;
    memset(&params, 0, sizeof(params));

    // The different pointers.
    params.dst = reinterpret_cast<Dst_type*>(dst);
    params.src = reinterpret_cast<Src_type const*>(src);
    params.softmax_sum = reinterpret_cast<float*>(softmax_sum);
    params.cu_q_seqlens = reinterpret_cast<int*>(cu_q_seqlens);
    params.mask = reinterpret_cast<int8_t const*>(mask);
    params.attention_sinks = reinterpret_cast<float const*>(attention_sinks);
    params.has_alibi = has_alibi;

    // The dimensions and precomputed values.
    params.b = b;
    params.h = h;
    params.bhs = b * h * s_inner;
    params.hs = h * s_inner;
    params.bs = b * s_inner;

    // The scaling factors for the int8 version to convert to/from float.
    params.scale_bmm1 = scale_bmm1;
    params.softcapping_scale_bmm1 = softcapping_scale_bmm1;
    params.scale_softmax = scale_softmax;
    // The number of warps_n used to identify the reduction strategy.
    params.warps_n = warps_n;

    // Compute the grid size.
    enum
    {
        WARPS_PER_CTA = 4
    };

    dim3 grid(s_outer, (h + WARPS_PER_CTA - 1) / WARPS_PER_CTA, b);
    dim3 threads_per_cta(THREADS_PER_WARP, WARPS_PER_CTA);

    // Launch the kernel.
    if (s_inner == 32)
    {
        softmax_kernel<Dst_type, Src_type, 32, WARPS_PER_CTA, 1><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 64)
    {
        softmax_kernel<Dst_type, Src_type, 64, WARPS_PER_CTA, 2><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 96)
    {
        softmax_kernel<Dst_type, Src_type, 96, WARPS_PER_CTA, 1><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 128)
    {
        softmax_kernel<Dst_type, Src_type, 128, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 192)
    {
        softmax_kernel<Dst_type, Src_type, 192, WARPS_PER_CTA, 2><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 256)
    {
        softmax_kernel<Dst_type, Src_type, 256, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 384)
    {
        softmax_kernel<Dst_type, Src_type, 384, WARPS_PER_CTA, 2><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 512)
    {
        softmax_kernel<Dst_type, Src_type, 512, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 1024)
    {
        softmax_kernel<Dst_type, Src_type, 1024, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 2048)
    {
        softmax_kernel<Dst_type, Src_type, 2048, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 4096)
    {
        softmax_kernel<Dst_type, Src_type, 4096, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 8192)
    {
        softmax_kernel<Dst_type, Src_type, 8192, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 16384)
    {
        softmax_kernel<Dst_type, Src_type, 16384, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 32768)
    {
        softmax_kernel<Dst_type, Src_type, 32768, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else if (s_inner == 65536)
    {
        softmax_kernel<Dst_type, Src_type, 65536, WARPS_PER_CTA><<<grid, threads_per_cta>>>(params);
    }
    else
    {
        assert(false);
    }
}

////////////////////////////////////////////////////////////////////////////////////////////////////