功能列表

各功能共存情况

Y表示两功能不冲突，N表示两功能冲突。

	控制可计算batch	全量/增量分离	高精度	clamp缩放	压缩mask	kv-bypass	logN缩放	qkv全量化	BNSD维度输入	kv tensorlist格式输入	MLA合并输入kvcache功能	SWA
控制可计算batch	-	Y	Y	Y	Y	Y	Y	N	Y	Y	N	N
全量/增量分离	-	-	Y	Y	Y	Y	Y	N	Y	Y	N	Y
高精度	-	-	-	Y	Y	Y	Y	Y	Y	Y	Y	N
clamp缩放	-	-	-	-	Y	Y	Y	Y	Y	Y	Y	N
压缩mask	-	-	-	-	-	Y	Y	Y	Y	Y	N	Y
kv-bypass	-	-	-	-	-	-	Y	N	Y	Y	N	Y
logN缩放	-	-	-	-	-	-	-	N	N	N	N	N
qkv全量化	-	-	-	-	-	-	-	-	N	N	Y	N
BNSD维度输入	-	-	-	-	-	-	-	-	-	N	N	N
kv tensorlist格式输入	-	-	-	-	-	-	-	-	-	-	N	N
MLA合并输入kvcache功能	-	-	-	-	-	-	-	-	-	-	-	N
SWA	-	-	-	-	-	-	-	-	-	-	-	-

控制可计算batch

功能
指定某几个batch参与attention计算。
开启方式
参数“batchRunStatusEnable”置为true，并传入batchRunStatus作为输入tensor。

batchRunStatus为0，1组成的tensor。0代表该位置的batch不参与计算，1代表参与计算。
特殊约束
- 不支持 Atlas 推理系列产品。
- 开启此功能时输入的cacheK, cacheV的维度为[batch, maxSeqLen, hiddenSize]。
- 不支持PA_ENCODER。

全量/增量分离

功能
当模型接入paged attention时，全量阶段选用SelfAttention且“calcType”置为PA_ENCODER，增量阶段选用PagedAttention。而当模型使用传统flash attention时，增量与全量阶段都使用SelfAttention；此时为提高计算效率，可在全量与增量阶段选用不同的calcType。
开启方式
全量阶段：参数“calcType”置为ENCODER。

增量阶段：参数“calcType”置为DECODER。
特殊约束
无。

高精度

功能
在进行attention计算时，Q∙KT的结果有可能溢出float16，导致算子输出tensor中间出现NAN值；此时可开启此功能，算子内部使用float32承载中间结果。
开启方式
参数“kernelType”置为KERNELTYPE_HIGH_PRECISION。
特殊约束
开启此功能时，传入的mask中需把-inf换成1。

高精度功能只在 Atlas 800I A2 推理产品 / Atlas A2 训练系列产品上才能生效。

clamp缩放

功能
相当于对传入的q做torch.clamp。
开启方式
参数“clampType”置为CLAMP_TYPE_MIN_MAX。

传入参数“clampMin”与参数“clampMax”以给出clamp的上下界。
特殊约束
不支持 Atlas 推理系列产品。

压缩mask

功能
又名mask-free。

在长序列场景下，由于seqLen较大，需要的mask的大小也会变大。此时可开启此功能，传入压缩后的mask，以减小显存占用。
开启方式
参数“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时生效。
- “maskType”为MASK_TYPE_ALIBI_COMPRESS或MASK_TYPE_ALIBI_COMPRESS_SQRT时，mask的维度：
  - 在 Atlas 800I A2 推理产品 / Atlas A2 训练系列产品上为[head_num, seqlen, 128]或[256, 256]。
  - 在 Atlas 推理系列产品上为[head_num,128//16,maxSeqlen,16]或[1,256//16,256,16]。
- “maskType”为MASK_TYPE_NORM_COMPRESS时，mask的维度：
  - 在 Atlas 800I A2 推理产品 / Atlas A2 训练系列产品上为[128,128]。
  - 在 Atlas 推理系列产品上为[1,128//16,128,16]。
MASK_TYPE_ALIBI_COMPRESS_LEFT_ALIGN仅支持 Atlas 800I A2 推理产品 / Atlas A2 训练系列产品，此时mask的维度为[256, 256]。
- Atlas 800I A2 推理产品 / Atlas A2 训练系列产品上，当数据类型为float16时，alibi压缩mask只有开启高精度才有效。
不同压缩mask的构造方法：
- 对于norm场景：
  对应maskType为MASK_TYPE_NORM_COMPRESS。
  
  mask为128 * 128的倒三角，其中：
  1. 当为fp16场景时，mask如下图所示。
    图1 fp16场景下的压缩norm mask
  2. 当为bf16/fp16高精度场景时，mask如下图所示。
    图2 bf16/fp16场景下的压缩norm mask
- 对于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 800I A2 推理产品 / Atlas A2 训练系列产品。

kv-bypass

功能
在某些场景下，客户想自行管理kvcache的过程，不想使用加速库做把kv刷新到kvcache的动作。
开启方式
参数“kvCacheCfg”置为K_BYPASS_V_BYPASS。同时输入tensor不传k, v。
特殊约束
此功能只有“calcType”为非PA_ENCODER时生效。

logN缩放

功能
在一般的情况下，上文计算公式中的Zoom函数，其缩放系数由参数qkScale给出。当用户想要使用logN形式的缩放函数时，可开启此功能。
开启方式
参数“scaleType”置为SCALE_TYPE_LOGN。

传入logN作为输入tensor。

“kernelType”置为KERNELTYPE_HIGH_PRECISION。
特殊约束
- 开启logN功能，“scaleType”需为SCALE_TYPE_LOGN，“calcType”需为DECODER或PA_ENCODER，分别对应增量阶段和全量阶段； Atlas 800I A2 推理产品 / Atlas A2 训练系列产品上“calcType”为PA_ENCODER时额外需要“kernelType”为KERNELTYPE_HIGH_PRECISION。
- logN功能与量化场景不支持同时开启。
- 当“inputLayout”为TYPE_BNSD时，“ScaleType”必须为SCALE_TYPE_TOR（不支持LogN缩放）。

qkv全量化

功能
支持量化好的q, k, v传入，降低显存占用。

开启方式

参数“quantType”置为TYPE_QUANT_QKV_OFFLINE或TYPE_QUANT_QKV_ONLINE, 分别为离线量化与在线量化。

在线全量化需要添加四个输入tensor，分别是qkDescale、qkOffset、vpvDescale、vpvOffset，当采用离线量化时，需要在在线全量化新增的输入tensor的基础上额外传入pScale作为输入tensor。

若干输入tensor需要传入指定的维度，如下表所示：

参数	维度	数据类型	格式	npu or cpu	描述
query	[nTokens, head_num, head_size]	int8	ND	npu	query矩阵。
key	[nTokens, head_num, head_size]	int8	ND	npu	key矩阵。
value	[nTokens, head_num, head_size]	int8	ND	npu	value矩阵。
mask	同mask类型, 开启mask压缩功能时与FA有所不同，见压缩mask	float16/bf16	ND/NZ	npu	四种shape分别对应： 1.所有batch相同，方阵。 2. batch不同时的方阵。 3. batch不同时的向量。 4. alibi场景。当maskType为undefined时不传此tensor。
seqLen	[batch]	int32	ND	cpu	等于1时，为增量或全量；大于1时，为全量。
slopes	[head_num]	Atlas 800I A2 推理产品 / Atlas A2 训练系列产品： float16	ND	npu	当maskType为alibi压缩，且mask为[256,256]时需传入此tensor，为alibi mask每个head的系数。
qkDescale	[head_num]	float	ND	npu	为Q*K^T的反量化scale参数。
qkOffset	[head_num]	int32	ND	npu	作为Q*K^T的反量化offset参数。预留tensor，需传任意非空tensor，实际暂未使用。
vpvDescale	[head_num]	float	ND	npu	quantType=2时，为P*V的反量化scale参数；quantType=3时，为V的反量化scale参数。
vpvOffset	[head_num]	int32	ND	npu	仅全量化场景传此tensor（即quantType=2或3）。quantType=2时，为P*V的反量化offset参数；quantType=3时，为V的反量化offset参数。预留tensor，需传任意非空tensor，实际暂未使用。
pScale	[head_num]	float	ND	npu	P的离线量化scale参数，当开启离线全量化时需要传此tensor（即quantType=2）。当开启在线全量化时不传此tensor（即quantType=3）。
output	[nTokens, head_num, head_size]	float16/bf16	ND	npu	输出。

其中pScale仅在离线量化场景传入，在线量化场景则不传此输入tensor。

特殊约束
使用全量化时需要指定输出tensor的数据类型，具体为使用参数outDataType，该参数只能是ACL_FLOAT16或ACL_BF16。

“calctype”须为PA_ENCODER。

不支持 Atlas 推理系列产品。

BNSD维度输入

功能
一般的，传入SelfAttention算子的q，k，v的维度为[batch, seqLen, head_num, head_dim]，即[b, s, n, d]，或者是它合轴后的变种。在某些场景下，传入[b, n, s, d]性能更好。

开启方式

参数“inputLayout”置为TYPE_BNSD。

输入tensor的维度如下：

当在 Atlas 800I A2 推理产品 / Atlas A2 训练系列产品上运行时：

输入tensor	维度
query	[batch, head_num, seq_len, head_size]
cacheK	[layer, batch, head_num, seq_len, head_size]
cacheV	[layer, batch, head_num, seq_len, head_size]

当在 Atlas 推理系列产品上运行时：

输入tensor	维度
query	[batch, head_num, seq_len, head_size]
cacheK	[layer, batch*head_num, head_size / 16, kv_max_seq, 16]
cacheV	[layer, batch*head_num, head_size / 16, kv_max_seq, 16]

特殊约束
- 此功能只有开启kv-bypass功能，即参数“kvcacheCfg”置为K_BYPASS_V_BYPASS时才可用。
- 使用BNSD维度输入时，“maskType”不能为MASK_TYPE_UNDEFINED。

kv tensorlist格式输入

功能
支持以tensorlist，即二级指针结构的形式而非整个tensor作为cacheK与cacheV的输入。

开启方式

如何构造该场景中的variantPack:

atb Tensor定义：

struct Tensor {
    //! \brief Tensor描述信息
    TensorDesc desc;
    //! \brief TensorNPU内存地址。
    void *deviceData = nullptr;
    //! \brief TensorCPU内存地址。
    void *hostData = nullptr;
    //! \brief “deviceData”或“hostData”指向内容的内存大小。
    uint64_t dataSize = 0;
};

variantPack中的输入tensors，在需要使能tensor list时，将对应tensor的deviceData置为nullptr，将hostData指向构造的tensorList即可。

特殊约束
此功能只在开启kv-bypass功能（即参数“kvcacheCfg”置为K_BYPASS_V_BYPASS），且“calcType”为DECODER时生效。

MLA合并输入kvcache功能

功能
过去attention算子会将k和v存储在显存中，随着模型越来越大，输入序列越来越长，kvcache也越来越大，对显存的容量造成很大压力，造成性能上的瓶颈。因此Multi-Head Latent Attention提出将kvcache压缩成一个一个较小的向量，代替原来的kvcache输入进fa算子进行计算。
开启方式
- calcType= PA_ENCODER
- mlaVHeadSize > 0
特殊约束
- MLA合并输入kvcache功能不支持 Atlas 推理系列产品，不支持alibi mask，多头自适应压缩（rope 场景），多头自适应压缩（alibi 场景），logN缩放，BNSD输入排布。
- 开启MLA合并输入kvcache功能后key和value的head_size范围为（0, 576]。
- mlaVHeadSize不能大于query和key的head_size。
- 开启MLA合并输入kvcache功能后query，key，value可传二维[nTokens, hiddenSize]或三维[nTokens, head_num, head_size]或四维[batch, seq_len, head_num, head_size]，且query，key的维度数可以不同。
- MLA合并kvcache功能支持全量化场景，一起开启时query和key的head_size范围为（0, 576]，mlaVHeadSize不能大于query和key的head_size。

Sliding window attention（SWA）功能

功能
Mistral 7B，Attention部分在GQA（Grouped-query attention）的基础上，叠加了SWA（Sliding window attention）的优化，可以进一步提高推理速度，并降低显存。

当前FA算子在进行推理时采用的是取所有KV进行计算，可以应用SWA提升推理速度，越远距离的信息，对当前位置的重要性越低，对于距离超过窗口大小的两个token不参与注意力分数计算。

通过设计mask实现n>m或m-n>window_len（n为kv_id, m为q_id）时，注意力分数跳过计算。如图7，若window_size = 3, “The”将不会计算与“the”之间的注意力分数。

图7 sliding window attention的mask功能示意

Decoder支持仅长度为window_len的最新KV历史参与计算。

KV Cache优化：比如窗口大小=4，则当第5个token需要缓存，直接替换掉第1个token，这样就可以保持kv缓存有一个最大值（为窗口大小），而不会无限增长。
开启方式
- windowSize > 0
- maskType必须为MASK_TYPE_SLIDING_WINDOW_NORM或MASK_TYPE_SLIDING_WINDOW_COMPRESS
特殊约束
- 开启特性后“cacheType”可以为CACHE_TYPE_NORM 或 CACHE_TYPE_SWA, 不开启特性“cacheType”只能为 CACHE_TYPE_NORM。
- Sliding Window Attention特性不支持动态batch，高精度，clamp缩放，qkv全量化，mla，logN缩放特性，BNSD数据排布。
- Sliding Window Attention特性在calcType=DECODER场景下，“maskType”不能为MASK_TYPE_SLIDING_WINDOW_COMPRESS，且不传入attentionMask。
- 开启特性必须满足以下两个条件。如果只满足一个条件，会校验报错。
  - windowSize > 0。
  - maskType必须为MASK_TYPE_SLIDING_WINDOW_NORM或MASK_TYPE_SLIDING_WINDOW_COMPRESS。

SWA mask样例生成参考

当windowSize >= seqlen时，mask和不开启SWA时一样，为上三角，否则mask生成可参考如下：

maskType=MASK_TYPE_SLIDING_WINDOW_NORM, dtype为float16。

以 max(seqLen)为5，windowSize为3为例。

图8 maskType=MASK_TYPE_SLIDING_WINDOW_NORM, dtype为float16时的mask样例

python代码示例：

          
               swa_mask = np.ones(shape=[maxseq,maxseq]) * -65536.0
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -window_size)
swa_mask = triu_mask + tril_mask

C++代码示例：

          
               static constexpr uint32_t  SWA_MASK_SIZE = maxSeq;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_MASK_SIZE * SWA_MASK_SIZE, -65536.0);

    for (uint32_t i = 0; i < SWA_MASK_SIZE; ++i) {
        uint32_t offset = i >= windowSize ? (i - windowSize + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

maskType=MASK_TYPE_SLIDING_WINDOW_NORM, dtype为bf16。

以 max(seqLen)为5，windowSize为3为例。

图9 maskType=MASK_TYPE_SLIDING_WINDOW_NORM, dtype为bf16时的mask样例

python代码示例：

          
               swa_mask = np.ones(shape=[maxseq,maxseq]) * 1.0
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -window_size)
swa_mask = triu_mask + tril_mask

C++代码示例：

          
               static constexpr uint32_t  SWA_MASK_SIZE = maxSeq;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_MASK_SIZE * SWA_MASK_SIZE, 1.0);

    for (uint32_t i = 0; i < SWA_MASK_SIZE; ++i) {
        uint32_t offset = i >= windowSize ? (i - windowSize + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

maskType=MASK_TYPE_SLIDING_WINDOW_COMPRESS, dtype为float16。

python代码示例：

          
               swa_mask = np.ones(shape=(1, 512, 512)) * -65536.0
pp_n = 128 if head_size <= 128 else 64 # head_size为每个注意力头的嵌入向量的大小
if window_size <= pp_n * 3:
    true_size = window_size
else:
    if window_size % pp_n == 0:
        true_size = pp_n * 3
    else:
        true_size = pp_n * 2 + window_size % pp_n
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -true_size)
swa_mask = triu_mask + tril_mask

C++代码示例：

          
               static constexpr uint32_t  SWA_COMPRESS_MASK_SIZE = 512;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_COMPRESS_MASK_SIZE * SWA_COMPRESS_MASK_SIZE, -65536.0);
    uint32_t blockSize = embeddim > 128 ? (16384 / embeddim / 16 * 16) : 128;
    uint32_t compressWindow = windowSize > 3 * blockSize ? (2 * blockSize + windowSize % blockSize) : windowSize;
    for (uint32_t i = 0; i < SWA_COMPRESS_MASK_SIZE; ++i) {
        uint32_t offset = i >= compressWindow ? (i - compressWindow + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

maskType=MASK_TYPE_SLIDING_WINDOW_COMPRESS, dtype为bf16。

python代码示例：

          
               swa_mask = np.ones(shape=(1, 512, 512)) * 1
pp_n = 128 if head_size <= 128 else 64 # head_size为每个注意力头的嵌入向量的大小
if window_size <= pp_n * 3:
    true_size = window_size
else:
    if window_size % pp_n == 0:
        true_size = pp_n * 3
    else:
        true_size = pp_n * 2 + window_size % pp_n
triu_mask = np.triu(swa_mask, 1)
tril_mask = np.tril(swa_mask, -true_size)
swa_mask = triu_mask + tril_mask

C++代码示例：

          
               static constexpr uint32_t  SWA_COMPRESS_MASK_SIZE = 512;
std::vector<float> create_attention_mask(uint32_t windowSize, uint32_t embeddim) {
    std::vector<float> attention_mask(SWA_COMPRESS_MASK_SIZE * SWA_COMPRESS_MASK_SIZE, 1.0);
    uint32_t blockSize = embeddim > 128 ? (16384 / embeddim / 16 * 16) : 128;
    uint32_t compressWindow = windowSize > 3 * blockSize ? (2 * blockSize + windowSize % blockSize) : windowSize;
    for (uint32_t i = 0; i < SWA_COMPRESS_MASK_SIZE; ++i) {
        uint32_t offset = i >= compressWindow ? (i - compressWindow + 1) : 0;
        for (uint32_t j = offset; j < i + 1; ++j) {
            attention_mask[i * SWA_COMPRESS_MASK_SIZE + j] = 0.0;
        }
    }
    return attention_mask;
}

父主题： SelfAttentionOperation