CANN / cann-recipes-infer: NPU DeepSeek-V3.2-Exp Ascend C 融合算子优化

NPU DeepSeek-V3.2-Exp Ascend C 融合算子优化【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer面向 DeepSeek-V3.2-Exp 架构本次发布两个全新的融合算子LightningIndexer(LI) 以及SparseFlashAttention(SFA)。其中LightningIndexer算子是本次 DeepSeek 新架构中最特别的部分之一用于对重要的上下文进行稀疏选择当前 Top-$k$ 计算已经融入LightningIndexer算子相较于 Top-$k$ 外置收益大约20%SparseFlashAttention则是针对稀疏 Attention 场景进行了特殊优化的 Attention 算子。A3 中的 Cube 核与 Vector 核数比为 $1:2$每1个 Cube 核与2个 Vector 核构成一个组核 (一个 AICore)其中一个组核的结构如下图所示其中 L1、L0A、L0B、L0C 为 Cube 核的片上 BufferUB 为 Vector 核的片上 Buffer。具体介绍可参考基本架构。在本篇技术报告中我们将对两个融合算子在 A3 上的具体优化进行细致地阐述包括具体的算子 Tiling 设计、Cube 核与 Vector 核间的流水优化、计算流程的实现优化、以及下一步即将开展的优化等内容同时将介绍当前融合算子在提升易用性上的进展。Highlights细致的算子 Tiling 及 Pipeline 设计发挥融合算子优势充分利用硬件特性基于 Ascend C API 高效实现聚合访存减少指令下发提高离散访存性能算子开发及调测的易用性提升LightningIndexerQuant量化支持Mat-Duce 算法充分利用 Cube 核OutlineNPU DeepSeek-V3.2-Exp Ascend C 融合算子优化HighlightsOutlineLightningIndexerLI 算子 Tiling 设计LI Pipeline 设计ReLU 随路实现Top-k 指令实现LI 量化支持Mat-Duce 算法SparseFlashAttentionSFA 算子 Tiling 设计SFA Pipeline 设计访存优化SFA Next Feature易用性提升TilingKey 模板化分核计算下沉LightningIndexerLightningIndexer基于一系列操作得到每一个 token 对应的 Top-$k$ 个位置。对于某个 token 对应的 Index Query $Q_{index}\in\R^{g\times d}$给定上下文 Index Key $K_{index}\in\R^{S_{k}\times d},W\in\R^{g\times 1}$其中 $g$ 为 GQA 对应的 group size$d$ 为每一个头的维度$S_{k}$ 是上下文的长度LightningIndexer的具体计算公式如下 $$ \text{Top-}k\left{[1]{1\times g}\left[(W[1]{1\times S_{k}})\odot\text{ReLU}\left(Q_{index}K_{index}^T\right)\right]\right} $$可拆分为如下计算流程计算矩阵乘法$S Q_{index}K_{index}^T$计算激活函数$S\text{ReLU}(S)$计算广播乘法$S_W(W[1]{1\times S{k}})\odot S$沿G轴进行Reduce操作$Score[1]_{1\times g} S_W$对 $Score$ 进行 $\text{Top-}k$ 计算即获取数值排序前 $k$ 个的结果并返回其对应的 Index在 Prefill 或开启 MTP 的 Decode 场景多个 token 对应的 $Q_{index}$ 会被合并统一计算本文后续部分不对单/多个 token 对应的 $Q_{index}$ 加以区分。实际开发中需要充分利用昇腾硬件的特征以实现更好的性能以下将详细介绍当前的解决方案及具体实现方式。LI 算子 Tiling 设计算子的 Tiling 设计决定了储存空间是否被充分利用进而影响整体的算力利用率以及算子性能。针对LightningIndexer本章分别阐述核间和核内的 Tiling 设计。由于片上 Buffer 空间有限任何计算本质上都需要进行切块因此需要确定一次迭代计算的基本块大小即核间的 Tiling 设计。当前LightningIndexer算子将基本块大小设置为$(512,512)$即一次迭代计算完成512行以减少右矩阵的重复搬运 (以$g64$为例一次会合并8个token统一计算)同时会完成 $S_{k}$ 方向上粒度为512的计算这个值的选取可以保证流水线头尾开销可控以及 Cube 核与 Vector 核之间同步开销可被计算掩盖。核内 Tiling 设计采用了如下方案L1 空间划分为如下部分$Q_{index}$ 矩阵常驻$512\times 128\times2\text{ Bytes}128\text{ KB}$L1层级基本块为 $(512,128)$在实际计算时可做更细致的搬运即一次只搬运$(256,128)$大小的数据到L1以降低搬运头开销$K_{index}$ 矩阵 3-Buffer 循环复用$3\times256\times 128\times2\text{ Bytes}192\text{ KB}$L1层级基本块为 $(256,128)$。L0AL0BL0C 使能 Double Buffer分别划分为$2\times32\text{ KB},2\times32\text{ KB},2\times64\text{ KB}$一次搬入 L0A 和 L0B 的矩阵块大小都为$(128,128)$这是一种对昇腾硬件较为亲和的设置可以提高算力利用率。这样的方式可以保证在计算过程中 $Q_{index}$ 只需要搬运一次并且降低 $K_{index}$ 矩阵的重复搬运次数假设 $Q_{index}$ 矩阵 L1 层级基本块大小为 $(128,128)$而行方向总的 token 数为 $S_q$那么 $K_{index}$ 矩阵需要重复搬运 $\lceil S_q/128\rceil$ 次选取基本块大小为 $(512,128)$ 可将重复搬运次数降低到 $\lceil S_q/512\rceil$。由于第一个计算过程矩阵乘的启动需要右矩阵的搬运为右矩阵分配多块 Buffer 可消除计算流程因为 Buffer 占用造成的阻塞所以为 $K_{index}$ 分配 3-Buffer 可以提升带宽爬坡期的搬运效率。LI Pipeline 设计LightningIndexer的计算流程只有两个阶段即矩阵乘和 ReLU 操作 ($C_1$) 以及后续的多个向量计算 ($V_1$)。为了防止多个阶段的串行执行需要提前下发 (Preload) 一次 $C_1$ 以实现流水的互相掩盖。ReLU 随路实现LightningIndexer首先需要进行矩阵乘法并对结果进行ReLU计算。由于矩阵乘法和 ReLU 计算分别发生在不同的计算单元朴素的实现需要将矩阵乘法的结果从 Cube 核搬运至 Vector 核进行。而昇腾NPU上 Cube 核的 FixP 组件提供随路 ReLU 能力即可以在矩阵乘法完成之后直接在数据搬运的过程中随路完成 ReLU 操作不需要显式进行 ReLU 计算。除 ReLU 之外FixP 组件还有更多强大功能具体可参考基础API: FixP相关。Top-k 指令实现LightningIndexer融合算子的核心是在长达数十万的序列中为每个 token 高效地筛选出分数最高的 $k$例如2048个索引。同时对于算子而言Top-$k$ 的计算必须是准确无误的不能采用近似算法求解。当前的实现方案基于昇腾支持的排序指令进行全量排序Top-$k$ 计算方案过程分为三步分组排序通过VBS32指令将每32个 token 按照其 $Score$ 进行稳定降序排列输出其排序向量以及对应索引向量直到将整个序列的 token 分组排序完毕。归并通过VMS4v2指令将至多4组可为2/3组长度为32、128、512的已排序向量进行归并直到合并后的向量长度达到2048。合并排列仍然按照其对应 $Score$ 进行稳定降序排列最终输出的向量长度为4组输入向量长度之和同时输出其对应索引向量。规约采用VMS4v2指令的耗尽模式 (当某个输入向量的可用数据已全部读取完毕排序指令立即暂停执行)将至多4组长度为2048已排序的向量进行归并。取出前 $k2048$ 个参数组成的向量与另外未参与合并排序的有序向量重复进行 Top-$k$ 计算直到将所有长度为2048的有序向量比较完毕得到最终所有向量中得分最高的2048个对应的 Index。假设上下文长度为 $S_k$ 分组排序与归并步骤的计算量为 $4S_k$。规约步骤的计算量无法精确评估采用正态分布结合蒙特卡洛法模拟结果显示第一次的计算量接近 $4k$ 从第二次开始由于已经有了 $k$ 个较大数的向量参与运算计算量开始大幅减少直到逼近耗尽极限 $k$ 。该部分计算量经过模拟分析后可表示为 $0.36S_k3.32k$ 。因此Top-$k$ 流程总计算量 $totalCount 4.36S_k 3.32k$。LI 量化支持当前LightningIndexer算子已支持 Per-Token-Head 量化 输入即输入的 $Q_{index}$ 和 $K_{index}$ 都被量化为 INT8 数据类型以此应对长序列场景下显存占用过大的问题。支持量化输入的LightningIndexerQuant对于某个 token 对应的 Index Query $Q_{index}^{INT8}\in\R^{g\times d}$ 及其反量化系数$Scale_{Q}\in\R^{g\times1}$给定上下文 Index Key $K_{index}^{INT8}\in\R^{S_{k}\times d}$ 及其反量化系数$Scale_{K}\in\R^{S_{k}\times1}$, $W\in\R^{g\times 1}$其中 $g$ 为 GQA 对应的 group size$d$ 为每一个头的维度$S_{k}$ 是上下文的长度LightningIndexerQuant的具体计算公式如下 $$ \text{Top-}k\left{[1]{1\times g}\left[(W[1]{1\times S_{k}})\odot\text{ReLU}\left(\left(Scale_QScale_K^T\right)\odot\left(Q_{index}^{INT8}{(K_{index}^{INT8})}^T\right)\right)\right]\right} $$Mat-Duce 算法LightningIndexerQuant算子采用了 Mat-Duce 算法利用数学变换调整计算逻辑用矩阵乘法实现规约计算。该算法充分利用 A3 芯片 Cube 核的随路能力将LightningIndexerQuant算子中的大部分计算都在 Cube 核内实现。由于 $Scale_Q,Scale_K\ge0$令 $\tilde{W} Scale_Q \odot W \in \mathbb{R}^{g\times1}$LightningIndexerQuant计算公式可等价变换为$$\text{Top-}k\left{\tilde{W}^T \text{ReLU}(Q_{index}^{INT8}(K_{index}^{INT8})^T) \odot Scale_K^T\right}.$$对应的具体计算流程为预处理在 UB 上计算 $\tilde{W} Scale_Q \odot W$CAST 到 FP16 写出到 GM;在 Cube 核内计算 $S^{INT32} Q_{index}^{INT8}(K_{index}^{INT8})^T$$S$ 数据类型为 INT32利用 FixP 的 DEQF16 指令计算 $S^{FP16}\text{FP16}(S^{INT32}/1024)$将 INT32 结果与 $1/1024$ 相乘并转化为 FP16 数据类型此处 $1/1024$ 可以防止数值溢出且不会影响最终 Top-$k$ 的结果继续利用 FixP 的随路 ReLU 功能计算 $S\text{ReLU}(S^{FP16})$ 并将结果搬运至 L1 以进行下一次矩阵乘法计算在 Cube 核内执行 $P\tilde{W}^TS$ 将结果写出到 GM 上。这一步骤同时完成了关于 $Scale_Q$ 的反量化、关于 $W$ 的行方向放缩以及不同头结果的求和操作将 $P$ 从 GM 读入到 Vector 核内并进行关于 $Scale_K$ 的反量化$Score P \odot Scale_K^T$在 Vector 核内对 $Score$ 计算 Top-$k$。计算流程如下图所示这一方案相比标准实现方案将大部分向量操作通过数学变换以 Cube 核实现 (步骤5)同时充分利用了 FixP 随路量化激活搬运 能力将量化矩阵乘法的结果从 INT32 数据类型转化为 FP16 数据类型 (步骤34) 得以进行第二次半精度的矩阵乘法操作显著提升了算子性能。SparseFlashAttention随着大模型上下文长度的增加Sparse Attention 的重要性与日俱增这一技术通过“只计算关键部分”大幅减少计算量然而会引入大量的离散访存造成数据搬运时间增加进而影响整体性能。Sparse Attention 计算可表示为$$ \text{softmax}(\frac{Q\tilde{K}^T}{\sqrt{d_k}})\tilde{V} $$ 其中 $\tilde{K},\tilde{V}$ 为基于某种选择算法 (如LightningIndexer) 得到的重要性较高的 Key 和 Value一般具有稀疏或分块稀疏的特征$d_k$ 为 $Q,\tilde{K}$ 每一个头的维度。本次公布的SparseFlashAttention是面向 Sparse Attention 的全新算子针对离散访存进行了指令缩减及搬运聚合的细致优化。SparseFlashAttention算子沿用 FlashAttention 的计算流程分为四个阶段$C_1$$QK^T$$V_1$online softmax$C_2$$PV$$V_2$rescaling $O$。SFA 算子 Tiling 设计SparseFlashAttention一次迭代计算的基本块大小为$(256,512)$。算子核内 Tiling 设计如下L1 空间划分为如下部分$Q/P$ 矩阵分时复用 ($P$矩阵是softmax计算的结果)$4\times128\times 288\times2\text{ Bytes}288\text{ KB}$即在一个 $C_1$ 或 $C_2$ 阶段中左矩阵分四次搬运到 L1 中最终在一次迭代内常驻于 L1其中 $Q$ 矩阵每次的搬运量大小为 $(128,288)$$P$ 矩阵每次的搬运量大小为 $(128,256)$$K/V$ 矩阵 3-Buffer 循环复用$3\times128\times 288\times2\text{ Bytes}216\text{ KB}$其中$K$矩阵每次的搬运大小为$(128,288)$$V$矩阵每次的搬运大小为$(256,128)$。L0AL0BL0C使能 Double Buffer分别划分为$2\times32\text{ KB},2\times32\text{ KB},2\times64\text{ KB}$在 $C_1$ 阶段一次搬入 L0A 和 L0B的矩阵块大小分别为 $(128,96),(96,128)$在 $C_2$ 阶段一次搬入 L0A 和 L0B 的矩阵块大小都为 $(128,128)$。SFA Pipeline 设计对于复杂的融合算子而言Cube 核和 Vector 核之间的高效协同是发挥硬件算力的核心挑战流水掩盖作为连接二者的核心调度机制其设计优劣直接决定了算子性能。由于沿用 AMLA 算法SparseFlashAttention除最终的归一化操作外只有三个计算阶段$C_1,V_1,C_2$。相较于LightningIndexer算子而言SparseFlashAttention算子的计算流程更为复杂依赖关系更多朴素的流水排布会导致计算过程中出现大量的空泡。如下图的上半部分所示由于计算流程中的不同阶段之间存在数据依赖不论是 Cube 流或 Vector 流上都出现了很多的空闲部分导致算力利用率较低。本算子采用 Preload 流水排布以消除依赖实现了除头尾以外的平台期部分 Cube 流完美掩盖 Vector 流效果如下图下半部分所示可以看到这种方式相较于朴素的流水排布有较大的性能收益。访存优化实际的访存行为可以分为两个步骤指令下发及数据搬运。对于非 Sparse 场景的 Attention 算子数据访存一般是连续的因此单次数据搬运耗时相较于指令下发耗时更长可实现连续的多次访存达到较高的访存带宽。然而对于 Sparse Attention 的离散访存而言单次访存的耗时大幅下降甚至可能小于指令下发的耗时这就会导致严重的指令下发阻塞如下图所示。Attention 计算从 Cube 核发起访存指令也由 Cube 核发射。基于 A3 芯片 Cube 核和 Vector 核数比为 $1:2$ 的特性利用 Vector 核发射访存指令将指令的下发效率提高一倍。另一方面单次访存的数据量本身也会影响实际达到的带宽而离散访存单次的访存量较小这会导致数据搬运本身的耗时高于预期。针对这一问题可以将离散的数据点两两聚合利用 srcGap 参数实现成对的数据访存提升访存带宽如图所示。SFA Next Feature后续将探索SparseFlashAttention算子的量化支持以应对长上下文场景下显存占用过大的挑战。易用性提升TilingKey 模板化TilingKey 是算子区分不同的实现而将代码进行区分的方法目标是以不同的分支处理不同的特性。然而 TilingKey 本身不易于记忆和理解因为它们通常是较长又没有明确含义的数字严重影响代码的可读性和可编程性。为了提高算子代码的可读性降低开发难度本次开源的两个全新算子均采用 Ascend C 的 Tiling 模板编程方法替代传统的 TilingKey 编程只需要定义模板参数及其对应的组合即可实现分支管理减少对 TilingKey 具体数值的依赖提高代码的管理效率。详细内容可参考Tiling模板编程。分核计算下沉Tiling 是算子执行前的预处理阶段在此阶段 Host 侧会计算 workspace 占用以及分核策略等算子执行需要的内容在单算子模式下可能会出现 Host Bound 的情况。通过静态图下沉技术我们规避了这种情况。然而对于 Attention 类算子而言实际输入的不同 Batch 对应的上下文长度可能不同因此 Tiling 中分核策略的计算对具体上下文长度 (actualSeqLengths) 存在值依赖这使得 Attention 类算子无法完全下沉到静态图执行导致整网产生过多小的静态子图引入大量静态子图调度开销同时进行分核计算时需要将对应的数据从 Device 侧向 Host 侧拷贝造成整网性能较差。针对这一问题Attention 类算子采用了 Tiling 下沉的解决方案将分核计算下沉到 Device 侧的 AICPU 进行。整网运行时分核计算依赖的数据就绪后即可启动且可以与算子执行流并行进行。然而对于分核计算依赖前序算子的情况这种方法依然无法避免执行断流。从具体使用的角度看这种处理方式引入了额外的 AICPU Tiling 逻辑在自定义算子流程中需要另外开发对应的 Tiling func并配合下沉图执行调度实现功能复杂、链条较长导致自定义开发和调测难度较高。本次发布的LightningIndexer及SparseFlashAttention算子进行了进一步的 Tiling 下沉优化将分核计算流程从 AICPU 下沉到 AICore 内部执行简化调用链条消除 Tiling 依赖前序算子引入的调度开销开发和调试流程得到了简化。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考