[object Object]

In the field of foundation models, sparseMode refers to the sparsity design of parameters or activations in the model architecture or computation formula, which is opposite to DenseMode.

This section describes the common sparseMode and the corresponding scenarios.

[object Object]undefined

The working principle of attenMask is to mask the values of the product between query (Q) and the transpose matrix of key (K) at the positions where the Mask is True, as shown in the following figure:

The $QK^T$ matrix is masked at the position where attenMask is True. The effect is as follows:

[object Object]

When sparseMode is set to 0, the defaultMask mode is used.

• If attenMask is not passed, the mask operation is not performed. In this case, attenMask is set to None, and the values of preTokens and nextTokens are ignored. The masked $QK^T$ matrix is as follows.

  

• If nextTokens is set to 0 and preTokens is greater than or equal to Sq, the sparse mode is causal. In this case, the lower triangular matrix should be passed to attenMask, and the part between preTokens and nextTokens needs to be computed. The masked $QK^T$ matrix is as follows:

  

  The lower triangular matrix should be passed to attenMask. The following figure shows the details.

  

• If preTokens is less than Sq, nextTokens is less than Skv, and both of them are greater than or equal to 0, the band mode is used. In this case, the part between preTokens and nextTokens needs to be computed. The masked $QK^T$ matrix is as follows.

  

  The attenMask should be a matrix in the shape of the band, as shown in the following figure.

  

• nextTokens is a negative number. For example, if preTokens is 9 and nextTokens is -3, the part between preTokens and nextTokens needs to be calculated. The masked $QK^T$ is as follows.

  Note: When nextTokens is a negative number, the value of preTokens must be greater than or equal to the absolute value of nextTokens, and the absolute value of nextTokens must be less than Skv.

  

• preTokens is a negative number. For example, if nextTokens is 7 and preTokens is -3, the part between preTokens and nextTokens needs to be calculated. The masked $QK^T$ is as follows.

  Note: When preTokens is a negative number, the value of nextTokens must be greater than or equal to the absolute value of preTokens, and the absolute value of preTokens must be less than Sq.

  

[object Object]

When sparseMode is set to 1, it indicates allMask, that is, the complete attenMask matrix is input.

In this scenario, the values of nextTokens and preTokens are ignored. The masked $QK^T$ matrix is shown as follows:

[object Object]

When sparseMode is set to 2, it indicates the mask of the leftUpCausal mode, which corresponds to the lower triangular scenario divided by the upper left vertex (the start point of the parameter is the upper left corner).

In this scenario, the values of preTokens and nextTokens are ignored. The masked $QK^T$ matrix is shown as follows:

The input attenMask is the optimized lower triangular matrix (2048 x 2048). The following figure shows the compressed lower triangular matrix.

[object Object]

When sparseMode is set to 3, it indicates the mask of the rightDownCausal mode, which corresponds to the lower triangular scenario divided by the lower right vertex (the start point of the parameter is the lower right corner).

In this scenario, the values of preTokens and nextTokens are ignored. attenMask is the optimized and compressed lower triangular matrix (2048 x 2048). The following figure shows the Masked $QK^T$ matrix.

[object Object]

When sparseMode is set to 4, it indicates the band scenario, that is, the part between preTokens and nextTokens is calculated. The start point of the parameter is the lower right corner, and there must be an intersection between preTokens and nextTokens. attenMask is the optimized and compressed lower triangular matrix (2048 x 2048). The masked $QK^T$ matrix is as follows.

[object Object]

When sparseMode is set to 5, it indicates the non-compressed prefix scenario. That is, a matrix with the length of Sq and the width of N is added to the left of the rightDownCausal. The value of N is obtained from the optional input prefix. For example, in the following figure, the prefix input array is [4, 5] when batch = 2. The value of N for each batch axis can be different. The start point of the parameter is the upper left corner.

In this scenario, the values of preTokens and nextTokens are ignored. The data format of the attenMask matrix must be BNSS or B1SS. The following figure shows the Masked $QK^T$ matrix.

The following figure shows the matrix to be transferred to attenMask.

[object Object]

When sparseMode is set to 6, it indicates the compressed prefix scenario. That is, in the prefix scenario, attenMask is the optimized and compressed lower triangular matrix and rectangular matrix (3072 x 2048). The upper part is a lower triangular matrix of [2048, 2048], and the lower part is a rectangular matrix of [1024, 2048]. The left part of the rectangular matrix is all zeros, and the right part is all ones. The following figure shows the matrix to be transferred to attenMask. In this scenario, the values of preTokens and nextTokens are ignored.

[object Object]

When sparseMode is set to 7, it indicates that the varlen mode is used and the long sequence is split externally (that is, the long sequence is split into multiple queries with different sequence lengths in the model script). You need to ensure that sparseMode 3 is used before the external splitting. In this mode, you need to set preTokens and nextTokens (the start point is the lower right vertex) and ensure that the parameters are correct. Otherwise, accuracy problems may occur.

The following figure shows the masked $QK^T$ matrix. The query is sliced in the second batch, and the key and value are not sliced. The 4 × 6 mask matrix is sliced into 2 × 6 and 2 × 6 masks, which are computed on device 1 and device 2, respectively.

• The last mask of card 1 is a band mask. In this case, preTokens is set to 6 (ensure that the value is greater than or equal to the last Skv), nextTokens is set to –2, actual_seq_qlen is set to {3,5}, and actual_seq_kvlen is set to {3,9}.
• The mask type of card 2 remains unchanged after the split. In this case, sparseMode is set to 3, actual_seq_qlen is set to {2,7,11}, and actual_seq_kvlen is set to {6,11,15}.

Note:

• When sparseMode is set to 7, band indicates the sparse type of the last non-empty tensor in the batch. If there is only one batch, you need to set parameters according to the requirements of the band mode. When sparseMode is set to 7, you need to input a 2048 x 2048 lower triangular mask as the input of the fusion operator.
• The sparse parameters of the band mode generated by external splitting when sparseMode=3 must meet the following conditions:
  • preTokens >= last_Skv.
  • last_Sq-last_Skv <= nextTokens <= 0.
  • In the current mode, the optional input pse is not supported.
• Batch in non-band mode must meet the following condition: Sq <= Skv.

[object Object]

When sparseMode is set to 8, it indicates that the varlen mode is used and the long sequence is split externally. You need to ensure that sparseMode 2 is used before the external splitting. In this mode, you need to set preTokens and nextTokens (the start point is the lower right vertex) and ensure that the parameters are correct. Otherwise, accuracy problems may occur.

The following figure shows the masked $QK^T$ matrix. The query is sliced in the second batch, the key and value are not sliced, and the 5 × 4 mask matrix is sliced into 2 × 4 and 3 × 4 masks, which are computed on device 1 and device 2, respectively.

• The mask type of card 1 remains unchanged after splitting. If sparseMode is 2, actual_seq_qlen should be set to {3,5} and actual_seq_kvlen should be set to {3,7}.
• The first mask of card 2 is a band mask. Set preTokens to 4 (ensure that the value is greater than or equal to the first Skv), nextTokens to 1, actual_seq_qlen to {3,8,12}, and actual_seq_kvlen to {4,9,13}.

Note:

• If sparseMode is 8, band indicates the sparse type of the batch of the first non-empty tensor. If there is only one batch, you need to set the parameters according to the requirements of the band mode. When sparseMode is 8, you need to input a 2048 x 2048 lower triangular mask as the input of the fusion operator.
• The parameters of the sparse in band mode generated by external splitting when sparseMode=2 must meet the following conditions:
  • preTokens >= first_Skv.
  • nextTokens >= first_Sq – first_Skv. Set this parameter based on the actual situation.
  • In the current mode, the optional input pse is not supported.