本文主要介绍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的一个算子耗时。

其中Importance score计算公式如下：

根据importanceScore公式及论文典型场景配置l=32，d=16，l’= 64下，importance score的计算需要在softmax结果的列上做累加(如下图)，

但是列的累加在硬件上不亲和，如果按照列累加，需要流程转变变为：转置->累加->转置，如下图所示：            

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内容，需要从GM搬运几乎重复的内容，效率非常低效。 参考下图，我们设计了一种固定小矩阵作为系数矩阵：

 上诉改造后，性能收益明显，性能收益来源分析如下：

  实际测试后发现，以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的复用及融合

          原论文中没有要求高精度计算重要性分数：由图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后断崖式下降。

复用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的计算结果

     前面我们提过vector计算重要性计算会导致需要所有的QK先计算完才能计算重要性分数，这就导致原先attention的QK过程不能采用online-attention计算，那么我们现在将重要性计算转成矩阵乘PW，并且将PW和PV的P进行了融合，那么是否可以重新转向online-attention计算呢？我们先简单看下online-attention计算的相关情况：

Online Softmax基本原理：分块（参见颜色）计算，依据rowmax逐块更新：            

Online Softmax流水：CV并行，cube bound掩盖下的vector空泡：

这里我们先提一个问题：相比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完整公式：

触发Flash decoding的业务场景，Aicore不满载场景，包括bs过小或者TP过大，单纯onlineSoftmax flash-attention会造成部分Aicore空闲，多batch下KV_seq_len长度不均，造成各核负载不均，所以当解决了flash-attention的问题后，也可以实现Flash-Decoding方案。    

因为阶段34工作量较大，实际穿刺未完成，前文提到的8.8加速收益比不包括阶段34的收益

  1. Select和sliding这两个算子没有依赖关系，是否可以“并行”、相互掩盖？在模型侧做“算子并行”依赖推理训练框架支持，难度过大，是否有其他并行方案。

   2. 从模型侧看，select是从全局选择的1k、16（*64）个重要性最高的文本段进行注意力计算，sliding是选末尾的0.5K的文本段进行注意力计算，二者业务上没有本质区别。

   3. 序列低于2K时，当前IFA算子模板不走入flash-decoding，说明短序列下单算子优化空间有限

   4.论文中sliding固定选末尾的0.5K，select选16*64.

   基于上诉分析和猜想，一个可能的优化方式为：将select和sliding在kernel内并行、各自减核运行，并寻找降低单个算子耗时的优化点。               

2.2.2.3：关键点-AiCore空闲分析：

      如图1分析，64K测试场景下单卡实际使用Aicore数量不足一半，只有10/24，而且每个Aicore只处理一个head，连pingpong都做不到。                     

   如图2实验可知， head_num=10时，减少单卡内可使用Aicore数量时， Aicore数量在5~10之间（

），select算子性能差异不大，低于5之后耗时明显上升（

每个Aicore处理3+个head，处理轮次上升 

） ，说明单Aicore处理一个或两个head性能最佳；

  如图3可知，head_num=20时，减少单卡内可使用Aicore数量至10个时，10个Aicore(

每个Aicore处理两个head，有doublebuffer

)性能约好（33.6us VS 41.7us， Aicore处理两个head比一个好），原因为：

20个Aicore同时copyin所有数据，触发了带宽争抢，带pingpong的10个Aicore某时刻搬运数据量比不带pingpong的20个Aicore要少一半。    

当单卡分到的head_num较少，可以减核、使单核处理两个head（如图4），既避免带宽争抢、也可以空出核做其他计算。

 只要单卡分配到的总head数低于20或24（参考当前TP=4，业务场景真实存在），都可以强制减核。

NSA推理在昇腾芯片上的一种亲和实现方案

人工智能

硬件

算法

本文介绍了NSA在昇腾平台的优化实现，通过矩阵计算优化compress算子，将耗时从200us降至80+us，并实现select和sliding算子并行融合，消除40us耗时。优化后推理加速比达8.8倍，理论加速比可超11.6倍。

图像处理

2026采购季 | AI焕新·智启新局

imageprocess

文章

问答

视频

教程

学习中心

腾讯云实验室

直播

竞赛

腾讯云代码分析专区

腾讯iOA零信任安全管理系统专区

腾讯云架构师技术同盟交流圈

腾讯云数据库专区

腾讯云智能顾问专区

腾讯云原生专区

腾讯混元专区

腾讯云TCE专区

腾讯云Lighthouse专区

腾讯云HAI专区

腾讯云Edgeone专区

腾讯云存储专区

腾讯云智能专区

腾讯轻联专区 

腾讯云开发专区

TAPD专区

腾讯轻量云游戏服专区

腾讯云最具价值专家

腾讯云架构师技术同盟

腾讯云创作之星

腾讯云开发者先锋

腾讯云代码助手

云原生构建

TAPD 敏捷项目管理

Cloud Studio

SDK中心

API中心

命令行工具

涵盖代码开发、场景应用、自动测试全流程，助你从零构建专属AI助手

一站式MCP教程库，解锁AI应用新玩法

聚焦“写作效率、视觉美观与运行性能”三方面进行全面升级，为您提供更高效、稳定的创作环境

社区富文本&Markdown编辑器全新改版上线，欢迎大家体验!

诚挚邀请您参与本次调研，分享您的真实使用感受与建议。您的反馈至关重要，感谢您的支持与参与！

NSA推理在昇腾芯片上的一种亲和实现方案

NSA推理在昇腾芯片上的一种亲和实现方案

社区

活动

圈层

关于

腾讯云开发者

热门产品

热门推荐

更多推荐