wehub-resource-sync
482 lines
17 KiB
C++
482 lines
17 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.
|
|
*/
|
|
|
|
#include "qkvToContextInt8InterleavedPlugin.h"
|
|
#include "NvInfer.h"
|
|
#include "common/bertCommon.h"
|
|
#include "common/plugin.h"
|
|
#include "common/serialize.hpp"
|
|
|
|
#include <cstring>
|
|
#include <cuda.h>
|
|
#include <iostream>
|
|
#include <optional>
|
|
#include <string_view>
|
|
#include <tuple>
|
|
#include <vector>
|
|
|
|
#include "bertQKVToContextPlugin/fused_multihead_attention_v2/fused_multihead_attention_v2.h"
|
|
|
|
using namespace nvinfer1;
|
|
using namespace nvinfer1::plugin;
|
|
using namespace nvinfer1::plugin::bert;
|
|
|
|
namespace
|
|
{
|
|
using namespace std::string_view_literals;
|
|
char const* const kQKV_TO_CONTEXT_INTERLEAVED_PLUGIN_VERSION{"6"};
|
|
char const* const kQKV_TO_CONTEXT_INTERLEAVED_PLUGIN_NAME{"CustomQKVToContextPluginDynamic"};
|
|
} // namespace
|
|
|
|
REGISTER_TENSORRT_PLUGIN(QKVToContextInterleavedPluginCreator);
|
|
|
|
constexpr uint32_t kIIDX = 0; // index of the input tensor
|
|
|
|
QKVToContextInterleavedPlugin::QKVToContextInterleavedPlugin(std::string const& name, int32_t hiddenSize,
|
|
int32_t numHeads, float dqProbs, bool useInt8ScaleMax, bool useExplicitInt8, float qkvScale, float ctxScale)
|
|
: mLayerName(name)
|
|
, mHeadSize(hiddenSize / numHeads)
|
|
, mHiddenSize(hiddenSize)
|
|
, mNumHeads(numHeads)
|
|
, mDqProbs(dqProbs)
|
|
, mQkvScale(qkvScale)
|
|
, mCtxScale(ctxScale)
|
|
{
|
|
mSM = getSmVersion();
|
|
mUseInt8ScaleMax = static_cast<int32_t>(useInt8ScaleMax);
|
|
mUseExplicitInt8 = static_cast<int32_t>(useExplicitInt8);
|
|
// variable sequence length is only supported with the fused MHA kernels
|
|
// we should not override mS!
|
|
bool isSMSupported = elem(mSM,
|
|
{kSM_AMPERE_100, kSM_AMPERE_10X, kSM_AMPERE_10B, kSM_TURING, kSM_XAVIER, kSM_ADA_10X, kSM_HOPPER_100,
|
|
kSM_BLACKWELL_100, kSM_BLACKWELL_120});
|
|
PLUGIN_VALIDATE(isSMSupported && "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
|
|
mXmmaKernel = getXMMAKernelsV2(DATA_TYPE_INT8, mSM);
|
|
}
|
|
|
|
QKVToContextInterleavedPlugin::~QKVToContextInterleavedPlugin() {}
|
|
|
|
IPluginV3* QKVToContextInterleavedPlugin::attachToContext(IPluginResourceContext* context) noexcept
|
|
{
|
|
return clone();
|
|
}
|
|
|
|
IPluginCapability* QKVToContextInterleavedPlugin::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;
|
|
}
|
|
|
|
IPluginV3* QKVToContextInterleavedPlugin::clone() noexcept
|
|
{
|
|
try
|
|
{
|
|
auto ret = std::make_unique<QKVToContextInterleavedPlugin>(
|
|
mLayerName, mHiddenSize, mNumHeads, mDqProbs, mUseInt8ScaleMax, mUseExplicitInt8, mQkvScale, mCtxScale);
|
|
|
|
ret->setPluginNamespace(mNamespace.c_str());
|
|
return ret.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
int32_t QKVToContextInterleavedPlugin::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
|
|
DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
|
|
IExprBuilder& exprBuilder) noexcept
|
|
{
|
|
try
|
|
{
|
|
// Input SHAPE is 1x(3*N*H)xTotalx1 (NCHW)
|
|
// Output SHAPE is 1x(N*H)xTotalx1
|
|
// In SupportsFormatCombination, we force the layout to be CHW, i.e.
|
|
// Input: 3xNx(H/32)xsumSx32, Output: 1xNx(H/32)xsumSx32
|
|
PLUGIN_ASSERT(inputs != nullptr);
|
|
PLUGIN_ASSERT(nbInputs == 3);
|
|
PLUGIN_ASSERT(nbShapeInputs == 0);
|
|
PLUGIN_ASSERT(outputs != nullptr);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
outputs[kIIDX] = inputs[kIIDX];
|
|
// Divide last dim by three
|
|
auto const* three = exprBuilder.constant(3);
|
|
outputs[kIIDX].d[1] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[kIIDX].d[1], *three);
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
bool QKVToContextInterleavedPlugin::supportsFormatCombination(
|
|
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t /*nbOutputs*/) noexcept
|
|
{
|
|
PLUGIN_ASSERT(pos >= 0);
|
|
PLUGIN_ASSERT(nbInputs == 3);
|
|
PLUGIN_ASSERT(pos <= nbInputs);
|
|
PLUGIN_ASSERT(inOut != nullptr);
|
|
// 3 inputs:
|
|
// 0: qkv
|
|
// 1: cu_seqlens
|
|
// 2: dummy
|
|
// 1 output
|
|
if (pos == 0 || pos == nbInputs)
|
|
{
|
|
return (inOut[pos].desc.type == DataType::kINT8) && (inOut[pos].desc.format == TensorFormat::kCHW32);
|
|
}
|
|
|
|
if (pos == 1)
|
|
{
|
|
// 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 == 2)
|
|
{
|
|
// this is the dummy input
|
|
return inOut[pos].desc.dims.nbDims == 1;
|
|
}
|
|
return false;
|
|
}
|
|
|
|
int32_t QKVToContextInterleavedPlugin::onShapeChange(
|
|
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
|
|
int32_t QKVToContextInterleavedPlugin::configurePlugin(
|
|
DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
|
|
size_t QKVToContextInterleavedPlugin::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
|
|
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
int32_t QKVToContextInterleavedPlugin::getOutputDataTypes(
|
|
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
outputTypes[0] = DataType::kINT8;
|
|
return pluginStatus_t::STATUS_SUCCESS;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return pluginStatus_t::STATUS_FAILURE;
|
|
}
|
|
|
|
char const* QKVToContextInterleavedPlugin::getPluginVersion() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_INTERLEAVED_PLUGIN_VERSION;
|
|
}
|
|
|
|
int32_t QKVToContextInterleavedPlugin::getNbOutputs() const noexcept
|
|
{
|
|
return 1;
|
|
}
|
|
|
|
void QKVToContextInterleavedPlugin::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* QKVToContextInterleavedPlugin::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
char const* QKVToContextInterleavedPlugin::getPluginName() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_INTERLEAVED_PLUGIN_NAME;
|
|
}
|
|
|
|
int32_t QKVToContextInterleavedPlugin::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);
|
|
|
|
int32_t const total = inputDesc[0].dims.d[2];
|
|
int32_t const B = inputDesc[1].dims.d[0] - 1;
|
|
int32_t const maxS = inputDesc[2].dims.d[0];
|
|
int32_t S = 384;
|
|
if (maxS <= 128)
|
|
{
|
|
S = 128;
|
|
}
|
|
else if (maxS <= 192)
|
|
{
|
|
S = 192;
|
|
}
|
|
else if (maxS <= 256)
|
|
{
|
|
S = 256;
|
|
}
|
|
Fused_multihead_attention_params_v2 params{};
|
|
params.b = B;
|
|
params.s = S;
|
|
params.h = mNumHeads;
|
|
params.d = mHeadSize;
|
|
|
|
params.interleaved = true;
|
|
|
|
params.o_ptr = outputs[0];
|
|
params.qkv_ptr = const_cast<void*>(inputs[0]);
|
|
params.cu_seqlens = static_cast<int32_t*>(const_cast<void*>(inputs[1]));
|
|
|
|
float scaleQkv = mUseExplicitInt8 ? mQkvScale : inputDesc[0].scale;
|
|
float scaleCtx = mUseExplicitInt8 ? mCtxScale : outputDesc[0].scale;
|
|
|
|
float scaleBmm1 = scaleQkv * scaleQkv * 0.125; // 1 / sqrt(64)
|
|
float scaleBmm2 = mDqProbs * scaleQkv / scaleCtx;
|
|
float scaleSoftmax = 1.F / mDqProbs;
|
|
|
|
params.scale_bmm1 = reinterpret_cast<uint32_t const&>(scaleBmm1);
|
|
params.scale_bmm2 = reinterpret_cast<uint32_t const&>(scaleBmm2);
|
|
params.scale_softmax = reinterpret_cast<uint32_t const&>(scaleSoftmax);
|
|
|
|
params.qkv_stride_in_bytes = total;
|
|
params.o_stride_in_bytes = total;
|
|
|
|
params.use_int8_scale_max = mUseInt8ScaleMax;
|
|
params.enable_i2f_trick
|
|
= -double(1 << 22) * double(scaleBmm2) <= -128.F && double(1 << 22) * double(scaleBmm2) >= 127.F;
|
|
|
|
try
|
|
{
|
|
mXmmaKernel->run(params, stream);
|
|
return cudaPeekAtLastError();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
return -1;
|
|
}
|
|
}
|
|
|
|
PluginFieldCollection const* QKVToContextInterleavedPlugin::getFieldsToSerialize() noexcept
|
|
{
|
|
mDataToSerialize.clear();
|
|
|
|
mDataToSerialize.emplace_back("hidden_size", &mHiddenSize, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("num_heads", &mNumHeads, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("use_int8_scale_max", &mUseInt8ScaleMax, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("use_explicit_int8", &mUseExplicitInt8, PluginFieldType::kINT32, 1);
|
|
mDataToSerialize.emplace_back("input_qkv_scale", &mQkvScale, PluginFieldType::kFLOAT32, 1);
|
|
mDataToSerialize.emplace_back("output_ctx_scale", &mCtxScale, PluginFieldType::kFLOAT32, 1);
|
|
|
|
if (mDqProbs >= 0)
|
|
{
|
|
mDataToSerialize.emplace_back("dq_probs", &mDqProbs, PluginFieldType::kFLOAT32, 1);
|
|
}
|
|
|
|
mFCToSerialize.nbFields = mDataToSerialize.size();
|
|
mFCToSerialize.fields = mDataToSerialize.data();
|
|
|
|
return &mFCToSerialize;
|
|
}
|
|
|
|
///////////////////////// Creator methods ////////////////////////
|
|
|
|
QKVToContextInterleavedPluginCreator::QKVToContextInterleavedPluginCreator()
|
|
{
|
|
static std::mutex sMutex;
|
|
std::lock_guard<std::mutex> lock(sMutex);
|
|
mPluginAttributes.clear();
|
|
mPluginAttributes.emplace_back(PluginField("hidden_size", nullptr, PluginFieldType::kINT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("num_heads", nullptr, PluginFieldType::kINT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("dq_probs", nullptr, PluginFieldType::kFLOAT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("use_int8_scale_max", nullptr, PluginFieldType::kINT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("use_explicit_int8", nullptr, PluginFieldType::kINT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("input_qkv_scale", nullptr, PluginFieldType::kFLOAT32, 1));
|
|
mPluginAttributes.emplace_back(PluginField("output_ctx_scale", nullptr, PluginFieldType::kFLOAT32, 1));
|
|
|
|
mFC.nbFields = mPluginAttributes.size();
|
|
mFC.fields = mPluginAttributes.data();
|
|
}
|
|
|
|
char const* QKVToContextInterleavedPluginCreator::getPluginName() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_INTERLEAVED_PLUGIN_NAME;
|
|
}
|
|
|
|
char const* QKVToContextInterleavedPluginCreator::getPluginVersion() const noexcept
|
|
{
|
|
return kQKV_TO_CONTEXT_INTERLEAVED_PLUGIN_VERSION;
|
|
}
|
|
|
|
PluginFieldCollection const* QKVToContextInterleavedPluginCreator::getFieldNames() noexcept
|
|
{
|
|
return &mFC;
|
|
}
|
|
|
|
IPluginV3* QKVToContextInterleavedPluginCreator::createPlugin(
|
|
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
|
{
|
|
try
|
|
{
|
|
// 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 QKVToContextInterleavedPlugin constructor.
|
|
int32_t numHeads{-1};
|
|
int32_t hiddenSize{0};
|
|
std::optional<int32_t> useInt8ScaleMax;
|
|
std::optional<int32_t> useExplicitInt8;
|
|
std::optional<float> qkvScale;
|
|
std::optional<float> ctxScale;
|
|
std::optional<float> dqProbs;
|
|
|
|
if (phase == TensorRTPhase::kBUILD)
|
|
{
|
|
|
|
plugin::validateRequiredAttributesExist({"hidden_size", "num_heads"}, fc);
|
|
}
|
|
else
|
|
{
|
|
PLUGIN_ASSERT(phase == TensorRTPhase::kRUNTIME);
|
|
plugin::validateRequiredAttributesExist({"hidden_size", "num_heads", "use_int8_scale_max",
|
|
"use_explicit_int8", "input_qkv_scale", "output_ctx_scale"},
|
|
fc);
|
|
}
|
|
|
|
for (int32_t i = 0; i < fc->nbFields; i++)
|
|
{
|
|
std::string_view const field_name = fc->fields[i].name;
|
|
|
|
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 == "dq_probs"sv)
|
|
{
|
|
dqProbs = *static_cast<float const*>(fc->fields[i].data);
|
|
PLUGIN_VALIDATE(
|
|
dqProbs.value() > 0.0F, ("QKV: Invalid dqProbs " + std::to_string(dqProbs.value())).c_str());
|
|
BERT_DEBUG_VALUE("Building dqProbs: ", dqProbs.value());
|
|
}
|
|
else if (field_name == "use_int8_scale_max"sv)
|
|
{
|
|
useInt8ScaleMax = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
PLUGIN_VALIDATE(useInt8ScaleMax.value() == 0 || useInt8ScaleMax.value() == 1,
|
|
("QKV: Invalid useInt8ScaleMax " + std::to_string(useInt8ScaleMax.value())).c_str());
|
|
BERT_DEBUG_VALUE("Building useInt8ScaleMax: ", useInt8ScaleMax.value());
|
|
}
|
|
else if (field_name == "use_explicit_int8"sv)
|
|
{
|
|
useExplicitInt8 = *static_cast<int32_t const*>(fc->fields[i].data);
|
|
PLUGIN_VALIDATE(useExplicitInt8.value() == 0 || useExplicitInt8.value() == 1,
|
|
("QKV: Invalid useExplicitInt8 " + std::to_string(useExplicitInt8.value())).c_str());
|
|
BERT_DEBUG_VALUE("Building use_explicit_int8: ", useExplicitInt8.value());
|
|
}
|
|
else if (field_name == "input_qkv_scale"sv)
|
|
{
|
|
qkvScale = *static_cast<float const*>(fc->fields[i].data);
|
|
PLUGIN_VALIDATE(
|
|
qkvScale.value() > 0, ("QKV: Invalid input_qkv_scale" + std::to_string(qkvScale.value())).c_str());
|
|
BERT_DEBUG_VALUE("Building input_qkv_scale: ", qkvScale.value());
|
|
}
|
|
else if (field_name == "output_ctx_scale"sv)
|
|
{
|
|
ctxScale = *static_cast<float const*>(fc->fields[i].data);
|
|
PLUGIN_VALIDATE(ctxScale.value() > 0,
|
|
("QKV: Invalid output_ctx_scale " + std::to_string(ctxScale.value())).c_str());
|
|
BERT_DEBUG_VALUE("Building output_ctx_scale: ", ctxScale.value());
|
|
}
|
|
}
|
|
|
|
if (!dqProbs.has_value())
|
|
{
|
|
gLogInfo << "Using default scale factor: 1.F/127.F" << std::endl;
|
|
dqProbs = 1.F / 127.F;
|
|
}
|
|
if (!useInt8ScaleMax.has_value())
|
|
{
|
|
gLogInfo << "Using default for use_int8_scale_max: 1" << std::endl;
|
|
useInt8ScaleMax = 1;
|
|
}
|
|
if (!useExplicitInt8.has_value())
|
|
{
|
|
gLogInfo << "Using default for use_explicit_int8: 0" << std::endl;
|
|
useExplicitInt8 = 0;
|
|
}
|
|
if (!qkvScale.has_value())
|
|
{
|
|
gLogInfo << "Using default for qkvScale: 1.F" << std::endl;
|
|
qkvScale = 1.F;
|
|
}
|
|
if (!ctxScale.has_value())
|
|
{
|
|
gLogInfo << "Using default for ctxScale: 1.F" << std::endl;
|
|
ctxScale = 1.F;
|
|
}
|
|
|
|
return new QKVToContextInterleavedPlugin(name, hiddenSize, numHeads, dqProbs.value(),
|
|
useInt8ScaleMax.value() != 0, useExplicitInt8.value() != 0, qkvScale.value(), ctxScale.value());
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
void QKVToContextInterleavedPluginCreator::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* QKVToContextInterleavedPluginCreator::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|