1.摘要
本文主要介绍NSA（主要是推理）在昇腾上的一种亲和实现，主要创新点包括：
(1):设计一种巧妙的系数矩阵，完成compress算子importanceScore的高效计算，并基于此成功将compress中attention计算转换为online方式，取得较大的优化(从200us下降到80+us)；
(2):利用长序列推理场景下，单卡分到的batchsize或groupnum不大的特点，将无相互依赖的select和sliding算子减核并行融合到一个算子中去，整体消除一个算子的运行时间（约40us）。由于工作量问题部分优化没有完成，但是叠加其小优化措施，推理在论文典型场景下推理加速收益比达到8.8（论文中理论收益11.6）。
此外值得一提的是，基于上面两个优化思路，理论上在昇腾实现NSA推理加速比有可能超过原论文中理论加速比11.6，其基本原理可以理解为相对IFA，compress中基于矩阵的重要性分数计算耗时，不应该超过被省掉的select或sliding的一个算子耗时。
2.方案详情
2.1场景问题：NAS算子情况介绍
NAS包括3个算子：
Compress 将KV块按照固定不长划分后通过网络进行压缩，捕捉粗粒度的语义信息，减少计算负担。同时输出重要性系数，用于TopK选择。
Select 对压缩后的KV块计算重要性系数，排序后选择Top-k重要的KV块，由此引入了稀疏Attention
sliding 按固定的窗口大小截取尾部的KV cache，弥补前两路分支只处理block-level信息的缺陷。
其中Importance score计算公式如下：
2.2方案描述
2.2.1：Compress算子优化
根据importanceScore公式及论文典型场景配置l=32，d=16，l’= 64下，importance score的计算需要在softmax结果的列上做累加(如下图)，
但是列的累加在硬件上不亲和，如果按照列累加，需要流程转变变为：转置->累加->转置，如下图所示：
当前思路存在以下问题：
importanceScore计算部分采用vector计算，涉及转置和大量小数据量指令运行，导致在推理decode阶段，耗时过长，拖累整个算子性能。如图所示：
importanceScore计算基于QK矩阵乘及softmax计算完之后得出P，才能进行运算，导致原本Attention计算部分不能采用onlineSoftmax方式（CB并行下的cube bound），只能CV串行，存在算力空泡。
2.2.1.1：compress的优化阶段1：重要性计算转向矩阵乘PW
我们经过分析和研究，发现可以用cube实现ImportanceScore部分，思路分两步，如下：
“带权重的滑窗求和”
经过对论文公式分析和计算过程的推导（如下两个图），可以得知ImportanceScore的第一个求和部分实质上是对左矩阵的“带权重的滑窗求和”，类似一维方向的卷积池化。如图1多个红框表示，滑窗的步幅和权重都是固定的，基于此可设计如图3矩阵乘方式，左矩阵（黄色）为softmax结果P，右矩阵(灰蓝色)为滑窗系数矩阵W。 只要构造与P等大的矩阵W，理论上即可完成所有计算过程，但是W矩阵极稀疏（90+%都是0）、极重复（只有1和2两个值），需要参考图像处理中卷积池化，设计一个很小的“卷积核”。
但是系数矩阵W在矩阵乘的过程中，由于硬件基本块的原因，会被切块，每次切块后的W内容不一样，导致数据计算不全，如下所示
此外还产生了系数矩阵不固定的问题，工程上落地还有困难，如下图所示：
“固定常驻的小矩阵W”
系数矩阵W在矩阵乘的过程中，由于硬件基本块的原因，会被切块，每次切块后的W内容不一样（如上所示）。此外如果不能将每次计算时W的内容统一，每次计算更换W内容，需要从GM搬运几乎重复的内容，效率非常低效。 参考下图，我们设计了一种固定小矩阵作为系数矩阵：
关键流程如下：
设计尾缺列首的系数矩阵W 如图中红字1所示，在常规W的第一列加入一列，内容为最后一列缺失的部分（尾缺列首），图中只有一个1，该矩阵可由模型侧预置。
PW矩阵乘时，多乘一列 如图中红字2，PW矩阵乘后，首列（橙色）部分 在本次计算中没有意义，但却是下一次计算需要的
首尾相加 如图中红字3，在前后两次矩阵乘结束后，在GM上 利用atomicAdd特性，避开整32byte对齐的要求， 进行第一次计算的尾列和第二次首列的相加，完成因切块导致缺失计算。
上诉改造后，性能收益明显，性能收益来源分析如下：
固定的右矩阵：左矩阵P不用变形（相对于原方案中转置），不停搬入计算；W矩阵固定不变，常驻L0B，带宽全用于左矩阵P的搬运；
单cycle计算量大：根据矩阵乘原理，一个cycle可以计算多行。实际开发时，结合L0B的大小，W矩阵shape可设计为128*33，即一次计算33个滑窗，相对原方案中一个vector指令计算一个滑窗有较大提高。
Cube算力大：当前cube利用率20/（1616）=5/64，CV算力比：C = 16个V ，原先方案实际只用一个vector ，实际为32倍，那么起码提升32 (5/64)=2.5倍，结合vector计算时有转置、频繁计算小数据量，实际收益估计会更高。
实际测试后发现，以S=64K，selectedBlockSize=64, compressBlockSize=32, compressStride=16，qHead=64，kvHead=1，630版本为274us，当前测试的结果为139us，274us->139us提升几乎一倍，其中importanceScore部分，大致是150us->20us，提升7.5倍，流水如下所示
2.2.1.2：compress的优化阶段2：PV和PW的复用及融合
P复用的出发点：
原论文中没有要求高精度计算重要性分数：由图1计算公式及论文上下文可知，重要性分数是基于softmax结果P进行的，之后同样的P进行PV计算，所以没有必须要用F32计算PW。630版本中之所以用f32计算，是因为在vector上softmax结果P是f32，继续进行vector的重要性计算时，需要cast成f16，省操作。
重复的P搬运：PW中的P为f32，PV中的P为f16，内容重复，数据量较大，且需要走MTE2带宽，可以减少这部分搬运（性能收益来源）
注意力的稀疏性：Q只与有限的历史token强相关，topk不会很多，重要性分数实际中会在某个K后断崖式下降。
P融合的具体措施：
复用PV中的P_f16计算PW，在原PV矩阵乘计算中，每轮迭代PV计算后，保持P不动，更换右矩阵为W（此时W从阶段1L0常驻退化为L1常驻）。以增加高带宽MTE1搬运为代价，减少低带宽MTE2搬运。
精度影响评估 ：
降低PW的精度，一定会引起重要性分数精度下降、进而影响topK选择，但是对业务影响可接受。主要理由为：如果误差发生，则影响K的排布，实际上会导致某个K的顺序发生竞争：
如果竞争发生在topK内(如图3中3/4块)，后续算子对topK内所有块均权计算attention，对业务完全不影响；
如果竞争发生在topK外，业务不需要，也没有影响；
如果竞争发生在topk的尾块(如图3中15/16块) ，可以相信topK尾部最低重要性的某个块发生近似变化，对整体影响是可以接受的；
此外需要补充说明的是，PV与PW融合是要求左矩阵P保持不动，先后乘W和V两个右矩阵，当groupSize * blockSize>= 128*128时，PV已经全占了L0C，无l0c空间存储importance score的PW结果。 当融合计算所需要的资源超长芯片资源时，则不能融合，只能复用P的搬运，计算时分开计算PW和PV，此时会多一次P的MTE1搬运（性能相对会有所下降），相对PW的W又可以变成L0常驻。 比如当groupSize（=Q_head/KV_head）128， blockSize（page_block_szie）也为128,128*128* 4(f32_len)*2(pingpong)=128K，刚好等于L0C的大小，已经没有空间再存放PW的计算结果
2.2.1.3：compress的优化阶段3:使用Online Softmax
前面我们提过vector计算重要性计算会导致需要所有的QK先计算完才能计算重要性分数，这就导致原先attention的QK过程不能采用online-attention计算，那么我们现在将重要性计算转成矩阵乘PW，并且将PW和PV的P进行了融合，那么是否可以重新转向online-attention计算呢？我们先简单看下online-attention计算的相关情况：
Online Softmax基本原理：分块（参见颜色）计算，依据rowmax逐块更新：
Online Softmax流水：CV并行，cube bound掩盖下的vector空泡：
onlineSoftmax基本公式：
这里我们先提一个问题：相比PV，基于阶段1、2，PW也是矩阵计算，为什么不能使用online模式？有什么区别？
我们通过下面这个图来分析这个问题：
其中黄色框是阶段1提到的系数矩阵。图中表达的核心观点为：PV的onlinesoftmax流程是相应颜色块完成矩阵乘后，带rowmax系数重叠相加而成（图上红黄绿紫重叠相加为一个黑块），PW则为带系数首尾相加（图上红黄绿紫，收尾相加为一个黑块），区别只在于求和方式，其他无差异。
补充一点细节说明：
Ø图中矩阵W并不全真正全部的矩阵，也不需要全部计算，参考阶段1的设计，只要保证W中每次黄色块是阶段1的带冗余头、“尾缺列首”的固定系数小矩阵W即可，减少大量重复搬运和计算。
Ø继承阶段2融合复用P的逻辑，复用onlineSoftmax现有的KV划块逻辑，复用依据rowmax逐块更新的逻辑
Ø首尾相加参照阶段1的设计，还是需要在GM上利用atomicAdd特性（稍微搓了点），避开整32byte对齐的要求，进行前一次计算的尾列和后一次首列的相加，完成因切块导致缺失计算。
基于onlineSoftmax的ImportanceScore+QKV attention完整公式：
2.2.1.4：compress的优化阶段4:使用Flash-Decoding对KV切核
触发Flash decoding的业务场景，Aicore不满载场景，包括bs过小或者TP过大，单纯onlineSoftmax flash-attention会造成部分Aicore空闲，多batch下KV_seq_len长度不均，造成各核负载不均，所以当解决了flash-attention的问题后，也可以实现Flash-Decoding方案。
因为阶段34工作量较大，实际穿刺未完成，前文提到的8.8加速收益比不包括阶段34的收益
2.2.2：Select和sliding推理优化
2.2.2.1：思路出发点：
Select和sliding这两个算子没有依赖关系，是否可以“并行”、相互掩盖？在模型侧做“算子并行”依赖推理训练框架支持，难度过大，是否有其他并行方案。
从模型侧看，select是从全局选择的1k、16（*64）个重要性最高的文本段进行注意力计算，sliding是选末尾的0.5K的文本段进行注意力计算，二者业务上没有本质区别。
序列低于2K时，当前IFA算子模板不走入flash-decoding，说明短序列下单算子优化空间有限
4.论文中sliding固定选末尾的0.5K，select选16*64.
