/*
 * SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 * SPDX-License-Identifier: Apache-2.0
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef _BERT_FMHA_FMHA
#define _BERT_FMHA_FMHA
#include "common/bertCommon.h"
#include "common/cudaDriverWrapper.h"
#include "common/plugin.h"
#include "cuda_runtime_api.h"
#include "fused_multihead_attention_common.h"
#include <memory>
#include <mutex>
#include <set>
#include <stdint.h>
#include <unordered_map>
#include <unordered_set>
#include <vector>

namespace nvinfer1
{

namespace pluginInternal
{
template <typename TKernelMeta, typename TKernelParam>
class TFusedMultiHeadAttentionXMMAKernel
{
public:
    using KernelMeta = TKernelMeta;
    using KernelParam = TKernelParam;
    inline uint64_t hashID(uint32_t s, uint32_t d) const
    {
        return (uint64_t) s << 32 | d;
    }
    virtual uint64_t hashID(const KernelMeta& kernelMeta) const
    {
        return hashID(kernelMeta.mS, kernelMeta.mD);
    }

    TFusedMultiHeadAttentionXMMAKernel(
        const TKernelMeta* pMetaStart, uint32_t nMetaCount, plugin::bert::Data_type type, uint32_t sm)
        : mDataType(type)
        , mKernelMeta(pMetaStart)
        , mKernelMetaCount(nMetaCount)
        , mSM(sm)
    {
        PLUGIN_ASSERT(mKernelMetaCount && "No kernels were loaded correctly.");
    }

    void loadXMMAKernels(uint32_t smVersion)
    {
        for (uint32_t i = 0; i < mKernelMetaCount; ++i)
        {
            const auto& kernelMeta = mKernelMeta[i];
            const auto kernelKey = hashID(kernelMeta);
            if (kernelMeta.mSM == smVersion && kernelMeta.mDataType == mDataType
                && mFunctions.find(kernelKey) == mFunctions.end())
            {
                const uint32_t DEFAULT_SMEM_SIZE{48 * 1024};
                if (kernelMeta.mSharedMemBytes >= DEFAULT_SMEM_SIZE)
                {
                    int32_t deviceID{0};
                    cudaGetDevice(&deviceID);
                    int32_t sharedMemPerMultiprocessor{0};
                    if (cudaDeviceGetAttribute(
                            &sharedMemPerMultiprocessor, cudaDevAttrMaxSharedMemoryPerBlockOptin, deviceID)
                            != cudaSuccess
                        || sharedMemPerMultiprocessor < static_cast<int32_t>(kernelMeta.mSharedMemBytes))
                    {
                        // skip load function because not enough shared memory to launch the kernel
                        continue;
                    }
                }

                CUmodule hmod{0};
                auto findModuleIter = mModules.find(kernelMeta.mCubin);
                if (findModuleIter != mModules.end())
                {
                    hmod = findModuleIter->second;
                }
                else
                {
                    cuErrCheck(mDriver.cuModuleLoadData(&hmod, kernelMeta.mCubin), mDriver);
                    mModules.insert(std::make_pair(kernelMeta.mCubin, hmod));
                }

                FusedMultiHeadAttentionKernelInfo funcInfo;
                funcInfo.mMetaInfoIndex = i;
                cuErrCheck(mDriver.cuModuleGetFunction(&funcInfo.mDeviceFunction, hmod, kernelMeta.mFuncName), mDriver);
                if (kernelMeta.mSharedMemBytes >= DEFAULT_SMEM_SIZE)
                {
                    if (mDriver.cuFuncSetAttribute(funcInfo.mDeviceFunction,
                            CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, kernelMeta.mSharedMemBytes)
                        != CUDA_SUCCESS)
                    {
                        // some chip may not have enough shared memory to launch the kernel
                        continue;
                    }
                }
                mFunctions.insert({kernelKey, funcInfo});
                uint64_t const s = kernelMeta.mS;
                uint64_t const headSize = kernelMeta.mD;
                uint64_t key = (headSize << 32 | s);
                if (mValidSequences.find(key) == mValidSequences.end())
                {
                    mValidSequences.insert(key);
                }
            }
        }
    }

    void loadXMMAKernels()
    {
        if (!mFunctions.empty())
        {
            return;
        }

        loadXMMAKernels(mSM);

        // sm_86 chips prefer sm_86 sass, but can also use sm_80 sass if sm_86 not exist.
        // sm_87 cannot run sm_80 sass
        if (mSM == kSM_86)
        {
            loadXMMAKernels(kSM_80);
        }

        // sm_89 will reuse sm_80 and sm_86 kernels
        if (mSM == kSM_89)
        {
            loadXMMAKernels(kSM_86);
            loadXMMAKernels(kSM_80);
        }
    }

    bool isValid(int32_t headSize, int32_t s) const
    {
        uint64_t key = (static_cast<uint64_t>(headSize) << 32 | static_cast<uint64_t>(s));
        return (mValidSequences.find(key) != mValidSequences.end());
    }

    virtual void run(TKernelParam& params, cudaStream_t ss) const
    {
        const auto findIter = mFunctions.find(hashID(params.s, params.d));
        std::stringstream errMsg;
        errMsg << "Could not find kernel for:\n"
               << "\t s: " << params.s << "\n"
               << "\t d: " << params.d << "\n"
               << "Was the plugin compiled on a compatible CUDA and SM version?\n"
               << "\t Compiled on CUDA " << CUDA_VERSION << "\n"
               << "\t Current SM version: " << mSM << "\n"
               << "\t SM versions enabled during compilation: "
#if defined(ENABLE_SM72)
               << "72 "
#endif
#if defined(ENABLE_SM75)
               << "75 "
#endif
#if defined(ENABLE_SM80)
               << "80 "
#endif
#if defined(ENABLE_SM86)
               << "86 "
#endif
#if defined(ENABLE_SM87)
               << "87 "
#endif
#if defined(ENABLE_SM89)
               << "89 "
#endif
#if defined(ENABLE_SM90)
               << "90 "
#endif
               << "\n";
        PLUGIN_VALIDATE(findIter != mFunctions.end(), errMsg.str().c_str());

        const auto& kernelMeta = mKernelMeta[findIter->second.mMetaInfoIndex];
        const CUfunction func = findIter->second.mDeviceFunction;

        void* kernelParams[] = {&params, nullptr};
        cuErrCheck(mDriver.cuLaunchKernel(func, params.h, params.b, 1, kernelMeta.mThreadsPerCTA, 1, 1,
                       kernelMeta.mSharedMemBytes, ss, kernelParams, nullptr),
            mDriver);
    }

    virtual ~TFusedMultiHeadAttentionXMMAKernel() = default;

protected:
    nvinfer1::CUDADriverWrapper mDriver;

    plugin::bert::Data_type mDataType;
    const TKernelMeta* mKernelMeta;
    uint32_t mKernelMetaCount;
    uint32_t mSM;
    std::unordered_map<const unsigned char*, CUmodule> mModules;
    struct FusedMultiHeadAttentionKernelInfo
    {
        uint32_t mMetaInfoIndex;
        CUfunction mDeviceFunction;
    };
    std::unordered_map<uint64_t, FusedMultiHeadAttentionKernelInfo> mFunctions;
    // Set of valid sequence and head size combination. We use (headSize << 32 | sequence) as key here.
    std::unordered_set<uint64_t> mValidSequences;
};
template <typename TFusedMHAKernelList>
class TFusedMHAKernelFactory
{
public:
    TFusedMHAKernelList const* getXMMAKernels(typename TFusedMHAKernelList::KernelMeta const* pKernelList,
        uint32_t nbKernels, plugin::bert::Data_type type, uint32_t sm)
    {
        static std::mutex s_mutex;
        std::lock_guard<std::mutex> lg(s_mutex);

        auto const id = hashID(type, sm);
        auto it = mKernels.find(id);
        if (it == mKernels.end())
        {
            auto newKernel = std::make_unique<TFusedMHAKernelList>(pKernelList, nbKernels, type, sm);
            newKernel->loadXMMAKernels();
            it = mKernels.emplace(id, std::move(newKernel)).first;
        }
        return it->second.get();
    }

    static TFusedMHAKernelFactory<TFusedMHAKernelList>& Get()
    {
        static TFusedMHAKernelFactory<TFusedMHAKernelList> s_factory;
        return s_factory;
    }

private:
    TFusedMHAKernelFactory() = default;

    inline uint64_t hashID(plugin::bert::Data_type type, uint32_t sm) const
    {
        // use deviceID in hasID for multi GPU support before driver support context-less loading of cubin
        int32_t deviceID{0};
        CSC(cudaGetDevice(&deviceID), STATUS_FAILURE);

        PLUGIN_ASSERT((deviceID & 0xFFFF) == deviceID);
        PLUGIN_ASSERT((type & 0xFFFF) == type);
        PLUGIN_ASSERT((sm & 0xFFFFFFFF) == sm);
        return (uint64_t) type << 48 | (uint64_t) deviceID << 32 | sm;
    }

    std::unordered_map<uint64_t, const std::unique_ptr<TFusedMHAKernelList>> mKernels;
};
} // namespace pluginInternal

namespace plugin
{
namespace bert
{
static inline size_t get_size_in_bytes(size_t n, Data_type dtype)
{
    switch (dtype)
    {
    case DATA_TYPE_E8M10: return n * 4;
    case DATA_TYPE_FP32: return n * 4;
    case DATA_TYPE_FP16: return n * 2;
    case DATA_TYPE_INT32: return n * 4;
    case DATA_TYPE_INT8: return n;
    case DATA_TYPE_INT4: return n / 2U;
    case DATA_TYPE_BOOL: return n / 8U;
    case DATA_TYPE_E8M7: return n * 2;
    default: PLUGIN_ASSERT(false); return 0;
    }
}

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

struct Fused_multihead_attention_params
{
    // The QKV matrices.
    void* qkv_ptr{};
    // The mask to implement drop-out.
    void* packed_mask_ptr{};
    // The O matrix (output).
    void* o_ptr{};

    // The stride between rows of the Q, K and V matrices.
    int64_t qkv_stride_in_bytes{};
    // The stride between matrices of packed mask.
    int64_t packed_mask_stride_in_bytes{};
    // The stride between rows of O.
    int64_t o_stride_in_bytes{};

#if defined(STORE_P)
    // The pointer to the P matrix (for debugging).
    void* p_ptr{};
    // The stride between rows of the P matrix (for debugging).
    int64_t p_stride_in_bytes{};
#endif // defined(STORE_P)

#if defined(STORE_S)
    // The pointer to the S matrix (for debugging).
    void* s_ptr{};
    // The stride between rows of the S matrix (for debugging).
    int64_t s_stride_in_bytes{};
#endif // defined(STORE_S)

    // The dimensions.
    int32_t b{};
    int32_t h{};
    int32_t s{};
    int32_t d{};
    // The scaling factors for the kernel.
    uint32_t scale_bmm1{};
    uint32_t scale_softmax{};
    uint32_t scale_bmm2{};

    // Do we use Niall's trick to avoid I2F/F2I in the INT8 kernel.
    // See https://confluence.nvidia.com/pages/viewpage.action?pageId=302779721 for details.
    bool enable_i2f_trick{};

    // The number of heads computed by one iteration of the wave.
    int32_t heads_per_wave{};
    // Buffers to perform a global sync and a critical section.
    int32_t* counters{};
    int32_t* max_barriers{};
    int32_t* sum_barriers{};
    int32_t* locks{};
    // Scratch buffers to finalize softmax.
    float* max_scratch_ptr{};
    float* sum_scratch_ptr{};
    // Scratch buffer to finalize the output (not needed for FP16).
    int32_t* o_scratch_ptr{};
};

////////////////////////////////////////////////////////////////////////////////////////////////////
#if defined(ENABLE_SM75)
extern unsigned char fused_multihead_attention_fp16_64_64_kernel_sm75_cu_o[];
extern unsigned char fused_multihead_attention_fp16_96_64_kernel_sm75_cu_o[];
extern unsigned char fused_multihead_attention_fp16_128_64_kernel_sm75_cu_o[];
extern unsigned char fused_multihead_attention_fp16_384_64_kernel_sm75_cu_o[];
extern unsigned char fused_multihead_attention_int8_128_64_kernel_sm75_cu_o[];
extern unsigned char fused_multihead_attention_int8_384_64_kernel_sm75_cu_o[];
#endif // defined(ENABLE_SM75)

#if defined(ENABLE_SM80) || defined(ENABLE_SM86) || defined(ENABLE_SM89)
extern unsigned char fused_multihead_attention_fp16_64_64_kernel_sm80_cu_o[];
extern unsigned char fused_multihead_attention_fp16_96_64_kernel_sm80_cu_o[];
extern unsigned char fused_multihead_attention_int8_384_64_kernel_sm80_cu_o[];
extern unsigned char fused_multihead_attention_int8_128_64_kernel_sm80_cu_o[];
extern unsigned char fused_multihead_attention_fp16_128_64_kernel_sm80_cu_o[];
extern unsigned char fused_multihead_attention_fp16_384_64_kernel_sm80_cu_o[];

extern unsigned char cubin_fmha_v1_int8_64_64_sm80_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_96_64_sm80_cu_cubin[];
#endif // defined(ENABLE_SM80) || defined(SM86) || defined(ENABLE_SM89)

#if defined(ENABLE_SM86) || defined(ENABLE_SM89)
extern unsigned char fused_multihead_attention_fp16_384_64_kernel_sm86_cu_o[];
#endif // defined(ENABLE_SM86) || defined(ENABLE_SM89)

#if defined(ENABLE_SM87)
extern unsigned char cubin_fmha_v1_int8_384_64_sm87_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_128_64_sm87_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_384_64_sm87_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_128_64_sm87_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_96_64_sm87_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_64_64_sm87_cu_cubin[];
#endif // defined(ENABLE_SM87)

#if defined(ENABLE_SM90)
extern unsigned char cubin_fmha_v1_int8_512_64_sm90_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_384_64_sm90_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_128_64_sm90_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_512_64_sm90_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_384_64_sm90_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_128_64_sm90_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_96_64_sm90_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_64_64_sm90_cu_cubin[];
#endif // defined(ENABLE_SM90)

#if defined(ENABLE_SM100)
extern unsigned char cubin_fmha_v1_int8_512_64_sm100_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_384_64_sm100_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_128_64_sm100_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_512_64_sm100_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_384_64_sm100_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_128_64_sm100_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_96_64_sm100_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_64_64_sm100_cu_cubin[];
#endif // defined(ENABLE_SM100)
#if defined(ENABLE_SM120)
extern unsigned char cubin_fmha_v1_int8_512_64_sm120_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_384_64_sm120_cu_cubin[];
extern unsigned char cubin_fmha_v1_int8_128_64_sm120_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_512_64_sm120_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_384_64_sm120_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_128_64_sm120_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_96_64_sm120_cu_cubin[];
extern unsigned char cubin_fmha_v1_fp16_64_64_sm120_cu_cubin[];
#endif // defined(ENABLE_SM120)

#if defined(ENABLE_SM75)
extern uint32_t fused_multihead_attention_fp16_64_64_kernel_sm75_cu_o_len;
extern uint32_t fused_multihead_attention_fp16_96_64_kernel_sm75_cu_o_len;
extern uint32_t fused_multihead_attention_fp16_128_64_kernel_sm75_cu_o_len;
extern uint32_t fused_multihead_attention_fp16_384_64_kernel_sm75_cu_o_len;
extern uint32_t fused_multihead_attention_int8_128_64_kernel_sm75_cu_o_len;
extern uint32_t fused_multihead_attention_int8_384_64_kernel_sm75_cu_o_len;
#endif // defined(ENABLE_SM75)

#if defined(ENABLE_SM80) || defined(ENABLE_SM86) || defined(ENABLE_SM89)
extern uint32_t fused_multihead_attention_fp16_64_64_kernel_sm80_cu_o_len;
extern uint32_t fused_multihead_attention_fp16_96_64_kernel_sm80_cu_o_len;
extern uint32_t fused_multihead_attention_int8_384_64_kernel_sm80_cu_o_len;
extern uint32_t fused_multihead_attention_int8_128_64_kernel_sm80_cu_o_len;
extern uint32_t fused_multihead_attention_fp16_128_64_kernel_sm80_cu_o_len;
extern uint32_t fused_multihead_attention_fp16_384_64_kernel_sm80_cu_o_len;
extern uint32_t cubin_fmha_v1_int8_64_64_sm80_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_96_64_sm80_cu_cubin_len;
#endif // defined(ENABLE_SM80) || defined(SM86) || defined(ENABLE_SM89)

#if defined(ENABLE_SM86) || defined(ENABLE_SM89)
extern uint32_t fused_multihead_attention_fp16_384_64_kernel_sm86_cu_o_len;
#endif // defined(ENABLE_SM86) || defined(ENABLE_SM89)

#if defined(ENABLE_SM87)
extern uint32_t cubin_fmha_v1_int8_384_64_sm87_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_128_64_sm87_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_384_64_sm87_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_128_64_sm87_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_96_64_sm87_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_64_64_sm87_cu_cubin_len;
#endif // defined(ENABLE_SM87)

#if defined(ENABLE_SM90)
extern uint32_t cubin_fmha_v1_int8_512_64_sm90_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_384_64_sm90_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_128_64_sm90_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_512_64_sm90_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_384_64_sm90_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_128_64_sm90_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_96_64_sm90_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_64_64_sm90_cu_cubin_len;
#endif // defined(ENABLE_SM90)

#if defined(ENABLE_SM100)
extern uint32_t cubin_fmha_v1_int8_512_64_sm100_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_384_64_sm100_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_128_64_sm100_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_512_64_sm100_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_384_64_sm100_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_128_64_sm100_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_96_64_sm100_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_64_64_sm100_cu_cubin_len;
#endif // defined(ENABLE_SM100)
#if defined(ENABLE_SM120)
extern uint32_t cubin_fmha_v1_int8_512_64_sm120_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_384_64_sm120_cu_cubin_len;
extern uint32_t cubin_fmha_v1_int8_128_64_sm120_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_512_64_sm120_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_384_64_sm120_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_128_64_sm120_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_96_64_sm120_cu_cubin_len;
extern uint32_t cubin_fmha_v1_fp16_64_64_sm120_cu_cubin_len;
#endif // defined(ENABLE_SM120)

#if !(defined(ENABLE_SM72) || defined(ENABLE_SM75) || defined(ENABLE_SM80) || defined(ENABLE_SM86)                     \
    || defined(ENABLE_SM87) || defined(ENABLE_SM89) || defined(ENABLE_SM90) || defined(ENABLE_SM100)                   \
    || defined(ENABLE_SM120))

#error This file can only be included if one of sm 72, 75, 80, 86, 87, 89, 90, 100 or 120 is defined.

#endif

static const struct FusedMultiHeadAttentionKernelMetaInfoV1
{
    Data_type mDataType;
    uint32_t mS;
    uint32_t mD;
    uint32_t mSM;
    const unsigned char* mCubin;
    uint32_t mCubinSize;
    const char* mFuncName;
    uint32_t mSharedMemBytes;
    uint32_t mThreadsPerCTA;
} sMhaKernelMetaInfos[] = {
#if defined(ENABLE_SM75)
    // Turing
    {DATA_TYPE_FP16, 64, 64, kSM_75, fused_multihead_attention_fp16_64_64_kernel_sm75_cu_o,
        fused_multihead_attention_fp16_64_64_kernel_sm75_cu_o_len,
        "fused_multihead_attention_v2_fp16_64_64_kernel_sm75", 24576, 128},
    {DATA_TYPE_FP16, 96, 64, kSM_75, fused_multihead_attention_fp16_96_64_kernel_sm75_cu_o,
        fused_multihead_attention_fp16_96_64_kernel_sm75_cu_o_len,
        "fused_multihead_attention_v2_fp16_96_64_kernel_sm75", 24576, 128},
    {DATA_TYPE_FP16, 128, 64, kSM_75, fused_multihead_attention_fp16_128_64_kernel_sm75_cu_o,
        fused_multihead_attention_fp16_128_64_kernel_sm75_cu_o_len, "fused_multihead_attention_fp16_128_64_kernel_sm75",
        32768, 128},
    {DATA_TYPE_FP16, 384, 64, kSM_75, fused_multihead_attention_fp16_384_64_kernel_sm75_cu_o,
        fused_multihead_attention_fp16_384_64_kernel_sm75_cu_o_len, "fused_multihead_attention_fp16_384_64_kernel_sm75",
        57344, 256},
    {DATA_TYPE_INT8, 128, 64, kSM_75, fused_multihead_attention_int8_128_64_kernel_sm75_cu_o,
        fused_multihead_attention_int8_128_64_kernel_sm75_cu_o_len, "fused_multihead_attention_int8_128_64_kernel_sm75",
        16384, 128},
    {DATA_TYPE_INT8, 384, 64, kSM_75, fused_multihead_attention_int8_384_64_kernel_sm75_cu_o,
        fused_multihead_attention_int8_384_64_kernel_sm75_cu_o_len, "fused_multihead_attention_int8_384_64_kernel_sm75",
        53284, 256},
#endif // defined(ENABLE_SM75)
#if defined(ENABLE_SM80) || defined(ENABLE_SM86) || defined(ENABLE_SM89)
    // Ampere
    {DATA_TYPE_FP16, 64, 64, kSM_80, fused_multihead_attention_fp16_64_64_kernel_sm80_cu_o,
        fused_multihead_attention_fp16_64_64_kernel_sm80_cu_o_len,
        "fused_multihead_attention_v2_fp16_64_64_kernel_sm80", 32768, 128},
    {DATA_TYPE_FP16, 96, 64, kSM_80, fused_multihead_attention_fp16_96_64_kernel_sm80_cu_o,
        fused_multihead_attention_fp16_96_64_kernel_sm80_cu_o_len,
        "fused_multihead_attention_v2_fp16_96_64_kernel_sm80", 49152, 128},
    {DATA_TYPE_FP16, 128, 64, kSM_80, fused_multihead_attention_fp16_128_64_kernel_sm80_cu_o,
        fused_multihead_attention_fp16_128_64_kernel_sm80_cu_o_len, "fused_multihead_attention_fp16_128_64_kernel_sm80",
        49152, 128},
    {DATA_TYPE_FP16, 384, 64, kSM_80, fused_multihead_attention_fp16_384_64_kernel_sm80_cu_o,
        fused_multihead_attention_fp16_384_64_kernel_sm80_cu_o_len, "fused_multihead_attention_fp16_384_64_kernel_sm80",
        114688, 256},
    {DATA_TYPE_INT8, 64, 64, kSM_80, cubin_fmha_v1_int8_64_64_sm80_cu_cubin, cubin_fmha_v1_int8_64_64_sm80_cu_cubin_len,
        "fmha_v1_int8_64_64_sm80_kernel", 24576, 128},
    {DATA_TYPE_INT8, 96, 64, kSM_80, cubin_fmha_v1_int8_96_64_sm80_cu_cubin, cubin_fmha_v1_int8_96_64_sm80_cu_cubin_len,
        "fmha_v1_int8_96_64_sm80_kernel", 28672, 128},
    {DATA_TYPE_INT8, 128, 64, kSM_80, fused_multihead_attention_int8_128_64_kernel_sm80_cu_o,
        fused_multihead_attention_int8_128_64_kernel_sm80_cu_o_len, "fused_multihead_attention_int8_128_64_kernel_sm80",
        24576, 128},
    {DATA_TYPE_INT8, 384, 64, kSM_80, fused_multihead_attention_int8_384_64_kernel_sm80_cu_o,
        fused_multihead_attention_int8_384_64_kernel_sm80_cu_o_len, "fused_multihead_attention_int8_384_64_kernel_sm80",
        57344, 256},
#endif // defined(ENABLE_SM80) || defined(SM86) || defined(ENABLE_SM89)
#if defined(ENABLE_SM86) || defined(ENABLE_SM89)
    // GA10x
    // Note: For GA10X keep only kernels whose sharedMemBytes < 100KiB
    {DATA_TYPE_FP16, 64, 64, kSM_86, fused_multihead_attention_fp16_64_64_kernel_sm80_cu_o,
        fused_multihead_attention_fp16_64_64_kernel_sm80_cu_o_len,
        "fused_multihead_attention_v2_fp16_64_64_kernel_sm80", 32768, 128},
    {DATA_TYPE_FP16, 96, 64, kSM_86, fused_multihead_attention_fp16_96_64_kernel_sm80_cu_o,
        fused_multihead_attention_fp16_96_64_kernel_sm80_cu_o_len,
        "fused_multihead_attention_v2_fp16_96_64_kernel_sm80", 49152, 128},
    {DATA_TYPE_FP16, 128, 64, kSM_86, fused_multihead_attention_fp16_128_64_kernel_sm80_cu_o,
        fused_multihead_attention_fp16_128_64_kernel_sm80_cu_o_len, "fused_multihead_attention_fp16_128_64_kernel_sm80",
        49152, 128},
    {DATA_TYPE_FP16, 384, 64, kSM_86, fused_multihead_attention_fp16_384_64_kernel_sm86_cu_o,
        fused_multihead_attention_fp16_384_64_kernel_sm86_cu_o_len, "fused_multihead_attention_fp16_384_64_kernel_sm80",
        65536, 256},
    {DATA_TYPE_INT8, 128, 64, kSM_86, fused_multihead_attention_int8_128_64_kernel_sm80_cu_o,
        fused_multihead_attention_int8_128_64_kernel_sm80_cu_o_len, "fused_multihead_attention_int8_128_64_kernel_sm80",
        24576, 128},
    {DATA_TYPE_INT8, 384, 64, kSM_86, fused_multihead_attention_int8_384_64_kernel_sm80_cu_o,
        fused_multihead_attention_int8_384_64_kernel_sm80_cu_o_len, "fused_multihead_attention_int8_384_64_kernel_sm80",
        57344, 256},
#endif // defined(ENABLE_SM86) || defined(ENABLE_SM89)
#if defined(ENABLE_SM87)
    // GA10b (Orin-Auto)
    {DATA_TYPE_INT8, 384, 64, kSM_87, cubin_fmha_v1_int8_384_64_sm87_cu_cubin,
        cubin_fmha_v1_int8_384_64_sm87_cu_cubin_len, "fmha_v1_int8_384_64_sm87_kernel", 40960, 256},
    {DATA_TYPE_INT8, 128, 64, kSM_87, cubin_fmha_v1_int8_128_64_sm87_cu_cubin,
        cubin_fmha_v1_int8_128_64_sm87_cu_cubin_len, "fmha_v1_int8_128_64_sm87_kernel", 24576, 128},
    {DATA_TYPE_FP16, 384, 64, kSM_87, cubin_fmha_v1_fp16_384_64_sm87_cu_cubin,
        cubin_fmha_v1_fp16_384_64_sm87_cu_cubin_len, "fmha_v1_fp16_384_64_sm87_kernel", 65536, 256},
    {DATA_TYPE_FP16, 128, 64, kSM_87, cubin_fmha_v1_fp16_128_64_sm87_cu_cubin,
        cubin_fmha_v1_fp16_128_64_sm87_cu_cubin_len, "fmha_v1_fp16_128_64_sm87_kernel", 49152, 128},
    {DATA_TYPE_FP16, 96, 64, kSM_87, cubin_fmha_v1_fp16_96_64_sm87_cu_cubin, cubin_fmha_v1_fp16_96_64_sm87_cu_cubin_len,
        "fmha_v1_fp16_96_64_sm87_kernel", 49152, 128},
    {DATA_TYPE_FP16, 64, 64, kSM_87, cubin_fmha_v1_fp16_64_64_sm87_cu_cubin, cubin_fmha_v1_fp16_64_64_sm87_cu_cubin_len,
        "fmha_v1_fp16_64_64_sm87_kernel", 32768, 128},
#endif // defined(ENABLE_SM87)
#if defined(ENABLE_SM90)
    // GH100 hopper
    {DATA_TYPE_INT8, 512, 64, kSM_90, cubin_fmha_v1_int8_512_64_sm90_cu_cubin,
        cubin_fmha_v1_int8_512_64_sm90_cu_cubin_len, "fmha_v1_int8_512_64_sm90_kernel", 73728, 256},
    {DATA_TYPE_INT8, 384, 64, kSM_90, cubin_fmha_v1_int8_384_64_sm90_cu_cubin,
        cubin_fmha_v1_int8_384_64_sm90_cu_cubin_len, "fmha_v1_int8_384_64_sm90_kernel", 73728, 256},
    {DATA_TYPE_INT8, 128, 64, kSM_90, cubin_fmha_v1_int8_128_64_sm90_cu_cubin,
        cubin_fmha_v1_int8_128_64_sm90_cu_cubin_len, "fmha_v1_int8_128_64_sm90_kernel", 24576, 128},
    {DATA_TYPE_FP16, 512, 64, kSM_90, cubin_fmha_v1_fp16_512_64_sm90_cu_cubin,
        cubin_fmha_v1_fp16_512_64_sm90_cu_cubin_len, "fmha_v1_fp16_512_64_sm90_kernel", 73728, 256},
    {DATA_TYPE_FP16, 384, 64, kSM_90, cubin_fmha_v1_fp16_384_64_sm90_cu_cubin,
        cubin_fmha_v1_fp16_384_64_sm90_cu_cubin_len, "fmha_v1_fp16_384_64_sm90_kernel", 65536, 256},
    {DATA_TYPE_FP16, 128, 64, kSM_90, cubin_fmha_v1_fp16_128_64_sm90_cu_cubin,
        cubin_fmha_v1_fp16_128_64_sm90_cu_cubin_len, "fmha_v1_fp16_128_64_sm90_kernel", 65536, 128},
    {DATA_TYPE_FP16, 96, 64, kSM_90, cubin_fmha_v1_fp16_96_64_sm90_cu_cubin, cubin_fmha_v1_fp16_96_64_sm90_cu_cubin_len,
        "fmha_v1_fp16_96_64_sm90_kernel", 49152, 128},
    {DATA_TYPE_FP16, 64, 64, kSM_90, cubin_fmha_v1_fp16_64_64_sm90_cu_cubin, cubin_fmha_v1_fp16_64_64_sm90_cu_cubin_len,
        "fmha_v1_fp16_64_64_sm90_kernel", 32768, 128},
#endif // defined(ENABLE_SM90)

#if defined(ENABLE_SM100)
    {DATA_TYPE_INT8, 512, 64, kSM_100, cubin_fmha_v1_int8_512_64_sm100_cu_cubin,
        cubin_fmha_v1_int8_512_64_sm100_cu_cubin_len, "fmha_v1_int8_512_64_sm100_kernel", 73728, 256},
    {DATA_TYPE_INT8, 384, 64, kSM_100, cubin_fmha_v1_int8_384_64_sm100_cu_cubin,
        cubin_fmha_v1_int8_384_64_sm100_cu_cubin_len, "fmha_v1_int8_384_64_sm100_kernel", 73728, 256},
    {DATA_TYPE_INT8, 128, 64, kSM_100, cubin_fmha_v1_int8_128_64_sm100_cu_cubin,
        cubin_fmha_v1_int8_128_64_sm100_cu_cubin_len, "fmha_v1_int8_128_64_sm100_kernel", 24576, 128},
    {DATA_TYPE_FP16, 512, 64, kSM_100, cubin_fmha_v1_fp16_512_64_sm100_cu_cubin,
        cubin_fmha_v1_fp16_512_64_sm100_cu_cubin_len, "fmha_v1_fp16_512_64_sm100_kernel", 73728, 256},
    {DATA_TYPE_FP16, 384, 64, kSM_100, cubin_fmha_v1_fp16_384_64_sm100_cu_cubin,
        cubin_fmha_v1_fp16_384_64_sm100_cu_cubin_len, "fmha_v1_fp16_384_64_sm100_kernel", 65536, 256},
    {DATA_TYPE_FP16, 128, 64, kSM_100, cubin_fmha_v1_fp16_128_64_sm100_cu_cubin,
        cubin_fmha_v1_fp16_128_64_sm100_cu_cubin_len, "fmha_v1_fp16_128_64_sm100_kernel", 65536, 128},
    {DATA_TYPE_FP16, 96, 64, kSM_100, cubin_fmha_v1_fp16_96_64_sm100_cu_cubin,
        cubin_fmha_v1_fp16_96_64_sm100_cu_cubin_len, "fmha_v1_fp16_96_64_sm100_kernel", 49152, 128},
    {DATA_TYPE_FP16, 64, 64, kSM_100, cubin_fmha_v1_fp16_64_64_sm100_cu_cubin,
        cubin_fmha_v1_fp16_64_64_sm100_cu_cubin_len, "fmha_v1_fp16_64_64_sm100_kernel", 32768, 128},
#endif // defined(ENABLE_SM100)
#if defined(ENABLE_SM120)
    {DATA_TYPE_INT8, 512, 64, kSM_120, cubin_fmha_v1_int8_512_64_sm120_cu_cubin,
        cubin_fmha_v1_int8_512_64_sm120_cu_cubin_len, "fmha_v1_int8_512_64_sm120_kernel", 73728, 256},
    {DATA_TYPE_INT8, 384, 64, kSM_120, cubin_fmha_v1_int8_384_64_sm120_cu_cubin,
        cubin_fmha_v1_int8_384_64_sm120_cu_cubin_len, "fmha_v1_int8_384_64_sm120_kernel", 73728, 256},
    {DATA_TYPE_INT8, 128, 64, kSM_120, cubin_fmha_v1_int8_128_64_sm120_cu_cubin,
        cubin_fmha_v1_int8_128_64_sm120_cu_cubin_len, "fmha_v1_int8_128_64_sm120_kernel", 24576, 128},
    {DATA_TYPE_FP16, 512, 64, kSM_120, cubin_fmha_v1_fp16_512_64_sm120_cu_cubin,
        cubin_fmha_v1_fp16_512_64_sm120_cu_cubin_len, "fmha_v1_fp16_512_64_sm120_kernel", 73728, 256},
    {DATA_TYPE_FP16, 384, 64, kSM_120, cubin_fmha_v1_fp16_384_64_sm120_cu_cubin,
        cubin_fmha_v1_fp16_384_64_sm120_cu_cubin_len, "fmha_v1_fp16_384_64_sm120_kernel", 65536, 256},
    {DATA_TYPE_FP16, 128, 64, kSM_120, cubin_fmha_v1_fp16_128_64_sm120_cu_cubin,
        cubin_fmha_v1_fp16_128_64_sm120_cu_cubin_len, "fmha_v1_fp16_128_64_sm120_kernel", 65536, 128},
    {DATA_TYPE_FP16, 96, 64, kSM_120, cubin_fmha_v1_fp16_96_64_sm120_cu_cubin,
        cubin_fmha_v1_fp16_96_64_sm120_cu_cubin_len, "fmha_v1_fp16_96_64_sm120_kernel", 49152, 128},
    {DATA_TYPE_FP16, 64, 64, kSM_120, cubin_fmha_v1_fp16_64_64_sm120_cu_cubin,
        cubin_fmha_v1_fp16_64_64_sm120_cu_cubin_len, "fmha_v1_fp16_64_64_sm120_kernel", 32768, 128},
#endif // defined(ENABLE_SM120)

};

using FusedMultiHeadAttentionXMMAKernel
    = pluginInternal::TFusedMultiHeadAttentionXMMAKernel<FusedMultiHeadAttentionKernelMetaInfoV1,
        Fused_multihead_attention_params>;
using FusedMHAKernelFactory = pluginInternal::TFusedMHAKernelFactory<FusedMultiHeadAttentionXMMAKernel>;

inline const FusedMultiHeadAttentionXMMAKernel* getXMMAKernels(Data_type type, uint32_t sm)
{
    return FusedMHAKernelFactory::Get().getXMMAKernels(
        sMhaKernelMetaInfos, sizeof(sMhaKernelMetaInfos) / sizeof(sMhaKernelMetaInfos[0]), type, sm);
}

} // namespace bert
} // namespace plugin
} // namespace nvinfer1
#endif // _BERT_FMHA_FMHA