FastSoftMaxOperation
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 unpad processing.
Softmax formula:
Application Scenarios
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.
Definition
struct FastSoftMaxParam {
int32_t headNum = 0;
std::vector<int32_t> qSeqLen;
uint8_t rsv[8] = {0};
};
Parameters
|
Member |
Type |
Default Value |
Description |
|---|---|---|---|
|
headNum |
int32_t |
0 |
Number of Attention heads. |
|
qSeqLen |
std::vector<int32_t> |
- |
Actual input length of each batch. The number of elements is batchSize. The maximum value is 32. |
|
rsv[8] |
uint8_t |
{0} |
Reserved |
Input
|
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
|
inTensor |
[nSquareTokens] |
float16 |
ND |
Input tensor, which is the 1D result of batch (headNum, qSeqLen[i], qSeqLen[i]) matrices in ND format. |
Output
|
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
|
outTensor |
[nSquareTokens] |
float16 |
ND |
Output tensor. The data range is [0, 1]. |
Restrictions
- The length of the qSeqLen array cannot exceed 32, and each element must be greater than 0.
- Currently, only the
Atlas A2 inference products is supported. - 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 qSeqLen = [qSeqLen0, qSeqLen1, ..., qSeqLenn], headNum = head is in the parameter list, the shape size of the corresponding input tensor is [nSquareTokens], nSquareTokens = ∑k = 0n(qSeqLenk)2∗head.
For example, if qSeqLen is set to [487, 16, 532, 413] and headNum is set to 8, the shape size of the input tensor is nSquareTokens = 4872∗8+162∗8+5322∗8+4132∗8 = 5528144,, and the shape of the input and output tensors is [5528144].