FlashAttention 昇腾优化：从 O(N²) 到 O(N) 的显存革命

前言7B 模型推理 seq4096 时Attention 的 KV Cache 占 1.2GB 显存。batch8 就爆显存吞吐只有 18 tokens/s。用 FlashAttention 后显存降到 180MB吞吐涨到 147 tokens/s。涨了 7 倍多不是算力变强了是 HBM 访问次数从 O(N²) 压到了 O(N)。很多人以为 FlashAttention 只是矩阵乘优化其实它的核心是显存访问模式的根本改变不存 QK^T 这个 N×N 的大矩阵改成分块算、边算边用中间结果走 L1 不落 HBM。Attention 的 O(N²) 显存问题标准 Attention 的计算公式Q x Wq K x Wk V x Wv S Q K^T / sqrt(d) P softmax(S) O P V问题出在 SQK^T这个矩阵。seq4096 时S 的大小是 4096×409616M 个元素。FP16 下占 32MB 显存。看起来不大但Prefill 阶段每个 batch 都要算一次 Sbatch8 就是 256MBDecode 阶段每生成一个 token 都要读整个 KV Cacheseq4096 时每个 token 读 57KB×4096233MB更严重的是S 矩阵要写回 HBMsoftmax 之后再读出来算 P×V。两次 HBM 读写延迟爆炸。实测数据Qwen2.5-7B910B 单卡FP16seq4096阶段HBM 访问量耗时占比QKV 投影38MB12%QK^T Softmax256MB35%P×V233MB28%其他-25%Attention 阶段占 63% 的时间其中 80% 是 HBM 读写。工程经验7B 模型推理时HBM 带宽利用率只有 35%。原因S 矩阵写 HBM 再读出来白白浪费带宽。开了 FlashAttention 后带宽利用率拉到 82%吞吐直接翻倍。FlashAttention 的 Tiling 策略FlashAttention 的核心思路不存 S 矩阵改成分块算。把 Q 分成 tile_q×d 的小块K 分成 tile_k×d 的小块V 分成 tile_k×d 的小块。每次只算一个 tile_q 的 Q 和一个 tile_k 的 K/Vfor each tile_q in Q: for each tile_k in K, V: S_tile Q_tile K_tile^T P_tile softmax(S_tile) O_tile P_tile V_tile关键点S_tile 和 P_tile 走 L1 缓存不落 HBM。为什么能省显存标准 Attention存整个 S 矩阵N×N显存 O(N²)FlashAttention只存 tile_q×tile_k 的 S_tile显存 O(tile_q×tile_k)seq4096 时标准 Attention 存 16M 元素FlashAttention 用 tile_q64、tile_k64只存 4K 元素。差了 4000 倍。为什么能省 HBM 访问标准 AttentionQ→HBMK→HBMS→HBMP→HBMV→HBMO→HBM6 次大矩阵读写FlashAttentionQ_tile 走 L1K_tile 走 L1S_tile 走 L1P_tile 走 L1V_tile 走 L1只有 O 最后写一次 HBMHBM 访问从 6 次降到 1 次。昇腾 NPU 的内存层次FlashAttention 在昇腾上的实现必须理解内存层次HBM高带宽内存 ↓ 带宽 1.2TB/s延迟 200ns L1 缓存每个 AI Core 独立 ↓ 容量 1MB带宽 ~10TB/s延迟 10ns L0A / L0BCube Unit 的输入缓冲区 ↓ 容量各 64KB Cube Unit矩阵乘单元 └─ MAC 阵列 16×16关键限制L1 只有 1MB。FlashAttention 的 tile_q×tile_k×dtype 必须小于 L1 容量否则溢出到 HBM性能暴跌。计算 tile 上限FP16S_tile tile_q × tile_k × 2 bytes P_tile tile_q × tile_k × 2 bytes Q_tile tile_q × d × 2 bytes K_tile tile_k × d × 2 bytes V_tile tile_k × d × 2 bytes O_tile tile_q × d × 2 bytes 总和 1MB 假设 d3584Qwen2.5-7B解出 tile_q × tile_k 16K实际选择tile_q64, tile_k64, block_size64tile_q×tile_k4096留足空间给其他 bufferCube/Vector 双缓冲流水线昇腾的 Cube Unit 算矩阵乘Vector Unit 算逐元素运算scale、mask、softmax。FlashAttention 要同时用 Cube 和 Vector关键是流水线设计。标准实现无流水线Cube: 算 Q×K^T → 等 Vector 算 softmax Vector: 等 Cube 算完 → 算 softmax → 等 Cube 算 P×V Cube: 等 Vector 算完 → 算 P×VCube 和 Vector 互相等空转 50% 时间。双缓冲流水线Cube: 算 Q1×K1^T → 算 Q1×K2^T → 算 P1×V1 → 算 P1×V2 → ... Vector: 等 Q1×K1^T → 算 softmax1 → 等 Q1×K2^T → 算 softmax2 → ...关键Cube 算下一个 tile 时Vector 在算上一个 tile 的 softmax。Cube 不等 VectorVector 不等 Cube。实测交叠率75%Cube 和 Vector 同时工作的时间占比工程经验Qwen2.5-7B 在 910B 上开双缓冲流水线后吞吐从 67 tokens/s 涨到 89 tokens/s33%。不开流水线Cube 等 Vector 占 40% 时间。tile_q / tile_k / block_size 参数调优这三个参数直接决定性能。tile_qQ 的分块大小太小32MAC 阵列填不满吞吐低太大128L1 装不下溢出到 HBM性能暴跌最优值64填满 MAC 阵列同时 L1 不溢出tile_kK/V 的分块大小太小32循环次数多调度开销大太大128L1 装不下最优值64block_sizesoftmax 的分块大小用于 Online Softmax影响 softmax 的数值稳定性通常等于 tile_k实测数据Qwen2.5-7B910B 单卡seq2048tile_qtile_k吞吐 (tokens/s)L1 溢出323272无646489无1286485轻微12812861严重tile_q64、tile_k64 最优。tile_q128 时 L1 开始溢出性能反而掉。工程经验tile_q 调优有个坑——不同 seq 的最优 tile_q 不同。seq1024 时 tile_q32 最快seq2048 时 tile_q64 最快。我们做了动态选择seq1024 用 tile_q32seq1024 用 tile_q64。与 GPU FlashAttention 的差异很多人以为FlashAttention 是通用的GPU 能跑 NPU 也能跑其实差异很大。维度GPUNVIDIANPU昇腾L2 缓存40-50MB全局共享无只有 L1L1 缓存128KB/SM1MB/AI Core执行单元SM统一Cube Vector分离最大 tile_q12864L1 限制流水线硬件自动调度软件显式编排核心差异 1GPU 的 L2 缓存大40-50MBtile_q 可以开到 128。昇腾 L1 只有 1MBtile_q 最大 64。核心差异 2GPU 的 SM 能同时跑矩阵乘和逐元素运算。昇腾 Cube 只能算矩阵乘Vector 只能算逐元素必须软件编排流水线。核心差异 3GPU 的 CUDA Stream 调度开销 1μs昇腾的 ACL 调用开销 12-15μs。FlashAttention 融合前 GPU 省 3 次 Kernel Launch~3μs昇腾省 3 次 ACL 调用~36μs。昇腾收益更大。性能收益总结模型优化前FlashAttention提升Qwen2.5-7B (seq2048)34 tokens/s89 tokens/s162%Qwen2.5-72B (seq4096, 4卡)320 TPS890 TPS178%DeepSeek-V3 (seq4096)580 TPS1420 TPS145%显存优化seq标准 AttentionFlashAttention节省2048580MB85MB-85%40961.2GB180MB-85%81924.8GB720MB-85%HBM 带宽利用率35%→82%省掉 S/P 矩阵的 HBM 读写踩坑实录坑 1短序列 FlashAttention 反而慢seq512 时FlashAttention 比 标准 Attention 慢 12%。原因Tiling 的调度开销比 HBM 访问省的时间还大。解决seq512 时不用 FlashAttention用标准 Attention。坑 2tile_q 开太大性能暴跌tile_q128 时L1 溢出到 HBM吞吐掉 30%。必须保证 tile_q×tile_k×dtype L1 容量的 80%。坑 3batch64 吞吐反而降FlashAttention 省了显存batch 能开到 64。但 KV Cache 太大开始 swap 到 Host 内存HBM 带宽利用率反而掉。解决batch 最大开到 32再大反而慢。https://atomgit.com/cann/ops-transformerhttps://atomgit.com/cann/ascend-transformer-boosthttps://atomgit.com/cann/cann-recipes-infer