去年11月帮一个团队做 DeepSeek-V3 的推理优化刚把模型权重搬到昇腾 910B 上跑出来 decode 吞吐 580 TPS。客户问GPU 上能到 2400差了 4 倍我第一反应是算子没写对翻 msprof 一看Attention 计算占 38%KV Cache 访存 27%剩下的全是算子调度开销。很多人以为 FlashAttention 只是矩阵乘的优化其实它的核心是把 O(N²) 显存访问压到 O(N)让中间结果不落 HBM。昇腾 NPU 的内存层次跟 GPU 完全不同直接移植 CUDA 的 FlashAttention kernel 根本跑不满。为什么GPU 的一个 SM 能同时干矩阵乘和逐元素运算昇腾是 Cube 干矩阵乘、Vector 干逐元素中间数据要走 L1 缓存。Attention 公式与显存访问瓶颈标准 Attention 的计算Attention(Q, K, V) softmax(QK^T / √d_k) × V拆开看Q × K^T → 输出 [batch, heads, seq_q, seq_k]除以 √d_kSoftmax逐 token 归一化乘 V → 输出 [batch, heads, seq_q, d_v]问题在第 1 步和第 4 步QK^T 的输出是 O(N²)N2048 时就是 4M 个 float16占 8MB。如果这个中间矩阵写回 HBM再读出来做 Softmax光这一来一回就吃掉 16GB/s 的带宽。序列长度 128K 时中间矩阵 512MHBM 带宽是瓶颈。FlashAttention 的做法不把 QK^T 的完整矩阵写出而是按 block 切分算完一个 block 的 Softmax 立刻乘 V中间结果塞在 SRAM昇腾上是 L1 缓存。工程经验Qwen2.5-7B 在 910B 上seq2048FlashAttention 融合前吞吐 34 tokens/s融合后 89 tokens/s涨了 162%。关键是少了 3 次 HBM 读写QK^T 写回、Softmax 读入、Attention 输出写回。昇腾 NPU 的内存层次昇腾 910B 的内存层次HBM ( High Bandwidth Memory ) ↓ 1.2 TB/s L1 缓存 ( 1MB / 核 ) ↓ 片上 SRAM L0A / L0B ( 128KB / Cube 单元 ) ↓ Cube Unit (矩阵乘) / Vector Unit (逐元素)跟 GPU 的区别GPU 的 L2 缓存是全局共享的所有 SM 都能直接访问昇腾的 L1 是每个 AI Core 独立的跨核通信走 HCCLGPU 的寄存器文件和 L1 是一体的昇腾的 L0A/L0B 是 Cube Unit 专用的这导致 Tiling 策略不能直接抄 GPU 的。GPU 上你可以把 QK^T 的中间结果存在 L2所有 block 共享昇腾上 L1 只有 1MB存不下 2048×2048 的半矩阵必须把 Q、K、V 都按 block 切分每个 AI Core 只算自己那份。CANN 的 Tiling 策略ops-transformer 里的 FlashAttention 算子Tiling 参数有三个参数含义典型值tile_qQ 的 token 维度切分粒度32 / 64 / 128tile_kK/V 的 token 维度切分粒度32 / 64 / 128block_sizeSoftmax 归一化的 block 大小256 / 512Tiling 的目标是让 QK^T 的中间结果刚好塞进 L1不溢出到 HBM。伪代码Ascend C 风格# Q: [batch, heads, seq_q, d] # K, V: [batch, heads, seq_k, d] tile_q 64 # 一次处理 64 个 query token tile_k 64 # 一次处理 64 个 key token for i in range(0, seq_q, tile_q): Q_tile Q[:, :, i:itile_q, :] # 从 HBM 读 Q 的这块 for j in range(0, seq_k, tile_k): K_tile K[:, :, j:jtile_k, :] # 从 HBM 读 K 的这块 # ---- Cube Unit 干活 ---- S_tile matmul(Q_tile, K_tile.T) # [tile_q, tile_k] # ---- 数据挪到 L1给 Vector Unit ---- S_tile_l1 copy_to_l1(S_tile) # 不落 HBM # ---- Vector Unit 干活 ---- S_tile_scaled S_tile_l1 / sqrt(d) S_tile_masked apply_causal_mask(S_tile_scaled) P_tile softmax(S_tile_masked, dim-1) # ---- 把 P 写回 L1等 V 来了直接算 ---- P_tile_l1 copy_to_l1(P_tile) V_tile V[:, :, j:jtile_k, :] # ---- Cube Unit 再干活 ---- O_tile matmul(P_tile_l1, V_tile) # [tile_q, d] # 累加多个 j 的结果在 L1 里完成 O_accum O_tile关键点Q × K^T 的结果 S_tile 算完之后不写 HBM直接留在 L1 里给 Vector Unit 算 Softmax。这就是 FlashAttention “显存访问 O(N)” 的本质。工程经验7B 模型推理时 batch1、seq2048FlashAttention 只用了 Cube 算力的 40%。原因tile_q32MAC 阵列 16×16 只填了一半。把 tile_q 调到 64吞吐直接涨了 22%。不是算力不够是 Tiling 参数没压满。Cube / Vector 协同机制达芬奇架构的核心设计Cube Unit 和 Vector Unit 是两个独立的执行单元中间靠 L1 缓存传数据。很多人以为 Cube 算完矩阵乘结果会自动传给 Vector其实不是。Cube 的输出默认写 L1Vector 从 L1 读数据这个接力需要程序员或编译器手动安排流水线。ops-transformer 的 FlashAttention 做了双缓冲Double Buffering时间线 Cycle 0-100: Cube 算 QK^T (tile 0) → 写 L1 Cycle 50-150: Vector 读 L1 (tile 0) → Softmax → 写 L1 Cycle 100-200: Cube 算 QK^T (tile 1) → 写 L1tile 0 的数据已经被 Vector 读完了L1 空间释放 Cycle 150-250: Vector 读 L1 (tile 1) → Softmax → 写 L1Cube 和 Vector 的交叠率能到 75%相当于两个单元都在跑没有等对方。如果 Tiling 没做好Cube 算完一块等 Vector或者 Vector 算完一块等 Cube交叠率掉到 30% 以下吞吐直接腰斩。AscendCL 调度流程没用 FlashAttention 融合之前Attention 的计算要走 7 次 AscendCL 调用PyTorch 代码 Q linear(x, W_q) # 1. 矩阵乘 K linear(x, W_k) # 2. 矩阵乘 V linear(x, W_v) # 3. 矩阵乘 S matmul(Q, K.T) / sqrt(d) # 4. 矩阵乘 逐元素 S causal_mask(S) # 5. 逐元素 P softmax(S, dim-1) # 6. 逐元素 O matmul(P, V) # 7. 矩阵乘7 次调用每次调用走 ACL → GE → Runtime 的调用链单次开销 12-15μs。30 层 Transformer 就是 210 次调用光调度开销就 2.5-3ms。用了 ops-transformer 的 FlashAttention 融合算子之后上面的 7 步变成 1 次 ACL 调用# 融合之后的调用 O ops_transformer_flash_attention(Q, K, V, mask, sqrt(d)) # 内部把 7 个算子融合成一个 Kernel # Cube/Vector 协同、Tiling、Double Buffering 全部在 Kernel 内部完成调度开销从 3ms 降到 0.05ms省出来的时间全用来算矩阵乘。工程经验DeepSeek-V3 的 MoE 层里有 15 个 expert每个 expert 都有独立的 Attention。没融合之前光 Attention 部分的调度开销就 18ms/token融合之后 4.2ms/token。省下来的 13.8ms 全变成了吞吐。Kernel Fusion 的边界在哪不是所有算子都能融合。Cube/Vector 的协同有边界能融合的矩阵乘 逐元素QK^T Scale Softmax P×V→ 中间结果走 L1LayerNorm 矩阵乘 → LayerNorm 的输出直接喂给 MatMul不落 HBM残差连接 激活函数 → Vector Unit 一次性干完不能融合的或者说融合了反而不快的两个独立的矩阵乘Q×W_q 和 K×W_k→ 它们没有数据依赖融合了反而抢 Cube 的算力需要跨 AI Core 通信的操作 → 融合之后通信开销比调度开销还大graph-autofusion 框架会自动判断哪些算子能融、哪些不能融。它的逻辑看算子之间的数据依赖关系如果 A 的输出是 B 的输入且中间结果的尺寸小于 L1 容量1MB就融否则不融。Attention 算子的瓶颈到底在哪很多人以为 Attention 的瓶颈在矩阵乘Cube Unit实测不是。模型Attention 占比瓶颈Qwen2.5-7B (seq512)Cube 72% / Vector 28%Cube矩阵乘Qwen2.5-7B (seq2048)Cube 38% / Vector 45% / 访存 17%VectorSoftmaxDeepSeek-V3 (seq4096)Cube 31% / Vector 52% / 访存 17%VectorSoftmax 归一化序列短的时候QK^T 的计算量大Cube 是瓶颈。序列一长Softmax 要遍历整个 seq_k 维度做归一化Vector Unit 反而成了瓶颈。很多人以为把 Cube 的 MAC 阵列填满就完事了其实序列长度超过 2048 之后优化重点应该从 Cube 转向 Vector——Softmax 的逐 token 归一化、KV Cache 的压缩/解压缩全是 Vector 的活。与 CUDA FlashAttention 的核心差异维度CUDA (GPU)Ascend C (昇腾 NPU)内存层次L2 全局共享 (几十MB)L1 每核独立 (1MB)融合粒度可以融整层 Transformer按 Cube/Vector 边界切分Tiling 策略tile 大小由 L2 容量决定tile 大小由 L1 容量决定双单元协同SM 同时干矩阵乘和逐元素Cube/Vector 靠 L1 传数据需要手动安排流水线调优工具nvFuser / PyTorch InductorAOE (AMCT) / graph-autofusion最直接的差异CUDA 上你可以把 QK^T、Softmax、P×V 全部融成一个巨型 Kernel中间结果全塞 L2昇腾上 L1 只有 1MB必须按 tile 切分算完一个 tile 的 Softmax 立刻乘 V不能等所有 tile 都算完再统一处理。这意味着同样的 FlashAttention 算法在昇腾上要写成多层嵌套循环外层遍历 tile_q内层遍历 tile_k而 CUDA 上可以是单层循环所有 block 并行算。工程经验从 GPU 迁移 FlashAttention 代码到昇腾千万别直接翻译 CUDA kernel。我们第一版就是直接翻的性能只有 GPU 的 40%。后来按达芬奇架构重写了 Tiling 和双缓冲逻辑性能追平了 CUDA 版本同算力档次对比。BatchSize 对性能的真实影响理论分析BatchSize 越大矩阵乘的并行度越高Cube 的 MAC 阵列利用率越高。实测数据Qwen2.5-7B910B 单卡FP16BatchSize吞吐 (tokens/s)TTFT (ms)显存 (GB)Cube 利用率1384514.223%41127816.851%817812020.168%1623121029.379%3226738039.285%64244720OOM81%BatchSize32 是拐点。再大KV Cache 占显存太多Attention 的 tile 数变多L1 缓存命中率下降吞吐反而掉。很多人以为 BatchSize 越大越好其实在线推理场景batch1-4和离线批处理batch16-32是完全不同的优化方向。在线要压 TTFT首 token 延迟离线要压总吞吐。KV Cache 压缩的收益DeepSeek-V4 的 CSA4倍压缩和 HCA128倍压缩不是免费的。压缩/解压缩需要 Vector Unit 的计算短序列4K反而不如不压缩。序列长度不压缩 (ms/token)CSA 4倍 (ms/token)HCA 128倍 (ms/token)5124.24.8 (14%)5.1 (21%)204812.69.3 (-26%)8.7 (-31%)819289.328.4 (-68%)18.2 (-80%)32768OOM89.731.4 (-65% vs CSA)131072OOMOOM89.3128K 序列HCA 把 KV Cache 从 7.3GB 压到 57MB显存不爆了吞吐才能跑起来。但压缩是有代价的——Vector Unit 要花 15-20% 的 cycle 做压缩/解压缩。很多人以为 KV Cache 压缩就是省显存其实真正的收益是省出的显存允许跑更大的 BatchSize更大的 BatchSize 才能让 Cube 的 MAC 阵列利用率从 30% 涨到 80%。仓库https://atomgit.com/cann/ops-transformer https://atomgit.com/cann/ascend-transformer-boost https://atomgit.com/cann/graph-autofusion
FlashAttention 在昇腾 NPU 上到底有多快？CANN 深度优化全解析
发布时间：2026/5/22 23:10:39
去年11月帮一个团队做 DeepSeek-V3 的推理优化刚把模型权重搬到昇腾 910B 上跑出来 decode 吞吐 580 TPS。客户问GPU 上能到 2400差了 4 倍我第一反应是算子没写对翻 msprof 一看Attention 计算占 38%KV Cache 访存 27%剩下的全是算子调度开销。很多人以为 FlashAttention 只是矩阵乘的优化其实它的核心是把 O(N²) 显存访问压到 O(N)让中间结果不落 HBM。昇腾 NPU 的内存层次跟 GPU 完全不同直接移植 CUDA 的 FlashAttention kernel 根本跑不满。为什么GPU 的一个 SM 能同时干矩阵乘和逐元素运算昇腾是 Cube 干矩阵乘、Vector 干逐元素中间数据要走 L1 缓存。Attention 公式与显存访问瓶颈标准 Attention 的计算Attention(Q, K, V) softmax(QK^T / √d_k) × V拆开看Q × K^T → 输出 [batch, heads, seq_q, seq_k]除以 √d_kSoftmax逐 token 归一化乘 V → 输出 [batch, heads, seq_q, d_v]问题在第 1 步和第 4 步QK^T 的输出是 O(N²)N2048 时就是 4M 个 float16占 8MB。如果这个中间矩阵写回 HBM再读出来做 Softmax光这一来一回就吃掉 16GB/s 的带宽。序列长度 128K 时中间矩阵 512MHBM 带宽是瓶颈。FlashAttention 的做法不把 QK^T 的完整矩阵写出而是按 block 切分算完一个 block 的 Softmax 立刻乘 V中间结果塞在 SRAM昇腾上是 L1 缓存。工程经验Qwen2.5-7B 在 910B 上seq2048FlashAttention 融合前吞吐 34 tokens/s融合后 89 tokens/s涨了 162%。关键是少了 3 次 HBM 读写QK^T 写回、Softmax 读入、Attention 输出写回。昇腾 NPU 的内存层次昇腾 910B 的内存层次HBM ( High Bandwidth Memory ) ↓ 1.2 TB/s L1 缓存 ( 1MB / 核 ) ↓ 片上 SRAM L0A / L0B ( 128KB / Cube 单元 ) ↓ Cube Unit (矩阵乘) / Vector Unit (逐元素)跟 GPU 的区别GPU 的 L2 缓存是全局共享的所有 SM 都能直接访问昇腾的 L1 是每个 AI Core 独立的跨核通信走 HCCLGPU 的寄存器文件和 L1 是一体的昇腾的 L0A/L0B 是 Cube Unit 专用的这导致 Tiling 策略不能直接抄 GPU 的。GPU 上你可以把 QK^T 的中间结果存在 L2所有 block 共享昇腾上 L1 只有 1MB存不下 2048×2048 的半矩阵必须把 Q、K、V 都按 block 切分每个 AI Core 只算自己那份。CANN 的 Tiling 策略ops-transformer 里的 FlashAttention 算子Tiling 参数有三个参数含义典型值tile_qQ 的 token 维度切分粒度32 / 64 / 128tile_kK/V 的 token 维度切分粒度32 / 64 / 128block_sizeSoftmax 归一化的 block 大小256 / 512Tiling 的目标是让 QK^T 的中间结果刚好塞进 L1不溢出到 HBM。伪代码Ascend C 风格# Q: [batch, heads, seq_q, d] # K, V: [batch, heads, seq_k, d] tile_q 64 # 一次处理 64 个 query token tile_k 64 # 一次处理 64 个 key token for i in range(0, seq_q, tile_q): Q_tile Q[:, :, i:itile_q, :] # 从 HBM 读 Q 的这块 for j in range(0, seq_k, tile_k): K_tile K[:, :, j:jtile_k, :] # 从 HBM 读 K 的这块 # ---- Cube Unit 干活 ---- S_tile matmul(Q_tile, K_tile.T) # [tile_q, tile_k] # ---- 数据挪到 L1给 Vector Unit ---- S_tile_l1 copy_to_l1(S_tile) # 不落 HBM # ---- Vector Unit 干活 ---- S_tile_scaled S_tile_l1 / sqrt(d) S_tile_masked apply_causal_mask(S_tile_scaled) P_tile softmax(S_tile_masked, dim-1) # ---- 把 P 写回 L1等 V 来了直接算 ---- P_tile_l1 copy_to_l1(P_tile) V_tile V[:, :, j:jtile_k, :] # ---- Cube Unit 再干活 ---- O_tile matmul(P_tile_l1, V_tile) # [tile_q, d] # 累加多个 j 的结果在 L1 里完成 O_accum O_tile关键点Q × K^T 的结果 S_tile 算完之后不写 HBM直接留在 L1 里给 Vector Unit 算 Softmax。这就是 FlashAttention “显存访问 O(N)” 的本质。工程经验7B 模型推理时 batch1、seq2048FlashAttention 只用了 Cube 算力的 40%。原因tile_q32MAC 阵列 16×16 只填了一半。把 tile_q 调到 64吞吐直接涨了 22%。不是算力不够是 Tiling 参数没压满。Cube / Vector 协同机制达芬奇架构的核心设计Cube Unit 和 Vector Unit 是两个独立的执行单元中间靠 L1 缓存传数据。很多人以为 Cube 算完矩阵乘结果会自动传给 Vector其实不是。Cube 的输出默认写 L1Vector 从 L1 读数据这个接力需要程序员或编译器手动安排流水线。ops-transformer 的 FlashAttention 做了双缓冲Double Buffering时间线 Cycle 0-100: Cube 算 QK^T (tile 0) → 写 L1 Cycle 50-150: Vector 读 L1 (tile 0) → Softmax → 写 L1 Cycle 100-200: Cube 算 QK^T (tile 1) → 写 L1tile 0 的数据已经被 Vector 读完了L1 空间释放 Cycle 150-250: Vector 读 L1 (tile 1) → Softmax → 写 L1Cube 和 Vector 的交叠率能到 75%相当于两个单元都在跑没有等对方。如果 Tiling 没做好Cube 算完一块等 Vector或者 Vector 算完一块等 Cube交叠率掉到 30% 以下吞吐直接腰斩。AscendCL 调度流程没用 FlashAttention 融合之前Attention 的计算要走 7 次 AscendCL 调用PyTorch 代码 Q linear(x, W_q) # 1. 矩阵乘 K linear(x, W_k) # 2. 矩阵乘 V linear(x, W_v) # 3. 矩阵乘 S matmul(Q, K.T) / sqrt(d) # 4. 矩阵乘 逐元素 S causal_mask(S) # 5. 逐元素 P softmax(S, dim-1) # 6. 逐元素 O matmul(P, V) # 7. 矩阵乘7 次调用每次调用走 ACL → GE → Runtime 的调用链单次开销 12-15μs。30 层 Transformer 就是 210 次调用光调度开销就 2.5-3ms。用了 ops-transformer 的 FlashAttention 融合算子之后上面的 7 步变成 1 次 ACL 调用# 融合之后的调用 O ops_transformer_flash_attention(Q, K, V, mask, sqrt(d)) # 内部把 7 个算子融合成一个 Kernel # Cube/Vector 协同、Tiling、Double Buffering 全部在 Kernel 内部完成调度开销从 3ms 降到 0.05ms省出来的时间全用来算矩阵乘。工程经验DeepSeek-V3 的 MoE 层里有 15 个 expert每个 expert 都有独立的 Attention。没融合之前光 Attention 部分的调度开销就 18ms/token融合之后 4.2ms/token。省下来的 13.8ms 全变成了吞吐。Kernel Fusion 的边界在哪不是所有算子都能融合。Cube/Vector 的协同有边界能融合的矩阵乘 逐元素QK^T Scale Softmax P×V→ 中间结果走 L1LayerNorm 矩阵乘 → LayerNorm 的输出直接喂给 MatMul不落 HBM残差连接 激活函数 → Vector Unit 一次性干完不能融合的或者说融合了反而不快的两个独立的矩阵乘Q×W_q 和 K×W_k→ 它们没有数据依赖融合了反而抢 Cube 的算力需要跨 AI Core 通信的操作 → 融合之后通信开销比调度开销还大graph-autofusion 框架会自动判断哪些算子能融、哪些不能融。它的逻辑看算子之间的数据依赖关系如果 A 的输出是 B 的输入且中间结果的尺寸小于 L1 容量1MB就融否则不融。Attention 算子的瓶颈到底在哪很多人以为 Attention 的瓶颈在矩阵乘Cube Unit实测不是。模型Attention 占比瓶颈Qwen2.5-7B (seq512)Cube 72% / Vector 28%Cube矩阵乘Qwen2.5-7B (seq2048)Cube 38% / Vector 45% / 访存 17%VectorSoftmaxDeepSeek-V3 (seq4096)Cube 31% / Vector 52% / 访存 17%VectorSoftmax 归一化序列短的时候QK^T 的计算量大Cube 是瓶颈。序列一长Softmax 要遍历整个 seq_k 维度做归一化Vector Unit 反而成了瓶颈。很多人以为把 Cube 的 MAC 阵列填满就完事了其实序列长度超过 2048 之后优化重点应该从 Cube 转向 Vector——Softmax 的逐 token 归一化、KV Cache 的压缩/解压缩全是 Vector 的活。与 CUDA FlashAttention 的核心差异维度CUDA (GPU)Ascend C (昇腾 NPU)内存层次L2 全局共享 (几十MB)L1 每核独立 (1MB)融合粒度可以融整层 Transformer按 Cube/Vector 边界切分Tiling 策略tile 大小由 L2 容量决定tile 大小由 L1 容量决定双单元协同SM 同时干矩阵乘和逐元素Cube/Vector 靠 L1 传数据需要手动安排流水线调优工具nvFuser / PyTorch InductorAOE (AMCT) / graph-autofusion最直接的差异CUDA 上你可以把 QK^T、Softmax、P×V 全部融成一个巨型 Kernel中间结果全塞 L2昇腾上 L1 只有 1MB必须按 tile 切分算完一个 tile 的 Softmax 立刻乘 V不能等所有 tile 都算完再统一处理。这意味着同样的 FlashAttention 算法在昇腾上要写成多层嵌套循环外层遍历 tile_q内层遍历 tile_k而 CUDA 上可以是单层循环所有 block 并行算。工程经验从 GPU 迁移 FlashAttention 代码到昇腾千万别直接翻译 CUDA kernel。我们第一版就是直接翻的性能只有 GPU 的 40%。后来按达芬奇架构重写了 Tiling 和双缓冲逻辑性能追平了 CUDA 版本同算力档次对比。BatchSize 对性能的真实影响理论分析BatchSize 越大矩阵乘的并行度越高Cube 的 MAC 阵列利用率越高。实测数据Qwen2.5-7B910B 单卡FP16BatchSize吞吐 (tokens/s)TTFT (ms)显存 (GB)Cube 利用率1384514.223%41127816.851%817812020.168%1623121029.379%3226738039.285%64244720OOM81%BatchSize32 是拐点。再大KV Cache 占显存太多Attention 的 tile 数变多L1 缓存命中率下降吞吐反而掉。很多人以为 BatchSize 越大越好其实在线推理场景batch1-4和离线批处理batch16-32是完全不同的优化方向。在线要压 TTFT首 token 延迟离线要压总吞吐。KV Cache 压缩的收益DeepSeek-V4 的 CSA4倍压缩和 HCA128倍压缩不是免费的。压缩/解压缩需要 Vector Unit 的计算短序列4K反而不如不压缩。序列长度不压缩 (ms/token)CSA 4倍 (ms/token)HCA 128倍 (ms/token)5124.24.8 (14%)5.1 (21%)204812.69.3 (-26%)8.7 (-31%)819289.328.4 (-68%)18.2 (-80%)32768OOM89.731.4 (-65% vs CSA)131072OOMOOM89.3128K 序列HCA 把 KV Cache 从 7.3GB 压到 57MB显存不爆了吞吐才能跑起来。但压缩是有代价的——Vector Unit 要花 15-20% 的 cycle 做压缩/解压缩。很多人以为 KV Cache 压缩就是省显存其实真正的收益是省出的显存允许跑更大的 BatchSize更大的 BatchSize 才能让 Cube 的 MAC 阵列利用率从 30% 涨到 80%。仓库https://atomgit.com/cann/ops-transformer https://atomgit.com/cann/ascend-transformer-boost https://atomgit.com/cann/graph-autofusion
:回调函数与sizeof/strlen深度对比)
