AttentionHelpers.cpp

/*
 * SPDX-License-Identifier: Apache-2.0
 *
 */

#include "AttentionHelpers.hpp"
#include "ImporterContext.hpp"
#include "NvInfer.h"
#include "ShapeTensor.hpp"
#include "errorHelpers.hpp"
#include "importerUtils.hpp"
#include <cmath>
#include <numeric>
#include <string>
#include <vector>

namespace
{
//!
//! \brief Return true if `divident` is divisible by `divisor`.
//!
bool isDivisible(int64_t const divident, int64_t const divisor)
{
    return (divisor != 0) && ((divident % divisor) == 0);
}
} // namespace

namespace onnx2trt
{

//!
//! \brief Reshape and return the Q, K, or V tensor from the input tensor.
//!
//! \param qkvInput The input tensor. This can either be a 4D tensor (batchSize, numHeads, sequenceLength, headSize) or
//!                 a 3D tensor (batchSize, sequenceLength, hiddenSize=numHeads*headSize). If it is a 3D tensor,
//!                 permute and reshape to the 4D shape before returning. Otherwise, return the input tensor.
//! \param ctx The importer context.
//! \param numHeadsValue The number of heads in the tensor.
//! \param needsReshape True if the tensor is 3D and needs to be reshaped to 4D.
//! \return nvinfer1::ITensor& The Q, K, or V tensor.
//!
nvinfer1::ITensor& reshapeQKVTensor(
    TensorOrWeights& qkvInput, ImporterContext* ctx, int64_t const numHeadsValue, bool const needsReshape)
{
    if (needsReshape)
    {
        // qkvInput is a 3D tensor (batchSize, sequenceLength, hiddenSize=numHeads * headSize).
        // Get relevant dimensions.
        ONNXTRT_CHECK(numHeadsValue != 0,
            "Number of attention heads is not specified, which is required for 3D Q/K/V tensors.",
            ErrorCode::kINVALID_NODE);
        ShapeTensor numHeads = shapeVector(numHeadsValue);

        ShapeTensor hiddenSize = gather(ctx, shapeOf(qkvInput), shapeVector(2));
        if (hiddenSize.allValuesKnown())
        {
            // Perform static check for divisibility.
            ONNXTRT_CHECK(isDivisible(hiddenSize[0], numHeads[0]),
                "hidden_size must be divisible by num_heads. Received hidden_size=" << hiddenSize[0]
                                                                                    << " and num_heads=" << numHeads,
                ErrorCode::kINVALID_NODE);
        }

        ShapeTensor headSize = floorDiv(ctx, hiddenSize, numHeads);

        // == Transform (batchSize, sequenceLength, hiddenSize) -> (batchSize, numHeads, sequenceLength, headSize) by ==
        // 1. Reshape to (batchSize, sequenceLength, numHeads, headSize).
        // Use (0, 0, numHeads, headSize) as a shorthand to propagate `batchSize` and `sequenceLength` from the input
        // tensor without instantiating them. Set `zeroIsPlaceholder` to enable this shorthand.
        ShapeTensor newShape = concat(ctx, fillShapeVector(ctx, 0, shapeVector(2)), concat(ctx, numHeads, headSize));
        nvinfer1::IShuffleLayer* shuffle
            = addShuffle(ctx, convertToTensor(qkvInput, ctx), newShape, /*zeroIsPlaceholder*/ true);

        // 2. Permute to (batchSize, numHeads, sequenceLength, headSize)
        shuffle->setSecondTranspose({0, 2, 1, 3});

        return *N_CHECK(shuffle->getOutput(0));
    }
    else
    {
        return convertToTensor(qkvInput, ctx);
    }
}

//!
//! \brief Scale the Q or K tensor by `sqrt(scale)`.
//!
//! `scale` is either provided as an attribute or set as the default value of `1/sqrt(headSize)`. `scale` is defined as
//! `QK^T -> QK^T * scale`, but we apply `Q -> Q * sqrt(scale)` and `K -> K * sqrt(scale)` for numerical stability.
//!
//! \param qkTensor The Q or K tensor to scale.
//! \param attrs The ONNX node attributes.
//! \param ctx The importer context.
//! \return nvinfer1::ITensor& The scaled Q or K tensor.
//!
nvinfer1::ITensor& scaleQKTensor(nvinfer1::ITensor& qkTensor, OnnxAttrs const& attrs, ImporterContext* ctx)
{
    nvinfer1::ITensor* sqrtScale = nullptr;

    // Use the tensor's actual rank so this works for both 4D padded BHND and 3D packed NHD tensors.
    int32_t const nbDims = qkTensor.getDimensions().nbDims;

    if (attrs.contains("scale"))
    {
        // Obtain the sqrt of scale as a constant (output of a constant layer).
        nvinfer1::IConstantLayer* constant = addConstantScalar(
            ctx, std::sqrt(attrs.get<float>("scale")), ::ONNX_NAMESPACE::TensorProto::FLOAT, nbDims);
        sqrtScale = castHelper(ctx, N_CHECK(constant)->getOutput(0), qkTensor.getType());
    }
    else
    {
        // headSize is always the last dimension: dim 3 for 4D BHND, dim 2 for 3D NHD.
        ShapeTensor headSize = gather(ctx, shapeOf(qkTensor), shapeScalar(nbDims - 1));
        nvinfer1::ITensor* headSizeF = castHelper(ctx, &headSize.tensor(ctx), qkTensor.getType());

        // By default, scale := 1/sqrt(headSize)
        nvinfer1::ITensor* sqrtHeadSize = getUnaryResult(ctx, *headSizeF, nvinfer1::UnaryOperation::kSQRT);
        nvinfer1::ITensor* scale = getUnaryResult(ctx, *sqrtHeadSize, nvinfer1::UnaryOperation::kRECIP);

        sqrtScale = getUnaryResult(ctx, *scale, nvinfer1::UnaryOperation::kSQRT);
        std::vector<int32_t> unsqueezeAxes(nbDims);
        std::iota(unsqueezeAxes.begin(), unsqueezeAxes.end(), 0);
        sqrtScale = unsqueezeTensor(ctx, *sqrtScale, unsqueezeAxes);
    }

    // Scale Q or K tensor by `sqrt(scale)`.
    return *getElementWiseResult(ctx, qkTensor, *sqrtScale, nvinfer1::ElementWiseOperation::kPROD);
}

nvinfer1::ITensor& convertToQTensor(
    TensorOrWeights& qInput, OnnxAttrs const& attrs, ImporterContext* ctx, bool const needsReshape)
{
    return convertToQTensor(qInput, attrs, ctx, attrs.get<int64_t>("q_num_heads", 0), needsReshape);
}

nvinfer1::ITensor& convertToQTensor(TensorOrWeights& qInput, OnnxAttrs const& attrs, ImporterContext* ctx,
    int64_t const numHeads, bool const needsReshape)
{
    return scaleQKTensor(reshapeQKVTensor(qInput, ctx, numHeads, needsReshape), attrs, ctx);
}

nvinfer1::ITensor& convertToKTensor(
    TensorOrWeights& kInput, OnnxAttrs const& attrs, ImporterContext* ctx, bool const needsReshape)
{
    return convertToKTensor(kInput, attrs, ctx, attrs.get<int64_t>("kv_num_heads", 0), needsReshape);
}

nvinfer1::ITensor& convertToKTensor(TensorOrWeights& kInput, OnnxAttrs const& attrs, ImporterContext* ctx,
    int64_t const numHeads, bool const needsReshape)
{
    return scaleQKTensor(reshapeQKVTensor(kInput, ctx, numHeads, needsReshape), attrs, ctx);
}

nvinfer1::ITensor& convertToVTensor(
    TensorOrWeights& vInput, OnnxAttrs const& attrs, ImporterContext* ctx, bool const needsReshape)
{
    return convertToVTensor(vInput, ctx, attrs.get<int64_t>("kv_num_heads", 0), needsReshape);
}

nvinfer1::ITensor& convertToVTensor(
    TensorOrWeights& vInput, ImporterContext* ctx, int64_t const numHeads, bool const needsReshape)
{
    return reshapeQKVTensor(vInput, ctx, numHeads, needsReshape);
}

nvinfer1::ITensor& convertToMaskTensor(TensorOrWeights& maskInput, ImporterContext* ctx)
{
    ONNXTRT_CHECK(maskInput.shape().nbDims <= 4,
        "Attention masks should have rank leq 4. Got mask with rank " << maskInput.shape().nbDims << ".",
        ErrorCode::kINVALID_NODE);

    if (maskInput.shape().nbDims == 4)
    {
        // Mask has rank 4. Directly return the mask tensor.
        return convertToTensor(maskInput, ctx);
    }
    else
    {
        // Mask has rank less than 4. Reshape to rank 4 by prepending dimensions.
        int32_t const numDimsToPrepend = 4 - maskInput.shape().nbDims;
        std::vector<int32_t> unsqueezeAxes(numDimsToPrepend);
        std::iota(unsqueezeAxes.begin(), unsqueezeAxes.end(), 0);

        return *unsqueezeTensor(ctx, convertToTensor(maskInput, ctx), unsqueezeAxes);
    }
}

nvinfer1::AttentionNormalizationOp parseNormalizationOp(OnnxAttrs const& attrs)
{
    std::string normalizationOp
        = attrs.get<std::string>("TRT_normalization_op", "softmax"); // Normalization op defaults to softmax.
    if (normalizationOp == "softmax")
    {
        return nvinfer1::AttentionNormalizationOp::kSOFTMAX;
    }
    else if (normalizationOp == "none")
    {
        return nvinfer1::AttentionNormalizationOp::kNONE;
    }
    else
    {
        ONNXTRT_CHECK(false, "Unsupported normalization op: " << normalizationOp, ErrorCode::kINVALID_NODE);
    }
}

nvinfer1::CausalMaskKind parseCausalKind(OnnxAttrs const& attrs)
{
    std::string const kind = attrs.get<std::string>("causal_kind", "none");
    if (kind == "none")
    {
        return nvinfer1::CausalMaskKind::kNONE;
    }
    else if (kind == "upper_left")
    {
        return nvinfer1::CausalMaskKind::kUPPER_LEFT;
    }
    else if (kind == "lower_right")
    {
        return nvinfer1::CausalMaskKind::kLOWER_RIGHT;
    }
    else
    {
        ONNXTRT_CHECK(false, "Unsupported causal_kind: " << kind, ErrorCode::kINVALID_NODE);
    }
}

nvinfer1::AttentionIOForm parseIOForm(OnnxAttrs const& attrs, std::string const& attrName)
{
    std::string form = attrs.get<std::string>(attrName, "padded_bhnd");
    if (form == "padded_bhnd")
    {
        return nvinfer1::AttentionIOForm::kPADDED_BHND;
    }
    else if (form == "packed_nhd")
    {
        return nvinfer1::AttentionIOForm::kPACKED_NHD;
    }
    else
    {
        ONNXTRT_CHECK(false, "Unsupported IO form: " << form, ErrorCode::kINVALID_NODE);
    }
}

nvinfer1::ITensor& reshapeOutputTensor(nvinfer1::ITensor& tensor, ImporterContext* ctx, bool const needsReshape)
{
    if (!needsReshape)
    {
        return tensor;
    }
    else
    {
        ShapeTensor numHeads = gather(ctx, shapeOf(tensor), shapeVector(1));
        ShapeTensor headSize = gather(ctx, shapeOf(tensor), shapeVector(3));
        ShapeTensor hiddenSize = mul(ctx, numHeads, headSize);

        // == Transform (batchSize, numHeads, sequenceLength, headSize) -> (batchSize, sequenceLength, hiddenSize) by ==
        // 1. Transpose the middle two dimensions: (batchSize, numHeads, sequenceLength, headSize) -> (batchSize,
        // sequenceLength, numHeads, headSize)
        // 2. Reshape to (batchSize, sequenceLength, hiddenSize).
        // Use (0, 0, hiddenSize) as a shorthand to propagate `batchSize` and `sequenceLength` from the input
        // tensor without instantiating them. Set `zeroIsPlaceholder` to enable this shorthand.
        ShapeTensor newShape = concat(ctx, fillShapeVector(ctx, 0, shapeVector(2)), hiddenSize);
        nvinfer1::IShuffleLayer* shuffle = addShuffle(ctx, tensor, newShape, /*zeroIsPlaceholder*/ true);
        shuffle->setFirstTranspose({0, 2, 1, 3});

        return *N_CHECK(shuffle->getOutput(0));
    }
}
} // namespace onnx2trt