昇腾CANN上FlashAttention的工程实践：catlass模板调优全记录

去年我们把一个13B参数的推理服务从GPU迁移到昇腾NPU，attention部分从标准实现换成catlass模板的FlashAttention，吞吐从1,200 tokens/s提到4,800 tokens/s。但这个过程不是"换个模板就完事"——数据布局、精度对齐、分块策略、算子融合，每一步都有坑。今天把整个调优过程记录下来，包含具体的配置参数和实测数据。

背景：为什么选catlass？

catlass不是CUTLASS的昇腾移植版，它是昇腾CANN体系内的算子模板库，定位是给开发者提供高性能算子的开发骨架。ops-nn、ops-math、ops-blas这些算子仓库底层都依赖catlass的模板。

选catlass而不是直接写Ascend C算子，原因很简单：手动写一个达芬奇架构上高性能的FlashAttention，你需要处理分块加载、Unified Buffer管理、bank conflict规避、流水线调度……一个人搞可能要一两个月。catlass模板把这些封装好了，你只需要调参数。

但"调参数"这三个字背后的事也不少。

精度选择：FP16还是BF16？

第一个决策点。昇腾910支持FP16和BF16两种半精度，catlass模板两种都支持。选择依据：

我们的场景是推理，softmax中间值容易爆FP16的范围。实测数据：

# FP16 FlashAttention，序列长度8192 [ERROR] softmax overflow detected, batch=2 head=15 tile_m=48 # 17个tile中有3个溢出，输出NaN # BF16 FlashAttention，同样配置 [PASS] no overflow, max softmax value = 1.2e+38 # 所有tile正常

所以长序列场景直接用BF16，省去溢出排查的麻烦。短序列（2048以内）FP16精度更好，推理结果跟FP32的误差更小。我们的折中方案：4K以内FP16，4K以上BF16。

FlashAttnConfig config; if (seq_len <= 4096) { config.use_fp16 = true; } else { config.use_fp16 = false; // 启用BF16 }

分块策略：不是越大越好

catlass模板的核心参数是block_m和block_n，控制Q和K/V的分块大小。直觉上block越大，并行度越高，性能越好。但达芬奇架构的约束不允许你无限加大：

约束1：Unified Buffer容量

达芬奇架构的Unified Buffer大约256KB（具体大小随芯片版本略有差异）。一个tile的数据量 = block_m × head_dim × sizeof(data_type) × 3（Q+K+V）。加上中间变量，实际占用大概是这个值的2-3倍。

block_m=128, head_dim=128, FP16: 单tile = 128 × 128 × 2 × 3 = 96KB 加上softmax统计量和O的累加buffer ≈ 200KB → 勉强能塞进去 block_m=256, head_dim=128, FP16: 单tile = 256 × 128 × 2 × 3 = 192KB 加上中间变量 ≈ 420KB → 超了

超了会怎样？catlass模板不会报错，而是自动降级——把一个tile拆成多次加载，性能反而比block_m=128更差。

约束2：K/V的复用模式

FlashAttention的outer loop是沿M方向（Q的序列方向）遍历，inner loop是沿N方向（K/V的序列方向）。每个Q的tile要跟所有K/V的tile做计算。所以K/V的tile会被反复加载，block_n越大，单次加载的数据量越大，但加载次数越少。

block_n=64比128快，因为小tile的cache命中率更高。block_n=32太碎了，调度开销吃掉了cache收益。block_m=256+block_n=64是最优组合，但要确认Unified Buffer够用。

数据布局：这步做错后面全白搭

catlass模板要求输入数据的layout是[batch, heads, seq_len, head_dim]，row-major存储，stride必须128字节对齐。PyTorch默认的tensor layout恰好满足，但如果你从其他框架（MindSpore、Paddle）传入数据，大概率layout不一样。

我们踩过的坑：MindSpore的attention输入layout是[batch, seq_len, heads, head_dim]，直接传给catlass模板，结果不对，但也不报错。数值偏了大概5%，肉眼不容易看出来，端到端推理结果就是差一截。

import torch_npu def ensure_layout(tensor, target_layout="BSHD"): """确保tensor的layout符合catlass要求""" current_layout = detect_layout(tensor) # 根据stride判断 if current_layout == "BSHD" and target_layout == "BHSD": # [batch, seq, heads, dim] -> [batch, heads, seq, dim] tensor = tensor.transpose(1, 2).contiguous() elif current_layout == "BHSD" and target_layout == "BSHD": tensor = tensor.transpose(1, 2).contiguous() # 128字节对齐检查 assert tensor.stride(0) % 128 == 0, f"stride未对齐: {tensor.stride(0)}" return tensor

另一个容易忽略的点：contiguous()。transpose之后tensor不再连续，必须调contiguous()才会真正重排内存。不调的话，catlass模板读到的数据是乱的。

Causal Mask的实现差异

自回归推理必须用causal mask，每个位置只能看到之前的token。catlass模板的causal实现有两种模式：

模式1：下三角mask矩阵

显式构造一个下三角bool矩阵，传入kernel。优点是通用，缺点是占用O(N²)显存——跟标准attention一样的毛病。

模式2：对角线跳过

kernel内部根据tile坐标判断哪些计算可以跳过。不需要额外显存，而且能跳过大量无效计算。

// catlass模板内部的对角线跳过逻辑（简化版） for (int tile_n = 0; tile_n < num_kv_tiles; tile_n++) { // 当前Q tile的行范围: [tile_m * block_m, (tile_m+1) * block_m) // 当前K tile的列范围: [tile_n * block_n, (tile_n+1) * block_n) if (causal && tile_n * block_n > (tile_m + 1) * block_m) { // 这个K tile完全在mask之外，跳过 continue; // 长序列时能跳过约50%的tile } // 加载K/V tile，做局部attention计算 load_kv_tile(k_tile, v_tile, tile_n); compute_local_attention(q_tile, k_tile, v_tile, o_tile); }

对角线跳过的收益跟序列长度正相关。序列越长，能跳过的tile越多：

收益随序列增长而增加，因为跳过计算的占比不变，但省下来的显存带宽可以用于有效计算。16384序列时，causal模式的吞吐比non-causal模式还高15%，这就是跳过无效计算的回报。

跟GE图引擎的融合优化

单算子调优到4,200 tokens/s之后，还有一档免费性能：算子融合。昇腾CANN的GE图引擎能自动把FlashAttention和相邻算子合并执行。

融合的前提是算子都走GE的图模式。如果你用AscendCL的单算子API调用FlashAttention，GE没法做融合。必须把整个模型编译成图：

import torch_npu from torch_npu.contrib import transfer_to_npu # 模型迁移到NPU，自动走GE图模式 model = model.npu() # GE日志确认融合 import os os.environ["GE_OPTYPE_BLACKLIST"] = "" # 清空黑名单，允许所有融合 os.environ["DUMP_GE_GRAPH"] = "1" # 导出GE图 # 推理一次，触发图编译 with torch.no_grad(): output = model(input_ids) # 检查融合结果 # 日志路径：/usr/local/Ascend/ascend-toolkit/latest/xx/dump/ # 搜索关键词："FlashAttention" "Fuse"

融合前后GE图的对比：

融合前（6个独立算子）： RMSNorm → MatMul(Q) → MatMul(K) → MatMul(V) → FlashAttention → MatMul(O) 融合后（2个融合算子）： FusedNormQKV(RMSNorm + MatMul Q/K/V) → FusedAttnProj(FlashAttention + MatMul O)

显存读写次数从12次降到4次，吞吐从4,200提到4,860 tokens/s。

反向传播的特殊处理

推理服务只跑前向，但如果你的场景是训练或finetune，FlashAttention的反向也需要catlass模板。反向有个额外参数：deterministic。

FlashAttnBwdConfig bwd_config; bwd_config.deterministic = false; // 非确定性模式，用atomic add bwd_config.deterministic = true; // 确定性模式，用排序累加

非确定性模式快15%左右，但梯度在多卡之间可能有微小差异（FP16的atomic add不满足交换律）。对训练来说这点差异通常不影响收敛，但如果你在做数值对比测试，建议开确定性模式。

完整调优结果

我们的13B模型在Ascend 910上的端到端性能变化：

从1,200到4,860，整体提升4倍。其中catlass模板贡献最大（3.5x），参数调优贡献7%，GE融合贡献8%。

想在自己的昇腾NPU上复现这些数据？去AtomGit拉catlass仓库：

https://atomgit.com/cann/catlass

建议先把examples目录下的FlashAttention示例跑通，确认环境没问题。然后对照本文的参数表逐步调优。如果遇到精度问题，先用BF16排除溢出，再逐步切回FP16。cann-recipes-train仓库里有FlashAttention在训练场景下的完整集成方案，包括反向传播和多卡并行。