/*
 * 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/alibi_params.h>
#include <fmha/fragment.h>
#include <fmha/gemm.h>
#include <fmha/gmem_tile_o.h>
#include <fmha/gmem_tile_o_packed.h>
#include <fmha/gmem_tile_ps.h>
#include <fmha/gmem_tile_qkv.h>
#include <fmha/gmem_tile_qkv_packed.h>
#include <fmha/smem_tile_o.h>
#include <fmha/smem_tile_qkv.h>
#include <fmha/smem_tile_v.h>
#include <fmha/softmax.h>
#include <fmha/traits.h>
#include <fmha/utils.h>

namespace fmha
{

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

// Ada hmma/imma reuses Ampere
template <typename Traits_>
struct Traits_reuse
{
    using Traits = Traits_;
};

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

template <>
struct Traits_reuse<fmha::Ada_hmma_fp16_traits>
{
    using Traits = fmha::Ampere_hmma_fp16_traits;
};

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

template <>
struct Traits_reuse<fmha::Ada_hmma_fp32_traits>
{
    using Traits = fmha::Ampere_hmma_fp32_traits;
};

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

template <>
struct Traits_reuse<fmha::Ada_imma_int8_int32_traits>
{
    using Traits = fmha::Ampere_imma_int8_int32_traits;
};

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

template <typename Traits_p, bool FORCE_EPILOGUE_FP16>
struct Traits_o_adapter
{
    using Traits = Traits_p;
};

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

template <bool FORCE_EPILOGUE_FP16>
struct Traits_o_adapter<fmha::Volta_hmma_fp16_traits, FORCE_EPILOGUE_FP16>
{
    using Traits = fmha::Volta_hmma_fp16_16x16x16_traits;
};

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

// convert to fp16 before smem_o store
template <>
struct Traits_o_adapter<fmha::Ampere_hmma_fp32_traits, true>
{
    using Traits = fmha::Ampere_hmma_fp16_traits;
};

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

// convert to fp16 before smem_o store
template <>
struct Traits_o_adapter<fmha::Turing_hmma_fp32_traits, true>
{
    using Traits = fmha::Turing_hmma_fp16_traits;
};

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

// convert to bf16 before smem_o store
template <>
struct Traits_o_adapter<fmha::Ampere_hmma_bf16_traits, true>
{
    using Traits = fmha::Ampere_hmma_bf16_bf16_traits;
};

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

template <
    // Instruction traits.
    typename Traits_,
    // The global memory tile for Q, K and V.
    template <typename, typename, int, int, int, int, bool, bool, int, bool> class Gmem_tile_q_,
    template <typename, typename, int, int, int, int, bool, bool, int, bool> class Gmem_tile_k_,
    template <typename, typename, int, int, int, int, bool, bool, int, bool> class Gmem_tile_v_,
    // The global memory tile for the output.
    template <typename, typename, int> class Gmem_tile_o_,
    // Sequence length.
    int S,
    // The valid hidden dimension.
    int VALID_D_,
    // The valid hidden dimension of V.
    int VALID_DV_,
    // The iteration step of the outer loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD_,
    // The flags to control the behaviour of LDGs.
    uint32_t FLAGS,
    // The version of the kernel.
    int VERSION_,
    // The mask version of the kernel
    int MASK_VERSION_,
    // Do we use half epilogue for the 2nd GEMM (hmma_fp32)
    bool BMM2_FP16_EPILOGUE = true,
    // non-positive means disabled
    int SAGE_BLOCK_SIZE_Q_ = 0, int SAGE_BLOCK_SIZE_K_ = 0, int SAGE_BLOCK_SIZE_V_ = 0>
struct Kernel_traits_
{

    // The instruction traits for the Q*K product.
    using Traits_p = typename Traits_reuse<Traits_>::Traits;
    // The instruction traits for the P*V product. Hack to change the traits for Volta HMMA.
    using Traits_o = typename Traits_o_adapter<Traits_p, false>::Traits;
    // The instruction traits for the epilogue of the 2nd GEMM. Always use FP16.
    using Traits_e = typename Traits_o_adapter<Traits_p, BMM2_FP16_EPILOGUE>::Traits;

    // The padded D dimension
    enum
    {
        VALID_D = VALID_D_
    };

    enum
    {
        D = Next_power_of_two<VALID_D>::VALUE
    };

    enum
    {
        VALID_DV = VALID_DV_ > 0 ? VALID_DV_ : VALID_D
    };

    enum
    {
        DV = Next_power_of_two<VALID_DV>::VALUE
    };

    enum
    {
        SAGE_ATTENTION = SAGE_BLOCK_SIZE_Q_ > 0 || SAGE_BLOCK_SIZE_K_ > 0 || SAGE_BLOCK_SIZE_V_ > 0
    };

    enum
    {
        SAGE_BLOCK_SIZE_Q = SAGE_BLOCK_SIZE_Q_
    };

    enum
    {
        SAGE_BLOCK_SIZE_K = SAGE_BLOCK_SIZE_K_
    };

    enum
    {
        SAGE_BLOCK_SIZE_V = SAGE_BLOCK_SIZE_V_
    };

    // TODO: expose these tiling params to the interface
    enum
    {
        USE_GRANULAR_TILING = (FLAGS & 0x1000) != 0u
    }; // TODO ANT: check FLAGS

    using Traits_tile_size = Traits_tile_size<(bool) USE_GRANULAR_TILING, STEP, S, D, DV, Traits_o::K_PER_MMA>;

    enum
    {
        CTA_P_TILE_M = Traits_tile_size::CTA_P_TILE_M
    };

    enum
    {
        CTA_P_TILE_N = Traits_tile_size::CTA_P_TILE_N
    };

    enum
    {
        CTA_P_TILE_K = Traits_tile_size::CTA_P_TILE_K
    };

    enum
    {
        CTA_O_TILE_M = Traits_tile_size::CTA_O_TILE_M
    };

    enum
    {
        CTA_O_TILE_N = Traits_tile_size::CTA_O_TILE_N
    };

    enum
    {
        CTA_O_TILE_K = Traits_tile_size::CTA_O_TILE_K
    };

    // Do we need to reload Q due to splitting the D ?
    enum
    {
        RELOAD_Q = static_cast<int>(CTA_P_TILE_K) != static_cast<int>(D)
    };

    // The CTA description for the 1st GEMM.
    using Cta_tile_p = typename Traits_p::template Cta_tile_extd<CTA_P_TILE_M, CTA_P_TILE_N, CTA_P_TILE_K, S, VALID_D,
        WARPS_M, WARPS_N, 1>;
    // The CTA description for the 2nd GEMM.
    using Cta_tile_o = typename Traits_o::template Cta_tile_extd<CTA_O_TILE_M, CTA_O_TILE_N, CTA_O_TILE_K, VALID_DV, S,
        WARPS_M, 1, WARPS_N>;

    // The MMA tile for the 1st GEMM.
    using Mma_tile_p = typename Traits_p::template Mma_tile<Cta_tile_p>;
    // The MMA tile for the 2nd GEMM.
    using Mma_tile_o = typename Traits_o::template Mma_tile<Cta_tile_o>;

    // Compute the total BMM2_MMAS_K (might not the same as Mma_tile_o::MMAS_K if the granular tiling is used).
    static_assert(S % CTA_O_TILE_K == 0, "");

    enum
    {
        TOTAL_BMM2_MMAS_K = Mma_tile_o::MMAS_K * (S / CTA_O_TILE_K)
    };

    // Constraints on the K dimension.
    static_assert(Mma_tile_p::K_PER_MMA <= static_cast<int>(D));
    static_assert(Mma_tile_o::K_PER_MMA <= S);

    // The version.
    enum
    {
        VERSION = VERSION_
    };

    // The mask version: padding (2), causal (3), sliding_window_causal (4), custom_mask (5).
    enum
    {
        MASK_VERSION = MASK_VERSION_
    };

    // Whether use causal mask or not.
    enum
    {
        CAUSAL_MASK = MASK_VERSION_ == 3 || MASK_VERSION_ == 4
    };

    // Whether use the sliding window attention or not.
    enum
    {
        SLIDING_WINDOW_ATTENTION = MASK_VERSION_ == 4
    };

    // Whether use the custom mask or not.
    enum
    {
        CUSTOM_MASK = MASK_VERSION_ == 5
    };

    // Do we use LDGSTS for Q, K or V.
    enum
    {
        USE_LDGSTS_Q = (FLAGS & 0x1u) != 0u
    };

    enum
    {
        USE_LDGSTS_K = (FLAGS & 0x2u) != 0u
    };

    enum
    {
        USE_LDGSTS_V = (FLAGS & 0x4u) != 0u
    };

    // Do we use one buffer for K and V.
    enum
    {
        SHARE_SMEM_FOR_K_AND_V = (FLAGS & 0x8u) != 0u
    };

    // Do we use the scale max trick.
    enum
    {
        USE_SCALE_MAX = (FLAGS & 0x10u) != 0u
    };

    // Are heads in QKV interleaved, i.e. total x h x 3 x d or total x 3 x h x d.
    enum
    {
        HEADS_INTERLEAVED = (FLAGS & 0x20u) == 0u
    };

    // Keep full K matrix in registers.
    enum
    {
        K_IN_REGS = (FLAGS & 0x40) == 0u
    };

    // Do we use only 2 fragments or full fragments for frag_q/k (only used by flash attention)
    enum
    {
        LIMIT_QK_FRAGMENTS = ((FLAGS & 0x80u) != 0u && !SHARE_SMEM_FOR_K_AND_V)
    };

    // Do we use only 2 fragments or full fragments for frag_v (only used by flash attention)
    enum
    {
        LIMIT_V_FRAGMENTS = ((FLAGS & 0x100u) != 0u && !SHARE_SMEM_FOR_K_AND_V)
    };

    // Limiting QK fragments implies SMEM_K has to reside in SMEM
    static_assert(!(LIMIT_QK_FRAGMENTS && SHARE_SMEM_FOR_K_AND_V), "");

    // Indicates that kernel does not loop over Q tensor, usually kernel name has _nl suffix
    enum
    {
        NO_LOOP = (FLAGS & 0x200u) != 0u
    };

    // Are sequences in one batch interleaved. i.e. s x b x ..., or b x s x ...
    enum
    {
        SEQUENCES_INTERLEAVED = (FLAGS & 0x400) != 0u
    };

    // Use BMM1 softcapping scale or not.
    enum
    {
        ENABLE_BMM1_SOFTCAPPING_SCALE = (FLAGS & 0x800) != 0u
    };

    // Use MTP (multi-token prediction for MLA kernels) or not.
    enum
    {
        IS_MTP = (FLAGS & 0x2000) != 0u
    };

    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    enum
    {
        CTAS_PER_HEAD = CTAS_PER_HEAD_
    };

    // The number of shared memory buffers to build a software pipeline for Q, K and V.
    enum
    {
        BUFFERS_PER_TILE_SMEM_Q = (USE_GRANULAR_TILING && D > 64) || (USE_LDGSTS_Q && !NO_LOOP) ? 2 : 1
    };

    enum
    {
        BUFFERS_PER_TILE_SMEM_K = USE_GRANULAR_TILING ? 2 : 1
    };

    enum
    {
        BUFFERS_PER_TILE_SMEM_V = USE_GRANULAR_TILING ? 2 : 1
    };

    // The global memory tile to load Q.
    using Gmem_tile_q = Gmem_tile_q_<Traits_p, Cta_tile_p, Traits_p::BITS_PER_ELEMENT_A, CTA_P_TILE_M, CTA_P_TILE_K,
        VALID_D, USE_LDGSTS_Q, HEADS_INTERLEAVED,
        3,                       // NUM_MATS
        SLIDING_WINDOW_ATTENTION // Not used.
        >;

    // The shared memory tile to swizzle Q.
    using Smem_tile_q
        = fmha::Smem_tile_a<Traits_p, Cta_tile_p, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, BUFFERS_PER_TILE_SMEM_Q>;

    // The global memory tile to load K.
    using Gmem_tile_k = Gmem_tile_k_<Traits_p, Cta_tile_p, Traits_p::BITS_PER_ELEMENT_B, CTA_P_TILE_N, CTA_P_TILE_K,
        VALID_D, USE_LDGSTS_K, HEADS_INTERLEAVED,
        3, // NUM_MATS
        SLIDING_WINDOW_ATTENTION>;

    // The shared memory tile to swizzle K.
    using Smem_tile_k
        = fmha::Smem_tile_b<Traits_p, Cta_tile_p, fmha::Col, Gmem_tile_k::BYTES_PER_LDG, BUFFERS_PER_TILE_SMEM_K>;

    // The global memory tile to load V.
    using Gmem_tile_v = Gmem_tile_v_<Traits_o, Cta_tile_o, Traits_o::BITS_PER_ELEMENT_B, CTA_O_TILE_K, CTA_O_TILE_N,
        VALID_DV, USE_LDGSTS_V, HEADS_INTERLEAVED,
        3, // NUM_MATS
        SLIDING_WINDOW_ATTENTION>;

    // The shared memory tile to swizzle V.
    using Smem_tile_v = fmha::Smem_tile_v<Traits_o, Cta_tile_o, BUFFERS_PER_TILE_SMEM_V>;

    // The global memory tile to store O.
    using Gmem_tile_o = Gmem_tile_o_<Traits_e, Cta_tile_o, CTAS_PER_HEAD>;
    // The shared memory tile for O.
    using Smem_tile_o = fmha::Smem_tile_o<Traits_e, Cta_tile_o>;

    // Make sure the number of threads match.
    static_assert((int) Gmem_tile_o::THREADS_PER_ROW == (int) Smem_tile_o::THREADS_PER_ROW, "");

    // The number of threads.
    enum
    {
        THREADS = Cta_tile_p::THREADS_PER_CTA
    };

    // Make sure the number of threads matches both CTAs.
    static_assert((int) THREADS == (int) Cta_tile_o::THREADS_PER_CTA, "");

    // The amount of shared memory needed to load Q and K.
    enum
    {
        BYTES_PER_SMEM_QK = Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE
    };

    // The extra amount of shared memory needed to load V.
    enum
    {
        BYTES_PER_SMEM_V = SHARE_SMEM_FOR_K_AND_V ? 0u : Smem_tile_v::BYTES_PER_TILE
    };

    // The amount of shared memory needed for Q, K and V..
    enum
    {
        BYTES_PER_SMEM_QKV = BYTES_PER_SMEM_QK + BYTES_PER_SMEM_V
    };

    // The amount of shared memory needed to load/store O.
    enum
    {
        BYTES_PER_SMEM_O = Smem_tile_o::BYTES_PER_TILE
    };

    // The amount of shared memory needed to load Q and store O.
    enum
    {
        BYTES_PER_SMEM_QO = NO_LOOP ? Smem_tile_o::BYTES_PER_TILE : Smem_tile_q::BYTES_PER_TILE + BYTES_PER_SMEM_O
    };

    // The amount of shared memory needed for Q, K, V and O.
    enum
    {
        BYTES_PER_SMEM = fmha::Max<BYTES_PER_SMEM_QKV, BYTES_PER_SMEM_QO>::VALUE
    };

    // Make sure we have enough shared memory.
    static_assert((NO_LOOP ? Smem_tile_o::BYTES_PER_TILE : Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE)
            <= BYTES_PER_SMEM,
        "");
};

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

template <
    // Instruction traits.
    typename Traits_,
    // The global memory tile for Q, K and V.
    template <typename, typename, int, int, int, bool, bool, int> class Gmem_tile_q_,
    // The global memory tile for the output.
    template <typename, typename, int> class Gmem_tile_o_,
    // Sequence length for K/V.
    int S_KV,
    // The hidden dimension.
    int D,
    // The iteration step of the outer loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD_,
    // The flags to control the behaviour of LDGs.
    uint32_t FLAGS,
    // The version of the kernel.
    int VERSION_,
    // Do we use half epilogue for the 2nd GEMM (hmma_fp32)
    bool BMM2_FP16_EPILOGUE = true>
struct Kernel_traits_fmhca_
{

    // The instruction traits for the Q*K product.
    using Traits_p = typename Traits_reuse<Traits_>::Traits;
    // The instruction traits for the P*V product. Hack to change the traits for Volta HMMA.
    using Traits_o = typename Traits_o_adapter<Traits_p, false>::Traits;
    // The instruction traits for the epilogue of the 2nd GEMM. Always use FP16.
    using Traits_e = typename Traits_o_adapter<Traits_p, BMM2_FP16_EPILOGUE>::Traits;

    // The CTA description for the 1st GEMM.
    using Cta_tile_p = typename Traits_p::template Cta_tile_extd<STEP, S_KV, D, S_KV, D, WARPS_M, WARPS_N, 1>;
    // The CTA description for the 2nd GEMM.
    using Cta_tile_o = typename Traits_o::template Cta_tile_extd<STEP, D, S_KV, D, S_KV, WARPS_M, 1, WARPS_N>;

    // The MMA tile for the 1st GEMM.
    using Mma_tile_p = typename Traits_p::template Mma_tile<Cta_tile_p>;
    // The MMA tile for the 2nd GEMM.
    using Mma_tile_o = typename Traits_o::template Mma_tile<Cta_tile_o>;

    // Constraints on the K dimension.
    static_assert(Mma_tile_p::K_PER_MMA <= D, "");
    static_assert(Mma_tile_o::K_PER_MMA <= S_KV, "");

    // The version.
    enum
    {
        VERSION = VERSION_
    };

    // The mask version
    enum
    {
        MASK_VERSION = VERSION_
    };

    // Whether use causal mask or not.
    enum
    {
        CAUSAL_MASK = MASK_VERSION >= 3
    };

    // Whether use the sliding window attention or not.
    enum
    {
        SLIDING_WINDOW_ATTENTION = MASK_VERSION == 4
    };

    // Do we use LDGSTS for Q, K or V.
    enum
    {
        USE_LDGSTS_Q = (FLAGS & 0x1u) != 0u
    };

    enum
    {
        USE_LDGSTS_K = (FLAGS & 0x2u) != 0u
    };

    enum
    {
        USE_LDGSTS_V = (FLAGS & 0x4u) != 0u
    };

    // Do we use one buffer for K and V.
    enum
    {
        SHARE_SMEM_FOR_K_AND_V = (FLAGS & 0x8u) != 0u
    };

    // Do we use the scale max trick.
    enum
    {
        USE_SCALE_MAX = (FLAGS & 0x10u) != 0u
    };

    // Are heads in QKV interleaved, i.e. total x h x 3 x d or total x 3 x h x d.
    enum
    {
        HEADS_INTERLEAVED = (FLAGS & 0x20u) == 0u
    };

    // Keep full K matrix in registers.
    enum
    {
        K_IN_REGS = (FLAGS & 0x40) == 0u
    };

    // Use BMM1 softcapping scale or not.
    enum
    {
        ENABLE_BMM1_SOFTCAPPING_SCALE = 0
    };

    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    enum
    {
        CTAS_PER_HEAD = CTAS_PER_HEAD_
    };

    // The global memory tile to load Q.
    using Gmem_tile_q
        = Gmem_tile_q_<Traits_p, Cta_tile_p, Traits_p::BITS_PER_ELEMENT_A, STEP, D, USE_LDGSTS_Q, HEADS_INTERLEAVED,
            1 // NUM_MATS
            >;

    // The shared memory tile to swizzle Q.
    using Smem_tile_q
        = fmha::Smem_tile_a<Traits_p, Cta_tile_p, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, USE_LDGSTS_Q ? 2 : 1>;

    // The global memory tile to load K.
    using Gmem_tile_k
        = Gmem_tile_q_<Traits_p, Cta_tile_p, Traits_p::BITS_PER_ELEMENT_B, S_KV, D, USE_LDGSTS_K, HEADS_INTERLEAVED,
            2 // NUM_MATS
            >;

    // The shared memory tile to swizzle K.
    using Smem_tile_k = fmha::Smem_tile_b<Traits_p, Cta_tile_p, fmha::Col>;

    // The global memory tile to load V.
    using Gmem_tile_v
        = Gmem_tile_q_<Traits_o, Cta_tile_o, Traits_o::BITS_PER_ELEMENT_B, S_KV, D, USE_LDGSTS_V, HEADS_INTERLEAVED,
            2 // NUM_MATS
            >;

    // The shared memory tile to swizzle V.
    using Smem_tile_v = fmha::Smem_tile_v<Traits_o, Cta_tile_o>;

    // The global memory tile to store O.
    using Gmem_tile_o = Gmem_tile_o_<Traits_e, Cta_tile_o, CTAS_PER_HEAD>;
    // The shared memory tile for O.
    using Smem_tile_o = fmha::Smem_tile_o<Traits_e, Cta_tile_o>;

    // Make sure the number of threads match.
    static_assert((int) Gmem_tile_o::THREADS_PER_ROW == (int) Smem_tile_o::THREADS_PER_ROW, "");

    // The number of threads.
    enum
    {
        THREADS = Cta_tile_p::THREADS_PER_CTA
    };

    // Make sure the number of threads matches both CTAs.
    static_assert((int) THREADS == (int) Cta_tile_o::THREADS_PER_CTA, "");

    // The amount of shared memory needed to load Q and K.
    enum
    {
        BYTES_PER_SMEM_QK = Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE
    };

    // The extra amount of shared memory needed to load V.
    enum
    {
        BYTES_PER_SMEM_V = SHARE_SMEM_FOR_K_AND_V ? 0u : Smem_tile_v::BYTES_PER_TILE
    };

    // The amount of shared memory needed for Q, K and V..
    enum
    {
        BYTES_PER_SMEM_QKV = BYTES_PER_SMEM_QK + BYTES_PER_SMEM_V
    };

    // The amount of shared memory needed to load Q and store O.
    enum
    {
        BYTES_PER_SMEM_QO = Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE
    };

    // The amount of shared memory needed for Q, K, V and O.
    enum
    {
        BYTES_PER_SMEM = fmha::Max<BYTES_PER_SMEM_QKV, BYTES_PER_SMEM_QO>::VALUE
    };

    // Make sure we have enough shared memory.
    static_assert(Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE <= BYTES_PER_SMEM, "");
};

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

template <
    // The instruction traits.
    typename Traits_,
    // The sequence length.
    int S,
    // The hidden size per head.
    int VALID_D,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD_,
    // The flags.
    uint32_t FLAGS = 0x8,
    // The mask version of the kernel
    int MASK_VERSION_ = 2>
struct Kernel_traits_interleaved_v2_
{

    // The instruction traits.
    using Traits = typename Traits_reuse<Traits_>::Traits;
    using Traits_p = Traits;
    using Traits_o = Traits;

    // The padded D dimension
    enum
    {
        D = Next_power_of_two<VALID_D>::VALUE
    };

    // The CTA description for the 1st GEMM.
    using Cta_tile_p = typename Traits::template Cta_tile_extd<STEP, S, D, S, VALID_D, WARPS_M, WARPS_N, 1>;
    // The CTA description for the 2nd GEMM.
    using Cta_tile_o = typename Traits::template Cta_tile_extd<STEP, D, S, VALID_D, S, WARPS_M, 1, WARPS_N>;

    // The version.
    enum
    {
        VERSION = 2
    };

    enum
    {
        MASK_VERSION = MASK_VERSION_
    };

    // Whether use causal mask or not.
    enum
    {
        CAUSAL_MASK = MASK_VERSION_ >= 3
    };

    // Whether use the sliding window attention or not.
    enum
    {
        SLIDING_WINDOW_ATTENTION = MASK_VERSION_ == 4
    };

    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    enum
    {
        CTAS_PER_HEAD = CTAS_PER_HEAD_
    };

    // Do we use LDGSTS for Q, K or V.
    enum
    {
        USE_LDGSTS_Q = (FLAGS & 0x1u) != 0u
    };

    enum
    {
        USE_LDGSTS_K = (FLAGS & 0x2u) != 0u
    };

    enum
    {
        USE_LDGSTS_V = (FLAGS & 0x4u) != 0u
    };

    // Do we use one buffer for K and V.
    enum
    {
        SHARE_SMEM_FOR_K_AND_V = (FLAGS & 0x8u) != 0u
    };

    // Do we use the scale max trick.
    enum
    {
        USE_SCALE_MAX = (FLAGS & 16) != 0u
    };

    // The global memory tile to load Q.
    using Gmem_tile_q
        = fmha::v2::Gmem_tile_qkv_interleaved<Traits, Cta_tile_p, Traits::BITS_PER_ELEMENT_A, STEP, D, USE_LDGSTS_Q>;
    // The shared memory tile to swizzle Q.
    using Smem_tile_q = fmha::Smem_tile_qk_interleaved_a<Traits, Cta_tile_p>;

    // The global memory tile to load K.
    using Gmem_tile_k
        = fmha::v2::Gmem_tile_qkv_interleaved<Traits, Cta_tile_p, Traits::BITS_PER_ELEMENT_B, S, D, USE_LDGSTS_K>;
    // The shared memory tile to swizzle K.
    using Smem_tile_k = fmha::Smem_tile_qk_interleaved_b<Traits, Cta_tile_p>;

    // The global memory tile to load V.
    using Gmem_tile_v
        = fmha::v2::Gmem_tile_qkv_interleaved<Traits, Cta_tile_o, Traits::BITS_PER_ELEMENT_B, S, D, USE_LDGSTS_V>;

    // The shared memory tile to swizzle V.
    using Smem_tile_v = fmha::Smem_tile_v_interleaved_b<Traits, Cta_tile_o>;

    // The global memory tile to store O.
    using Gmem_tile_o = fmha::v2::Imma_gmem_tile_o_interleaved<Traits, Cta_tile_o, CTAS_PER_HEAD>;
    // The shared memory tile for O.
    using Smem_tile_o = fmha::Smem_tile_o_interleaved<Traits, Cta_tile_o>;

    // Make sure the number of threads match.
    static_assert((int) Gmem_tile_o::THREADS_PER_ROW == (int) Smem_tile_o::THREADS_PER_ROW, "");

    // The number of threads.
    enum
    {
        THREADS = Cta_tile_p::THREADS_PER_CTA
    };

    // Make sure the number of threads matches both CTAs.
    static_assert((int) THREADS == (int) Cta_tile_o::THREADS_PER_CTA, "");

    // The amount of shared memory needed to load Q and K.
    enum
    {
        BYTES_PER_SMEM_QK = Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE
    };

    // The extra amount of shared memory needed to load V.
    enum
    {
        BYTES_PER_SMEM_V = SHARE_SMEM_FOR_K_AND_V ? 0u : Smem_tile_v::BYTES_PER_TILE
    };

    // The amount of shared memory needed for Q, K and V..
    enum
    {
        BYTES_PER_SMEM_QKV = BYTES_PER_SMEM_QK + BYTES_PER_SMEM_V
    };

    // The amount of shared memory needed to load Q and store O.
    enum
    {
        BYTES_PER_SMEM_QO = Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE
    };

    // The amount of shared memory needed for Q, K, V and O.
    enum
    {
        BYTES_PER_SMEM = fmha::Max<BYTES_PER_SMEM_QKV, BYTES_PER_SMEM_QO>::VALUE
    };

    // Make sure we have enough shared memory.
    static_assert(Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE <= BYTES_PER_SMEM, "");
};

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

template <
    // The instruction traits.
    typename Traits_,
    // The sequence length.
    int S,
    // The hidden size per head.
    int VALID_D,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD_,
    // The flags.
    uint32_t FLAGS = 0x8,
    // The mask version of the kernel
    int MASK_VERSION_ = 2>
using Kernel_traits_interleaved_v2
    = Kernel_traits_interleaved_v2_<Traits_, S, VALID_D, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD_, FLAGS, MASK_VERSION_>;

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

template <
    // The instruction traits.
    typename Traits,
    // The sequence length.
    int S,
    // The hidden size per head.
    int D,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD,
    // The flags.
    uint32_t FLAGS = 0x8>
using Kernel_traits_v1 = Kernel_traits_<Traits, fmha::v1::Gmem_tile_qkv, fmha::v1::Gmem_tile_qkv,
    fmha::v1::Gmem_tile_qkv, fmha::v1::Gmem_tile_o, S, D, 0, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS, 1, 1>;

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

template <
    // The instruction traits.
    typename Traits,
    // The sequence length.
    int S,
    // The hidden size per head.
    int D,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD,
    // The flags.
    uint32_t FLAGS = 0x8>
using Kernel_traits_v1_causal_mask = Kernel_traits_<Traits, fmha::v1::Gmem_tile_qkv, fmha::v1::Gmem_tile_qkv,
    fmha::v1::Gmem_tile_qkv, fmha::v1::Gmem_tile_o, S, D, 0, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS,
    1,  // VERSION_
    3>; // MASK_VERSION_

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

template <typename Traits_, typename OutputType>
struct Gmem_tile_o_dispatcher
{

    template <typename Traits, typename Cta_tile, int CTAS_PER_HEAD>
    using Gmem_tile_o = fmha::v2::Gmem_tile_o<Traits, Cta_tile, CTAS_PER_HEAD>;
};

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

template <>
struct Gmem_tile_o_dispatcher<fmha::Ada_qmma_e4m3_fp32_traits, uint16_t>
{

    template <typename Traits, typename Cta_tile, int CTAS_PER_HEAD>
    using Gmem_tile_o = fmha::v2::Gmem_tile_o_uint16<Traits, Cta_tile, CTAS_PER_HEAD>;
};

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

template <>
struct Gmem_tile_o_dispatcher<fmha::Ada_qmma_e4m3_fp32_traits, nv_bfloat16>
{

    template <typename Traits, typename Cta_tile, int CTAS_PER_HEAD>
    using Gmem_tile_o = fmha::v2::Gmem_tile_o_bfloat16<Traits, Cta_tile, CTAS_PER_HEAD>;
};

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

template <
    // The instruction traits.
    typename Traits,
    // The sequence length.
    int S,
    // The hidden size per head.
    int D,
    // The hidden dimension of V.
    int DV,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD,
    // The flags.
    uint32_t FLAGS = 0x8,
    // The attention mask version (see src/mask.h).
    int MASK_VERSION = 2,
    // Do we use half epilogue for the 2nd GEMM (hmma_fp32)
    bool BMM2_FP16_EPILOGUE = true,
    // The output type.
    typename OutputType = typename Traits::A_type,
    // The sage attention block size for Q, K and V
    int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0>
using Kernel_traits_v2 = Kernel_traits_<Traits, fmha::v2::Gmem_tile_qkv, fmha::v2::Gmem_tile_qkv,
    fmha::v2::Gmem_tile_qkv, Gmem_tile_o_dispatcher<Traits, OutputType>::Gmem_tile_o, S, D, DV, STEP, WARPS_M, WARPS_N,
    CTAS_PER_HEAD, FLAGS, 2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>;

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

template <
    // The instruction traits.
    typename Traits,
    // The sequence length.
    int S,
    // The hidden size per head.
    int D,
    // The hidden dimension of V.
    int DV,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD,
    // The flags.
    uint32_t FLAGS = 0x8,
    // The attention mask version (see src/mask.h).
    int MASK_VERSION = 2,
    // Do we use half epilogue for the 2nd GEMM (hmma_fp32)
    bool BMM2_FP16_EPILOGUE = true,
    // The output type.
    typename OutputType = typename Traits::A_type,
    // The sage attention block size for Q, K and V
    int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0>
using Kernel_traits_v2_q_k_v
    = Kernel_traits_<Traits, fmha::v2::Gmem_tile_q_k_v, fmha::v2::Gmem_tile_q_k_v, fmha::v2::Gmem_tile_q_k_v,
        Gmem_tile_o_dispatcher<Traits, OutputType>::Gmem_tile_o, S, D, DV, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS,
        2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>;

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

template <
    // The instruction traits.
    typename Traits,
    // The sequence length.
    int S,
    // The hidden size per head.
    int D,
    // The hidden dimension of V.
    int DV,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD,
    // The flags.
    uint32_t FLAGS = 0x8,
    // The attention mask version (see src/mask.h).
    int MASK_VERSION = 2,
    // Do we use half epilogue for the 2nd GEMM (hmma_fp32)
    bool BMM2_FP16_EPILOGUE = true,
    // The output type.
    typename OutputType = typename Traits::A_type,
    // The sage attention block size for Q, K and V
    int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0>
using Kernel_traits_v2_paged_kv_cache
    = Kernel_traits_<Traits, fmha::v2::Gmem_tile_q_k_v, fmha::v2::Gmem_tile_paged_kv, fmha::v2::Gmem_tile_paged_kv,
        Gmem_tile_o_dispatcher<Traits, OutputType>::Gmem_tile_o, S, D, DV, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS,
        2, MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>;

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

template <
    // The instruction traits.
    typename Traits,
    // The sequence length.
    int S,
    // The hidden size per head.
    int D,
    // The hidden dimension of V.
    int DV,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD,
    // The flags.
    uint32_t FLAGS = 0x8,
    // The attention mask version (see src/mask.h).
    int MASK_VERSION = 2,
    // Do we use half epilogue for the 2nd GEMM (hmma_fp32)
    bool BMM2_FP16_EPILOGUE = true,
    // The output type.
    typename OutputType = typename Traits::A_type,
    // The sage attention block size for Q, K and V
    int SAGE_BLOCK_SIZE_Q = 0, int SAGE_BLOCK_SIZE_K = 0, int SAGE_BLOCK_SIZE_V = 0>
using Kernel_traits_v2_contiguous_kv_cache = Kernel_traits_<Traits, fmha::v2::Gmem_tile_q_k_v,
    fmha::v2::Gmem_tile_contiguous_kv, fmha::v2::Gmem_tile_contiguous_kv,
    Gmem_tile_o_dispatcher<Traits, OutputType>::Gmem_tile_o, S, D, 0, STEP, WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS, 2,
    MASK_VERSION, BMM2_FP16_EPILOGUE, SAGE_BLOCK_SIZE_Q, SAGE_BLOCK_SIZE_K, SAGE_BLOCK_SIZE_V>;

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

template <
    // The instruction traits.
    typename Traits,
    // The sequence length for K and V.
    int S_KV,
    // The hidden size per head.
    int D,
    // The number of timesteps per iteration of the main loop.
    int STEP,
    // The number of vertical warps.
    int WARPS_M,
    // The number of horizontal warps.
    int WARPS_N,
    // The number of CTAs per head for Cta_tile_p; equivalent to BMM1 split-K
    int CTAS_PER_HEAD,
    // The flags.
    uint32_t FLAGS = 0x8>
using Kernel_traits_fmhca = Kernel_traits_fmhca_<Traits, fmha::v2::Gmem_tile_q_kv, fmha::v2::Gmem_tile_o, S_KV, D, STEP,
    WARPS_M, WARPS_N, CTAS_PER_HEAD, FLAGS, 2>;

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

} // namespace fmha