CANN 8.5 在 2024 Q4 发布ops-transformer 仓库跟进了三个重要更新FlashAttention V2 的反向传播融合、MC2 通算融合的多卡拓扑适配、以及新增的 GroupedAttention 算子。如果你已经用 CANN 8.0 在跑大模型推理这篇帮你判断要不要升级。FlashAttention V2反向传播终于融合了CANN 8.0 的 FlashAttention 只融合了前向传播。训练场景下反向传播还是要拆成三个独立 kerneldQ、dK、dV中间结果落显存。CANN 8.5 把反向传播也融合进去了——一次 kernel 完成 dQ/dK/dV 的计算。训练场景的收益配置训练吞吐 (tokens/s/p)显存占用CANN 8.0 FlashAttention V11,82056 GBCANN 8.5 FlashAttention V22,41044 GB训练吞吐提升 32%显存省了 21%。显存省下来意味着可以开更大的 batch 或更长的序列。前向推理不受影响——如果你只做推理这个更新对你没用不用特意升级。MC2 通算融合多卡拓扑适配CANN 8.0 的 MC2 只支持同一台服务器内的卡间通信。8 卡的 Atlas 800I A2 没问题但如果你要用两台服务器做 16 卡的 MoE 训练跨机 All-to-All 通信走的是 RoCEMC2 没法把它和计算重叠。CANN 8.5 加了 RoCE 通算融合支持。MC2 可以同时管理 HCCL 的卡间通信和 RoCE 的跨机通信让两者都跟计算流水线重叠。实测数据Mixtral 8x7B 的 16 卡训练配置通信占比吞吐 (tokens/s/p)CANN 8.0跨机不通算融合42%680CANN 8.5跨机通算融合23%1,050通信占比从 42% 降到 23%跨机场景的 MC2 终于能用了。新增GroupedAttention 算子Grouped-Query AttentionGQA在 Llama2、Mistral 等模型里广泛使用。CANN 8.0 需要把 GQA 展开 MHA 来跑CANN 8.5 新增了原生 GQA 支持importtorch_npu# GQA: num_q_heads32, num_kv_heads8qtorch.randn(1,32,4096,128,devicenpu,dtypetorch.float16)ktorch.randn(1,8,4096,128,devicenpu,dtypetorch.float16)vtorch.randn(1,8,4096,128,devicenpu,dtypetorch.float16)# CANN 8.5 直接支持 KV heads Q headsouttorch_npu.npu.flash_attention(q,k,v)CANN 8.0 要手动把 K/V repeat 到 32 个 head显存和计算都浪费。原生 GQA 省掉了 repeat 操作显存占用降低 75%延迟降 15-20%。升级建议场景是否建议升级到 8.5只做推理单机不急8.0 够用推理 GQA 模型建议升级原生 GQA 收益大训练单机建议升级FlashAttention V2 反向融合省显存训练多机 MoE必须升级跨机 MC2 是刚需升级方式# 更新 CANN toolkit./Ascend-cann-toolkit_8.5.run--install# 重新编译 ops-transformercdops-transformergitpullbashbuild.sh# 更新 torch_npupipinstalltorch_npu2.3.0# CANN 8.5 对应版本兼容性注意CANN 8.5 的 FlashAttention V2 API 跟 8.0 的 V1 有个不兼容变更npu.flash_attention的scale参数从位置参数改成了关键字参数。如果你之前的代码是flash_attention(q, k, v, 1.0/math.sqrt(dim))需要改成flash_attention(q, k, v, scale1.0/math.sqrt(dim))。不改的话会报参数类型错误。如果你的 MoE 训练要上多机CANN 8.5 的跨机 MC2 是硬需求不升级就是浪费卡。单机推理用户可以观望等下一个大版本再看。仓库在这里https://atomgit.com/cann/ops-transformer
