LaserAttentionGradOperation
Description
In training scenarios, the LaserAttention algorithm is used to implement self-attention computation (reverse).
Definition
struct LaserAttentionGradParam {
int headNum = 0;
std::string inputLayout = "BNSD";
float scaleValue = 0.08838834764831843;
float keepProb = 1.0f;
int preTokens = 2147483647;
int nextTokens = 1;
int sparseMode = 0;
int innerPrecise = 1;
uint8_t rsv[8] = {0};
};
Parameters
Member |
Type |
Default Value |
Value Range |
Description |
|---|---|---|---|---|
headNum |
int |
0 |
(0, INT_MAX] |
Number of heads |
inputLayout |
string |
"BNSD" |
"BNSD" |
Input layout |
scaleValue |
float |
0.08838834764831843 |
(0, 1] |
Scaling coefficient |
keepProb |
float |
1.0f |
Currently, this parameter is not used. The configured value does not take effect. You can set this parameter to any value. |
Percentage of resources to be reserved, which is related to dropMask. |
preTokens |
int |
2147483647 |
[1, INT_MAX], an integer multiple of 256 |
Size of the sparse vertical sliding window. This parameter is related to seqSize. For details about the constraints, see Restrictions. |
nextTokens |
int |
1 |
Currently, this parameter is not used. The configured value does not take effect. You can set this parameter to any value. |
Size of the sparse horizontal sliding window. |
sparseMode |
int |
0 |
Currently, this parameter is not used. The configured value does not take effect. You can set this parameter to any value. |
Whether the lower triangle/trapezoid mask exists in sparse mode. |
innerPrecise |
int |
1 |
1 |
Precision mode. 1: high-precision mode. |
rsv[8] |
uint8_t |
{0} |
- |
Reserved |
Input
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
query |
[batch, qHeadNum, seqSize, qHeadDim] |
bf16 |
ND |
Q matrix of the attention mechanism |
key |
[batch, kvHeadNum, kvSize, kHeadDim] |
bf16 |
ND |
K matrix of the attention mechanism |
value |
[batch, kvHeadNum, kvSize, vHeadDim] |
bf16 |
ND |
V matrix of the attention mechanism |
attentionOutGrad |
[batch, qHeadNum, seqSize, vHeadDim] |
bf16 |
ND |
Gradient of the forward output (attentionOut) |
pseShift |
- |
- |
- |
Currently, this API is reserved. The input is an empty tensor. |
dropMask |
- |
- |
- |
Currently, this API is reserved. The input is an empty tensor. |
paddingMask |
- |
- |
- |
Currently, this API is reserved. The input is an empty tensor. |
attenMask |
[seqSize, kvSize] |
float16 |
ND |
Attention range. This API is optional. You can input a normal tensor or input an empty tensor that is not involved in computation. |
softmaxMax |
[batch, qHeadNum, seqSize] |
float |
ND |
Corresponding to the forward output softmaxMax |
softmaxSum |
[batch, qHeadNum, seqSize] |
float |
ND |
Corresponding to the forward output softmaxSum |
softmaxIn |
[...] |
bf16 |
ND |
Corresponding to the forward output softmaxOut. The dimension can be configured randomly. |
attentionIn |
[batch, qHeadNum, seqSize, vHeadDim] |
bf16 |
ND |
Corresponding to the forward output attentionOut |
prefix |
- |
- |
- |
Currently, this API is reserved. The input is an empty tensor. |
actualSeqQLen |
- |
- |
- |
Currently, this API is reserved. The input is an empty tensor. |
actualSeqKVLen |
- |
- |
- |
Currently, this API is reserved. The input is an empty tensor. |
Output
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
queryGrad |
[batch, qHeadNum, seqSize, qHeadDim] |
bf16 |
ND |
Output tensor. |
keyGrad |
[batch, kvHeadNum, kvSize, kHeadDim] |
bf16 |
ND |
Output tensor. |
valueGrad |
[batch, kvHeadNum, kvSize, vHeadDim] |
bf16 |
ND |
Output tensor. |
dpse |
[...] |
bf16 |
ND |
Currently, this API is reserved. The dimension can be configured as required. |
Description
- The value of qHeadNum is the same as that of headNum.
- The value of qHeadNum is an integer multiple of the value of kvHeadNum.
- The values of seqSize and kvSize are integer multiples of 256.
- When a query, key, or value is constructed, if the value range is within [-100, 100], data is generated in uniform distribution mode. Otherwise, data is generated in normal distribution mode. The average value is randomly selected from [-100, 100], and the standard deviation is randomly selected from [1, 25].
- When attenMask is not an empty tensor:
- The values of seqSize and kvSize are the same.
- In lower triangle form, when preTokens < seqSize, the operator is treated as a trapezoid.
- The value range of attentionOutGrad is [-0.5, 0.5].
- The values of qHeadDim, kHeadDim, and vHeadDim are as follows:
qHeadDim
kHeadDim
vHeadDim
Constraints
128
128
128
The value of seqSize must be greater than or equal to that of preTokens.
192
256
128
The value of seqSize is the same as that of preTokens.