chore: import upstream snapshot with attribution

plugin/embLayerNormPlugin/embLayerNormVarSeqlenPlugin.cpp

@@ -0,0 +1,832 @@
+/*
+ * 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 <cstring>
+#include <cuda.h>
+#include <memory>
+#include <set>
+#include <string_view>
+#include <vector>
+
+#include "NvInfer.h"
+#include "common/serialize.hpp"
+#include "embLayerNormVarSeqlenPlugin.h"
+
+using namespace nvinfer1;
+using namespace nvinfer1::plugin;
+using namespace nvinfer1::plugin::bert;
+
+namespace
+{
+using namespace std::string_view_literals;
+constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE{"4"};
+constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON{"5"};
+constexpr char const* kEMB_LAYER_NORM_VAR_SEQLEN_NAME{"CustomEmbLayerNormPluginDynamic"};
+
+void checkConfigurationInputs(
+    PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
+{
+    // Validate input arguments
+    PLUGIN_ASSERT(nbInputs == 4);
+    PLUGIN_ASSERT(nbOutputs == 2);
+
+    PLUGIN_ASSERT(inputs[0].dims.nbDims == 1);
+    PLUGIN_ASSERT(inputs[1].dims.nbDims == 1);
+
+    PLUGIN_ASSERT(inputs[1].dims.d[0] == inputs[0].dims.d[0]);
+
+    PLUGIN_ASSERT(inputs[2].dims.nbDims == 1);
+
+    PLUGIN_ASSERT(outputs[0].dims.nbDims == 4);
+    PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[0]) == static_cast<size_t>(inputs[0].dims.d[0]));
+    PLUGIN_ASSERT(outputs[0].dims.d[2] == 1);
+    PLUGIN_ASSERT(outputs[0].dims.d[3] == 1);
+
+    PLUGIN_ASSERT(inputs[0].type == DataType::kINT32);
+    PLUGIN_ASSERT(inputs[1].type == DataType::kINT32);
+    PLUGIN_ASSERT(inputs[2].type == DataType::kINT32);
+}
+
+bool initializeFields(char const* name, PluginFieldCollection const* fc, Weights& beta, Weights& gamma,
+    Weights& word_emb, Weights& pos_emb, Weights& tok_emb)
+{
+    bool output_fp16 = false;
+    std::set<std::string> const requiredAttributes{
+        "bert_embeddings_layernorm_beta",
+        "bert_embeddings_layernorm_gamma",
+        "bert_embeddings_word_embeddings",
+        "bert_embeddings_token_type_embeddings",
+        "bert_embeddings_position_embeddings",
+    };
+    plugin::validateRequiredAttributesExist(requiredAttributes, fc);
+
+    for (int32_t i = 0; i < fc->nbFields; i++)
+    {
+        std::string_view const field_name = fc->fields[i].name;
+        if (field_name == "bert_embeddings_layernorm_beta"sv)
+        {
+            BERT_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
+            beta.values = fc->fields[i].data;
+            beta.count = fc->fields[i].length;
+            beta.type = fieldTypeToDataType(fc->fields[i].type);
+        }
+
+        else if (field_name == "bert_embeddings_layernorm_gamma"sv)
+        {
+            BERT_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
+            gamma.values = fc->fields[i].data;
+            gamma.count = fc->fields[i].length;
+            gamma.type = fieldTypeToDataType(fc->fields[i].type);
+        }
+
+        else if (field_name == "bert_embeddings_word_embeddings"sv)
+        {
+            BERT_DEBUG_MSG("Building bert_embeddings_word_embeddings...");
+            word_emb.values = fc->fields[i].data;
+            word_emb.count = fc->fields[i].length;
+            word_emb.type = fieldTypeToDataType(fc->fields[i].type);
+        }
+
+        else if (field_name == "bert_embeddings_token_type_embeddings"sv)
+        {
+            BERT_DEBUG_MSG("Building bert_embeddings_token_type_embeddings...");
+            tok_emb.values = fc->fields[i].data;
+            tok_emb.count = fc->fields[i].length;
+            tok_emb.type = fieldTypeToDataType(fc->fields[i].type);
+        }
+
+        else if (field_name == "bert_embeddings_position_embeddings"sv)
+        {
+            BERT_DEBUG_MSG("Building bert_embeddings_position_embeddings...");
+            pos_emb.values = fc->fields[i].data;
+            pos_emb.count = fc->fields[i].length;
+            pos_emb.type = fieldTypeToDataType(fc->fields[i].type);
+        }
+        else if (field_name == "output_fp16"sv)
+        {
+            BERT_DEBUG_MSG("Building output_fp16...");
+            PLUGIN_VALIDATE(fc->fields[i].type == PluginFieldType::kINT32);
+            output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
+        }
+    }
+    return output_fp16;
+}
+
+} // namespace
+
+REGISTER_TENSORRT_PLUGIN(EmbLayerNormVarSeqlenPluginHFaceCreator);
+REGISTER_TENSORRT_PLUGIN(EmbLayerNormVarSeqlenPluginMTronCreator);
+
+EmbLayerNormVarSeqlenPluginBase::EmbLayerNormVarSeqlenPluginBase(std::string const& name, DataType type,
+    Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb, Weights const& tokEmb,
+    DataType maskType)
+    : mLayerName(name)
+    , mLd(beta.count)
+    , mType(type)
+    , mMaskType(maskType)
+{
+    // Assuming Weights.count is the number of elements and not bytes
+    PLUGIN_VALIDATE(beta.count == gamma.count);
+    PLUGIN_VALIDATE(mLd > 0U);
+    PLUGIN_VALIDATE(wordEmb.count % mLd == 0);
+    PLUGIN_VALIDATE(posEmb.count % mLd == 0);
+    PLUGIN_VALIDATE(tokEmb.count % mLd == 0);
+    mWordVocabSize = wordEmb.count / mLd;
+    mPosVocabSize = posEmb.count / mLd;
+    mTokVocabSize = tokEmb.count / mLd;
+
+    mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
+    mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
+
+    mWordEmb.convertAndCopy(wordEmb, mType);
+    mTokEmb.convertAndCopy(tokEmb, mType);
+    mPosEmb.convertAndCopy(posEmb, mType);
+
+    copyToDevice(mGamma, sizeof(float) * mGamma.count, mGammaDev);
+    copyToDevice(mBeta, sizeof(float) * mBeta.count, mBetaDev);
+
+    copyToDevice(mWordEmb, getWeightsSize(mWordEmb, mType), mWordEmbDev);
+    copyToDevice(mPosEmb, getWeightsSize(mPosEmb, mType), mPosEmbDev);
+    copyToDevice(mTokEmb, getWeightsSize(mTokEmb, mType), mTokEmbDev);
+}
+
+EmbLayerNormVarSeqlenPluginHFace::EmbLayerNormVarSeqlenPluginHFace(std::string const& name, DataType const type,
+    Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb, Weights const& tokEmb)
+    : EmbLayerNormVarSeqlenPluginBase(name, type, beta, gamma, wordEmb, posEmb, tokEmb, DataType::kINT32)
+{
+    BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace creation");
+}
+
+EmbLayerNormVarSeqlenPluginMTron::EmbLayerNormVarSeqlenPluginMTron(std::string const& name, DataType const type,
+    Weights const& beta, Weights const& gamma, Weights const& wordEmb, Weights const& posEmb, Weights const& tokEmb)
+    : EmbLayerNormVarSeqlenPluginBase(name, type, beta, gamma, wordEmb, posEmb, tokEmb, type)
+{
+    BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron creation");
+}
+
+EmbLayerNormVarSeqlenPluginBase::~EmbLayerNormVarSeqlenPluginBase()
+{
+    try
+    {
+        // This gets called when the network containing plugin is destroyed
+        mGammaDev.reset(nullptr);
+        mBetaDev.reset(nullptr);
+        mWordEmbDev.reset(nullptr);
+        mPosEmbDev.reset(nullptr);
+        mTokEmbDev.reset(nullptr);
+        // delete this; (TRT will delete this plugin object)
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+EmbLayerNormVarSeqlenPluginHFace::~EmbLayerNormVarSeqlenPluginHFace()
+{
+    BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace destruction");
+}
+
+EmbLayerNormVarSeqlenPluginMTron::~EmbLayerNormVarSeqlenPluginMTron()
+{
+    BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron destruction");
+}
+
+//////
+// IPluginV3 method definitions:
+// - getCapabilityInterface() (Base)
+// - clone() (HFace, MTron)
+//////
+IPluginCapability* EmbLayerNormVarSeqlenPluginBase::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* EmbLayerNormVarSeqlenPluginHFace::clone() noexcept
+{
+    try
+    {
+        BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace clone");
+
+        auto p = std::make_unique<EmbLayerNormVarSeqlenPluginHFace>(
+            mLayerName, mType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb);
+        p->setPluginNamespace(mNamespace.c_str());
+        return p.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+IPluginV3* EmbLayerNormVarSeqlenPluginMTron::clone() noexcept
+{
+    try
+    {
+        BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron clone");
+
+        auto p = std::make_unique<EmbLayerNormVarSeqlenPluginMTron>(
+            mLayerName, mType, mBeta, mGamma, mWordEmb, mPosEmb, mTokEmb);
+        p->setPluginNamespace(mNamespace.c_str());
+
+        return p.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+// End IPluginV3 method definitions
+
+//////
+// IPluginV3OneRuntime method definitions:
+// - getFieldsToSerialize() (Base)
+// - onShapeChange() (Base)
+// - attachToContext() (Base)
+// - enqueue() (HFace, MTron)
+/////
+
+PluginFieldCollection const* EmbLayerNormVarSeqlenPluginBase::getFieldsToSerialize() noexcept
+{
+    mDataToSerialize.clear();
+    bool output_fp16 = mType == DataType::kHALF;
+    mDataToSerialize.emplace_back("output_fp16", &output_fp16, PluginFieldType::kINT32, 1);
+    mDataToSerialize.emplace_back("bert_embeddings_layernorm_beta", static_cast<float const*>(mBeta.values),
+        PluginFieldType::kFLOAT32, mBeta.count);
+    mDataToSerialize.emplace_back("bert_embeddings_layernorm_gamma", static_cast<float const*>(mGamma.values),
+        PluginFieldType::kFLOAT32, mGamma.count);
+    if (output_fp16)
+    {
+        mDataToSerialize.emplace_back("bert_embeddings_word_embeddings", static_cast<half const*>(mWordEmb.values),
+            PluginFieldType::kFLOAT16, mWordEmb.count);
+        mDataToSerialize.emplace_back("bert_embeddings_token_type_embeddings", static_cast<half const*>(mTokEmb.values),
+            PluginFieldType::kFLOAT16, mTokEmb.count);
+        mDataToSerialize.emplace_back("bert_embeddings_position_embeddings", static_cast<half const*>(mPosEmb.values),
+            PluginFieldType::kFLOAT16, mPosEmb.count);
+    }
+    else
+    {
+        mDataToSerialize.emplace_back("bert_embeddings_word_embeddings", static_cast<float const*>(mWordEmb.values),
+            PluginFieldType::kFLOAT32, mWordEmb.count);
+        mDataToSerialize.emplace_back("bert_embeddings_token_type_embeddings",
+            static_cast<float const*>(mTokEmb.values), PluginFieldType::kFLOAT32, mTokEmb.count);
+        mDataToSerialize.emplace_back("bert_embeddings_position_embeddings", static_cast<float const*>(mPosEmb.values),
+            PluginFieldType::kFLOAT32, mPosEmb.count);
+    }
+    mFCToSerialize.nbFields = mDataToSerialize.size();
+    mFCToSerialize.fields = mDataToSerialize.data();
+    return &mFCToSerialize;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginHFace::onShapeChange(
+    PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
+{
+    try
+    {
+        BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace onShapeChange");
+        checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
+
+        // output 0 is the embedding
+        PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[1]) == static_cast<size_t>(mLd));
+        PLUGIN_ASSERT(outputs[0].type == mType);
+        // output 1 is the mask indices (empty for HFace variant)
+        PLUGIN_ASSERT(outputs[1].dims.nbDims == 0);
+        PLUGIN_ASSERT(outputs[1].type == mMaskType);
+        return pluginStatus_t::STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return pluginStatus_t::STATUS_FAILURE;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginMTron::onShapeChange(
+    PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
+{
+    try
+    {
+        // Validate input arguments
+        BERT_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron onShapeChange");
+        checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
+        PLUGIN_ASSERT(static_cast<size_t>(outputs[0].dims.d[1]) == static_cast<size_t>(mLd));
+
+        PLUGIN_ASSERT(outputs[1].dims.nbDims == 4);
+        PLUGIN_ASSERT(static_cast<size_t>(outputs[1].dims.d[0]) == static_cast<size_t>(inputs[0].dims.d[0]));
+        PLUGIN_ASSERT(static_cast<size_t>(outputs[1].dims.d[1]) == static_cast<size_t>(mLd));
+        PLUGIN_ASSERT(outputs[1].dims.d[2] == 1);
+        PLUGIN_ASSERT(outputs[1].dims.d[3] == 1);
+
+        PLUGIN_ASSERT(outputs[0].type == mType);
+        PLUGIN_ASSERT(outputs[1].type == mMaskType);
+        return pluginStatus_t::STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return pluginStatus_t::STATUS_FAILURE;
+}
+
+IPluginV3* EmbLayerNormVarSeqlenPluginBase::attachToContext(IPluginResourceContext* context) noexcept
+{
+    return clone();
+}
+
+int32_t EmbLayerNormVarSeqlenPluginHFace::enqueue(PluginTensorDesc const* inputDesc,
+    PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs, void* /* workspace */,
+    cudaStream_t stream) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
+
+        int32_t const batchSize = inputDesc[2].dims.d[0] - 1;
+        // read out the maximum sequence length from the dummy input
+        int32_t const maxSeqlen = inputDesc[3].dims.d[0];
+
+        // There are four versions of the kernel which are optimized for sequence lengths 384, 256, 192 and 128.
+        // Find the closest sequence length bigger than the max seq length in this batch.
+        int32_t S = 384;
+        if (maxSeqlen <= 128)
+        {
+            S = 128;
+        }
+        else if (maxSeqlen <= 192)
+        {
+            S = 192;
+        }
+        else if (maxSeqlen <= 256)
+        {
+            S = 256;
+        }
+
+        // Our plugin outputs only one tensor
+        auto const inputIds = static_cast<int32_t const*>(inputs[0]);
+        auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
+        int32_t const* cuSeqlens = static_cast<int32_t const*>(inputs[2]);
+
+        float const* beta = mBetaDev.get();
+        float const* gamma = mGammaDev.get();
+        if (mType == DataType::kFLOAT)
+        {
+            auto output = static_cast<float*>(outputs[0]);
+            auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
+            auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
+            auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
+
+            return embSkipLayerNormHFace<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
+                cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
+        }
+        if (mType == DataType::kHALF)
+        {
+            auto output = static_cast<half*>(outputs[0]);
+            auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
+            auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
+            auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
+
+            return embSkipLayerNormHFace<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
+                cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output);
+        }
+        else
+        {
+            gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
+                      << std::endl;
+
+            return STATUS_NOT_SUPPORTED;
+        }
+
+        return STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return STATUS_FAILURE;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginMTron::enqueue(PluginTensorDesc const* inputDesc,
+    PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs, void* /* workspace */,
+    cudaStream_t stream) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
+
+        int32_t const batchSize = inputDesc[2].dims.d[0] - 1;
+        // read out the maximum sequence length from the dummy input
+        int32_t const maxSeqlen = inputDesc[3].dims.d[0];
+
+        // There are four versions of the kernel which are optimized for sequence lengths 384, 256, 192 and 128.
+        // Find the closest sequence length bigger than the max seq length in this batch.
+        int32_t S = 384;
+        if (maxSeqlen <= 128)
+        {
+            S = 128;
+        }
+        else if (maxSeqlen <= 192)
+        {
+            S = 192;
+        }
+        else if (maxSeqlen <= 256)
+        {
+            S = 256;
+        }
+
+        // Our plugin outputs only one tensor
+        auto const inputIds = static_cast<int32_t const*>(inputs[0]);
+        auto const segmentIds = static_cast<int32_t const*>(inputs[1]);
+        int32_t const* cuSeqlens = static_cast<int32_t const*>(inputs[2]);
+
+        float const* beta = mBetaDev.get();
+        float const* gamma = mGammaDev.get();
+        if (mType == DataType::kFLOAT)
+        {
+            auto output = static_cast<float*>(outputs[0]);
+            auto skip = static_cast<float*>(outputs[1]);
+            auto const wordEmb = static_cast<float const*>(mWordEmbDev.get());
+            auto const tokEmb = static_cast<float const*>(mTokEmbDev.get());
+            auto const posEmb = static_cast<float const*>(mPosEmbDev.get());
+
+            return embSkipLayerNormMTron<float>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
+                cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output, skip);
+        }
+        if (mType == DataType::kHALF)
+        {
+            auto output = static_cast<half*>(outputs[0]);
+            auto skip = static_cast<half*>(outputs[1]);
+            auto const wordEmb = static_cast<half const*>(mWordEmbDev.get());
+            auto const tokEmb = static_cast<half const*>(mTokEmbDev.get());
+            auto const posEmb = static_cast<half const*>(mPosEmbDev.get());
+
+            return embSkipLayerNormMTron<half>(stream, static_cast<int32_t>(mLd), batchSize, S, inputIds, segmentIds,
+                cuSeqlens, beta, gamma, wordEmb, posEmb, tokEmb, mWordVocabSize, mTokVocabSize, output, skip);
+        }
+        else
+        {
+            gLogError << "Unsupported type error, expected [kHALF,kFLOAT], but received " << static_cast<int32_t>(mType)
+                      << std::endl;
+
+            return STATUS_NOT_SUPPORTED;
+        }
+
+        return STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return STATUS_FAILURE;
+}
+
+// end IPluginV3OneRuntime method definitions
+
+///////
+// IPluginV3OneBuild method definitions
+// - getNbOutputs() (Base)
+// - supportsFormatCombination() (Base)
+// - getOutputShapes (HFace, MTron)
+// - getOutputDataTypes() (Base)
+// - configurePlugin() (Base)
+// - getWorkSpaceSize() (Base)
+//////
+
+int32_t EmbLayerNormVarSeqlenPluginBase::getNbOutputs() const noexcept
+{
+    return 2;
+}
+
+bool EmbLayerNormVarSeqlenPluginBase::supportsFormatCombination(
+    int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
+{
+    // The four inputs to this plugin input_ids, segment_ids, cu_seqlens and a dummy input with the
+    // size of the max seq length in that order
+    PLUGIN_ASSERT(nbInputs == 4);
+    // The two outputs of the plugin are embedding and the mask
+    PLUGIN_ASSERT(nbOutputs == 2);
+
+    PluginTensorDesc const& desc = inOut[pos].desc;
+    if (desc.format != TensorFormat::kLINEAR)
+    {
+        return false;
+    }
+    if (pos == 0 || pos == 2) // input_ids and cu_seqlens
+    {
+        return desc.type == DataType::kINT32 && desc.dims.nbDims == 1;
+    }
+
+    PluginTensorDesc const& prev = inOut[pos - 1].desc;
+    if (pos == 1) // segment ids: check it's the same as input_ids
+    {
+        return desc.type == DataType::kINT32 && desc.dims.nbDims == 1 && desc.dims.d[0] == prev.dims.d[0];
+    }
+
+    if (pos == 3)
+    {
+        return desc.dims.nbDims == 1;
+    }
+
+    // embedded sequence
+    if (pos == nbInputs)
+    {
+        return desc.type == mType && desc.dims.nbDims == 4 && desc.dims.d[0] == inOut[0].desc.dims.d[0]
+            && desc.dims.d[2] == 1 && desc.dims.d[3] == 1;
+    }
+    // mask
+    return desc.type == mMaskType;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginHFace::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
+    DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
+    IExprBuilder& exprBuilder) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputs != nullptr);
+        PLUGIN_VALIDATE(outputs != nullptr);
+
+        // Input should be input ids and token ids and cumulative seqlens
+        // Output should be the embeddings tensor and mask indices
+        PLUGIN_ASSERT(nbInputs == 4);
+        PLUGIN_ASSERT(nbOutputs == 2);
+
+        PLUGIN_ASSERT(inputs[0].nbDims == 1); // sum of all s
+        PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
+
+        PLUGIN_ASSERT(inputs[2].nbDims == 1); // B+1
+
+        // output 0 : embedded input
+        outputs[0].nbDims = 4;
+        outputs[0].d[0] = inputs[0].d[0];
+        outputs[0].d[1] = exprBuilder.constant(mLd);
+        outputs[0].d[2] = exprBuilder.constant(1);
+        outputs[0].d[3] = exprBuilder.constant(1);
+
+        // Output 1 : maskIdx
+        // Return empty tensor since this is dummy output, we do not delete it for backward compatibility.
+        outputs[1].nbDims = 0;
+        return pluginStatus_t::STATUS_SUCCESS;
+    }
+    catch (const std::exception& e)
+    {
+        caughtError(e);
+    }
+    return pluginStatus_t::STATUS_FAILURE;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginMTron::getOutputShapes(DimsExprs const* inputs, int32_t nbInputs,
+    DimsExprs const* shapeInputs, int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs,
+    IExprBuilder& exprBuilder) noexcept
+{
+    try
+    {
+        PLUGIN_VALIDATE(inputs != nullptr);
+        PLUGIN_VALIDATE(outputs != nullptr);
+        // Input should be input ids and token ids and cumulative seqlens
+        // Output should be the embeddings tensor and mask indices
+        PLUGIN_ASSERT(nbInputs == 4);
+        PLUGIN_ASSERT(nbOutputs == 2);
+
+        PLUGIN_ASSERT(inputs[0].nbDims == 1); // sum of all s
+        PLUGIN_ASSERT(inputs[0].nbDims == inputs[1].nbDims);
+        PLUGIN_ASSERT(inputs[2].nbDims == 1); // B+1
+
+        // Output 0 : embedded input
+        outputs[0].nbDims = 4;
+        outputs[0].d[0] = inputs[0].d[0];
+        outputs[0].d[1] = exprBuilder.constant(mLd);
+        outputs[0].d[2] = exprBuilder.constant(1);
+        outputs[0].d[3] = exprBuilder.constant(1);
+
+        // Output 1 : maskIdx
+        outputs[1].nbDims = 4;
+        outputs[1].d[0] = inputs[0].d[0];
+        outputs[1].d[1] = exprBuilder.constant(mLd);
+        outputs[1].d[2] = exprBuilder.constant(1);
+        outputs[1].d[3] = exprBuilder.constant(1);
+
+        return pluginStatus_t::STATUS_SUCCESS;
+    }
+    catch (const std::exception& e)
+    {
+        caughtError(e);
+    }
+    return pluginStatus_t::STATUS_FAILURE;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginBase::getOutputDataTypes(
+    DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
+{
+    try
+    {
+        PLUGIN_ASSERT(mType == DataType::kHALF || mType == DataType::kFLOAT);
+        outputTypes[0] = mType;
+        outputTypes[1] = mMaskType;
+        return pluginStatus_t::STATUS_SUCCESS;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return pluginStatus_t::STATUS_FAILURE;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginBase::configurePlugin(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) noexcept
+{
+    return pluginStatus_t::STATUS_SUCCESS;
+}
+
+size_t EmbLayerNormVarSeqlenPluginBase::getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
+    DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
+{
+    return 0;
+}
+// End IPluginV3OneBuild method definitions
+
+//////
+// IPluginV3OneCore method definitions
+// - getPluginVersion() (MTron, HFace)
+// - getPluginName() (Base)
+// - getPluginNamespace() (Base)
+// - setPluginNamespace() (Base)
+//////
+char const* EmbLayerNormVarSeqlenPluginHFace::getPluginVersion() const noexcept
+{
+    return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
+}
+
+char const* EmbLayerNormVarSeqlenPluginMTron::getPluginVersion() const noexcept
+{
+    return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
+}
+
+char const* EmbLayerNormVarSeqlenPluginBase::getPluginName() const noexcept
+{
+    return kEMB_LAYER_NORM_VAR_SEQLEN_NAME;
+}
+
+char const* EmbLayerNormVarSeqlenPluginBase::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}
+
+void EmbLayerNormVarSeqlenPluginBase::setPluginNamespace(char const* libNamespace) noexcept
+{
+    mNamespace = libNamespace;
+}
+// End IPluginV3OneCore method definitions
+
+//////////////////////////// Plugin Creator member definitions /////////////////////////////
+
+EmbLayerNormVarSeqlenPluginBaseCreator::EmbLayerNormVarSeqlenPluginBaseCreator()
+{
+    static std::mutex sMutex;
+    std::lock_guard<std::mutex> lock(sMutex);
+    mPluginAttributes.clear();
+    mPluginAttributes.emplace_back(PluginField("output_fp16", nullptr, PluginFieldType::kINT32, 1));
+    // the length of beta, gamma, word_emb, pos_emb, and tok_emb will only be known at the time of plugin creation
+    // so we set it to 0 here
+    mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_beta", nullptr, PluginFieldType::kFLOAT32, 0));
+    mPluginAttributes.emplace_back(PluginField("bert_embeddings_layernorm_gamma", nullptr, PluginFieldType::kFLOAT32, 0));
+    // the embeddings datatype is determined by the output_fp16 attribute known at runtime
+    // so we set it to kUNKNOWN here
+    mPluginAttributes.emplace_back(PluginField("bert_embeddings_word_embeddings", nullptr, PluginFieldType::kUNKNOWN, 0));
+    mPluginAttributes.emplace_back(PluginField("bert_embeddings_token_type_embeddings", nullptr, PluginFieldType::kUNKNOWN, 0));
+    mPluginAttributes.emplace_back(PluginField("bert_embeddings_position_embeddings", nullptr, PluginFieldType::kUNKNOWN, 0));
+    mFC.nbFields = mPluginAttributes.size();
+    mFC.fields = mPluginAttributes.data();
+}
+
+char const* EmbLayerNormVarSeqlenPluginBaseCreator::getPluginName() const noexcept
+{
+    return kEMB_LAYER_NORM_VAR_SEQLEN_NAME;
+}
+
+char const* EmbLayerNormVarSeqlenPluginHFaceCreator::getPluginVersion() const noexcept
+{
+    return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
+}
+
+char const* EmbLayerNormVarSeqlenPluginMTronCreator::getPluginVersion() const noexcept
+{
+    return kEMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
+}
+
+PluginFieldCollection const* EmbLayerNormVarSeqlenPluginBaseCreator::getFieldNames() noexcept
+{
+    return &mFC;
+}
+
+IPluginV3* EmbLayerNormVarSeqlenPluginHFaceCreator::createPlugin(
+    char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
+{
+    try
+    {
+        BERT_DEBUG_MSG("EmbLayerNormVarSeqlenHFace createPlugin");
+
+        Weights beta{};  // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
+                         // existence
+        Weights gamma{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
+                         // existence
+        Weights word_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
+                            // verify existence
+        Weights pos_emb{};  // required attribute: validateRequiredAttributesExist() call in initializeFields() will
+                            // verify existence
+        Weights tok_emb{};  // required attribute: validateRequiredAttributesExist() call in initializeFields() will
+                            // verify existence
+        bool output_fp16 = initializeFields(name, fc, beta, gamma, word_emb, pos_emb, tok_emb);
+
+        BERT_DEBUG_MSG("Building the Plugin...");
+        auto p = std::make_unique<EmbLayerNormVarSeqlenPluginHFace>(
+            name, output_fp16 ? DataType::kHALF : DataType::kFLOAT, beta, gamma, word_emb, pos_emb, tok_emb);
+        return p.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+IPluginV3* EmbLayerNormVarSeqlenPluginMTronCreator::createPlugin(
+    char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
+{
+    try
+    {
+        BERT_DEBUG_MSG("EmbLayerNormVarSeqlenMTron createPlugin");
+
+        Weights beta{};  // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
+                         // existence
+        Weights gamma{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will verify
+                         // existence
+        Weights word_emb{}; // required attribute: validateRequiredAttributesExist() call in initializeFields() will
+                            // verify existence
+        Weights pos_emb{};  // required attribute: validateRequiredAttributesExist() call in initializeFields() will
+                            // verify existence
+        Weights tok_emb{};  // required attribute: validateRequiredAttributesExist() call in initializeFields() will
+                            // verify existence
+        bool output_fp16 = initializeFields(name, fc, beta, gamma, word_emb, pos_emb, tok_emb);
+
+        BERT_DEBUG_MSG("Building the Plugin...");
+        auto p = std::make_unique<EmbLayerNormVarSeqlenPluginMTron>(
+            name, output_fp16 ? DataType::kHALF : DataType::kFLOAT, beta, gamma, word_emb, pos_emb, tok_emb);
+        return p.release();
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+    return nullptr;
+}
+
+void EmbLayerNormVarSeqlenPluginBaseCreator::setPluginNamespace(char const* libNamespace) noexcept
+{
+    try
+    {
+        mNamespace = libNamespace;
+    }
+    catch (std::exception const& e)
+    {
+        caughtError(e);
+    }
+}
+
+char const* EmbLayerNormVarSeqlenPluginBaseCreator::getPluginNamespace() const noexcept
+{
+    return mNamespace.c_str();
+}