/*
 * 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.
 */

#pragma once

#include <fmha/softmax.h>
#include <fmha/traits.h>
#include <fmha/utils.h>

namespace fmha
{
namespace ws
{

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

// Special Softmax struct to handle optimization tricks on Hopper Warp-Specialized Kernels.
template <template <int, int, int, bool, bool> class Traits, typename Kernel_traits>
struct Softmax_base
{

    // The instruction traits for BMM1.
    using Traits_p = typename Kernel_traits::Traits_p;
    // The instruction traits for BMM2.
    using Traits_o = typename Kernel_traits::Traits_o;

    // The CTA description for BMM1.
    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
    // The CTA description for BMM2.
    using Cta_tile_o = typename Kernel_traits::Cta_tile_o;

    // The GMMA compute tile for BMM1.
    using Compute_tile_p = typename Kernel_traits::Compute_tile_p;
    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Kernel_traits::Compute_tile_o;

    // The MMA tile for the BMM1.
    using Mma_tile_p = typename Kernel_traits::Mma_tile_p;
    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Kernel_traits::Mma_tile_o;

    // The fragment of BMM1 output.
    using Fragment_p = typename Compute_tile_o::Fragment;

    // The step size of KV loop.
    enum
    {
        STEP_KV = Kernel_traits::STEP_KV
    };

    // Whether apply causal mask or not.
    enum
    {
        CAUSAL_MASK = Kernel_traits::CAUSAL_MASK
    };

    // Whether do we attend to the specific sliding window or chunk ?
    enum
    {
        SLIDING_OR_CHUNKED_ATTENTION = Kernel_traits::SLIDING_OR_CHUNKED_ATTENTION
    };

    // Are we applying alibi bias (drop FMA optimizations for accuracy reasons).
    enum
    {
        APPLY_ALIBI = Kernel_traits::APPLY_ALIBI
    };

    // Are we applying softcapping scale for qk products ?
    enum
    {
        ENABLE_BMM1_SOFTCAPPING_SCALE = Kernel_traits::ENABLE_BMM1_SOFTCAPPING_SCALE
    };

    // Do we use custom mask input ?
    enum
    {
        USE_CUSTOM_MASK = Kernel_traits::USE_CUSTOM_MASK
    };

    // Apply the exp2f optimization (fuse bmm1_scale and -max into FMAs).
    enum
    {
        EXP2F_OPTIMIZATION = Kernel_traits::EXP2F_OPTIMIZATION
    };

    // Whether we need to check if local_max could be -inf or not.
    enum
    {
        CHECK_IF_NEG_INF_EXISTS = SLIDING_OR_CHUNKED_ATTENTION || USE_CUSTOM_MASK
    };

    // Ctor.
    template <typename Params>
    inline __device__ Softmax_base(Params params, int tidx)
        : tidx_(tidx)
        , scale_bmm1_(params.scale_bmm1_d ? *params.scale_bmm1_d : params.scale_bmm1)
        , softcapping_scale_bmm1_(params.softcapping_scale_bmm1)
        , sliding_window_size_(params.sliding_window_size)
        , log2_chunked_attention_size_(params.log2_chunked_attention_size)
        , packed_mask_ptr_{reinterpret_cast<uint32_t*>(params.packed_mask_ptr)}
        , params_packed_mask_stride_in_bytes_{params.packed_mask_stride_in_bytes}
    {

        int warp = tidx / 32;
        int lane = tidx % 32;
        // The corresponding row/col for each thread after MMA.
        // fixed 4x1 warp layout.
        quad_col_ = lane % 4;
        if (CAUSAL_MASK)
        {
            quad_row_ = warp * 16 + lane / 4;
        }
    }

    // Compute the sliding window or chunk start.
    inline __device__ int compute_sliding_window_or_chunk_start(int row)
    {
        // If the chunked atteniton is used.
        if (log2_chunked_attention_size_ > 0)
        {
            // The attention chunk start.
            return (row >> log2_chunked_attention_size_) << log2_chunked_attention_size_;
        }
        else
        {
            // The sliding window start is the max of 0 and row - sliding_window_size.
            return max(0, row - sliding_window_size_);
        }
    }

    // Load the packed mask in global memory.
    inline __device__ void load_packed_mask(int row_offset, int col_offset)
    {
        if constexpr (USE_CUSTOM_MASK)
        {
            static_assert(Mma_tile_p::CORES_M == 2, "Not implemented!");
            // Note that row_offset takes sum_s into consideration.
            // row_offset % 64 == 0.
            // 32 bits per thread, so 128 bits (32 bytes) per mma row (4 threads).
            int64_t mask_row_offset_in_bytes = row_offset * params_packed_mask_stride_in_bytes_;
            // offset_in_bytes = (tidx_ * 32 + (col_offset / (16 * 4)) * 128 * 32) / 8.
            // note that col_offset % 64 == 0 here.
            int64_t mask_col_offset_in_bytes = tidx_ * 4 + col_offset * Cta_tile_p::THREADS_PER_CTA / 16;
            // add the two offsets for uint32 packed mask.
            int64_t mask_offset = (mask_row_offset_in_bytes + mask_col_offset_in_bytes) / 4;
            if constexpr (STEP_KV == 64)
            {
                // 32 bits (2 rows, 16 cols) are needed.
                packed_mask_.x = packed_mask_ptr_[mask_offset];
            }
            else if constexpr (STEP_KV == 128)
            {
                // 2 x 32 bits (4 rows, 16 cols) are needed.
                packed_mask_.x = packed_mask_ptr_[mask_offset];
                packed_mask_.y = packed_mask_ptr_[mask_offset + 128];
            }
            else if constexpr (STEP_KV == 256)
            {
                // 4 x 32 bits (4 rows, 16 cols) are needed.
                packed_mask_.x = packed_mask_ptr_[mask_offset];
                packed_mask_.y = packed_mask_ptr_[mask_offset + 128];
                packed_mask_.z = packed_mask_ptr_[mask_offset + 256];
                packed_mask_.w = packed_mask_ptr_[mask_offset + 384];
            }
        }
    }

    // Check if the two positions are valid or not.
    inline __device__ void valid_positions(int mi, int ni, bool& v0, bool& v1)
    {
        // Only need one uint32_t packed mask in this case.
        if constexpr (STEP_KV == 64)
        {
            // Packed mask input.
            v0 = packed_mask_.x & (1 << (ni * 4 + mi * 2 + 0));
            v1 = packed_mask_.x & (1 << (ni * 4 + mi * 2 + 1));
        }
        else if constexpr (STEP_KV == 128)
        {
            // Packed mask input.
            if (ni < 8)
            {
                v0 = packed_mask_.x & (1 << (ni * 4 + mi * 2 + 0));
                v1 = packed_mask_.x & (1 << (ni * 4 + mi * 2 + 1));
            }
            else
            {
                v0 = packed_mask_.y & (1 << (ni * 4 + mi * 2 + 0 - 32));
                v1 = packed_mask_.y & (1 << (ni * 4 + mi * 2 + 1 - 32));
            }
        }
        else if constexpr (STEP_KV == 256)
        {
            // KV step size is 256 in this case (i.e CORES_N = 32).
            if (ni < 8)
            {
                v0 = packed_mask_.x & (1 << (ni * 4 + mi * 2 + 0));
                v1 = packed_mask_.x & (1 << (ni * 4 + mi * 2 + 1));
            }
            else if (ni < 16)
            {
                v0 = packed_mask_.y & (1 << (ni * 4 + mi * 2 + 0 - 32));
                v1 = packed_mask_.y & (1 << (ni * 4 + mi * 2 + 1 - 32));
            }
            else if (ni < 24)
            {
                v0 = packed_mask_.z & (1 << (ni * 4 + mi * 2 + 0 - 64));
                v1 = packed_mask_.z & (1 << (ni * 4 + mi * 2 + 1 - 64));
            }
            else
            {
                v0 = packed_mask_.w & (1 << (ni * 4 + mi * 2 + 0 - 96));
                v1 = packed_mask_.w & (1 << (ni * 4 + mi * 2 + 1 - 96));
            }
        }
    }

    // Convert from bmm1 output fragments to floats.
    inline __device__ void unpack(Compute_tile_p& ctile_p)
    {
#pragma unroll
        for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
        {
#pragma unroll
            for (int ni = 0; ni < Mma_tile_p::CORES_N; ni++)
            {
                // Satfinite in case of overflow (due to fp16 accumulation).
                // When there is no alibi bias, we fuse bmm1_scale with -max by FMAs.
                uint32_t scaled_h2 = EXP2F_OPTIMIZATION
                    ? satfinite_h2(ctile_p.acc_[0][0].reg(ni * Mma_tile_p::CORES_M + mi))
                    : satfinite_h2(hmul2(ctile_p.acc_[0][0].reg(ni * Mma_tile_p::CORES_M + mi), scale_bmm1_));
                // Convert from half2 to float2.
                reinterpret_cast<float2*>(&elt_[mi][2 * ni])[0] = half2_to_float2(scaled_h2);
            }
        }
    }

    // Convert from bmm1 output fragments to floats.
    template <bool APPLY_MASK, typename AlibiParams>
    inline __device__ void apply_alibi_and_mask(Compute_tile_p& ctile_p, AlibiParams const& alibi_params,
        float const alibi_head_scale, int actual_seqlen, int row_offset, int col_offset)
    {
#pragma unroll
        for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
        {
#pragma unroll
            for (int ni = 0; ni < Mma_tile_p::CORES_N; ni++)
            {

                bool v0 = true, v1 = true;
                int col = 0;
                if constexpr (APPLY_MASK)
                {
                    // Custom mask input.
                    if constexpr (USE_CUSTOM_MASK)
                    {
                        // Packed mask input.
                        valid_positions(mi, ni, v0, v1);
                        // Causal mask.
                    }
                    else if constexpr (CAUSAL_MASK)
                    {
                        // Causal Mask: we have to apply mask before getting max.
                        int row = row_offset + quad_row_ + mi * 8;
                        col = col_offset + quad_col_ * 2 + ni * 8;
                        // Mask for the two N elements.
                        v0 = (col <= row);
                        v1 = (col + 1 <= row);

                        // Attend to the specific sliding window or chunk.
                        if constexpr (SLIDING_OR_CHUNKED_ATTENTION)
                        {
                            int sliding_window_or_chunk_start = compute_sliding_window_or_chunk_start(row);
                            v0 &= (col >= sliding_window_or_chunk_start);
                            v1 &= (col + 1 >= sliding_window_or_chunk_start);
                        }
                        // Dense(padding) mask.
                    }
                    else
                    {
                        col = col_offset + quad_col_ * 2 + ni * 8;
                        v0 = (col < actual_seqlen);
                        v1 = (col + 1 < actual_seqlen);
                    }
                }

                // The unpacked floats from the array.
                float2& f2 = reinterpret_cast<float2*>(&elt_[mi][2 * ni])[0];

                // Attention logit softcapping scale.
                // 1.0f / softcapping_scale has been fused into scale_bmm1.
                if constexpr (ENABLE_BMM1_SOFTCAPPING_SCALE)
                {
                    f2.x = softcapping_scale_bmm1_ * fmha::__tanhf(f2.x);
                    f2.y = softcapping_scale_bmm1_ * fmha::__tanhf(f2.y);
                }

                // Use minimum value of float here to avoid generating NANs with expf.
                if constexpr (APPLY_ALIBI)
                {
                    f2.x = v0 ? (f2.x * alibi_params.scale_after_alibi
                               + (col + alibi_params.sequence_pos_offset) * alibi_head_scale)
                              : -FLT_MAX;
                    f2.y = v1 ? (f2.y * alibi_params.scale_after_alibi
                               + (col + 1 + alibi_params.sequence_pos_offset) * alibi_head_scale)
                              : -FLT_MAX;
                }
                else
                {
                    f2.x = v0 ? f2.x : -FLT_MAX;
                    f2.y = v1 ? f2.y : -FLT_MAX;
                }
            }
        }
    }

    // Calculate max/sum, and update flash-attention scales.
    template <bool IS_FIRST_COL>
    inline __device__ void compute_and_update_scale(
        float (&global_max)[Mma_tile_p::CORES_M], float (&global_sum)[Mma_tile_p::CORES_M])
    {
        float const scale = reinterpret_cast<float const&>(scale_bmm1_);

// Row-wise max of current tile.
#pragma unroll
        for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
        {
            if (IS_FIRST_COL)
            {
                local_max_[mi] = elt_[mi][0];
            }
            else
            {
                local_max_[mi] = fmaxf(global_max[mi], elt_[mi][0]);
            }
#pragma unroll
            for (int ni = 1; ni < Mma_tile_p::CORES_N * 2; ni++)
            {
                local_max_[mi] = fmaxf(local_max_[mi], elt_[mi][ni]);
            }
            local_max_[mi] = fmaxf(__shfl_xor_sync(uint32_t(-1), local_max_[mi], 1), local_max_[mi]);
            local_max_[mi] = fmaxf(__shfl_xor_sync(uint32_t(-1), local_max_[mi], 2), local_max_[mi]);
        }

// Softmax Exp.
#pragma unroll
        for (int ni = 0; ni < Mma_tile_p::CORES_N; ni++)
        {
#pragma unroll
            for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
            {

                float& p0 = elt_[mi][2 * ni + 0];
                float& p1 = elt_[mi][2 * ni + 1];

                // When all elts of the tile are -FLT_MAX, we have to make sure
                //  expf generates 0 for all values instead of 1.
                if constexpr (!EXP2F_OPTIMIZATION)
                {
                    float masked_max = (!CHECK_IF_NEG_INF_EXISTS || local_max_[mi] != -FLT_MAX) ? local_max_[mi] : 0.f;
                    p0 = expf(p0 - masked_max);
                    p1 = expf(p1 - masked_max);
                }
                else
                {
                    // Use exp2f optimization for cases without alibi.
                    float masked_max
                        = (!CHECK_IF_NEG_INF_EXISTS || local_max_[mi] != -FLT_MAX) ? local_max_[mi] * scale : 0.f;
                    p0 = custom_exp2f(p0, scale, masked_max);
                    p1 = custom_exp2f(p1, scale, masked_max);
                }
            }
        }

// Row-wise sum of current tile.
#pragma unroll
        for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
        {
            local_sum_[mi] = elt_[mi][0];
#pragma unroll
            for (int ni = 1; ni < Mma_tile_p::CORES_N * 2; ni++)
            {
                local_sum_[mi] += elt_[mi][ni];
            }
            local_sum_[mi] += __shfl_xor_sync(uint32_t(-1), local_sum_[mi], 1);
            local_sum_[mi] += __shfl_xor_sync(uint32_t(-1), local_sum_[mi], 2);
        }

        // Initialize or update the global sum and max.
        if (IS_FIRST_COL)
        {
#pragma unroll
            for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
            {
                global_sum[mi] = local_sum_[mi];
                global_max[mi] = local_max_[mi];
            }
        }
        else
        {
#pragma unroll
            for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
            {
                float max_old = global_max[mi];
                float max_new = local_max_[mi];
                float sum_old = global_sum[mi];
                float sum_new = local_sum_[mi];
                // Remove the old max and replace by the new one.
                if constexpr (!EXP2F_OPTIMIZATION)
                {
                    correction_[mi] = expf(max_old - max_new);
                }
                else
                {
                    // Use exp2f optimization for cases without alibi.
                    correction_[mi] = exp2f((max_old - max_new) * scale);
                }
                global_sum[mi] = sum_old * correction_[mi] + sum_new;

                // New max already takes into account old max.
                global_max[mi] = max_new;
            }
        }
    }

    // Update flash attention scales and pack elements for BMM2.
    template <bool IS_FIRST_COL>
    inline __device__ void pack(Compute_tile_o& ctile_o, Fragment_p (&frag_p)[Mma_tile_o::MMAS_K])
    {

// Pack 4 cols for BMM2 A tile.
#pragma unroll
        for (int ni = 0; ni < Mma_tile_o::MMAS_K; ni++)
        {
            frag_p[ni].reg(0) = float2_to_half2(elt_[0][4 * ni + 0], elt_[0][4 * ni + 1]);
            frag_p[ni].reg(1) = float2_to_half2(elt_[1][4 * ni + 0], elt_[1][4 * ni + 1]);
            frag_p[ni].reg(2) = float2_to_half2(elt_[0][4 * ni + 2], elt_[0][4 * ni + 3]);
            frag_p[ni].reg(3) = float2_to_half2(elt_[1][4 * ni + 2], elt_[1][4 * ni + 3]);
        }

        if (!IS_FIRST_COL)
        {
// Correct accumulators to current max.
#pragma unroll
            for (int mi = 0; mi < Mma_tile_o::CORES_M; mi++)
            {
                uint32_t const scale = float_to_half2(correction_[mi]);

// Assume only N has multiple MMAs (MMAS_M = 1).
// MMAS_N > 1 when N dimension is split.
#pragma unroll
                for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
                {
#pragma unroll
                    for (int ni = 0; ni < Mma_tile_o::CORES_N; ni++)
                    {
                        uint32_t& reg = ctile_o.acc_[0][mma_ni].reg(ni * Mma_tile_o::CORES_M + mi);
                        reg = hmul2(reg, scale);
                    }
                }
            }
        }
        else
        {
            ctile_o.clear();
        }
    }

    // BMM1 scale.
    uint32_t const scale_bmm1_;
    // BMM1 softcapping scale.
    float const softcapping_scale_bmm1_;
    // The thread idx in the warp group.
    int tidx_;
    // The col index for the mma thread layout.
    int quad_col_;
    // The row index for the mma thread layout.
    int quad_row_;

    // The sliding window size.
    int const sliding_window_size_;
    // The log2 attention chunk size.
    int const log2_chunked_attention_size_;

    // The packed mask ptr.
    uint32_t const* packed_mask_ptr_;
    // The packed mask k-dim stride in bytes;
    int64_t const params_packed_mask_stride_in_bytes_;

    // Unpacked BMM1 output buffer.
    float elt_[Mma_tile_p::CORES_M][Mma_tile_p::CORES_N * 2];
    // Local max.
    float local_max_[Mma_tile_p::CORES_M];
    // Local sum.
    float local_sum_[Mma_tile_p::CORES_M];
    // Correction_ scales for ctil_o.
    float correction_[Mma_tile_p::CORES_M];
    // The packed mask.
    uint4 packed_mask_;
};

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

template <template <int, int, int, bool, bool> class Traits, typename Kernel_traits>
struct Softmax
{
};

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

// Hopper_hgmma_fp16_traits
template <typename Kernel_traits>
struct Softmax<Hopper_hgmma_fp16_traits, Kernel_traits> : public Softmax_base<Hopper_hgmma_fp16_traits, Kernel_traits>
{

    // The Base class.
    using Base = Softmax_base<Hopper_hgmma_fp16_traits, Kernel_traits>;

    // Ctor.
    template <typename Params>
    inline __device__ Softmax(Params const& params, int tidx)
        : Base(params, tidx)
    {
    }
};

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

// Fp32 accumulation traits.
template <template <int, int, int, bool, bool> class Traits, typename Kernel_traits>
struct Softmax_fp32_base : public Softmax_base<Traits, Kernel_traits>
{

    // The Base class.
    using Base = Softmax_base<Traits, Kernel_traits>;

    // The instruction traits for BMM1.
    using Traits_p = typename Base::Traits_p;
    // The instruction traits for BMM2.
    using Traits_o = typename Base::Traits_o;

    // The CTA description for BMM1.
    using Cta_tile_p = typename Base::Cta_tile_p;
    // The CTA description for BMM2.
    using Cta_tile_o = typename Base::Cta_tile_o;

    // The GMMA compute tile for BMM1.
    using Compute_tile_p = typename Base::Compute_tile_p;
    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Base::Compute_tile_o;

    // The MMA tile for the BMM1.
    using Mma_tile_p = typename Base::Mma_tile_p;
    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Base::Mma_tile_o;

    // The fragment of BMM1 output.
    using Fragment_p = typename Compute_tile_o::Fragment;

    // Whether apply causal mask or not.
    enum
    {
        CAUSAL_MASK = Base::CAUSAL_MASK
    };

    // Do we use custom mask input ?
    enum
    {
        USE_CUSTOM_MASK = Base::USE_CUSTOM_MASK
    };

    // Whether we attend to the specific sliding window or chunk ?
    enum
    {
        SLIDING_OR_CHUNKED_ATTENTION = Base::SLIDING_OR_CHUNKED_ATTENTION
    };

    // Are we applying alibi bias (drop FMA optimizations for accuracy reasons).
    enum
    {
        APPLY_ALIBI = Base::APPLY_ALIBI
    };

    // Are we applying softcapping_scale for qk products ?
    enum
    {
        ENABLE_BMM1_SOFTCAPPING_SCALE = Base::ENABLE_BMM1_SOFTCAPPING_SCALE
    };

    // Apply the exp2f optimization (fuse bmm1_scale and -max into FMAs).
    enum
    {
        EXP2F_OPTIMIZATION = Base::EXP2F_OPTIMIZATION
    };

    // Whether we need to check if local_max could be -inf or not.
    enum
    {
        CHECK_IF_NEG_INF_EXISTS = Base::CHECK_IF_NEG_INF_EXISTS
    };

    // Ctor.
    template <typename Params>
    inline __device__ Softmax_fp32_base(Params const& params, int tidx)
        : Base(params, tidx)
    {
    }

    // Convert from bmm1 output fragments to floats.
    inline __device__ void unpack(Compute_tile_p& ctile_p)
    {
#pragma unroll
        for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
        {
#pragma unroll
            for (int ni = 0; ni < Mma_tile_p::CORES_N; ni++)
            {
                float2 f2;
                f2.x = ctile_p.acc_[0][0].elt(2 * ni * Mma_tile_p::CORES_M + 2 * mi);
                f2.y = ctile_p.acc_[0][0].elt(2 * ni * Mma_tile_p::CORES_M + 2 * mi + 1);

                float const scale = reinterpret_cast<float const&>(this->scale_bmm1_);
                f2.x = EXP2F_OPTIMIZATION ? f2.x : f2.x * scale;
                f2.y = EXP2F_OPTIMIZATION ? f2.y : f2.y * scale;

                // Store to elt array.
                reinterpret_cast<float2*>(&this->elt_[mi][2 * ni])[0] = f2;
            }
        }
    }
};

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

// Hopper_hgmma_fp32_traits
template <typename Kernel_traits>
struct Softmax<Hopper_hgmma_fp32_traits, Kernel_traits>
    : public Softmax_fp32_base<Hopper_hgmma_fp32_traits, Kernel_traits>
{

    // The Base class.
    using Base = Softmax_fp32_base<Hopper_hgmma_fp32_traits, Kernel_traits>;

    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Base::Compute_tile_o;

    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Base::Mma_tile_o;

    // The fragment of BMM1 output.
    using Fragment_p = typename Compute_tile_o::Fragment;

    // Ctor.
    template <typename Params>
    inline __device__ Softmax(Params const& params, int tidx)
        : Base(params, tidx)
    {
    }

    // Update flash attention scales and pack elements for BMM2.
    template <bool IS_FIRST_COL>
    inline __device__ void pack(Compute_tile_o& ctile_o, Fragment_p (&frag_p)[Mma_tile_o::MMAS_K])
    {

// Pack 4 cols for BMM2 A tile.
#pragma unroll
        for (int ni = 0; ni < Mma_tile_o::MMAS_K; ni++)
        {
            frag_p[ni].reg(0) = float2_to_half2(this->elt_[0][4 * ni + 0], this->elt_[0][4 * ni + 1]);
            frag_p[ni].reg(1) = float2_to_half2(this->elt_[1][4 * ni + 0], this->elt_[1][4 * ni + 1]);
            frag_p[ni].reg(2) = float2_to_half2(this->elt_[0][4 * ni + 2], this->elt_[0][4 * ni + 3]);
            frag_p[ni].reg(3) = float2_to_half2(this->elt_[1][4 * ni + 2], this->elt_[1][4 * ni + 3]);
        }

        if (!IS_FIRST_COL)
        {
// Correct accumulators to current max.
#pragma unroll
            for (int mi = 0; mi < Mma_tile_o::CORES_M; mi++)
            {
                // Assume only N has multiple MMAs (MMAS_M = 1).
                // MMAS_N > 1 when N dimension is split.
                float correction = this->correction_[mi];
#pragma unroll
                for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
                {
#pragma unroll
                    for (int ni = 0; ni < Mma_tile_o::CORES_N; ni++)
                    {
                        uint32_t& reg0 = ctile_o.acc_[0][mma_ni].reg(2 * ni * Mma_tile_o::CORES_M + 2 * mi);
                        uint32_t& reg1 = ctile_o.acc_[0][mma_ni].reg(2 * ni * Mma_tile_o::CORES_M + 2 * mi + 1);
                        asm volatile("mul.f32 %0, %0, %1;\n" : "+r"(reg0) : "f"(correction));
                        asm volatile("mul.f32 %0, %0, %1;\n" : "+r"(reg1) : "f"(correction));
                    }
                }
            }
        }
        else
        {
            ctile_o.clear();
        }
    }
};

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

// Hopper_hgmma_bf16_traits
template <typename Kernel_traits>
struct Softmax<Hopper_hgmma_bf16_traits, Kernel_traits>
    : public Softmax_fp32_base<Hopper_hgmma_bf16_traits, Kernel_traits>
{

    // The Base class.
    using Base = Softmax_fp32_base<Hopper_hgmma_bf16_traits, Kernel_traits>;

    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Base::Compute_tile_o;

    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Base::Mma_tile_o;

    // The fragment of BMM1 output.
    using Fragment_p = typename Compute_tile_o::Fragment;

    // Ctor.
    template <typename Params>
    inline __device__ Softmax(Params const& params, int tidx)
        : Base(params, tidx)
    {
    }

    // Update flash attention scales and pack elements for BMM2.
    template <bool IS_FIRST_COL>
    inline __device__ void pack(Compute_tile_o& ctile_o, Fragment_p (&frag_p)[Mma_tile_o::MMAS_K])
    {

// Pack 4 cols for BMM2 A tile.
#pragma unroll
        for (int ni = 0; ni < Mma_tile_o::MMAS_K; ni++)
        {
            frag_p[ni].reg(0) = float2_to_bf16_x2(this->elt_[0][4 * ni + 0], this->elt_[0][4 * ni + 1]);
            frag_p[ni].reg(1) = float2_to_bf16_x2(this->elt_[1][4 * ni + 0], this->elt_[1][4 * ni + 1]);
            frag_p[ni].reg(2) = float2_to_bf16_x2(this->elt_[0][4 * ni + 2], this->elt_[0][4 * ni + 3]);
            frag_p[ni].reg(3) = float2_to_bf16_x2(this->elt_[1][4 * ni + 2], this->elt_[1][4 * ni + 3]);
        }

        if (!IS_FIRST_COL)
        {
// Correct accumulators to current max.
#pragma unroll
            for (int mi = 0; mi < Mma_tile_o::CORES_M; mi++)
            {
                // Assume only N has multiple MMAs (MMAS_M = 1).
                // MMAS_N > 1 when N dimension is split.
                float correction = this->correction_[mi];
#pragma unroll
                for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
                {
#pragma unroll
                    for (int ni = 0; ni < Mma_tile_o::CORES_N; ni++)
                    {
                        uint32_t& reg0 = ctile_o.acc_[0][mma_ni].reg(2 * ni * Mma_tile_o::CORES_M + 2 * mi);
                        uint32_t& reg1 = ctile_o.acc_[0][mma_ni].reg(2 * ni * Mma_tile_o::CORES_M + 2 * mi + 1);
                        asm volatile("mul.f32 %0, %0, %1;\n" : "+r"(reg0) : "f"(correction));
                        asm volatile("mul.f32 %0, %0, %1;\n" : "+r"(reg1) : "f"(correction));
                    }
                }
            }
        }
        else
        {
            ctile_o.clear();
        }
    }
};

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

// Hopper_qgmma_e4m3_fp32_traits
template <typename Kernel_traits>
struct Softmax<Hopper_qgmma_e4m3_fp32_traits, Kernel_traits>
    : public Softmax_fp32_base<Hopper_qgmma_e4m3_fp32_traits, Kernel_traits>
{

    // The Base class.
    using Base = Softmax_fp32_base<Hopper_qgmma_e4m3_fp32_traits, Kernel_traits>;

    // The instruction traits for BMM1.
    using Traits_p = typename Base::Traits_p;
    // The instruction traits for BMM2.
    using Traits_o = typename Base::Traits_o;

    // The CTA description for BMM1.
    using Cta_tile_p = typename Base::Cta_tile_p;
    // The CTA description for BMM2.
    using Cta_tile_o = typename Base::Cta_tile_o;

    // The GMMA compute tile for BMM1.
    using Compute_tile_p = typename Base::Compute_tile_p;
    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Base::Compute_tile_o;

    // The MMA tile for the BMM1.
    using Mma_tile_p = typename Base::Mma_tile_p;
    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Base::Mma_tile_o;

    // The fragment of BMM1 output.
    using Fragment_p = typename Compute_tile_o::Fragment;

    // Whether apply causal mask or not.
    enum
    {
        CAUSAL_MASK = Base::CAUSAL_MASK
    };

    // Whether we attend to the specific sliding window or chunk ?
    enum
    {
        SLIDING_OR_CHUNKED_ATTENTION = Base::SLIDING_OR_CHUNKED_ATTENTION
    };

    // Are we applying alibi bias (drop FMA optimizations for accuracy reasons).
    enum
    {
        APPLY_ALIBI = Base::APPLY_ALIBI
    };

    // Apply the exp2f optimization (fuse bmm1_scale and -max into FMAs).
    enum
    {
        EXP2F_OPTIMIZATION = Base::EXP2F_OPTIMIZATION
    };

    // Are we applying softcapping_scale for qk products ?
    enum
    {
        ENABLE_BMM1_SOFTCAPPING_SCALE = Base::ENABLE_BMM1_SOFTCAPPING_SCALE
    };

    // Whether we need to check if local_max could be -inf or not.
    enum
    {
        CHECK_IF_NEG_INF_EXISTS = Base::CHECK_IF_NEG_INF_EXISTS
    };

    // Ctor.
    template <typename Params>
    inline __device__ Softmax(Params const& params, int tidx)
        : Base(params, tidx)
    {
    }

    // Calculate max/sum, and update flash-attention scales.
    template <bool IS_FIRST_COL>
    inline __device__ void compute_and_update_scale(
        float (&global_max)[Mma_tile_p::CORES_M], float (&global_sum)[Mma_tile_p::CORES_M])
    {
        float const scale = reinterpret_cast<float const&>(this->scale_bmm1_);
        float(&local_max_)[Mma_tile_p::CORES_M] = this->local_max_;
        float(&local_sum_)[Mma_tile_p::CORES_M] = this->local_sum_;
        float(&correction_)[Mma_tile_p::CORES_M] = this->correction_;
        float(&elt_)[Mma_tile_p::CORES_M][Mma_tile_p::CORES_N * 2] = this->elt_;

// Row-wise max of current tile.
#pragma unroll
        for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
        {
            if (IS_FIRST_COL)
            {
                local_max_[mi] = elt_[mi][0];
            }
            else
            {
                local_max_[mi] = fmaxf(global_max[mi], elt_[mi][0]);
            }
#pragma unroll
            for (int ni = 1; ni < Mma_tile_p::CORES_N * 2; ni++)
            {
                local_max_[mi] = fmaxf(local_max_[mi], elt_[mi][ni]);
            }
            local_max_[mi] = fmaxf(__shfl_xor_sync(uint32_t(-1), local_max_[mi], 1), local_max_[mi]);
            local_max_[mi] = fmaxf(__shfl_xor_sync(uint32_t(-1), local_max_[mi], 2), local_max_[mi]);
        }

// Softmax Exp.
#pragma unroll
        for (int ni = 0; ni < Mma_tile_p::CORES_N; ni++)
        {
#pragma unroll
            for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
            {
                // The equation exp2(scale * x - max) * q_scale_s
                // equals to:   exp2(scale * x - max) * exp2(log2(q_scale_s))
                // equals to:   exp2(scale * x - (max - log2(q_scale_s)))
                //                   ^^^^^   ^   ^^^^^^^^^^^^^^^^^^^^^^^
                // So instead of per-accumulator muls, we can do per-row subs which saves FP cycles.
                // As we scale the softmax output early, before doing the local_sum, we have to unscale
                // the local_sum afterwards.
                float& p0 = elt_[mi][2 * ni + 0];
                float& p1 = elt_[mi][2 * ni + 1];

                // When all elts of the tile are -FLT_MAX, we have to make sure
                //  expf generates 0 for all values instead of 1.
                if constexpr (!EXP2F_OPTIMIZATION)
                {
                    float masked_max = (!CHECK_IF_NEG_INF_EXISTS || local_max_[mi] != -FLT_MAX)
                        ? local_max_[mi] - logf(q_scale_s_)
                        : 0.f;
                    p0 = expf(p0 - masked_max);
                    p1 = expf(p1 - masked_max);
                }
                else
                {
                    // Use exp2f optimization for cases without alibi.
                    float masked_max = (!CHECK_IF_NEG_INF_EXISTS || local_max_[mi] != -FLT_MAX)
                        ? local_max_[mi] * scale - log2f(q_scale_s_)
                        : 0.f;
                    p0 = custom_exp2f(p0, scale, masked_max);
                    p1 = custom_exp2f(p1, scale, masked_max);
                }
            }
        }

// Row-wise sum of current tile.
#pragma unroll
        for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
        {
            local_sum_[mi] = elt_[mi][0];
#pragma unroll
            for (int ni = 1; ni < Mma_tile_p::CORES_N * 2; ni++)
            {
                local_sum_[mi] += elt_[mi][ni];
            }
            local_sum_[mi] += __shfl_xor_sync(uint32_t(-1), local_sum_[mi], 1);
            local_sum_[mi] += __shfl_xor_sync(uint32_t(-1), local_sum_[mi], 2);
        }

        // Initialize or update the global sum and max.
        if (IS_FIRST_COL)
        {
#pragma unroll
            for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
            {
                global_sum[mi] = local_sum_[mi];
                global_max[mi] = local_max_[mi];
            }
        }
        else
        {
#pragma unroll
            for (int mi = 0; mi < Mma_tile_p::CORES_M; mi++)
            {
                float max_old = global_max[mi];
                float max_new = local_max_[mi];
                float sum_old = global_sum[mi];
                float sum_new = local_sum_[mi];
                // Remove the old max and replace by the new one.
                if constexpr (!EXP2F_OPTIMIZATION)
                {
                    correction_[mi] = expf(max_old - max_new);
                }
                else
                {
                    // Use exp2f optimization for cases without alibi.
                    correction_[mi] = exp2f((max_old - max_new) * scale);
                }
                global_sum[mi] = sum_old * correction_[mi] + sum_new;

                // New max already takes into account old max.
                global_max[mi] = max_new;
            }
        }
    }

    // Update flash attention scales and pack elements for BMM2.
    template <bool IS_FIRST_COL>
    inline __device__ void pack(Compute_tile_o& ctile_o, Fragment_p (&frag_p)[Mma_tile_o::MMAS_K])
    {

// Pack 4 cols for BMM2 A tile.
#pragma unroll
        for (int ni = 0; ni < Mma_tile_o::MMAS_K; ni++)
        {

            // 1st row - 8 elements per row.
            float tmp_00 = this->elt_[0][8 * ni + 0]; // + 0
            float tmp_01 = this->elt_[0][8 * ni + 1]; // + 1
            float tmp_02 = this->elt_[0][8 * ni + 2]; // + 8
            float tmp_03 = this->elt_[0][8 * ni + 3]; // + 9
            float tmp_04 = this->elt_[0][8 * ni + 4]; // +16
            float tmp_05 = this->elt_[0][8 * ni + 5]; // +17
            float tmp_06 = this->elt_[0][8 * ni + 6]; // +24
            float tmp_07 = this->elt_[0][8 * ni + 7]; // +25

            // 2nd row - 8 elements per row.
            float tmp_10 = this->elt_[1][8 * ni + 0]; // + 0
            float tmp_11 = this->elt_[1][8 * ni + 1]; // + 1
            float tmp_12 = this->elt_[1][8 * ni + 2]; // + 8
            float tmp_13 = this->elt_[1][8 * ni + 3]; // + 9
            float tmp_14 = this->elt_[1][8 * ni + 4]; // +16
            float tmp_15 = this->elt_[1][8 * ni + 5]; // +17
            float tmp_16 = this->elt_[1][8 * ni + 6]; // +24
            float tmp_17 = this->elt_[1][8 * ni + 7]; // +25

            // Pack to 4 registers.
            frag_p[ni].reg(0) = fmha::float4_to_fp8x4<Kernel_traits::Element_data_type>(tmp_00, tmp_01, tmp_02, tmp_03);
            frag_p[ni].reg(1) = fmha::float4_to_fp8x4<Kernel_traits::Element_data_type>(tmp_10, tmp_11, tmp_12, tmp_13);
            frag_p[ni].reg(2) = fmha::float4_to_fp8x4<Kernel_traits::Element_data_type>(tmp_04, tmp_05, tmp_06, tmp_07);
            frag_p[ni].reg(3) = fmha::float4_to_fp8x4<Kernel_traits::Element_data_type>(tmp_14, tmp_15, tmp_16, tmp_17);
        }

        if (!IS_FIRST_COL)
        {
// Correct accumulators to current max.
#pragma unroll
            for (int mi = 0; mi < Mma_tile_o::CORES_M; mi++)
            {
                // Assume only N has multiple MMAs (MMAS_M = 1).
                // MMAS_N > 1 when N dimension is split.
                float correction = this->correction_[mi];
#pragma unroll
                for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
                {
#pragma unroll
                    for (int ni = 0; ni < Mma_tile_o::CORES_N; ni++)
                    {
                        uint32_t& reg0 = ctile_o.acc_[0][mma_ni].reg(2 * ni * Mma_tile_o::CORES_M + 2 * mi);
                        uint32_t& reg1 = ctile_o.acc_[0][mma_ni].reg(2 * ni * Mma_tile_o::CORES_M + 2 * mi + 1);
                        asm volatile("mul.f32 %0, %0, %1;\n" : "+r"(reg0) : "f"(correction));
                        asm volatile("mul.f32 %0, %0, %1;\n" : "+r"(reg1) : "f"(correction));
                    }
                }
            }
        }
        else
        {
            ctile_o.clear();
        }
    }

    static constexpr float q_scale_s_ = Traits_o::SOFTMAX_FP_QUANT_SCALE;
};

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

// BMM2 epilogue to apply scales (flash attention).
// FP32 accumulation as default.
template <template <int, int, int, bool, bool> class Traits, typename Kernel_traits>
struct Tile_o_epilogue_base
{

    // The instruction traits for BMM2.
    using Traits_o = typename Kernel_traits::Traits_o;

    // The CTA description for BMM2.
    using Cta_tile_o = typename Kernel_traits::Cta_tile_o;

    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Kernel_traits::Compute_tile_o;

    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Kernel_traits::Mma_tile_o;

    template <typename Params>
    inline __device__ Tile_o_epilogue_base(Params const& params)
    {
        ; // nothing to construct.
    };

    // Scale ctile_o output by 1/sum
    inline __device__ void scale(Compute_tile_o& ctile_o, float (&global_sum)[Mma_tile_o::CORES_M])
    {
// Final step's update.
#pragma unroll
        for (int mi = 0; mi < Mma_tile_o::CORES_M; mi++)
        {
            global_sum[mi] = global_sum[mi] == 0.f ? 1.f : 1.0f / global_sum[mi];

// Assume only N has multiple MMAs (MMAS_M = 1).
#pragma unroll
            for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
            {
#pragma unroll
                for (int ni = 0; ni < Mma_tile_o::CORES_N; ni++)
                {
                    float& reg0 = ctile_o.acc_[0][mma_ni].elt(2 * ni * Mma_tile_o::CORES_M + 2 * mi);
                    float& reg1 = ctile_o.acc_[0][mma_ni].elt(2 * ni * Mma_tile_o::CORES_M + 2 * mi + 1);
                    reg0 *= global_sum[mi];
                    reg1 *= global_sum[mi];
                }
            }
        }
    }
};

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

template <template <int, int, int, bool, bool> class Traits, typename Kernel_traits>
struct Tile_o_epilogue
{
};

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

// Hopper_hgmma_fp16_traits
template <typename Kernel_traits>
struct Tile_o_epilogue<Hopper_hgmma_fp16_traits, Kernel_traits>
    : public Tile_o_epilogue_base<Hopper_hgmma_fp16_traits, Kernel_traits>
{

    // The Base class.
    using Base = Tile_o_epilogue_base<Hopper_hgmma_fp16_traits, Kernel_traits>;

    // The instruction traits for BMM2.
    using Traits_o = typename Base::Traits_o;

    // The CTA description for BMM2.
    using Cta_tile_o = typename Base::Cta_tile_o;

    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Base::Compute_tile_o;

    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Base::Mma_tile_o;

    // Base constructor.
    using Base::Tile_o_epilogue_base;

    // Scale ctile_o output by 1/sum
    inline __device__ void scale(Compute_tile_o& ctile_o, float (&global_sum)[Mma_tile_o::CORES_M])
    {
// Final step's update.
#pragma unroll
        for (int mi = 0; mi < Mma_tile_o::CORES_M; mi++)
        {
            global_sum[mi] = global_sum[mi] == 0.f ? 1.f : 1.0f / global_sum[mi];
            uint32_t const scale = float_to_half2(global_sum[mi]);

// Assume only N has multiple MMAs (MMAS_M = 1).
#pragma unroll
            for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
            {
#pragma unroll
                for (int ni = 0; ni < Mma_tile_o::CORES_N; ni++)
                {
                    uint32_t& reg = ctile_o.acc_[0][mma_ni].reg(ni * Mma_tile_o::CORES_M + mi);
                    reg = hmul2(reg, scale);
                }
            }
        }
    }
};

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

// Hopper_hgmma_fp32_traits
template <typename Kernel_traits>
struct Tile_o_epilogue<Hopper_hgmma_fp32_traits, Kernel_traits>
    : public Tile_o_epilogue_base<Hopper_hgmma_fp32_traits, Kernel_traits>
{

    // The Base class.
    using Base = Tile_o_epilogue_base<Hopper_hgmma_fp32_traits, Kernel_traits>;

    // Base constructor.
    using Base::Tile_o_epilogue_base;
};

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

// Hopper_hgmma_bf16_traits
template <typename Kernel_traits>
struct Tile_o_epilogue<Hopper_hgmma_bf16_traits, Kernel_traits>
    : public Tile_o_epilogue_base<Hopper_hgmma_bf16_traits, Kernel_traits>
{

    // The Base class.
    using Base = Tile_o_epilogue_base<Hopper_hgmma_bf16_traits, Kernel_traits>;

    // Base constructor.
    using Base::Tile_o_epilogue_base;
};

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

// Hopper_qgmma_e4m3_fp32_traits
template <typename Kernel_traits>
struct Tile_o_epilogue<Hopper_qgmma_e4m3_fp32_traits, Kernel_traits>
    : public Tile_o_epilogue_base<Hopper_qgmma_e4m3_fp32_traits, Kernel_traits>
{

    // The Base class.
    using Base = Tile_o_epilogue_base<Hopper_qgmma_e4m3_fp32_traits, Kernel_traits>;

    // The instruction traits for BMM2.
    using Traits_o = typename Base::Traits_o;

    // The CTA description for BMM2.
    using Cta_tile_o = typename Base::Cta_tile_o;

    // The GMMA compute tile for BMM2.
    using Compute_tile_o = typename Base::Compute_tile_o;

    // The MMA tile for the BMM2.
    using Mma_tile_o = typename Base::Mma_tile_o;

    // Ctor.
    template <typename Params>
    inline __device__ Tile_o_epilogue(Params const& params)
        : Base(params)
        , scale_bmm2_(*params.scale_bmm2_d)
    {
    }

    // Scale ctile_o output by 1/sum
    inline __device__ void scale(Compute_tile_o& ctile_o, float (&global_sum)[Mma_tile_o::CORES_M])
    {
// Final step's update.
#pragma unroll
        for (int mi = 0; mi < Mma_tile_o::CORES_M; mi++)
        {
#ifdef UNIFIED_EPILOGUE_SCALE
            // Descaling factor
            float const scale_bmm2_f_ = reinterpret_cast<float&>(scale_bmm2_);
            global_sum[mi] = global_sum[mi] == 0.f ? scale_bmm2_f_ : scale_bmm2_f_ / global_sum[mi];
#else
            global_sum[mi] = global_sum[mi] == 0.f ? 1.0f : 1.0f / global_sum[mi];
#endif
// Assume only N has multiple MMAs (MMAS_M = 1).
#pragma unroll
            for (int mma_ni = 0; mma_ni < Mma_tile_o::MMAS_N; mma_ni++)
            {
#pragma unroll
                for (int ni = 0; ni < Mma_tile_o::CORES_N; ni++)
                {
                    float& reg0 = ctile_o.acc_[0][mma_ni].elt(2 * ni * Mma_tile_o::CORES_M + 2 * mi);
                    float& reg1 = ctile_o.acc_[0][mma_ni].elt(2 * ni * Mma_tile_o::CORES_M + 2 * mi + 1);
                    reg0 *= global_sum[mi];
                    reg1 *= global_sum[mi];
                }
            }
        }
    }

    // Scale_bmm2.
    uint32_t scale_bmm2_;
};

} // namespace ws
} // namespace fmha
