NPU DeepSeek-V4 Ascend C 融合算子优化【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer面向 DeepSeek-V4 架构本次发布两个全新的融合算子SparseAttnSharedKV(SAS)以及Compressor并且对LightningIndexer(LI) 进行了能力增强。其中SparseAttnSharedKV算子为本次 DeepSeek 新架构中的 Attention 模块设计可覆盖 SharedKV 情形下多种不同的 FlashAttention 应用场景包括Sliding Window AttentionCompressed AttentionSparse Attention以及它们的组合LightningIndexer算子的计算流程与核心功能保持不变针对 $KV$ 压缩场景适配了对应的mask计算方式并且在 950PR/DT 上利用直方图 API 加速了 Top-$k$ 计算Compressor算子面向新架构中较为复杂的压缩算法将原本分散的一系列操作实现为高效的融合算子。本次的三个算子均同步支持 Atlas-A3 及 950PR/DT 平台并且在 950PR/DT 上适配了 FP8 量化模式以进一步提高 CANN 算子对于不同数据类型的支持。Highlights融合算子在 Atlas-A3 和 950PR/DT 同步开源强化昇腾芯片对低精度浮点数 FP8 的类型支持融合复杂的Compressor计算流程兼顾不同压缩方式强化LightningIndexer算子利用 950PR/DT 全新指令提升算子性能统一接口设计SparseAttnSharedKV算子覆盖多种不同 Attention 计算场景将值依赖的 Scheduler 计算部分独立为 AICPU Scheduler 算子使用更为灵活且适配不同图模式分阶段的分核策略充分利用硬件计算资源OutlineNPU DeepSeek-V4 Ascend C 融合算子优化HighlightsOutlineDeepSeek-V4 Attention 整体结构Compressor计算公式Tiling 设计错位写入LightningIndexer计算公式Tiling 设计Top-k 950PR/DT 方案SparseAttnSharedKV计算公式Tiling 设计多场景覆盖AICPU SchedulerAICPU scheduler 算子DeepSeek-V4 Attention 整体结构DeepSeek-V4 Attention 模块本次 DeepSeek-V4 将 $K,V$ 统一为同一份即 Shared Key-Value Attention并且采用了 Attention Sink同时引入了复杂的交错 Hybrid Attention 结构不同层的 Attention 模块的计算模式不同基于每层不同的 compress ratio 分为三类C1A: 前两层 C1A 层只涉及关于原始$KV$的 Sliding Window Attention 计算每个 token 只与相邻 128 个 token 对应的原始$KV$cache进行计算C1A-AttentionC4A除第一层外的奇数层为 C4A 层。C4A 层的 Attention 除了会与原始$KV$(ori_KV) 进行 window size 为 128 的 Sliding Window Attention 计算外还涉及到与压缩$KV$(cmp_KV) 的稀疏计算其中 cmp_KV 会基于每 4 个 token 及其前 4 个 token 的语义进行压缩。C4A-AttentionC128A除第二层外的偶数层为 C128A 层。C128A 层的 Attention 除了会与原始$KV$(ori_KV) 进行 window size 为 128 的 Sliding Window Attention 计算外还涉及到与压缩$KV$(cmp_KV) 的稠密计算其中 cmp_KV 会基于每 128 个 token 的语义进行压缩。C128A-Attention在 C4A 及 C128A 层输入的 hidden state 会通过一系列计算分别产生 ori_KV 和 cmp_KV其中计算 cmp_KV 的流程较为复杂而且在 C4A 和 C128A 的场景下有所不同。我们将这一部分计算设计为 Atlas-A3 和 950PR/DT 上的Compressor融合算子兼顾不同情况下的计算流程并且实现了较好的性能。C4A 层还会涉及到对 cmp_KV 的进一步稀疏化延续 DeepSeek-V3.2 的设计这一部分的稀疏化基于LightningIndexer进行。本次发布的LightningIndexer融合算子在 Atlas-A3 原本的算子基础上适配了 compress ratio 对应的阶梯形状 mask 计算逻辑并且在 950PR/DT 上适配 FP8 全量化版本。由于不同层的 Attention 计算不尽相同涉及到多种场景。针对 DeepSeek-V4 的 Attention 核心计算部分我们分别在 Atlas-A3 和 950PR/DT 上设计了支持 Attention Sink 的全新非量化SparseAttnSharedKV及伪量化SparseAttnSharedKV算子通过一套接口覆盖不同的功能需求。CompressorDeepSeek-V4 采用了一个新的attention架构Compress-4-Attention (C4A) 和Compress-128-Attention (C128A). 具体来说是将每 4 或 128 个 token 的$KV$cache 压缩成一个然后每个token与这些压缩的$KV$cache去进行 Attention 计算。在长序列的情况下C4A 和 C128A 可以有效地减少计算开销。计算公式C4A层C4A层的 Compressor 包含如下一系列计算操作给定输入$X\in \mathbb{R}^{s\times h}$其中 $s$是 序列长度$h$ 是hidden size首先计算其对应的 $2$ 个$KV$输入 $C^a, C^b \in \mathbb{R}^{s \times d}$,与压缩权重 $Z^a, Z^b \in \mathbb{R}^{s \times d}$其中 $d$ 是 head dimension. 具体公式如下$$ \begin{aligned} C^aX\cdot W^{aKV}, C^bX\cdot W^{bKV}, \ Z^aX\cdot W^{aGate}, Z^bX\cdot W^{bGate}, \end{aligned} $$其中$W^{aKV}, W^{bKV}, W^{aGate}, W^{bGate} \in \mathbb{R}^{h \times d}$是C4A对应KV和压缩权重的权重参数。长度为$s$的KV序列$C^a, C^b$中的每4个KV会被压缩成1个 $C^{\text{Comp}} \in \mathbb{R}^{\frac{s}{4} \times d}$其第 $i$ 行 $C_{i}^{\text{Comp}} \in \mathbb{R}^{1 \times d}$ 的计算公式如下$$ \begin{aligned} C_{i}^{\text{Comp}} \frac{\sum_{j4i}^{4(i1)-1}{e^{Z_j^aB_{j-4i}}} \odot C^a_j \sum_{j4(i1)}^{4(i2)-1}{e^{Z_j^bB_{j-4i}}} \odot C^b_j}{\sum_{j4i}^{4(i1)-1}{e^{Z_j^aB_{j-4i}}} \sum_{j4(i1)}^{4(i2)-1}{e^{Z_j^bB_{j-4i}}}} \ \left[1\right]{1\times8} (softmax(\left[Z^a{\left[4(i-1)1:4i,:\right]} ; Z^b_{\left[4i1:4(i1),:\right]}\right] B) \odot \left[C^a_{\left[4(i-1)1:4i,:\right]} ; C^b_{\left[4i1:4(i1),:\right]}\right]) \end{aligned} $$其中 $B \in \mathbb{R}^{8 \times d}$ 为 $C^a, C^b$ 对应的positional biases。C4A 示意图C128A层相比C4A, C128A层的 Compressor 会以更大的压缩比率对KV序列去进行压缩并且只依赖单一KV序列$C \in \mathbb{R}^{s \times d}$和压缩权重$Z \in \mathbb{R}^{s \times d}$其中 $$ \begin{aligned} CX\cdot W^{KV}, \ ZX\cdot W^{Gate} \end{aligned} $$ $W^{KV}, W^{Z}$ 是C128A对应KV和压缩权重的权重参数。然后长度为$s$的KV序列$C$中的每 $128$ 个$KV$会被压缩成 $1$ 个 $C^{\text{Comp}} \in \mathbb{R}^{\frac{s}{128} \times d}$其第 $i$ 行$$ \begin{aligned} C_{i}^{\text{Comp}} \frac{\sum_{j128i}^{128(i1)-1}{e^{Z_jB_{j-128i}}} \odot C_j}{\sum_{j128i}^{128(i1)-1}{e^{Z_jB_{j-128i}}}} \ \left[1\right]{1\times128} (softmax(Z{\left[128(i-1)1:128i,:\right]} B) \odot C_{\left[128(i-1)1:128i,:\right]}) \end{aligned} $$$i 1, \cdots, \frac{s}{128},$ 其中 $B \in \mathbb{R}^{128 \times d}$ 为 $C$ 对应的positional biases.C128A 示意图具体计算流程Compressor会先把 $r (r4 \text{ or } 128)$ 个$KV$压缩成 $1$ 个再对压缩后的 $KV$ 做 RMSNorm 与 Partial Rope 的后处理。各阶段具体计算如下压缩阶段计算矩阵乘法C4A: $\left[C^a, Z^a\right] X \left[W^{aKV}, W^{aGate}\right], \left[C^b, Z^b\right] X \left[W^{bKV}, W^{bGate}\right];$C128A: $\left[C, Z\right] X \left[W^{KV}, W^{Gate}\right]$计算分组加法C4A: $Z_i^\prime \left[Z_{\left[4(i-1)1:4i,:\right]}^a; Z_{\left[4i1:4(i1),:\right]}^b\right]B,~i1,2,\cdots, \frac{s}{4};$C128A: $Z_i^\prime Z_{\left[128(i-1)1:128i,:\right]} B,~i1,2,\cdots, \frac{s}{128};$计算分组SoftmaxC4A: $S_i^\prime softmax(Z_i^\prime),~i1,2,\cdots, \frac{s}{4};$C128A: $S_i^\prime softmax(Z_i^\prime),~i1,2,\cdots, \frac{s}{128};$计算Hadamard乘积:C4A: $(S_H)i S_i^\prime \odot \left[C^a{\left[4(i-1)1:4i,:\right]} ; C^b_{\left[4i1:4(i1),:\right]}\right],~i1,2,\cdots, \frac{s}{4};$C128A: $S_H S^\prime \odot C;$沿着压缩轴分组求和C4A: $C_{i}^{\text{Comp}} \left[1\right]_{1\times8} (S_H)_i, ~i1,2,\cdots, \frac{s}{4};$C128A: $C_{i}^{\text{Comp}} \left[1\right]_{1\times128} (S_H)_i, ~i1,2,\cdots, \frac{s}{128};$后处理阶段计算RMSNorm;计算Partial RoPE。融合价值Compressor 计算流程从上述计算流程图可以看出Compressor 中存在大量的非计算向量操作 (如图中split,slice等)若以小算子实现则会引入大量无意义的数据重复搬入搬出导致性能极差。以融合算子实现可以将这一类操作在数据搬运的过程中随路实现大幅降低 Compressor 部分的耗时。图中的coff在 C128A 层为1在 C4A 层为2。coff2意味着需要用相邻的两组 token 语义进行压缩因此 C4A 层的 Compressor 会引入一个较为复杂的数据重排操作C4A Compressor 数据重排如上图所示C4A 层的 Compressor 会引入一个较为复杂的数据重排操作即将每个 token 映射为两个 $d$ 大小的向量后 (图中的红色方块与黄色方块)将每 4 个 token 对应的黄色方块向量与前 4 个 token 对应的红色方块向量重排为连续存储。该操作如果用小算子单独实现同样会引入大量非计算的向量操作。而在该融合算子中我们巧妙地完成了这一部分的数据重排可参考后续章节错位写入。Tiling 设计Atlas-A3Atlas-A3 上算子 Tiling 如下对于 C128A:L1 空间划分为如下部分$X$ 矩阵$2 \times 256 \times 256 \times 2 \text{ Bytes}256 \text{ KB}$, L1 层级基本块为 $(256, 256)$$W^{KV}, W^{Gate}$ 矩阵$2 \times 256 \times 64 \times 2 \text{ Bytes}64 \text{ KB}$, L1 层级基本块为 $(256, 128)$一轮迭代会搬入 $W^{KV}, W^{Gate}$ 各 $2$ 个 $(256, 64)$ 的矩阵块L0A, L0B 使能 double buffer, 分别划分为 $2\times32\text{ KB},2\times32\text{ KB}$一次搬入 L0A 和 L0B 的矩阵块大小都为 $(128,128)$ L0C 使能 double buffer划分成 $2\times64\text{ KB}$分别存放 $1$ 个 $X$ 分块与 $W^{KV}, W^{Gate}$ 矩阵相乘的结果并且在 L0C 上累加。对于 C4A:L1 空间划分为如下部分$X$ 矩阵$128 \times 256 \times 2 \text{ Bytes}64 \text{ KB}$, L1 层级基本块为 $(128, 256)$$W^{KV}, W^{Gate}$ 矩阵$4 \times 256 \times 64 \times 2 \text{ Bytes}128 \text{ KB}$, L1 层级基本块为 $(256, 128)$一轮迭代会搬入 $W^{aKV}, W^{bKV}, W^{aGate}, W^{bGate}$ 各 $1$ 个 $(256, 64)$ 的矩阵块。L0A, L0B 使能 double buffer, 分别划分为 $2\times32\text{ KB},2\times32\text{ KB}$一次搬入 L0A 和 L0B 的矩阵块大小都为 $(128,128)$ L0C 使能 double buffer划分成 $2\times64\text{ KB}$分别存放 $2$ 个 $X$ 分块与 $1$ 个 KVGate 矩阵相乘的结果并且在 L0C 上累加。Atlas-A3 上 Compressor 核内 Tiling以此基本块做完矩阵乘法并按 reduce 轴累加后均分给 2 个 Vector 核, 进行后续Vector操作。Decode在 Decode 阶段为了解决 $s$ 轴方向无法分核的问题进一步切分 $d$ 轴从 $64$ 降低成 $32$来提高启动的核数与带宽利用率。基本块 Cost Model基本块的选取是算子设计中的关键环节需依据硬件参数通过理论计算确定合适的 Tiling 策略为使一轮基本块迭代中搬运左矩阵或右矩阵数据的时间能被对应计算时间掩盖需满足搬运耗时不超过计算耗时即 $$ \frac{(baseM baseN) \times baseK \times 2 \text{ Bytes}}{\text{Bandwidth}} \leq \frac{baseM \times baseN \times baseK}{4096} $$ 约去 $baseK$对于 Atlas-A3, C128A 中选取 $(baseM, baseN, baseK)(256, 128, 512)$, C4A 中选取 $(baseM, baseN, baseK)(128, 256, 512)$这里 $baseK$ 的选取主要是考虑昇腾亲和 Cacheline 对齐提高搬运效率。对于 950PR/DT考虑 L0AL0B 容量限制可取 $baseMbaseNbaseK256$ 达成条件。错位写入注意到在计算 C4A 时$Z^b, C^b$ 计算 $C_{i}^{\text{Comp}}$ 的下标与 $Z^a, C^a$ 计算 $C_{i1}^{\text{Comp}}$ 的下标有重叠在这里我们用错位写入来实现这个处理。如下图所示完成前序的矩阵乘计算后红色部分的计算结果在写入目标地址时会制定一个额外的地址偏移保证红色部分和黄色部分的计算结果对齐将 state_cache 里所需要的数据读入 UB 即可直接进行后续的向量计算。错位写入LightningIndexer计算公式LightningIndexer基于一系列操作得到每一个 token 对应的 Top-$k$ 个位置。不考虑 batch 轴对于 Query 对应的 $Q_{index}\in\R^{n \times S_q \times d}$给定上下文 Index Key $K_{index}\in\R^{S_{k}\times d},W\in\R^{n \times S_q\times 1}$其中 $n$ 为头的个数$d$ 为每一个头的维度$S_{q}$ 是请求的个数$S_{k}$ 是上下文的长度LightningIndexer的具体计算公式如下 $$ \text{Top-}k\left{[1]{1\times n}n\left[(W{1}[1]{1\times S_{k}})\odot\text{ReLU}\left(Q_{index}{d}K{index}^T\right)\right]\right} $$此时默认 $Q_{index}$ 和 $K_{index}$ 是量化后的FP8(E4M3)矩阵。量化过程中已对每个头和每个token得到反量化系数记为矩阵形式 $Scale_q \in\R^{n \times S_q\times 1}$, $Scale_k \in \R^{S_k \times 1}$。LightningIndexer可拆分为如下计算流程 (其中 ${d}$${1}$ 和 $_{n}$ 分别代表 $d,1,n$ 三个维度上相乘, $\odot$ 代表逐元素相乘)对量化后的矩阵计算矩阵乘法$S Q_{index}{d} K{index}^T$计算激活函数$S\text{ReLU}(S)$W 吸收反量化系数 $Scale_q$$W Scale_q \odot W$;计算广播乘法$S_W(W {1}[1]{1\times S_{k}})\odot S$关于头进行ReduceSum操作$Score[1]{1\times n}{n} S_W$广播乘反量化系数$Score (Scale_k {1}[1]{1\times S_{q}})^T \odot Score$对 $Score$ 进行 $\text{Top-}k$ 计算即获取数值排序前 $k$ 个的结果并返回其对应的 Index。在 Prefill 或开启 MTP 的 Decode 场景多个 token 对应的 $Q_{index}$ 会被合并统一计算本文后续部分不对单/多个 token 对应的 $Q_{index}$ 加以区分。实际开发中需要充分利用昇腾硬件的特征以实现更好的性能以下将详细介绍当前的解决方案及具体实现方式。Lightning Indexer 计算流程Tiling 设计Atlas-A3本次 Atlas-A3 上LightningIndexer的 Tiling 沿用之前的设计可参考Atlas-A3 LightningIndexer。Top-k 950PR/DT 方案LightningIndexer融合算子的核心是在长达数十万 (序列长度记为$S_{k}$) 的序列中为每个 token 高效地筛选出分数最高的 $k$ (例如 512) 个索引。同时对于算子而言Top-$k$ 的计算必须是准确无误的不能采用近似算法求解。一种昇腾亲和的解决方案是使用 950PR/DT 的 Histograms 指令。该指令可对若干个 uint8 数字进行直方图统计一个 cycle 可将 64 个数字分到 0-127 或 128-255 桶中。基于此思路以一个 Query 为例使用 Histograms 指令在 $S_{k}$ 个 float 类型的 Score 中选取 Top-$k$ 的核心算法总结如下保序变换设 $x$ 为浮点数 $f$ 的二进制表示, $\text{convert}(f)$ 为变换后的 uint32 值。可采用如下的保序变换算法uint32_t mask (x0x80000000)?0xffffffff:0x80000000; return (ff)?(x^mask):0x0;这等价于以下分段函数: $$ \text{convert}(f) \begin{cases} 0 \text{if}\ f\ \text{is nan}\ 2^{32}-1-x (\equiv \sim x) \text{if}\ x\ \geq 2^{31} (\text{i.e.}\ f 0) \ x 2^{31} (\equiv x \oplus 2^{31}) \text{if}\ x 2^{31} (\text{i.e.}\ f \geq 0) \end{cases} $$ 若 Score 的输入为 bfloat16 类型则可用类似方法转化为 uint16。多级直方图分桶以 uint32 类型为例介绍一种昇腾亲和的多级直方图分桶方法。首先将 uint32 的 32 位由高到低切分为 4 个 8 位分别视为 uint8 类型。已知对于 uint32 类型的两个数 a 和 b记两者的高8位构成的 uint8 分别为 a_8 和 b_8则 a_8 b_8 一定可以推导出 a b。基于这一思想可分别定位出第 Top-$k$ 值的每一个 8 位。以 uint32 为例若为 uint16 则只需二级直方图如图所示多级直方图筛选第 Top-$k$ 的值流程如下对所有的 uint32 元素提取从高到低的第一个 8 位用hist指令绘制直方图返回的结果是一个长度为256的数组 Array_1其中第 $i$ 个值代表 $0$-$i$ 个桶的个数和。将 Array_1 中的值与临界值 $S_k-k1$ 比较大于等于 $S_k-k1$ 的位置为 true。从小到大第一个为 true 的桶即为第 Top-$k$ 元素所在的桶 (图中蓝色部分)记为 $a_1$。此时可将前8位小于 $a_1$ 的元素用 mask 掩去 (即为图中的灰色部分)。这部分元素将被丢弃设本次丢弃的个数为 $C_1$ ($C_1 S_k-k$)。对剩余的 $S_k - C_1$ 个元素提取第 2 个 8 位绘制直方图重复 1-2 步的过程选取出临界桶 $a_2$并丢弃第2个8位小于 $a_2$ 的元素统计本次丢弃个数 $C_2$。对剩余的 $S_k - C_1 - C_2$ 个元素提取第 3 个 8 位绘制直方图重复 1-2 步的过程选取出临界桶 $a_3$并丢弃第 3 个 8 位小于 $a_3$ 的元素统计本次丢弃的个数 $C_3$。对剩余的 $S_k - C_1 - C_2 - C_3$ 个元素提取第 4 个 8 位绘制直方图重复 1-2 步的过程选取出临界桶 $a_4$。由此得到第 $k$ 大的值为 $$ value_k a_1 \times 2^{24} a_2 \times 2^{16} a_3 \times 2^{8} a_4 $$向量化筛选此时已拿到第 $k$ 大的值为 $value_k$只需考虑大于等于 $ value_k $ 的值总体流程如下大于 $value_k$ 的部分以向量化的方式256B 为单位将原序列中的值与 $value_k$ 比较得到一个 mask大于 $value_k$ 的位置为 true。以此 mask 筛选出所有大于 $value_k$ 的值对应的指标并聚集到 $result$ 中此时个数小于 $k$。等于 $value_k$ 的部分以 256B 为单位循环遍历所有 $Score$找到等于 $value_k$ 的值并筛选出来直到 $result$ 的总数等于 $k$ 时退出循环。实现指令和理论开销以 uint32 为例设序列长度 $S_k16384$, 选取的 Top-$k$ 为 $512$Top-$k$ 部分具体的实现指令和理论开销如下:hist1: 用DeInterleave解交织提取第一个8位并用直方图指令hist分桶。 一条hist指令每个cycle处理64个数字且需要两条指令分别处理0-127和128-255。实现过程中以 $256$ 个数字为一组进行循环处理每组需8个cycle处理直方图。加上初始化等其他开销约 10 cycle总计 $(S_k/256 * 9 10)$ cycle。find1: 对hist1的结果每 64 个桶一组与临界值$S_k-k1$做比较将比较的结果用squeeze的gather模式拼到一起并计算新的临界值开销约 $(4*7 10)$ cycle。hist2: 相比于hist1每轮循环增加一次比较Compare总开销约 $(S_k/256 * 10 10)$ cycle。find2: 和find1相同 开销 $(4 * 7 10)$ cycle。hist3: 相比于hist1除了每轮循环增加两次比较Compare和一次And还需要增加一次DeInterleave解交织提取第三个8位故总开销约 $(S_k/256 * 13 10)$ cycle。find3: 和find1相同开销 $(4 * 7 10)$ cycle。hist4: 相比于hist1增加一次DeInterleave三次Compare和两次And开销估计为 $(S_k/256 * 15 10)$ cycle。findKth: 计算临界 Top-$k$ 的值 $value_k$开销和 和find1基本相同共计 $(4 * 7 10)$ cycle。findGTO: 64 个 uint32 为一组做Compare和squeeze将大于 $value_k$ 的值对应的指标取出共计约 $S_k/64*7$ cycle。findEQ: 64 个 uint32 为一组做Compare和squeeze将等于 $value_k$ 的值对应的指标依次取出直到取满 $k$ 个共计 $S_k/64*9$ cycle。累计以上所有开销, 在 950PR/DT 上采用多级直方图和向量化筛选的方式理论开销估计为 $(S_k/256*111192)$ cycle, 在 $S_k16384$时理论开销计算为 7296 cycle。相比之下对于相同的 $S_k$ 和 $k$ Atlas-A3 上基于 sort32 和 mrgsort 的实现方式 (Atlas-A3 LI Top-k 指令实现) 需要 $ (1.36S_k3.32k\lceil\frac{max(log_2(k)-5,0)}{2}S_k \rceil)/2 $ cycle 的理论开销。在 $S_k16384, k512$ 时开销约 28375 cycle。SparseAttnSharedKV计算公式SparseAttnSharedKV算子旨在完成以下形式的 Attention 计算 $$O \text{softmax}(Q\tilde{K}^T \cdot \text{softmax_scale})\tilde{V}$$ 其中 $\tilde{K}\tilde{V}$ 为基于入参控制的实际参与计算的 $KV$。适配DeepSeek-V4 模型的全新 Attention该算子支持 Sliding Window Attention、Compressed Attention (以及在此基础上的稀疏化) 和两者混合。本次的SparseAttnSharedKV算子在 Atlas-A3 和 950PR/DT 上同步发布Atlas-A3 上为非量化版本即输入的 $Q,KV$ 都是半精度算子计算流程分为下面四个阶段$C_1$$QK^T$$V_1$online softmax$C_2$$PV$$V_2$rescaling $O$。950PR/DT 上为伪量化版本即输入的 $KV$ 已被量化到 8bit算子计算流程在最初新增一个 $V_0$ 阶段用于对 $KV$ 进行反量化整体计算流程为五个阶段。Tiling 设计由于 Atlas-A3 和 950PR/DT 的芯片架构不同我们采用了不同的 Tiling 策略。Atlas-A3一次迭代计算的基本块大小为(256,512)即完成 $Q$ 序列方向 256 和 $KV$ 序列方向 512 长度的计算。考虑到实际情况例如 C4A 层中存在的稀疏计算无法将不同 $Q$ 进行合轴计算只能一个 Group 内统一计算因此一次迭代实际计算的基本块大小可能为(64,512)但不影响整体 Cache 划分。算子核内 Tiling 设计如下L1 空间划分为如下部分$Q/P$ 矩阵分时复用 ($P$矩阵是softmax计算的结果)4×128×256×2 Bytes256 KB即在一个 $C_1$ ​或 $C_2$ ​ 阶段中左矩阵分四次搬运到 L1 中最终在一次迭代内常驻于 L1其中 $Q$ 矩阵每次的搬运量大小为 (128,256)$P$ 矩阵每次的搬运量大小为 (128,256)$K/V$ 矩阵 3-Buffer 循环复用3×128×256×2 Bytes192 KB其中 $K$ 矩阵每次的搬运大小为(128,256)$V$ 矩阵每次的搬运大小为(256,128)。L0AL0BL0C使能 Double Buffer分别划分为$2\times32\text{ KB}$ KB,$2\times32\text{ KB}$,$2\times64\text{ KB}$在 $C_1$ / $C_2$ 阶段一次搬入 L0A 和 L0B的矩阵块大小分别为 (128,128),(128,128)。多场景覆盖SparseAttnSharedKV算子通过多种入参组合实现不同场景的 Attention 计算场景其中关键入参为ori_kv: SWA对应的原始 $KV$ cacheori_win_left, ori_win_right: 表示 sliding window attention 中参与计算的 key-value 在序列轴的范围。假设 $t$ 为 query 的sequence id 则参与计算的key/value 的sequence id 范围为 $t - \text{win_left}, t - \text{win_left} 1, ..., t, ..., t \text{win_right}$cmp_kv: C4A 层及 C128A 层压缩后 $KV$ cachecmp_ratio: cmp_kv 中 $KV$ 的压缩比例。上述入参可通过如下组合方式完成 DeepSeek-V4 中不同场景的 Attention 计算AICPU SchedulerScheduler 核心职能在于执行流的静态预演与管控。通过预分配显存空间、锁定执行分支并优化分核策略该机制能够有效消除运行时开销确保指令流水线始终处于饱和状态极大提升硬件吞吐。AICPU scheduler 算子对于 Attention 及 LI 这类“值依赖算子”而言仅依靠输入的 shape 信息无法实现好的分核策略还需要在运行时获取实际的上下文长度。传统 Scheduler 实现方式需要回到 HOST 执行会导致整体执行流水被打断而直接在算子 kernel 内计算分核策略会导致这一部分的计算耗时无法被掩盖引入极大的头开销。昇腾芯片内置若干个 ARM 计算核AICPU把这一部分 Scheduler 计算放到 AICPU 上执行既可以保证算子入图又可以使得这一部分耗时被掩盖。本次同步发布了SparseAttnSharedKV及LightningIndexer对应的AICPU Scheduler 算子: SAS AICPU Scheduler, LI AICPU Scheduler。现可以将整个 Scheduler 过程分为两个部分只需要输入的 shape 即可计算最大显存占用 max workspace用于提前给算子分配足够的device空间具体执行路径 tiling key确定实际编译的算子模板启用的核数 block dim用于确定开核情况tiling data提取实际的 shape 信息进行常量折叠减少算子运行时的标量计算。需要实际的上下文长度才可以计算具体的分核策略每个核要执行的实际计算部分提高计算负载的均衡性以实现高算力利用率。第一个部分的 Scheduler 计算由 Host 执行而本次算子将第二个值依赖部分的 Scheduler 单独定义为一个 AICPU 算子其输出蕴含分核信息的 metadata 供实际的 NPU 算子使用# run in aicpu metadata torch_npu.npu_sparse_attn_sharedkv_metadata(......) # run in aicore npu_result, softmax_lse torch.ops.custom.npu_sparse_attn_sharedkv(...,metadatametadata,...)这样的方式有很多优势更为灵活不需要框架侧额外识别这类特殊的“值依赖”算子也减轻了算子侧额外适配框架的工作量用户通过脚本直接调用 AICPU 算子即可对于 aclgraph 等图模式更为友好对于单次推理而言每种计算模式对应的 AICPU Scheduler 算子只需要执行一次可以复用。对于同样有实时 Scheduler 需求的算子也可以自定义开发 AICPU Scheduler 算子实现更好的整网性能可参考AICPU算子开发了解开发方式获取更多信息。【免费下载链接】cann-recipes-infer本项目针对LLM与多模态模型推理业务中的典型模型、加速算法提供基于CANN平台的优化样例项目地址: https://gitcode.com/cann/cann-recipes-infer创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考