FastSoftMaxOperation

Applicable Products

Description

Training operator, which is used together with unpadOperation to perform high-performance SoftMax processing on the result of multiplying matrix Q and matrix K after unpadding.

SoftMax formula:

Use Cases

Figure 1 Model flowchart

The Multi-Head Attention calculation process is shown in Figure 1. After unpad processing, the result of multiplying matrix Q and matrix K is actually concatenated by batchSize matrices whose sizes are (seqLen[i], seqLen[i]) (as shown in Figure 2). The SoftMax operation needs to be performed on the last axis of the result. The FastSoftMax operator is used to receive the seqLen and headNum information and implement the SoftMax operation on the data in this arrangement mode.

Figure 2 Arrangement of input data

Definition

1
2
3
4
5

struct FastSoftMaxParam {
    int32_t headNum = 0;
    std::vector<int32_t> qSeqLen;
    uint8_t rsv[8] = {0};
};

Parameters

Input

Output

Constraints

• The length of the qSeqLen array cannot exceed 32, and each element must be greater than 0.
• The dimension size nSquareTokens of the input tensor (inTensor) and output tensor (outTensor) is related to headNum and qSeqLen in the parameters.

  

  : indicates the kth element in qSeqLen.

Functions

When the parameter list contains

,

the shape size of the corresponding input tensor is .

For example,

In this case, the shape size of the input tensor is

The shape of the input and output tensors is [5528144].