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

// Need 10.1 for cublasGemmStridedBatchedEx
#include <cuda.h>
#if CUDA_VERSION >= 10010

#include "NvInfer.h"
#include "bertQKVToContextPlugin/fused_multihead_attention/fused_multihead_attention.h"
#include "bertQKVToContextPlugin/fused_multihead_attention_v2/fused_multihead_attention_v2.h"
#include "common/bertCommon.h"
#include "common/serialize.hpp"
#include "mhaRunner.h"
#include "qkvToContextPluginLegacy.h"

#include <cstdint>
#include <cstring>
#include <iostream>
#include <memory>
#include <string_view>
#include <tuple>
#include <vector>

using namespace nvinfer1;
using namespace nvinfer1::plugin;
using namespace nvinfer1::plugin::bert;
using namespace nvinfer1::pluginInternal;

namespace
{
using namespace std::string_view_literals;
char const* const kQKV_TO_CONTEXT_PLUGIN_LEGACY_VERSION{"1"};
char const* const kQKV_TO_CONTEXT_VAR_SEQLEN_LEGACY_PLUGIN_VERSION{"2"};
char const* const kQKV_TO_CONTEXT_PLUGIN_NAME{"CustomQKVToContextPluginDynamic"};
} // namespace

REGISTER_TENSORRT_PLUGIN(QKVToContextPluginDynamicLegacyCreator);

constexpr uint32_t kIIDX = 0; // index of the input tensor
constexpr uint32_t kMIDX = 1; // index of the mask

REGISTER_TENSORRT_PLUGIN(QKVToContextVarSeqlenPluginLegacyCreator);

QKVToContextPluginDynamicLegacy::QKVToContextPluginDynamicLegacy(std::string const name, DataType const type,
    int32_t const hiddenSize, int32_t const numHeads, float const dqProbs, bool hasImask)
    : mLayerName(name)
    , mS(0)
    , mB(0)
    , mHeadSize(hiddenSize / numHeads)
    , mHiddenSize(hiddenSize)
    , mNumHeads(numHeads)
    , mHasImask(hasImask)
    , mType(type)
    , mDqProbs(dqProbs)

{
    mSM = getSmVersion();
}

QKVToContextPluginDynamicLegacy::QKVToContextPluginDynamicLegacy(
    std::string const name, void const* data, size_t length)
    : mLayerName(name)
{
    BERT_DEBUG_MSG("QKV Deser Start");
    deserialize_value(&data, &length, &mType);
    deserialize_value(&data, &length, &mNumHeads);
    deserialize_value(&data, &length, &mHeadSize);
    deserialize_value(&data, &length, &mHasImask);
    deserialize_value(&data, &length, &mHiddenSize);
    deserialize_value(&data, &length, &mSM);
    deserialize_value(&data, &length, &mS);
    deserialize_value(&data, &length, &mB);

    deserialize_value(&data, &length, &mDqProbs);

    createMHARunner();

    int32_t hasUnfusedRunner = 0;
    deserialize_value(&data, &length, &hasUnfusedRunner);
    if (hasUnfusedRunner)
    {
        PLUGIN_ASSERT(unfusedDispatcher.get());
        unfusedDispatcher->deserialize(data, length);
    }

    BERT_DEBUG_MSG("QKV Deser done");
}

void QKVToContextPluginDynamicLegacy::createMHARunner()
{
    if (!fusedDispatcher.get())
    {
        if (mType == DataType::kHALF)
        {
            fusedDispatcher.reset(new FusedMHARunnerFP16(mNumHeads, mSM));
        }
        else if (mType == DataType::kINT8)
        {
            fusedDispatcher.reset(new FusedMHARunnerInt8(mNumHeads, mSM, mDqProbs));
        }
    }

    if (!unfusedDispatcher.get())
    {
        unfusedDispatcher.reset(new UnfusedMHARunner(mType, mNumHeads, mSM));
    }
}

// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt* QKVToContextPluginDynamicLegacy::clone() const noexcept
{
    BERT_DEBUG_MSG("QKV Clone");

    std::unique_ptr<QKVToContextPluginDynamicLegacy> ret;
    // the workspacesize is 0 if we have not call setup the dispatcher yet.
    if (unfusedDispatcher.get() && unfusedDispatcher->getWorkspaceSize())
    {
        auto buff = std::make_unique<std::byte[]>(getSerializationSize());
        serialize(buff.get());

        ret = std::make_unique<QKVToContextPluginDynamicLegacy>(mLayerName, buff.get(), getSerializationSize());
    }
    else
    {
        ret = std::make_unique<QKVToContextPluginDynamicLegacy>(
            mLayerName, mType, mHiddenSize, mNumHeads, mDqProbs, mHasImask);
    }

    ret->setPluginNamespace(mNamespace.c_str());
    BERT_DEBUG_MSG("QKV Clone done");
    return ret.release();
}

DimsExprs QKVToContextPluginDynamicLegacy::getOutputDimensions(
    int32_t outputIndex, DimsExprs const* inputs, int32_t /*nbInputs*/, IExprBuilder& exprBuilder) noexcept
{
    // Input is BxSx3*N*H, output should be BxSxN*H
    PLUGIN_ASSERT(outputIndex == 0);
    // Copy over everything
    DimsExprs output(inputs[kIIDX]);
    // Divide last dim by three
    auto const* three = exprBuilder.constant(3);
    output.d[HDIM] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[kIIDX].d[HDIM], *three);
    return output;
}
// NOLINTNEXTLINE(readability-function-cognitive-complexity)
bool QKVToContextPluginDynamicLegacy::supportsFormatCombination(
    int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t /*nbOutputs*/) noexcept
{
    PLUGIN_ASSERT(pos >= 0);
    PLUGIN_ASSERT(pos < 2 + mHasImask);
    PLUGIN_ASSERT(nbInputs == 1 + mHasImask);
    auto const* in = inOut;
    auto const* out = inOut + nbInputs;
    int32_t packedSize = getMHAMaskPackedSize(mSM, mType, in->dims.d[SDIM]);

    // we only support int8 IO in fused mha runner, and we only support fused mha runner on Xavier, Turing and Ampere
    if (mType == DataType::kINT8)
    {
        if (!elem(mSM, {kSM_75, kSM_80, kSM_86, kSM_87, kSM_89, kSM_90, kSM_100, kSM_120}))
        {
            gLogError << "INT8 IO is only supported on Turing, Ampere, Hopper and Blackwell for plugin "
                      << kQKV_TO_CONTEXT_PLUGIN_NAME << std::endl;
            return false;
        }
        if (in->dims.d[SDIM] == -1)
        {
            gLogError << "INT8 IO not support dynamic shape in sequence dimension for plugin "
                      << kQKV_TO_CONTEXT_PLUGIN_NAME << std::endl;
            return false;
        }
        if (packedSize == unfusedMaskSize)
        {
            gLogError << "INT8 IO only support sequence length 128,384 for plugin " << kQKV_TO_CONTEXT_PLUGIN_NAME
                      << std::endl;
            return false;
        }
    }

    if (pos == 0)
    {
        bool isFormatSupported = in->format == TensorFormat::kLINEAR;
        if (mType == DataType::kINT8)
        {
            if (in->dims.d[HDIM] % 32U == 0)
            {
                isFormatSupported = in->format == TensorFormat::kCHW32;
            }
            else
            {
                isFormatSupported = in->format == TensorFormat::kCHW4;
            }
        }

        // must not check descriptions > pos
        return (in->type == mType) &&         // precision
            isFormatSupported &&              // format
            (in->dims.nbDims == 5) &&         // num dims
            ((in->dims.d[HDIM] % 3U) == 0) && // see getOutputDimensions
            ((in->dims.d[3]) == 1) &&         // for fc
            ((in->dims.d[4]) == 1)            // for fc
            ;
    }

    // pos==1
    if ((mHasImask && pos == 1)) // pos 1 is the mask
    {
        auto const* inMask = &inOut[1];
        if (inMask->dims.d[1] != -1 && inMask->dims.d[1] != packedSize)
        {
            gLogError << "CustomEmbLayerNormPluginDynamic returned mask with pack size " << inMask->dims.d[1]
                      << ", but " << kQKV_TO_CONTEXT_PLUGIN_NAME << " expects mask pack size " << packedSize
                      << std::endl;
            return false;
        }

        // detect full mask and check that it was produced
        return (inMask->type == DataType::kINT32) &&     // precision
            (inMask->format == TensorFormat::kLINEAR) && // format
            (inMask->dims.nbDims == 2) &&                // Bx2*maskSize
            (inMask->dims.d[0] == in->dims.d[BDIM]);
    }

    if (!mHasImask || pos == 2) // output pos
    {
        bool isFormatSupported = out->format == TensorFormat::kLINEAR;
        if (mType == DataType::kINT8)
        {
            if (out->dims.d[HDIM] % 32U == 0)
            {
                isFormatSupported = out->format == TensorFormat::kCHW32;
            }
            else
            {
                isFormatSupported = out->format == TensorFormat::kCHW4;
            }
        }

        return (in->type == out->type) &&                      // precision
            isFormatSupported &&                               // format
            (out->dims.nbDims == 5) &&                         // num dims
            ((in->dims.d[HDIM] / 3) == (out->dims.d[HDIM])) && // div 3
            ((out->dims.d[3]) == 1) &&                         // for fc
            ((out->dims.d[4]) == 1) &&                         // for fc
            ((out->dims.d[BDIM]) == in->dims.d[BDIM]) &&       // check B
            ((out->dims.d[SDIM]) == in->dims.d[SDIM])          // check S
            ;
    }

    return false;
}
void QKVToContextPluginDynamicLegacy::configurePlugin(
    DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
{
    PLUGIN_ASSERT(nbInputs == 1 + mHasImask);
    PLUGIN_ASSERT(nbOutputs == 1);
    PluginTensorDesc const& inDesc = in[kIIDX].desc;
    TRT_UNUSED inDesc;
    PluginTensorDesc const& outDesc = out->desc;
    TRT_UNUSED outDesc;
    PLUGIN_ASSERT(mType == inDesc.type);
    PLUGIN_ASSERT(mType == outDesc.type);
    PLUGIN_ASSERT(inDesc.dims.d[BDIM] == outDesc.dims.d[BDIM]);
    PLUGIN_ASSERT(inDesc.dims.d[SDIM] == outDesc.dims.d[SDIM]);
    PLUGIN_ASSERT(inDesc.dims.d[HDIM] == 3 * outDesc.dims.d[HDIM]);
    if (mHasImask)
    {
        PluginTensorDesc const& maskDesc = in[kMIDX].desc;
        TRT_UNUSED maskDesc;
        PLUGIN_ASSERT(maskDesc.dims.d[0] == inDesc.dims.d[BDIM]);
    }

    createMHARunner();

    int32_t const S = inDesc.dims.d[SDIM];
    int32_t const B = inDesc.dims.d[BDIM] <= 0 ? in->max.d[BDIM] : inDesc.dims.d[BDIM];
    if (S <= 0)
    {
        // in dynamic shape build stage, we setup with max sequence that cannot fused
        int32_t const Smin = in->min.d[SDIM];
        int32_t const Smax = in->max.d[SDIM];

        if (fusedDispatcher.get())
        {
            for (int32_t i = Smax; i >= Smin; --i)
            {
                if (!fusedDispatcher->isValid(mHeadSize, i))
                {
                    unfusedDispatcher->setup(i, B, mHeadSize);
                    mS = i;
                    mB = B;
                    break;
                }
            }
        }
        else
        {
            unfusedDispatcher->setup(Smax, B, mHeadSize);
            mS = Smax;
            mB = B;
        }
    }
    else
    {
        // in inference stage or in static shape build stage
        if (fusedDispatcher.get() && fusedDispatcher->isValid(mHeadSize, S))
        {
            fusedDispatcher->setup(S, B, mHeadSize);
        }
        else
        {
            unfusedDispatcher->setup(S, B, mHeadSize);
        }
        mS = S;
        mB = B;
    }
}

size_t QKVToContextPluginDynamicLegacy::getWorkspaceSize(PluginTensorDesc const* /*inputs*/, int32_t /*nbInputs*/,
    PluginTensorDesc const* /*outputs*/, int32_t /*nbOutputs*/) const noexcept
{
    // only unfused kernel need workspace, and we need larger workspace for larger sequence length
    // we have already setup unfusedDispatcher with max sequence in configurePlugin
    // if unfusedDispatcher is not initialized in configurePlugin
    PLUGIN_ASSERT(unfusedDispatcher.get());
    return unfusedDispatcher->getWorkspaceSize();
}

// IPluginV2Ext Methods
DataType QKVToContextPluginDynamicLegacy::getOutputDataType(
    int32_t index, nvinfer1::DataType const* inputTypes, int32_t /*nbInputs*/) const noexcept
{
    PLUGIN_ASSERT(index == 0);
    PLUGIN_ASSERT(
        inputTypes[0] == DataType::kFLOAT || inputTypes[0] == DataType::kHALF || inputTypes[0] == DataType::kINT8);
    return inputTypes[0];
}

void QKVToContextPluginDynamicLegacy::attachToContext(
    cudnnContext* cudnn, cublasContext* cublas, nvinfer1::IGpuAllocator* allocator) noexcept
{
    try
    {
        mCublasWrapper = createPluginCublasWrapper(allocator);
        mCublas = mCublasWrapper->getCublasHandle();
        PLUGIN_VALIDATE(mCublas != nullptr);
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
}

// IPluginV2 Methods
char const* QKVToContextPluginDynamicLegacy::getPluginType() const noexcept
{
    return kQKV_TO_CONTEXT_PLUGIN_NAME;
}

char const* QKVToContextPluginDynamicLegacy::getPluginVersion() const noexcept
{
    return kQKV_TO_CONTEXT_PLUGIN_LEGACY_VERSION;
}

int32_t QKVToContextPluginDynamicLegacy::getNbOutputs() const noexcept
{
    return 1;
}

int32_t QKVToContextPluginDynamicLegacy::initialize() noexcept
{
    return 0;
}

void QKVToContextPluginDynamicLegacy::terminate() noexcept {}

size_t QKVToContextPluginDynamicLegacy::getSerializationSize() const noexcept
{
    PLUGIN_ASSERT(unfusedDispatcher.get());
    return sizeof(mNumHeads) + sizeof(mHeadSize) + sizeof(DataType) + sizeof(mHasImask) + sizeof(mHiddenSize)
        + sizeof(mSM) + sizeof(mS) + sizeof(mB) + sizeof(mDqProbs) + sizeof(int32_t)
        + unfusedDispatcher->getSerializationSize();
}

void QKVToContextPluginDynamicLegacy::serialize(void* buffer) const noexcept
{
    serialize_value(&buffer, mType);
    serialize_value(&buffer, mNumHeads);
    serialize_value(&buffer, mHeadSize);
    serialize_value(&buffer, mHasImask);
    serialize_value(&buffer, mHiddenSize);
    serialize_value(&buffer, mSM);
    serialize_value(&buffer, mS);
    serialize_value(&buffer, mB);

    serialize_value(&buffer, mDqProbs);
    if (unfusedDispatcher.get() && unfusedDispatcher->getWorkspaceSize())
    {
        int32_t hasUnfusedRunner = 1;
        serialize_value(&buffer, hasUnfusedRunner);
        unfusedDispatcher->serialize(buffer);
    }
    else
    {
        int32_t hasUnfusedRunner = 0;
        serialize_value(&buffer, hasUnfusedRunner);
    }
}

void QKVToContextPluginDynamicLegacy::destroy() noexcept
{
    delete this;
}

void QKVToContextPluginDynamicLegacy::setPluginNamespace(char const* libNamespace) noexcept
{
    mNamespace = libNamespace;
}

char const* QKVToContextPluginDynamicLegacy::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

int32_t QKVToContextPluginDynamicLegacy::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
    void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
    PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr);
    PLUGIN_ASSERT(mS == inputDesc->dims.d[SDIM]);
    PLUGIN_ASSERT(mB == inputDesc->dims.d[BDIM]);

    try
    {
        void const* const maskPtr = mHasImask ? inputs[1] : nullptr;
        if (mHasImask && fusedDispatcher.get() && fusedDispatcher->isValid(mHeadSize, inputDesc->dims.d[SDIM]))
        {
            fusedDispatcher->run(
                inputDesc[0], outputDesc[0], inputs[0], maskPtr, outputs[0], workspace, stream, mCublas);
        }
        else
        {
            PLUGIN_VALIDATE(unfusedDispatcher.get(), "The Unfused MHARunner is uninitialized, no MHARunner available!");
            PLUGIN_VALIDATE(mType != DataType::kINT8, "The Unfused MHARunner does not support INT8!");
            unfusedDispatcher->run(
                inputDesc[0], outputDesc[0], inputs[0], maskPtr, outputs[0], workspace, stream, mCublas);
        }
    }
    catch (std::exception const& e)
    {
        caughtError(e);
        return -1;
    }
    return 0;
}

QKVToContextPluginDynamicLegacyCreator::QKVToContextPluginDynamicLegacyCreator()
{
    mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("hidden_size", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("has_mask", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("dq_probs", nullptr, PluginFieldType::kFLOAT32, 1));

    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

char const* QKVToContextPluginDynamicLegacyCreator::getPluginName() const noexcept
{
    return kQKV_TO_CONTEXT_PLUGIN_NAME;
}

char const* QKVToContextPluginDynamicLegacyCreator::getPluginVersion() const noexcept
{
    return kQKV_TO_CONTEXT_PLUGIN_LEGACY_VERSION;
}

PluginFieldCollection const* QKVToContextPluginDynamicLegacyCreator::getFieldNames() noexcept
{
    return &mFC;
}

IPluginV2* QKVToContextPluginDynamicLegacyCreator::createPlugin(
    char const* name, PluginFieldCollection const* fc) noexcept
{
    try
    {
        BERT_DEBUG_MSG("Creating QKV2ContextPlugin...");
        PLUGIN_VALIDATE(fc != nullptr);
        int32_t hiddenSize = 0;
        // Since numHeads must always exist or validateRequiredAttributes will fail,
        // we can set numHeads to -1 so that static analysis tools don't warn about
        // a division by zero in QKVToContextPluginDynamicLegacy constructor.
        int32_t numHeads{-1};
        bool hasMask = false;
        int32_t typeId = -1;

        float dqProbs = -1.0F;

        PLUGIN_VALIDATE(fc->fields != nullptr);
        plugin::validateRequiredAttributesExist({"type_id", "hidden_size", "num_heads", "has_mask"}, fc);

        for (int32_t i = 0; i < fc->nbFields; i++)
        {
            PLUGIN_VALIDATE(fc->fields[i].name != nullptr);
            PLUGIN_VALIDATE(fc->fields[i].data != nullptr);
            std::string_view const field_name = fc->fields[i].name;

            if (field_name == "type_id"sv)
            {
                typeId = *static_cast<int32_t const*>(fc->fields[i].data);
                PLUGIN_VALIDATE(typeId >= 0 && typeId <= 2, ("QKV: Invalid TypeId " + std::to_string(typeId)).c_str());
                BERT_DEBUG_VALUE("Building typeId: ", typeId);
            }
            if (field_name == "hidden_size"sv)
            {
                hiddenSize = *static_cast<int32_t const*>(fc->fields[i].data);
                PLUGIN_VALIDATE(hiddenSize > 0, ("QKV: Invalid hiddenSize " + std::to_string(hiddenSize)).c_str());
                BERT_DEBUG_VALUE("Building hiddenSize: ", hiddenSize);
            }
            if (field_name == "num_heads"sv)
            {
                numHeads = *static_cast<int32_t const*>(fc->fields[i].data);
                PLUGIN_VALIDATE(numHeads > 0, ("QKV: Invalid numHeads " + std::to_string(numHeads)).c_str());
                BERT_DEBUG_VALUE("Building numHeads: ", numHeads);
            }
            if (field_name == "has_mask"sv)
            {
                auto hasMaskValue = *static_cast<int32_t const*>(fc->fields[i].data);
                PLUGIN_VALIDATE(hasMaskValue == 0 || hasMaskValue == 1,
                    ("QKV: Invalid hasMask " + std::to_string(hasMaskValue)).c_str());
                hasMask = static_cast<bool>(hasMaskValue);
                BERT_DEBUG_VALUE("Building hasMask: ", hasMask);
            }

            if (field_name == "dq_probs"sv)
            {
                dqProbs = *static_cast<float const*>(fc->fields[i].data);
                PLUGIN_VALIDATE(dqProbs > 0.0F, ("QKV: Invalid dqProbs " + std::to_string(dqProbs)).c_str());
                BERT_DEBUG_VALUE("Building dqProbs: ", dqProbs);
            }
        }

        BERT_DEBUG_MSG("Building the Plugin...");
        auto type = static_cast<DataType>(typeId);
        if (type == DataType::kINT8 && dqProbs < 0)
        {
            BERT_DEBUG_MSG("Using default scale factor");
            dqProbs = 1.F / 127.F;
        }

        auto p = std::make_unique<QKVToContextPluginDynamicLegacy>(name, type, hiddenSize, numHeads, dqProbs, hasMask);
        return p.release();
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
    return nullptr;
}

IPluginV2* QKVToContextPluginDynamicLegacyCreator::deserializePlugin(
    char const* name, void const* serialData, size_t serialLength) noexcept
{
    // This object will be deleted when the network is destroyed, which will
    // call QKVToContextPluginDynamicLegacy::destroy()
    return new QKVToContextPluginDynamicLegacy(name, serialData, serialLength);
}

void QKVToContextPluginDynamicLegacyCreator::setPluginNamespace(char const* libNamespace) noexcept
{
    mNamespace = libNamespace;
}

char const* QKVToContextPluginDynamicLegacyCreator::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

QKVToContextVarSeqlenPluginLegacy::QKVToContextVarSeqlenPluginLegacy(std::string const name, DataType const type,
    int32_t const hiddenSize, int32_t const numHeads, float const dqProbs, bool hasImask, bool varSeqlen,
    bool useInt8ScaleMax)
    : mLayerName(name)
    , mS(0)
    , mB(0)
    , mHeadSize(hiddenSize / numHeads)
    , mHiddenSize(hiddenSize)
    , mNumHeads(numHeads)
    , mHasImask(hasImask)
    , mType(type)
    , mDqProbs(dqProbs)
    , mHdim(HDIM)
    , mUseVarSeqlen(varSeqlen)
    , mUseInt8ScaleMax(useInt8ScaleMax)
{
    mSM = getSmVersion();

    if (varSeqlen)
    {
        // variable sequence length is only supported with the fused MHA kernels
        // we should not override mS!
        bool isSMSupported = elem(mSM, {kSM_75, kSM_80, kSM_86, kSM_87, kSM_89, kSM_90, kSM_100, kSM_120});
        PLUGIN_ASSERT(isSMSupported && (type == DataType::kINT8 || type == DataType::kHALF)
            && "requesting maxSeqlen not compatible with GPU arch");
        // the layout changes: SxB will be a combined \sum_i s_i and hdim will be the 2nd dimension instead of the third
        mHdim = 1;
    }
}

QKVToContextVarSeqlenPluginLegacy::QKVToContextVarSeqlenPluginLegacy(
    std::string const name, void const* data, size_t length)
    : mLayerName(name)
{
    BERT_DEBUG_MSG("QKV Deser Start");
    deserialize_value(&data, &length, &mType);
    deserialize_value(&data, &length, &mNumHeads);
    deserialize_value(&data, &length, &mHeadSize);
    deserialize_value(&data, &length, &mHasImask);
    deserialize_value(&data, &length, &mHiddenSize);
    deserialize_value(&data, &length, &mSM);
    deserialize_value(&data, &length, &mS);
    deserialize_value(&data, &length, &mB);

    deserialize_value(&data, &length, &mDqProbs);

    deserialize_value(&data, &length, &mUseVarSeqlen);
    deserialize_value(&data, &length, &mHdim);
    deserialize_value(&data, &length, &mUseInt8ScaleMax);

    createMHARunner();
    mDispatcher->deserialize(data, length);

    BERT_DEBUG_MSG("QKV Deser done");
}

void QKVToContextVarSeqlenPluginLegacy::createMHARunner()
{
    if (mDispatcher.get())
    {
        return;
    }

    if (mUseVarSeqlen)
    {
        PLUGIN_ASSERT(mHeadSize <= 64);
        {
            if (mHeadSize != 64)
            {
                mPatcher.reset(new QkvPaddingRunner(mType));
            }
            if (mType == DataType::kHALF)
            {
                mDispatcher.reset(new FusedMHARunnerFP16v2(mNumHeads, mSM));
            }
            else if (mType == DataType::kINT8)
            {
                mDispatcher.reset(new FusedMHARunnerInt8v2(mNumHeads, mSM, mDqProbs, mUseInt8ScaleMax));
            }
        }
    }
    else
    {
        PLUGIN_ASSERT(mType != DataType::kINT8);
        mDispatcher.reset(new UnfusedMHARunner(mType, mNumHeads, mSM));
    }
}

// IPluginV2DynamicExt Methods
nvinfer1::IPluginV2DynamicExt* QKVToContextVarSeqlenPluginLegacy::clone() const noexcept
{
    BERT_DEBUG_MSG("QKV Clone");

    std::unique_ptr<QKVToContextVarSeqlenPluginLegacy> ret;
    if (mDispatcher.get())
    {
        auto buff = std::make_unique<std::byte[]>(getSerializationSize());
        serialize(buff.get());

        ret = std::make_unique<QKVToContextVarSeqlenPluginLegacy>(mLayerName, buff.get(), getSerializationSize());
    }
    else
    {
        ret = std::make_unique<QKVToContextVarSeqlenPluginLegacy>(
            mLayerName, mType, mHiddenSize, mNumHeads, mDqProbs, mHasImask, mUseVarSeqlen, mUseInt8ScaleMax);
    }

    ret->setPluginNamespace(mNamespace.c_str());
    BERT_DEBUG_MSG("QKV Clone done");
    return ret.release();
}

DimsExprs QKVToContextVarSeqlenPluginLegacy::getOutputDimensions(
    int32_t outputIndex, DimsExprs const* inputs, int32_t /*nbInputs*/, IExprBuilder& exprBuilder) noexcept
{
    // Input is BxSx3*N*H, output should be BxSxN*H
    PLUGIN_ASSERT(outputIndex == 0);
    // Copy over everything
    DimsExprs output(inputs[kIIDX]);
    // Divide last dim by three
    auto const* three = exprBuilder.constant(3);
    output.d[mHdim] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[kIIDX].d[mHdim], *three);
    return output;
}

bool QKVToContextVarSeqlenPluginLegacy::supportsFormatCombination(
    int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
    // we only support variable sequence and int8 IO in fused mha runner, and we only support fused mha runner on
    // Turing, Ampere, Hopper and Blackwell
    bool const hasV2Kernels = elem(mSM, {kSM_75, kSM_80, kSM_86, kSM_87, kSM_89, kSM_90, kSM_100, kSM_120});
    PLUGIN_ASSERT((mType != DataType::kINT8 || hasV2Kernels)
        && "INT8 IO is only supported on Xavier, Turing, Ampere, Hopper and Blackwell!");
    PLUGIN_ASSERT((!mUseVarSeqlen || hasV2Kernels)
        && "Variable sequence is only supported on Xavier, Turing, Ampere, Hopper and Blackwell!");

    PLUGIN_ASSERT(pos >= 0);
    PLUGIN_ASSERT(pos < 2 + mHasImask + 2 * mUseVarSeqlen);
    PLUGIN_ASSERT(nbInputs == 1 + mHasImask + 2 * mUseVarSeqlen);
    PLUGIN_ASSERT(nbOutputs == 1);
    auto const* in = inOut;
    auto const* out = inOut + nbInputs;
    if (mUseVarSeqlen)
    {
        PLUGIN_ASSERT((mType == DataType::kHALF || mType == DataType::kINT8)
            && "Conditions for variable seqlen support not fulfilled");
        // qkv, mask, cu_seqlens, dummy
        PLUGIN_ASSERT(nbInputs == 4 && "for varseqlen, expected 4 inputs");
    }

    auto const inDims = in->dims;
    auto const outDims = out->dims;

    auto supportedFormat = TensorFormat::kLINEAR;
    if (mType == DataType::kINT8)
    {
        supportedFormat = (inDims.d[mHdim] % 32U == 0) ? TensorFormat::kCHW32 : TensorFormat::kCHW4;
    }

    int32_t supportedNbDims = 5;
    if (mUseVarSeqlen)
    {
        supportedNbDims = 4;
    }

    bool supportedHdim = (pos == 0) ? (inDims.d[mHdim] % 3U == 0) : (inDims.d[mHdim] / 3 == outDims.d[mHdim]);

    if (pos == 0 || pos == nbInputs)
    { // check input and output
        auto const& desc = inOut[pos];
        return (desc.type == mType) &&               // check type
            (desc.format == supportedFormat) &&      // check format
            (desc.dims.nbDims == supportedNbDims) && // check dims:
            (supportedHdim) &&                       // - hidden dims multiple of 3 for qkv
            (desc.dims.d[mHdim + 1] == 1) &&         // - dummy 1 or h
            (desc.dims.d[mHdim + 2] == 1)            // - dummy 1 or w
            ;
    }

    PLUGIN_ASSERT(mHasImask);
    if (pos == 1)
    { // must be input mask
        auto const* mask = &inOut[pos];
        if (mUseVarSeqlen)
        {
            // dummy input
            return true;
        }

        return mask->format == TensorFormat::kLINEAR && (mask->type == DataType::kINT32) && // precision
            (mask->dims.nbDims == 1)                                                        // num dims
            ;
    }
    PLUGIN_ASSERT(mUseVarSeqlen);
    if (pos == 2)
    { // must be cuSeqlens
        // cuSeqlens is a int32_t array of size B+1
        auto const* seqlens = &inOut[pos];
        return (seqlens->type == DataType::kINT32) && (seqlens->format == TensorFormat::kLINEAR);
    }
    if (pos == 3)
    {
        // this is the dummy input
        return inOut[pos].dims.nbDims == 1;
    }
    return false;
}

void QKVToContextVarSeqlenPluginLegacy::configurePlugin(
    DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
{
    PLUGIN_ASSERT(nbInputs == 1 + mHasImask + 2 * mUseVarSeqlen);
    PLUGIN_ASSERT(nbOutputs == 1);
    PluginTensorDesc const& inDesc = in[kIIDX].desc;
    TRT_UNUSED inDesc;
    PluginTensorDesc const& outDesc = out->desc;
    TRT_UNUSED outDesc;
    PLUGIN_ASSERT(mType == inDesc.type);
    PLUGIN_ASSERT(mType == outDesc.type);
    if (!mUseVarSeqlen)
    {
        PLUGIN_ASSERT(inDesc.dims.d[BDIM] == outDesc.dims.d[BDIM]);
        PLUGIN_ASSERT(inDesc.dims.d[SDIM] == outDesc.dims.d[SDIM]);
        PLUGIN_ASSERT(inDesc.dims.d[mHdim] == 3 * outDesc.dims.d[mHdim]);
        if (mHasImask)
        {
            PluginTensorDesc const& maskDesc = in[kMIDX].desc;
            TRT_UNUSED maskDesc;
            PLUGIN_ASSERT(maskDesc.dims.d[0] == inDesc.dims.d[BDIM]);
        }

        int32_t const S = inDesc.dims.d[SDIM] <= 0 ? in->max.d[SDIM] : inDesc.dims.d[SDIM];
        int32_t const B = inDesc.dims.d[BDIM] <= 0 ? in->max.d[BDIM] : inDesc.dims.d[BDIM];

        if (S != mS || B != mB)
        {
            BERT_DEBUG_MSG("setting up MHA runner for single sequence length");
            createMHARunner();
            this->mDispatcher->setup(S, B, mHeadSize);
            mS = S;
            mB = B;
        }
    }
    else
    {
        BERT_DEBUG_MSG("setting up MHA runner for variable sequence length");
        createMHARunner();
        // need to initialize S and B with somewhat useful values, they will be reset at enqueue for the actual
        // batchsize
        this->mDispatcher->setup(256, 1, mHeadSize);
    }
}

size_t QKVToContextVarSeqlenPluginLegacy::getWorkspaceSize(PluginTensorDesc const* inputs, int32_t /* nbInputs */,
    PluginTensorDesc const* /* outputs */, int32_t /* nbOutputs */) const noexcept
{
    size_t paddingWorkpaceSize = mPatcher ? mPatcher->getWorkspaceSize(inputs[0].dims.d[0], mNumHeads) : 0;
    return mDispatcher->getWorkspaceSize() + paddingWorkpaceSize;
}

// IPluginV2Ext Methods
DataType QKVToContextVarSeqlenPluginLegacy::getOutputDataType(
    int32_t index, nvinfer1::DataType const* inputTypes, int32_t /*nbInputs*/) const noexcept
{
    PLUGIN_ASSERT(index == 0);
    PLUGIN_ASSERT(
        inputTypes[0] == DataType::kFLOAT || inputTypes[0] == DataType::kHALF || inputTypes[0] == DataType::kINT8);
    return inputTypes[0];
}

void QKVToContextVarSeqlenPluginLegacy::attachToContext(
    cudnnContext* cudnn, cublasContext* cublas, nvinfer1::IGpuAllocator* allocator) noexcept
{
    try
    {
        mCublasWrapper = createPluginCublasWrapper(allocator);
        mCublas = mCublasWrapper->getCublasHandle();
        PLUGIN_VALIDATE(mCublas != nullptr);
    }
    catch (std::exception const& e)
    {
        caughtError(e);
    }
}

// IPluginV2 Methods
char const* QKVToContextVarSeqlenPluginLegacy::getPluginType() const noexcept
{
    return kQKV_TO_CONTEXT_PLUGIN_NAME;
}

char const* QKVToContextVarSeqlenPluginLegacy::getPluginVersion() const noexcept
{
    return kQKV_TO_CONTEXT_VAR_SEQLEN_LEGACY_PLUGIN_VERSION;
}

int32_t QKVToContextVarSeqlenPluginLegacy::getNbOutputs() const noexcept
{
    return 1;
}

int32_t QKVToContextVarSeqlenPluginLegacy::initialize() noexcept
{
    return 0;
}

void QKVToContextVarSeqlenPluginLegacy::terminate() noexcept {}

size_t QKVToContextVarSeqlenPluginLegacy::getSerializationSize() const noexcept
{
    return sizeof(mNumHeads) + sizeof(mHeadSize) + sizeof(DataType) + sizeof(mHasImask) + sizeof(mHiddenSize)
        + sizeof(mSM) + sizeof(mS) + sizeof(mB) + sizeof(mDqProbs) + mDispatcher->getSerializationSize()
        + sizeof(mUseVarSeqlen) + sizeof(mHdim) + sizeof(mUseInt8ScaleMax);
}

void QKVToContextVarSeqlenPluginLegacy::serialize(void* buffer) const noexcept
{
    serialize_value(&buffer, mType);
    serialize_value(&buffer, mNumHeads);
    serialize_value(&buffer, mHeadSize);
    serialize_value(&buffer, mHasImask);
    serialize_value(&buffer, mHiddenSize);
    serialize_value(&buffer, mSM);
    serialize_value(&buffer, mS);
    serialize_value(&buffer, mB);

    serialize_value(&buffer, mDqProbs);
    serialize_value(&buffer, mUseVarSeqlen);
    serialize_value(&buffer, mHdim);
    serialize_value(&buffer, mUseInt8ScaleMax);
    mDispatcher->serialize(buffer);
}

void QKVToContextVarSeqlenPluginLegacy::destroy() noexcept
{
    delete this;
}

void QKVToContextVarSeqlenPluginLegacy::setPluginNamespace(char const* libNamespace) noexcept
{
    mNamespace = libNamespace;
}

char const* QKVToContextVarSeqlenPluginLegacy::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

// NOLINTNEXTLINE(readability-function-cognitive-complexity)
int32_t QKVToContextVarSeqlenPluginLegacy::enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
    nvinfer1::PluginTensorDesc const* outputDesc, void const* const* inputs, void* const* outputs, void* workspace,
    cudaStream_t stream) noexcept
{
    PLUGIN_VALIDATE(inputDesc != nullptr && outputDesc != nullptr && inputs != nullptr && outputs != nullptr);

    if (mUseVarSeqlen)
    {
        int32_t const B = inputDesc[2].dims.d[0] - 1;
        int32_t const maxS = inputDesc[3].dims.d[0];
        PLUGIN_ASSERT((maxS <= 512)
            && "No implementation for variable sequence length multi-head attention plugin with sequence > 512.");

        int32_t S = 512;
        if (DataType::kHALF == mType && maxS <= 64)
        {
            S = 64;
        }
        else if (DataType::kHALF == mType && maxS <= 96)
        {
            S = 96;
        }
        else if (maxS <= 128)
        {
            S = 128;
        }
        else if (maxS <= 192)
        {
            S = 192;
            if (mType == DataType::kHALF)
            {
                S = 256;
            }
        }
        else if (maxS <= 256)
        {
            S = 256;
        }
        else if (maxS <= 384)
        {
            S = 384;
        }

        auto runV2Kernel = [this, &S, &B, &workspace, &inputDesc, &outputDesc, &stream, &inputs, &outputs](
                               MHARunner* dispatcher, QkvPaddingRunner* patcher, int32_t padSize) {
            PLUGIN_ASSERT(dispatcher);
            // Validate that we can padding to the dispatch required head size also there is kernel exist for this
            // sequence length.
            if (mHeadSize > padSize || !dispatcher->isValid(padSize, S))
            {
                return false;
            }
            dispatcher->setup(S, B, padSize);

            // Need pad and unpad to run the V2 kernel.
            if (mHeadSize < padSize)
            {
                PLUGIN_ASSERT(patcher);
                PLUGIN_ASSERT(padSize <= patcher->getMaxPaddingHeadSize());
                auto sumSeqLen = inputDesc[0].dims.d[0];
                auto paddingWorkspace = patcher->get16BytesAlignedPointer(workspace, dispatcher->getWorkspaceSize());
                auto ret = mPatcher->pad(inputs[0], paddingWorkspace, sumSeqLen, mNumHeads, mHeadSize, padSize, stream);
                if (ret != cudaSuccess)
                {
                    return false;
                }

                MhaRunParameter paddingArgs = patcher->patchMhaArgs(
                    inputDesc, outputDesc, inputs, outputs, paddingWorkspace, sumSeqLen, mNumHeads, padSize);
                try
                {
                    dispatcher->run(paddingArgs.inputDesc, paddingArgs.outputDesc, paddingArgs.inputs,
                        paddingArgs.outputs, workspace, stream, mCublas);
                }
                catch (std::exception const& e)
                {
                    caughtError(e);
                    return false;
                }

                ret = patcher->unpad(
                    paddingArgs.outputs[0], outputs[0], sumSeqLen, mNumHeads, mHeadSize, padSize, stream);
                return ret == cudaSuccess;
            }
            else
            {
                // No pad/unpad is needed.
                try
                {
                    dispatcher->run(inputDesc, outputDesc, inputs, outputs, workspace, stream, mCublas);
                }
                catch (std::exception const& e)
                {
                    caughtError(e);
                    return false;
                }
                return true;
            }
        };
        // Try pad head size to 32 first, if it failed, then try to pad head size to 64.
        if (!runV2Kernel(mDispatcher.get(), mPatcher.get(), 32) && !runV2Kernel(mDispatcher.get(), mPatcher.get(), 64))
        {
            return false;
        }

        return cudaGetLastError();
    }

    PLUGIN_ASSERT(mS == inputDesc->dims.d[SDIM]);
    PLUGIN_ASSERT(mB == inputDesc->dims.d[BDIM]);

    void const* maskPtr = mHasImask ? inputs[1] : nullptr;
    mDispatcher->run(inputDesc[0], outputDesc[0], inputs[0], maskPtr, outputs[0], workspace, stream, mCublas);
    return cudaGetLastError();
}

QKVToContextVarSeqlenPluginLegacyCreator::QKVToContextVarSeqlenPluginLegacyCreator()
{
    mPluginAttributes.clear();
    mPluginAttributes.emplace_back(PluginField("type_id", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("hidden_size", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("has_mask", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("dq_probs", nullptr, PluginFieldType::kFLOAT32, 1));
    mPluginAttributes.emplace_back(PluginField("var_seqlen", nullptr, PluginFieldType::kINT32, 1));
    mPluginAttributes.emplace_back(PluginField("use_int8_scale_max", nullptr, PluginFieldType::kINT32, 1));

    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

char const* QKVToContextVarSeqlenPluginLegacyCreator::getPluginName() const noexcept
{
    return kQKV_TO_CONTEXT_PLUGIN_NAME;
}

char const* QKVToContextVarSeqlenPluginLegacyCreator::getPluginVersion() const noexcept
{
    return kQKV_TO_CONTEXT_VAR_SEQLEN_LEGACY_PLUGIN_VERSION;
}

PluginFieldCollection const* QKVToContextVarSeqlenPluginLegacyCreator::getFieldNames() noexcept
{
    return &mFC;
}

IPluginV2* QKVToContextVarSeqlenPluginLegacyCreator::createPlugin(
    char const* name, PluginFieldCollection const* fc) noexcept
{
    BERT_DEBUG_MSG("Creating QKV2ContextPlugin...");

    int32_t hiddenSize = 0;
    // Since numHeads must always exist or validateRequiredAttributes will fail,
    // we can set numHeads to -1 so that static analysis tools don't warn about
    // a division by zero in QKVToContextVarSeqlenPlugin constructor.
    int32_t numHeads{-1};
    bool hasMask = false;
    int32_t typeId = -1;

    int32_t varSeqlen = 0;

    float dqProbs = -1;
    int32_t useInt8ScaleMax{-1};

    plugin::validateRequiredAttributesExist({"type_id", "hidden_size", "num_heads", "has_mask"}, fc);
    for (int32_t i = 0; i < fc->nbFields; i++)
    {
        std::string_view const field_name = fc->fields[i].name;

        if (field_name == "type_id"sv)
        {
            typeId = *static_cast<int32_t const*>(fc->fields[i].data);
            PLUGIN_VALIDATE(typeId >= 0 && typeId <= 2, ("QKV: Invalid TypeId " + std::to_string(typeId)).c_str());
            BERT_DEBUG_VALUE("Building typeId: ", typeId);
        }
        if (field_name == "hidden_size"sv)
        {
            hiddenSize = *static_cast<int32_t const*>(fc->fields[i].data);
            PLUGIN_VALIDATE(hiddenSize > 0, ("QKV: Invalid hiddenSize " + std::to_string(hiddenSize)).c_str());
            BERT_DEBUG_VALUE("Building hiddenSize: ", hiddenSize);
        }
        if (field_name == "num_heads"sv)
        {
            numHeads = *static_cast<int32_t const*>(fc->fields[i].data);
            PLUGIN_VALIDATE(numHeads > 0, ("QKV: Invalid numHeads " + std::to_string(numHeads)).c_str());
            BERT_DEBUG_VALUE("Building numHeads: ", numHeads);
        }
        if (field_name == "has_mask"sv)
        {
            hasMask = *static_cast<bool const*>(fc->fields[i].data);
            PLUGIN_VALIDATE(hasMask == 0 || hasMask == 1, ("QKV: Invalid hasMask " + std::to_string(hasMask)).c_str());
            BERT_DEBUG_VALUE("Building hasMask: ", hasMask);
        }

        if (field_name == "dq_probs"sv)
        {
            dqProbs = *static_cast<float const*>(fc->fields[i].data);
            PLUGIN_VALIDATE(dqProbs > 0.0F, ("QKV: Invalid dqProbs " + std::to_string(dqProbs)).c_str());
            BERT_DEBUG_VALUE("Building dqProbs: ", dqProbs);
        }
        if (field_name == "var_seqlen"sv)
        {
            varSeqlen = *static_cast<int32_t const*>(fc->fields[i].data);
            BERT_DEBUG_VALUE("Building var_seqlen: ", varSeqlen);
        }
        if (field_name == "use_int8_scale_max"sv)
        {
            useInt8ScaleMax = *static_cast<int32_t const*>(fc->fields[i].data);
            PLUGIN_VALIDATE(useInt8ScaleMax == 0 || useInt8ScaleMax == 1,
                ("QKV: Invalid useInt8ScaleMax " + std::to_string(useInt8ScaleMax)).c_str());
            BERT_DEBUG_VALUE("Building useInt8ScaleMax: ", useInt8ScaleMax);
        }
    }

    if (useInt8ScaleMax < 0)
    {
        gLogInfo << "Using default for use_int8_scale_max: true" << std::endl;
        useInt8ScaleMax = 1;
    }

    BERT_DEBUG_MSG("Building the Plugin...");
    DataType type = static_cast<DataType>(typeId);
    if (type == DataType::kINT8 && dqProbs < 0)
    {
        gLogInfo << "Using default scale factor\n";
        dqProbs = 1.F / 127.F;
    }

    auto const useInt8ScaleMaxFlag = static_cast<bool>(useInt8ScaleMax);

    auto p = std::make_unique<QKVToContextVarSeqlenPluginLegacy>(
        name, type, hiddenSize, numHeads, dqProbs, hasMask, varSeqlen, useInt8ScaleMaxFlag);
    return p.release();
}

IPluginV2* QKVToContextVarSeqlenPluginLegacyCreator::deserializePlugin(
    char const* name, void const* serialData, size_t serialLength) noexcept
{
    // This object will be deleted when the network is destroyed, which will
    // call QKVToContextVarSeqlenPluginLegacy::destroy()
    return new QKVToContextVarSeqlenPluginLegacy(name, serialData, serialLength);
}

void QKVToContextVarSeqlenPluginLegacyCreator::setPluginNamespace(char const* libNamespace) noexcept
{
    mNamespace = libNamespace;
}

char const* QKVToContextVarSeqlenPluginLegacyCreator::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

#endif // CUDA_VERSION >= 10010
