Compression Mask

Description

Also called mask-free.

In the long sequence scenario, because seqLen is large, the size of the required mask also increases. In this case, you can enable this function and pass the compressed mask to reduce the display memory usage.

How to Enable

Set isTriuMask to 1.

• In the rope scenario, set maskType to MASK_TYPE_NORM_COMPRESS.
• In the alibi scenario, set maskType to MASK_TYPE_ALIBI_COMPRESS, MASK_TYPE_ALIBI_COMPRESS_SQRT, or MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN as required.

In the alibi compression mask scenario, slopes need to be passed.

Constraints

• The ALiBi compression mask takes effect only when calcType is set to PA_ENCODER or PREFIX_ENCODER.
• When maskType is set to MASK_TYPE_ALIBI_COMPRESS or MASK_TYPE_ALIBI_COMPRESS_SQRT, the mask dimension is as follows:
  • On the Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products : [headNum, seqLen, 128] or [256, 256]
  • On the Atlas inference products : [headNum,128//16,maxSeqLen,16] or [1,256//16,256,16]
• When maskType is set to MASK_TYPE_NORM_COMPRESS, the mask dimension is as follows:
  • On the Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products : [128,128]
  • On the Atlas inference products : [1,128//16,128,16]
• When the ALiBi compression mask is used, headSize of q, k, and v must be less than or equal to 128.
• MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN is supported only by the Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products . In this case, the mask dimension is [256, 256].
• Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products : When the data type is float16, the ALiBi compression mask is valid only when high precision is enabled.
• Methods for constructing different compression masks:
  • In the norm scenario:

    maskType = MASK_TYPE_NORM_COMPRESS

    The mask is the inverted triangle of 128 x 128.

    1. The following figure shows the mask in the float16 scenario.
       Figure 1 Compression norm mask in the float16 scenario
       
    2. The following figure shows the mask in the bf16/float16 high-precision scenario.
       Figure 2 Compression norm mask in the bf16/float16 scenario
       
  • In the alibi scenario:

    The alibi mask can be disassembled as follows:

    

    The alibi coefficient is different for each head, and triangularMask represents an inverted triangle mask. In the compressed mask scenario, the mask in the input tensor is the compressed alibi bias, and slopes is the alibi coefficient. There are three compression scenarios for alibi bias.

    1. Right-aligned alibi bias

       maskType = MASK_TYPE_ALIBI_COMPRESS

       The following figure shows the 512 x 512 alibi bias before compression.

       Figure 3 512 x 512 alibi bias before compression
       

       The following figure shows the 256 x 256 compressed alibi bias.

       Figure 4 256 x 256 compressed alibi bias
       
    2. Right-aligned alibi bias for square extraction

       maskType = MASK_TYPE_ALIBI_COMPRESS_SQRT

       On the basis of right-aligned ALiBi bias, the root sign of bias is changed. The bias of 1 is changed to , and other biases are the same as the basis of 1.

    3. Left-aligned alibi bias

       maskType = MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN

       The following figure shows the alibi bias before compression.

       Figure 5 Alibi bias before compression
       

       The following figure shows the compressed alibi bias.

       Figure 6 Compressed alibi bias
       

       Note: The left-aligned compression mask is supported only by Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products .