压缩mask

功能说明

又名mask-free。

在长序列场景下，由于seqLen较大，需要的mask的大小也会变大。此时可开启此功能，传入压缩后的mask，以减小显存占用。

开启方式

参数“isTriuMask”置为1。在 Atlas 推理系列产品上使用MASK_TYPE_NORM_COMPRESS时，源码在算子构造函数中自动强制isTriuMask=1，用户无需手动设置。

在rope场景下：将变量“maskType”置为MASK_TYPE_NORM_COMPRESS。
在alibi场景下：依据具体需要，将变量“maskType”置为MASK_TYPE_ALIBI_COMPRESS或MASK_TYPE_ALIBI_COMPRESS_SQRT或MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN。

alibi压缩mask场景需传入slopes。

约束说明

alibi压缩mask只有“calcType”置为PA_ENCODER和PREFIX_ENCODER时生效。
“maskType”为MASK_TYPE_ALIBI_COMPRESS或MASK_TYPE_ALIBI_COMPRESS_SQRT时，mask的维度：
- 在 Atlas A2 训练系列产品 / Atlas A2 推理系列产品和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品上为[headNum, seqLen, 128]或[256, 256]。
- 在 Atlas 推理系列产品上为[headNum,128//16,maxSeqLen,16]或[1,256//16,256,16]。
“maskType”为MASK_TYPE_NORM_COMPRESS时，mask的维度：
- 在 Atlas A2 训练系列产品 / Atlas A2 推理系列产品和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品上为[128,128]。
- 在 Atlas 推理系列产品上为[1,128//16,128,16]或 [1,2048//16,2048,16]（即[1,128,2048,16]，2048长窗口）。
使用alibi压缩mask时，q、k、v的headSize要小于等于128。
MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN仅支持 Atlas A2 训练系列产品 / Atlas A2 推理系列产品和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品。此时mask的维度为[256, 256]。
Atlas A2 训练系列产品 / Atlas A2 推理系列产品和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品上，当数据类型为float16时，alibi压缩mask只有开启高精度才有效。
Atlas 推理系列产品上，使用MASK_TYPE_NORM_COMPRESS的2048长窗口（mask维度为[1,128,2048,16]或[2048,2048]）时，需满足以下约束：
- 仅 Atlas 推理系列产品支持，其他推理型号不支持。
- mask数据类型仅支持float16。
- 不支持量化（quantType必须为TYPE_QUANT_UNDEFINED）。
- scaleType必须为SCALE_TYPE_TOR（不支持LOGN）。
- isTriuMask由源码强制为1（构造函数层，当maskType为MASK_TYPE_NORM_COMPRESS且310P时），外部传入非必需（倒三角mask，即下三角及对角线为0/可见，上三角为负大值/不可见，推荐-60000或-10000）。
- maskStride为2048。
- calcType仅支持PA_ENCODER时，q_seqlen需等于kv_seqlen。
不同压缩mask的构造方法：
- 对于norm场景：
  对应maskType为MASK_TYPE_NORM_COMPRESS。
  
  mask为128 * 128的倒三角，其中：
  1. 当为float16场景时，mask如下图所示。
    图1 float16场景下的压缩norm mask
  2. 当为bf16/float16高精度场景时，mask如下图所示。
    图2 bf16/float16场景下的压缩norm mask
  3. 当 Atlas 推理系列产品为2048长窗口场景时：
    mask为2048 * 2048的倒三角，其中上三角（不含对角线）填充负大值（float16下等效于不可见，推荐-60000或-10000），下三角及对角线为0（可见）。该mask为一张共享的2048×2048画布，所有batch/head共用同一块mask。不论query/kv序列长度如何，均使用同一块2048×2048 mask，算子根据kvSeqLen与qSeqLen自动计算前缀长度用于mask对齐，用户无需传入额外参数。
    
    ND形状为[2048, 2048]，NZ形状为[1, 128, 2048, 16]（即[1, 2048/16, 2048, 16]，NZ分块大小为16）。 Atlas 推理系列产品 上需使用NZ格式输入，ND格式需在用户侧通过npu_format_cast转换为NZ。
- 对于alibi场景：
  alibi mask可以拆解为
  
  $\text{[math]}$
  
  其中，alibi coefficient为每个head各不相同的系数，triangularMask代表倒三角mask。压缩mask的场景下，输入tensor中的mask即为压缩后的alibi bias，slopes即为alibi coefficient。针对alibi bias的压缩有如下三种情况：
  1. 右对齐alibi bias
    对应“maskType”为MASK_TYPE_ALIBI_COMPRESS。
    
    如下图所示，为512 * 512的压缩前的alibi bias。
    
    图3 512 * 512压缩前的alibi bias
    
    对应的压缩后的256 * 256的压缩后的alibi bias如下图所示。
    
    图4 256 * 256压缩后的alibi bias
  2. 开平方的右对齐alibi bias
    对应“maskType”为MASK_TYPE_ALIBI_COMPRESS_SQRT。
    
    在右对齐alibi bias的基础上对于bias本身做了开根号处理，1的bias变为..., 其余与1一致。
  3. 左对齐的alibi bias
    对应“maskType”为MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN。
    
    压缩前的alibi bias如下图所示。
    
    图5 压缩前的alibi bias
    
    压缩后的alibi bias如下图所示。
    
    图6 压缩后的alibi bias
    
    注意：左对齐的压缩mask只支持 Atlas A2 训练系列产品 / Atlas A2 推理系列产品和 Atlas A3 推理系列产品 / Atlas A3 训练系列产品。

父主题： mask