倒三角稀疏attention优化

基本原理

LLM是基于only-decoder的模式，在进行self-attention的计算时会采用sequence-mask的方法，目的是为了使得decoder不能看见未来的信息。也就是对于一个序列中的第i个token，解码的时候只能够依靠i时刻之前(包括i)的的输出，而不能依赖于i时刻之后的输出。因此通常采取一个遮盖的方法(Mask)使得其在计算self-attention的时候只用i个时刻之前的token进行计算。self-attention的计算会将mask的部分也计算进去，针对这一场景，提出了倒三角attention加速算法，避免attention mask中上三角位置中无需计算的block参与计算，从而减少计算量。

Transformer-Decoder的Attention中相似矩阵与概率矩阵的右上三角部分为无效数据，结合矩阵乘法公式和Softmax计算公式发现，Q矩阵的第i行只需要与K和V的前i行做Attention计算，不需要与i+1及其后面的行做Attention计算。基于此特点，在矩阵的行方向（文本长度维度所在轴）将Q分成了无交叠的t份；将K和V分成有交叠的t份，每份的起始idx均为0，结束与对应Q分块的结束idx相同；将attn_mask分成无交叠的t份，每份的行idx与Q的idx一致，列idx与K的idx一致。然后再用对应的分块去做Attention计算，得到每一个输出分块，最后再将所有的输出分块拼接得到最终的Attention输出。

使用场景

倒三角特性在大模型中self-attention占比高的场景下会取得较大的性能收益，也适用于由于batch size大，sequence length过长导致的memory bound和计算bound的场景。而在一些由于计算和内存都不饱和的场景下，使用倒三角会增加小算子的数量，从而导致倒三角特性性能收益不明显。在已验证的场景中，字节GPT-150M：吞吐提升18%移动llama-70B：吞吐提升8%，电信智源Aquila-7B：吞吐提升6%。

操作步骤

使用AscendSpeed，配置--triangle attention标志，默认每次计算的block size为512，block size可根据sequence length的大小来进行调整，例如sequence length=4096，可使用--triangle block size 1024 标志将block size设置为1024。分块大小带来的影响：小分块无效数据和无效计算少，但调度开销大，批处理的效率低。大分块无效数据和无效计算多，但调度开销小，批处理的效率高。