Mask Compression

Description

Also called mask-free.

For long sequences, the size of the required mask increases because of a large seqLen. In this case, you can enable this function and pass the compressed mask to reduce the graphics memory usage.

How to Enable

Set isTriuMask to 1. When MASK_TYPE_NORM_COMPRESS is used on Atlas inference products , the source code automatically forces isTriuMask to become 1 in the operator constructor, and you do not need to manually set it.

• RoPE: Set maskType to MASK_TYPE_NORM_COMPRESS.
• ALiBi: Set maskType to MASK_TYPE_ALIBI_COMPRESS, MASK_TYPE_ALIBI_COMPRESS_SQRT, or MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN as needed.

Pass slopes for ALiBi-compressed masks.

Constraints

• ALiBi-compressed masks take 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:
  • Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products : [headNum, seqLen, 128] or [256, 256]
  • 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:
  • Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products : [128,128]
  • Atlas inference products : [1,128//16,128,16] or [1,2048//16,2048,16] (that is, [1,128,2048,16], 2048-long window)
• When an ALiBi-compressed 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 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-compressed mask is valid only when high precision is enabled.
• For Atlas inference products, follow the constraints below when using the 2048-long window of MASK_TYPE_NORM_COMPRESS (mask dimension: [1,128,2048,16] or [2048,2048]):
  • Only Atlas inference products are supported.
  • The mask data type can only be float16.
  • Quantization is not supported. quantType must be TYPE_QUANT_UNDEFINED.
  • scaleType must be SCALE_TYPE_TOR (LOGN is not supported).
  • isTriuMask is forcibly set to 1 in the source code (at the constructor layer, when maskType is MASK_TYPE_NORM_COMPRESS and 310P is used). It is not necessary to pass this parameter externally (a triangular mask, where the lower triangle and diagonal are 0 (visible), and the upper triangle is a large negative value (invisible); -60000 or -10000 is recommended).
  • maskStride = 2048.
  • When calcType supports only PA_ENCODER, q_seqlen must be equal to kv_seqlen.
• Methods for constructing different compressed masks:
  • NORM:

    maskType = MASK_TYPE_NORM_COMPRESS.

    The mask is an inverted triangle of 128 x 128.

    1. The following figure shows a mask with float16 in use.
       Figure 1 Compressed norm mask with float16 in use
       
    2. The following figure shows a mask for bf16/float16 high precision.
       Figure 2 Compressed norm mask for bf16/float16
       
    3. 2048-long window for Atlas inference products :

       The mask is a 2048 x 2048 inverted triangle. The upper triangle (excluding the diagonal) is filled with a large negative value (invisible in float16; -60000 or -10000 is recommended), and the lower triangle and diagonal are 0 (visible). The mask is a shared 2048 x 2048 canvas. All batches/heads share the same mask. Regardless of the query/KV sequence length, the same 2048 x 2048 mask is used. The operator automatically calculates the prefix length based on kvSeqLen and qSeqLen for mask alignment. You do not need to pass any additional parameters.

       The ND shape is [2048, 2048], and the NZ shape is [1, 128, 2048, 16] (that is, [1, 2048/16, 2048, 16], where the NZ block size is 16). Atlas inference products require inputs in NZ format. You need to convert ND to NZ using npu_format_cast.

  • ALiBi:

    alibiMask can be understood as follows:

    

    alibiCoefficient is different for each head, and triangularMask represents an inverted triangle mask. For mask compression, the mask in the input tensor is the compressed alibiBias, and slopes is the alibiCoefficient. There are three compression scenarios for alibiBias.

    1. Right-aligned alibiBias

       maskType = MASK_TYPE_ALIBI_COMPRESS.

       The following figure shows the 512 x 512 alibiBias before compression.

       Figure 3 512 x 512 alibiBias before compression
       

       The following figure shows the 256 x 256 compressed alibiBias.

       Figure 4 256 x 256 compressed alibiBias
       
    2. Right-aligned alibiBias with square root

       maskType = MASK_TYPE_ALIBI_COMPRESS_SQRT.

       The bias in 1 is changed to , and the rest remains the same as in 1.

    3. Left-aligned alibiBias

       maskType = MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN.

       The following figure shows the alibiBias before compression.

       Figure 5 alibiBias before compression
       

       The following figure shows the compressed alibiBias.

       Figure 6 Compressed alibiBias
       

       Note: Left-aligned compressed masks are supported only by Atlas A2 training products / Atlas A2 inference products and Atlas A3 inference products / Atlas A3 training products .