wehub-resource-sync
1535 lines
56 KiB
C++
1535 lines
56 KiB
C++
/*
|
|
* 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 "qkvToContextPlugin.h"
|
|
|
|
#include <cstdint>
|
|
#include <cstring>
|
|
#include <iostream>
|
|
#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_VERSION{"4"};
|
|
char const* const kQKV_TO_CONTEXT_VAR_SEQLEN_PLUGIN_VERSION{"5"};
|
|
char const* const kQKV_TO_CONTEXT_PLUGIN_NAME{"CustomQKVToContextPluginDynamic"};
|
|
} // namespace
|
|
|
|
REGISTER_TENSORRT_PLUGIN(QKVToContextPluginDynamicCreator);
|
|
|
|
constexpr uint32_t kIIDX = 0; // index of the input tensor
|
|
constexpr uint32_t kMIDX = 1; // index of the mask
|
|
|
|
REGISTER_TENSORRT_PLUGIN(QKVToContextVarSeqlenPluginCreator);
|
|
QKVToContextPluginDynamic::~QKVToContextPluginDynamic() {}
|
|
|
|
QKVToContextPluginDynamic::QKVToContextPluginDynamic(const std::string name, const DataType type,
|
|
const int32_t hiddenSize, const int32_t numHeads, float const dqProbs, bool hasImask)
|
|
: mLayerName(name)
|
|
, mS(0)
|
|
, mB(0)
|
|
, mHeadSize(hiddenSize / numHeads)
|
|
, mHiddenSize(hiddenSize)
|
|
, mNumHeads(numHeads)
|
|
, mType(type)
|
|
, mDqProbs(dqProbs)
|
|
|
|
{
|
|
mHasImask = static_cast<int32_t>(hasImask);
|
|
mSM = getSmVersion();
|
|
}
|
|
|
|
QKVToContextPluginDynamic::QKVToContextPluginDynamic(const std::string name, const DataType type, const int32_t S,
|
|
const int32_t B, const int32_t SM, const int32_t hiddenSize, const int32_t numHeads, float const dqProbs,
|
|
bool hasImask, bool hasUnfusedDispatcher, void const* runnerStateBuffer)
|
|
: mLayerName(name)
|
|
, mS(S)
|
|
, mB(B)
|
|
, mSM(SM)
|
|
, mHeadSize(hiddenSize / numHeads)
|
|
, mHiddenSize(hiddenSize)
|
|
, mNumHeads(numHeads)
|
|
, mType(type)
|
|
, mDqProbs(dqProbs)
|
|
|
|
{
|
|
BERT_DEBUG_MSG("MHA Runner Deser");
|
|
mHasImask = static_cast<int32_t>(hasImask);
|
|
mHasUnfusedDispatcher = static_cast<int32_t>(hasUnfusedDispatcher);
|
|
|
|
createMHARunner();
|
|
|
|
if (hasUnfusedDispatcher)
|
|
{
|
|
PLUGIN_ASSERT(unfusedDispatcher.get());
|
|
PLUGIN_ASSERT(runnerStateBuffer != nullptr);
|
|
auto length = unfusedDispatcher->getSerializationSize();
|
|
unfusedDispatcher->deserialize(runnerStateBuffer, length);
|
|
}
|
|
|
|
BERT_DEBUG_MSG("MHA Runner Deser Done");
|
|
}
|
|
|
|
|
|
IPluginCapability* QKVToContextPluginDynamic::getCapabilityInterface(PluginCapabilityType type) noexcept
|
|
{
|
|
try
|
|
{
|
|
if (type == PluginCapabilityType::kBUILD)
|
|
{
|
|
return static_cast<IPluginV3OneBuild*>(this);
|
|
}
|
|
if (type == PluginCapabilityType::kRUNTIME)
|
|
{
|
|
return static_cast<IPluginV3OneRuntime*>(this);
|
|
}
|
|
PLUGIN_ASSERT(type == PluginCapabilityType::kCORE);
|
|
return static_cast<IPluginV3OneCore*>(this);
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
void QKVToContextPluginDynamic::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));
|
|
}
|
|
}
|
|
|
|
IPluginV3* QKVToContextPluginDynamic::clone() noexcept
|
|
{
|
|
BERT_DEBUG_MSG("QKV Clone");
|
|
|
|
mHasUnfusedDispatcher = 0;
|
|
void const* bufferData = nullptr;
|
|
// the workspacesize is 0 if we have not call setup the dispatcher yet.
|
|
if (unfusedDispatcher.get() && unfusedDispatcher->getWorkspaceSize())
|
|
{
|
|
mHasUnfusedDispatcher = 1;
|
|
mRunnerStateBuffer.resize(unfusedDispatcher->getSerializationSize());
|
|
unfusedDispatcher->serialize(mRunnerStateBuffer.data());
|
|
bufferData = mRunnerStateBuffer.data();
|
|
}
|
|
|
|
auto ret = std::make_unique<QKVToContextPluginDynamic>(mLayerName, mType, mS, mB, mSM, mHiddenSize, mNumHeads,
|
|
mDqProbs, static_cast<bool>(mHasImask), mHasUnfusedDispatcher, bufferData);
|
|
ret->setPluginNamespace(mNamespace.c_str());
|
|
BERT_DEBUG_MSG("QKV Clone done");
|
|
return ret.release();
|
|
}
|
|
|
|
int32_t QKVToContextPluginDynamic::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
|
|
DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
|
|
IExprBuilder& exprBuilder) noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(inputs != nullptr);
|
|
PLUGIN_ASSERT(nbInputs == 1 + mHasImask);
|
|
PLUGIN_ASSERT(nbShapeInputs == 0);
|
|
PLUGIN_ASSERT(outputs != nullptr);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
// Input is BxSx3*N*H, output should be BxSxN*H
|
|
// Copy over everything
|
|
outputs[kIIDX] = inputs[kIIDX];
|
|
// Divide last dim by three
|
|
auto const* three = exprBuilder.constant(3);
|
|
outputs[kIIDX].d[HDIM] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[kIIDX].d[HDIM], *three);
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
// NOLINTNEXTLINE(readability-function-cognitive-complexity)
|
|
bool QKVToContextPluginDynamic::supportsFormatCombination(
|
|
int32_t pos, DynamicPluginTensorDesc 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->desc.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->desc.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->desc.format == TensorFormat::kLINEAR;
|
|
if (mType == DataType::kINT8)
|
|
{
|
|
if (in->desc.dims.d[HDIM] % 32U == 0)
|
|
{
|
|
isFormatSupported = in->desc.format == TensorFormat::kCHW32;
|
|
}
|
|
else
|
|
{
|
|
isFormatSupported = in->desc.format == TensorFormat::kCHW4;
|
|
}
|
|
}
|
|
|
|
// must not check descriptions > pos
|
|
return (in->desc.type == mType) && // precision
|
|
isFormatSupported && // format
|
|
(in->desc.dims.nbDims == 5) && // num dims
|
|
((in->desc.dims.d[HDIM] % 3U) == 0) && // see getOutputDimensions
|
|
((in->desc.dims.d[3]) == 1) && // for fc
|
|
((in->desc.dims.d[4]) == 1) // for fc
|
|
;
|
|
}
|
|
|
|
// pos==1
|
|
if ((mHasImask && pos == 1)) // pos 1 is the mask
|
|
{
|
|
auto const* inMask = &inOut[1].desc;
|
|
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->desc.dims.d[BDIM]);
|
|
}
|
|
|
|
if (!mHasImask || pos == 2) // output pos
|
|
{
|
|
bool isFormatSupported = out->desc.format == TensorFormat::kLINEAR;
|
|
if (mType == DataType::kINT8)
|
|
{
|
|
if (out->desc.dims.d[HDIM] % 32U == 0)
|
|
{
|
|
isFormatSupported = out->desc.format == TensorFormat::kCHW32;
|
|
}
|
|
else
|
|
{
|
|
isFormatSupported = out->desc.format == TensorFormat::kCHW4;
|
|
}
|
|
}
|
|
|
|
return (in->desc.type == out->desc.type) && // precision
|
|
isFormatSupported && // format
|
|
(out->desc.dims.nbDims == 5) && // num dims
|
|
((in->desc.dims.d[HDIM] / 3) == (out->desc.dims.d[HDIM])) && // div 3
|
|
((out->desc.dims.d[3]) == 1) && // for fc
|
|
((out->desc.dims.d[4]) == 1) && // for fc
|
|
((out->desc.dims.d[BDIM]) == in->desc.dims.d[BDIM]) && // check B
|
|
((out->desc.dims.d[SDIM]) == in->desc.dims.d[SDIM]) // check S
|
|
;
|
|
}
|
|
|
|
return false;
|
|
}
|
|
|
|
int32_t QKVToContextPluginDynamic::onShapeChange(
|
|
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(in != nullptr);
|
|
PLUGIN_ASSERT(nbInputs == 1 + mHasImask);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
PluginTensorDesc const& inDesc = in[kIIDX];
|
|
TRT_UNUSED inDesc;
|
|
PLUGIN_ASSERT(out != nullptr);
|
|
PluginTensorDesc const& outDesc = out[0];
|
|
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];
|
|
TRT_UNUSED maskDesc;
|
|
PLUGIN_ASSERT(maskDesc.dims.d[0] == inDesc.dims.d[BDIM]);
|
|
}
|
|
|
|
createMHARunner();
|
|
|
|
// mS and mB that are set by configurePlugin() may be stale
|
|
mS = inDesc.dims.d[SDIM];
|
|
mB = inDesc.dims.d[BDIM];
|
|
PLUGIN_ASSERT(mS);
|
|
PLUGIN_ASSERT(mB);
|
|
if (fusedDispatcher.get() && fusedDispatcher->isValid(mHeadSize, mS))
|
|
{
|
|
fusedDispatcher->setup(mS, mB, mHeadSize);
|
|
}
|
|
else
|
|
{
|
|
unfusedDispatcher->setup(mS, mB, mHeadSize);
|
|
}
|
|
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
int32_t QKVToContextPluginDynamic::configurePlugin(
|
|
DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(in != nullptr);
|
|
PLUGIN_ASSERT(nbInputs == 1 + mHasImask);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
PluginTensorDesc const& inDesc = in[kIIDX].desc;
|
|
TRT_UNUSED inDesc;
|
|
PLUGIN_ASSERT(out != nullptr);
|
|
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();
|
|
|
|
const int32_t S = inDesc.dims.d[SDIM];
|
|
const int32_t 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
|
|
const int32_t Smin = in->min.d[SDIM];
|
|
const int32_t 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;
|
|
}
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
size_t QKVToContextPluginDynamic::getWorkspaceSize(DynamicPluginTensorDesc const* /*inputs*/, int32_t /*nbInputs*/,
|
|
DynamicPluginTensorDesc 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
|
|
int32_t QKVToContextPluginDynamic::getOutputDataTypes(
|
|
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(
|
|
inputTypes[0] == DataType::kFLOAT || inputTypes[0] == DataType::kHALF || inputTypes[0] == DataType::kINT8);
|
|
outputTypes[0] = inputTypes[0];
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
void QKVToContextPluginDynamic::setCublasResources(std::shared_ptr<CublasWrapper> cublasWrapper)
|
|
{
|
|
mCublasWrapper = cublasWrapper;
|
|
// The shared cublasWrapper resource owns the handle.
|
|
// but `this` instance has a non-owning pointer to the handle.
|
|
// Note that the cublasWrapper inits the handle and checks for nullptr
|
|
// so we don't have to do that here.
|
|
mCublasHandle = mCublasWrapper->getCublasHandle();
|
|
}
|
|
|
|
IPluginV3* QKVToContextPluginDynamic::attachToContext(IPluginResourceContext* context) noexcept
|
|
{
|
|
try
|
|
{
|
|
auto p = static_cast<QKVToContextPluginDynamic*>(clone());
|
|
// the clone has shared ownership of underling cublasWrapper instance
|
|
// that is mapped to current context
|
|
p->setCublasResources(createPluginCublasWrapper(context));
|
|
return p;
|
|
}
|
|
catch (const std::exception& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
char const* QKVToContextPluginDynamic::getPluginVersion() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_PLUGIN_VERSION;
|
|
}
|
|
|
|
int32_t QKVToContextPluginDynamic::getNbOutputs() const noexcept
|
|
{
|
|
return 1;
|
|
}
|
|
|
|
char const* QKVToContextPluginDynamic::getPluginName() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_PLUGIN_NAME;
|
|
}
|
|
|
|
|
|
void QKVToContextPluginDynamic::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* QKVToContextPluginDynamic::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
int32_t QKVToContextPluginDynamic::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, mCublasHandle);
|
|
}
|
|
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, mCublasHandle);
|
|
}
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
return -1;
|
|
}
|
|
return 0;
|
|
}
|
|
|
|
PluginFieldCollection const* QKVToContextPluginDynamic::getFieldsToSerialize() noexcept
|
|
{
|
|
mDataToSerialize.clear();
|
|
|
|
mDataToSerialize.emplace_back("type_id", &mType, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("hidden_size", &mHiddenSize, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("num_heads", &mNumHeads, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("has_mask", &mHasImask, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("S", &mS, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("B", &mB, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("SM", &mSM, PluginFieldType::kINT32, 1);
|
|
|
|
if (unfusedDispatcher.get() && unfusedDispatcher->getWorkspaceSize())
|
|
{
|
|
mHasUnfusedDispatcher = 1;
|
|
mRunnerStateBuffer.resize(unfusedDispatcher->getSerializationSize());
|
|
unfusedDispatcher->serialize(mRunnerStateBuffer.data());
|
|
mDataToSerialize.emplace_back("runnerStateBuffer", (void const*) mRunnerStateBuffer.data(),
|
|
PluginFieldType::kUNKNOWN, mRunnerStateBuffer.size());
|
|
}
|
|
else
|
|
{
|
|
mHasUnfusedDispatcher = 0;
|
|
}
|
|
|
|
mDataToSerialize.emplace_back("hasUnfusedDispatcher", &mHasUnfusedDispatcher, PluginFieldType::kINT32, 1);
|
|
|
|
if (mDqProbs >= 0)
|
|
{
|
|
mDataToSerialize.emplace_back("dq_probs", &mDqProbs, PluginFieldType::kFLOAT32, 1);
|
|
}
|
|
|
|
mFCToSerialize.nbFields = mDataToSerialize.size();
|
|
mFCToSerialize.fields = mDataToSerialize.data();
|
|
|
|
return &mFCToSerialize;
|
|
}
|
|
|
|
QKVToContextPluginDynamicCreator::QKVToContextPluginDynamicCreator()
|
|
{
|
|
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* QKVToContextPluginDynamicCreator::getPluginName() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_PLUGIN_NAME;
|
|
}
|
|
|
|
char const* QKVToContextPluginDynamicCreator::getPluginVersion() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_PLUGIN_VERSION;
|
|
}
|
|
|
|
PluginFieldCollection const* QKVToContextPluginDynamicCreator::getFieldNames() noexcept
|
|
{
|
|
return &mFC;
|
|
}
|
|
|
|
IPluginV3* QKVToContextPluginDynamicCreator::createPlugin(
|
|
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) 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 QKVToContextPluginDynamic constructor.
|
|
int32_t numHeads{-1};
|
|
bool hasMask = false;
|
|
int32_t typeId = -1;
|
|
int32_t s = -1;
|
|
int32_t b = -1;
|
|
int32_t sm = -1;
|
|
bool hasUnfusedDispatcher = false;
|
|
void const* runnerStateBuffer = nullptr;
|
|
float dqProbs = -1.0F;
|
|
|
|
PLUGIN_VALIDATE(fc->fields != nullptr);
|
|
if (phase == TensorRTPhase::kBUILD)
|
|
{
|
|
plugin::validateRequiredAttributesExist({"type_id", "hidden_size", "num_heads", "has_mask"}, fc);
|
|
}
|
|
else
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
plugin::validateRequiredAttributesExist(
|
|
{"type_id", "S", "B", "hidden_size", "num_heads", "has_mask", "SM", "hasUnfusedDispatcher"}, 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);
|
|
}
|
|
else 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);
|
|
}
|
|
else 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);
|
|
}
|
|
else 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);
|
|
}
|
|
else 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);
|
|
}
|
|
else if (field_name == "S"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
s = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
BERT_DEBUG_VALUE("Building S: ", s);
|
|
}
|
|
else if (field_name == "B"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
b = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
BERT_DEBUG_VALUE("Building B: ", b);
|
|
}
|
|
else if (field_name == "SM"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
sm = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
BERT_DEBUG_VALUE("Building SM: ", sm);
|
|
}
|
|
else if (field_name == "hasUnfusedDispatcher"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
auto hasUnfusedDispatcherValue = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
PLUGIN_VALIDATE(hasUnfusedDispatcherValue == 0 || hasUnfusedDispatcherValue == 1,
|
|
("QKV: Invalid hasUnfusedDispatcher " + std::to_string(hasUnfusedDispatcherValue)).c_str());
|
|
hasUnfusedDispatcher = static_cast<bool>(hasUnfusedDispatcherValue);
|
|
BERT_DEBUG_VALUE("Building hasUnfusedDispatcher: ", hasUnfusedDispatcher);
|
|
}
|
|
else if (field_name == "runnerStateBuffer"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
runnerStateBuffer = static_cast<void const*>(fc->fields[i].data);
|
|
}
|
|
}
|
|
|
|
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;
|
|
}
|
|
|
|
if (phase == TensorRTPhase::kBUILD)
|
|
{
|
|
return new QKVToContextPluginDynamic(name, type, hiddenSize, numHeads, dqProbs, hasMask);
|
|
}
|
|
|
|
PLUGIN_VALIDATE(s != -1, "invalid S during runtime plugin creation");
|
|
PLUGIN_VALIDATE(b != -1, "invalid B during runtime plugin creation");
|
|
PLUGIN_VALIDATE(sm != -1, "invalid SM during runtime plugin creation");
|
|
if (hasUnfusedDispatcher == 1)
|
|
{
|
|
PLUGIN_VALIDATE(runnerStateBuffer != nullptr, "invalid runnerStateBuffer during runtime plugin creation");
|
|
}
|
|
else
|
|
{
|
|
PLUGIN_VALIDATE(runnerStateBuffer == nullptr, "invalid runnerStateBuffer during runtime plugin creation");
|
|
}
|
|
|
|
return new QKVToContextPluginDynamic(
|
|
name, type, s, b, sm, hiddenSize, numHeads, dqProbs, hasMask, hasUnfusedDispatcher, runnerStateBuffer);
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
|
|
void QKVToContextPluginDynamicCreator::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* QKVToContextPluginDynamicCreator::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
|
|
///// QKVToContextVarSeqlenPlugin (CustomQKVToContextPluginDynamic v5) ////
|
|
|
|
QKVToContextVarSeqlenPlugin::~QKVToContextVarSeqlenPlugin() {}
|
|
|
|
QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(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)
|
|
, mType(type)
|
|
, mDqProbs(dqProbs)
|
|
, mHdim(HDIM)
|
|
{
|
|
mSM = getSmVersion();
|
|
mUseVarSeqlen = static_cast<int32_t>(varSeqlen);
|
|
mUseInt8ScaleMax = static_cast<int32_t>(useInt8ScaleMax);
|
|
mHasImask = static_cast<int32_t>(hasImask);
|
|
|
|
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;
|
|
}
|
|
}
|
|
|
|
QKVToContextVarSeqlenPlugin::QKVToContextVarSeqlenPlugin(std::string const name, int32_t const S, int32_t const B,
|
|
DataType const type, int32_t const hiddenSize, int32_t const numHeads, float const dqProbs, bool hasImask,
|
|
bool varSeqlen, bool useInt8ScaleMax, void const* runnerStateBuffer)
|
|
: mLayerName(name)
|
|
, mS(S)
|
|
, mB(B)
|
|
, mHeadSize(hiddenSize / numHeads)
|
|
, mHiddenSize(hiddenSize)
|
|
, mNumHeads(numHeads)
|
|
, mType(type)
|
|
, mDqProbs(dqProbs)
|
|
, mHdim(HDIM)
|
|
{
|
|
mSM = getSmVersion();
|
|
mUseVarSeqlen = static_cast<int32_t>(varSeqlen);
|
|
mUseInt8ScaleMax = static_cast<int32_t>(useInt8ScaleMax);
|
|
mHasImask = static_cast<int32_t>(hasImask);
|
|
|
|
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;
|
|
}
|
|
|
|
createMHARunner();
|
|
|
|
PLUGIN_ASSERT(runnerStateBuffer != nullptr);
|
|
auto length = mDispatcher->getSerializationSize();
|
|
mDispatcher->deserialize(runnerStateBuffer, length);
|
|
}
|
|
|
|
IPluginCapability* QKVToContextVarSeqlenPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept
|
|
{
|
|
try
|
|
{
|
|
if (type == PluginCapabilityType::kBUILD)
|
|
{
|
|
return static_cast<IPluginV3OneBuild*>(this);
|
|
}
|
|
if (type == PluginCapabilityType::kRUNTIME)
|
|
{
|
|
return static_cast<IPluginV3OneRuntime*>(this);
|
|
}
|
|
PLUGIN_ASSERT(type == PluginCapabilityType::kCORE);
|
|
return static_cast<IPluginV3OneCore*>(this);
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
void QKVToContextVarSeqlenPlugin::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));
|
|
}
|
|
}
|
|
|
|
|
|
IPluginV3* QKVToContextVarSeqlenPlugin::clone() noexcept
|
|
{
|
|
BERT_DEBUG_MSG("QKV Clone");
|
|
|
|
std::unique_ptr<QKVToContextVarSeqlenPlugin> ret;
|
|
|
|
// the workspacesize is 0 if we have not call setup the dispatcher yet.
|
|
if (mDispatcher.get())
|
|
{
|
|
mRunnerStateBuffer.resize(mDispatcher->getSerializationSize());
|
|
void* const bufferData = mRunnerStateBuffer.data();
|
|
mDispatcher->serialize(bufferData);
|
|
|
|
ret = std::make_unique<QKVToContextVarSeqlenPlugin>(mLayerName, mS, mB, mType, mHiddenSize, mNumHeads, mDqProbs,
|
|
mHasImask, mUseVarSeqlen, mUseInt8ScaleMax, bufferData);
|
|
}
|
|
else
|
|
{
|
|
// dispatcher not setup yet, use type 1 constructor
|
|
ret = std::make_unique<QKVToContextVarSeqlenPlugin>(
|
|
mLayerName, mType, mHiddenSize, mNumHeads, mDqProbs, mHasImask, mUseVarSeqlen, mUseInt8ScaleMax);
|
|
}
|
|
|
|
ret->setPluginNamespace(mNamespace.c_str());
|
|
BERT_DEBUG_MSG("QKV Clone done");
|
|
return ret.release();
|
|
}
|
|
|
|
int32_t QKVToContextVarSeqlenPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
|
|
DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
|
|
IExprBuilder& exprBuilder) noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(inputs != nullptr);
|
|
PLUGIN_ASSERT(nbInputs == 1 + mHasImask + 2 * mUseVarSeqlen);
|
|
PLUGIN_ASSERT(nbShapeInputs == 0);
|
|
PLUGIN_ASSERT(outputs != nullptr);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
// Input is BxSx3*N*H, output should be BxSxN*H
|
|
// Copy over everything
|
|
outputs[kIIDX] = inputs[kIIDX];
|
|
// Divide last dim by three
|
|
auto const* three = exprBuilder.constant(3);
|
|
// mHdim is 2 for fixed seqlen and is 1 for varseqlen
|
|
outputs[kIIDX].d[mHdim]
|
|
= exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[kIIDX].d[mHdim], *three);
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
bool QKVToContextVarSeqlenPlugin::supportsFormatCombination(
|
|
int32_t pos, DynamicPluginTensorDesc 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->desc.dims;
|
|
auto const outDims = out->desc.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].desc;
|
|
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].desc;
|
|
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].desc;
|
|
return (seqlens->type == DataType::kINT32) && (seqlens->format == TensorFormat::kLINEAR);
|
|
}
|
|
if (pos == 3)
|
|
{
|
|
// this is the dummy input
|
|
return inOut[pos].desc.dims.nbDims == 1;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
int32_t QKVToContextVarSeqlenPlugin::onShapeChange(
|
|
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(in != nullptr);
|
|
PLUGIN_ASSERT(nbInputs == 1 + mHasImask + 2 * mUseVarSeqlen);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
PluginTensorDesc const& inDesc = in[kIIDX];
|
|
TRT_UNUSED inDesc;
|
|
PluginTensorDesc const& outDesc = out[0];
|
|
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];
|
|
TRT_UNUSED maskDesc;
|
|
PLUGIN_ASSERT(maskDesc.dims.d[0] == inDesc.dims.d[BDIM]);
|
|
}
|
|
|
|
// during build, configurePlugin() should have set mS, mB appropriately
|
|
// during inference, the engine should have mS, mB information
|
|
PLUGIN_ASSERT(mS);
|
|
PLUGIN_ASSERT(mB);
|
|
|
|
BERT_DEBUG_MSG("setting up MHA runner for single sequence length");
|
|
createMHARunner();
|
|
this->mDispatcher->setup(mS, mB, mHeadSize);
|
|
}
|
|
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);
|
|
}
|
|
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
int32_t QKVToContextVarSeqlenPlugin::configurePlugin(
|
|
DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(in != nullptr);
|
|
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]);
|
|
}
|
|
|
|
const int32_t S = inDesc.dims.d[SDIM] <= 0 ? in->max.d[SDIM] : inDesc.dims.d[SDIM];
|
|
const int32_t 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);
|
|
}
|
|
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
size_t QKVToContextVarSeqlenPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t /* nbInputs */,
|
|
DynamicPluginTensorDesc const* /* outputs */, int32_t /* nbOutputs */) const noexcept
|
|
{
|
|
size_t paddingWorkpaceSize = mPatcher ? mPatcher->getWorkspaceSize(inputs[0].desc.dims.d[0], mNumHeads) : 0;
|
|
return mDispatcher->getWorkspaceSize() + paddingWorkpaceSize;
|
|
}
|
|
|
|
int32_t QKVToContextVarSeqlenPlugin::getOutputDataTypes(
|
|
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(
|
|
inputTypes[0] == DataType::kFLOAT || inputTypes[0] == DataType::kHALF || inputTypes[0] == DataType::kINT8);
|
|
outputTypes[0] = inputTypes[0];
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
void QKVToContextVarSeqlenPlugin::setCublasResources(std::shared_ptr<CublasWrapper> cublasWrapper)
|
|
{
|
|
mCublasWrapper = cublasWrapper;
|
|
// The shared cublasWrapper resource owns the handle.
|
|
// but `this` instance has a non-owning pointer to the handle.
|
|
// Note that the cublasWrapper inits the handle and checks for nullptr
|
|
// so we don't have to do that here.
|
|
mCublasHandle = mCublasWrapper->getCublasHandle();
|
|
}
|
|
|
|
IPluginV3* QKVToContextVarSeqlenPlugin::attachToContext(IPluginResourceContext* context) noexcept
|
|
{
|
|
try
|
|
{
|
|
auto p = static_cast<QKVToContextVarSeqlenPlugin*>(clone());
|
|
// the clone has shared ownership of underling cublasWrapper instance
|
|
// that is mapped to current context
|
|
p->setCublasResources(createPluginCublasWrapper(context));
|
|
return p;
|
|
}
|
|
catch (const std::exception& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
char const* QKVToContextVarSeqlenPlugin::getPluginVersion() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_VAR_SEQLEN_PLUGIN_VERSION;
|
|
}
|
|
|
|
int32_t QKVToContextVarSeqlenPlugin::getNbOutputs() const noexcept
|
|
{
|
|
return 1;
|
|
}
|
|
|
|
char const* QKVToContextVarSeqlenPlugin::getPluginName() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_PLUGIN_NAME;
|
|
}
|
|
|
|
|
|
void QKVToContextVarSeqlenPlugin::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* QKVToContextVarSeqlenPlugin::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
// NOLINTNEXTLINE(readability-function-cognitive-complexity)
|
|
int32_t QKVToContextVarSeqlenPlugin::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)
|
|
{
|
|
const int32_t B = inputDesc[2].dims.d[0] - 1;
|
|
const int32_t 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, mCublasHandle);
|
|
}
|
|
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, mCublasHandle);
|
|
}
|
|
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, mCublasHandle);
|
|
return cudaGetLastError();
|
|
}
|
|
|
|
PluginFieldCollection const* QKVToContextVarSeqlenPlugin::getFieldsToSerialize() noexcept
|
|
{
|
|
mDataToSerialize.clear();
|
|
|
|
mDataToSerialize.emplace_back("type_id", &mType, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("hidden_size", &mHiddenSize, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("num_heads", &mNumHeads, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("has_mask", &mHasImask, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("var_seqlen", &mUseVarSeqlen, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("use_int8_scale_max", &mUseInt8ScaleMax, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("S", &mS, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("B", &mB, PluginFieldType::kINT32, 1);
|
|
|
|
mRunnerStateBuffer.resize(mDispatcher->getSerializationSize());
|
|
mDispatcher->serialize(mRunnerStateBuffer.data());
|
|
mDataToSerialize.emplace_back("runnerStateBuffer", (void const*) mRunnerStateBuffer.data(),
|
|
PluginFieldType::kUNKNOWN, mRunnerStateBuffer.size());
|
|
|
|
if (mDqProbs >= 0)
|
|
{
|
|
mDataToSerialize.emplace_back("dq_probs", &mDqProbs, PluginFieldType::kFLOAT32, 1);
|
|
}
|
|
|
|
mFCToSerialize.nbFields = mDataToSerialize.size();
|
|
mFCToSerialize.fields = mDataToSerialize.data();
|
|
|
|
return &mFCToSerialize;
|
|
}
|
|
|
|
QKVToContextVarSeqlenPluginCreator::QKVToContextVarSeqlenPluginCreator()
|
|
{
|
|
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* QKVToContextVarSeqlenPluginCreator::getPluginName() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_PLUGIN_NAME;
|
|
}
|
|
|
|
char const* QKVToContextVarSeqlenPluginCreator::getPluginVersion() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_VAR_SEQLEN_PLUGIN_VERSION;
|
|
}
|
|
|
|
PluginFieldCollection const* QKVToContextVarSeqlenPluginCreator::getFieldNames() noexcept
|
|
{
|
|
return &mFC;
|
|
}
|
|
|
|
IPluginV3* QKVToContextVarSeqlenPluginCreator::createPlugin(
|
|
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) 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 QKVToContextVarSeqelnPlugin constructor.
|
|
int32_t numHeads = -1;
|
|
bool hasMask = false;
|
|
int32_t typeId = -1;
|
|
int32_t s = -1;
|
|
int32_t b = -1;
|
|
void const* runnerStateBuffer = nullptr;
|
|
int32_t varSeqlen = 0;
|
|
float dqProbs = -1.0F;
|
|
int32_t useInt8ScaleMax = -1;
|
|
|
|
PLUGIN_VALIDATE(fc->fields != nullptr);
|
|
|
|
if (phase == TensorRTPhase::kBUILD)
|
|
{
|
|
plugin::validateRequiredAttributesExist({"type_id", "hidden_size", "num_heads", "has_mask"}, fc);
|
|
}
|
|
else
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
// since fc is from a deserialized engine,
|
|
// we expect all attributes (except dq_probs) to be present during runtime
|
|
plugin::validateRequiredAttributesExist({"type_id", "S", "B", "hidden_size", "num_heads", "has_mask",
|
|
"var_seqlen", "use_int8_scale_max", "runnerStateBuffer"},
|
|
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);
|
|
}
|
|
else 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);
|
|
}
|
|
else 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);
|
|
}
|
|
else 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);
|
|
}
|
|
|
|
else 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);
|
|
}
|
|
else if (field_name == "var_seqlen"sv)
|
|
{
|
|
varSeqlen = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
BERT_DEBUG_VALUE("Building var_seqlen: ", varSeqlen);
|
|
}
|
|
else 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);
|
|
}
|
|
else if (field_name == "S"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
s = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
BERT_DEBUG_VALUE("Building S: ", s);
|
|
}
|
|
else if (field_name == "B"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
b = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
BERT_DEBUG_VALUE("Building B: ", b);
|
|
}
|
|
else if (field_name == "runnerStateBuffer"sv)
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
runnerStateBuffer = static_cast<void const*>(fc->fields[i].data);
|
|
}
|
|
}
|
|
|
|
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);
|
|
|
|
if (phase == TensorRTPhase::kBUILD)
|
|
{
|
|
return new QKVToContextVarSeqlenPlugin(
|
|
name, type, hiddenSize, numHeads, dqProbs, hasMask, varSeqlen, useInt8ScaleMaxFlag);
|
|
}
|
|
|
|
PLUGIN_VALIDATE(s != -1, "invalid S during runtime plugin creation");
|
|
PLUGIN_VALIDATE(b != -1, "invalid B during runtime plugin creation");
|
|
PLUGIN_VALIDATE(runnerStateBuffer != nullptr, "invalid runnerStateBuffer during runtime plugin creation");
|
|
|
|
return new QKVToContextVarSeqlenPlugin(name, s, b, type, hiddenSize, numHeads, dqProbs, hasMask, varSeqlen,
|
|
useInt8ScaleMaxFlag, runnerStateBuffer);
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
void QKVToContextVarSeqlenPluginCreator::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* QKVToContextVarSeqlenPluginCreator::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
#endif // CUDA_VERSION >= 10010
|