企业官网建设流程全解析

oil-gas-ops-prospect融合注意力算子FusedBiasSoftmax揭秘：3次算子发射合成1个kernel

【免费下载链接】oil-gas-ops-prospect面向油气勘探（oil & gas exploration）领域的昇腾自定义算子库。仓名即 oil-gas-ops（油气算子）+ prospect（勘探目标），面向地震成像、全波形反演等勘探计算场景项目地址: https://gitcode.com/cann/oil-gas-ops-prospect

oil-gas-ops-prospect是面向油气勘探领域的昇腾（Ascend）自定义算子库，其中的FusedBiasSoftmax融合算子把注意力前向里常见的"两次 bias 累加 + 一次 softmax"共 3 次算子发射，合成了1 个 kernel，原地写回结果，显著减少中间访存与 kernel 发射开销。本文带你读懂它的原理、实现与调用方式。

🎯 一分钟看懂：FusedBiasSoftmax 算什么？

数学定义很简单：

$$y = \mathrm{softmax}(x + b_1 + b_2,\ \mathrm{dim}=-1)$$

这正是 openfold 结构预测中注意力核心的标准形态：logits 先加上mask bias（可广播）和triangle bias（全尺寸），再沿最后一维做 softmax。

⚡ 为什么要把 3 次发射合成 1 个 kernel？

在原版 openfold 的attn_core_inplace_cuda.forward_中，同样的语义要发射 3 次 kernel：

1. logits += bias1（add 算子）
2. logits += bias2（add 算子）
3. logits = softmax(logits)（in-place softmax）

每次发射都意味着：中间结果写回 HBM（主机内存）→ 下一次 kernel 再读出来，外加 kernel 发射本身的调度开销。而 bias 累加与 softmax 都是逐行独立、计算强度低的"访存密集"操作，单独发射时性能会被访存和 launch 开销主导。

FusedBiasSoftmax 的做法是：数据进片上 UB 后就不落地——bias 累加、fp32 域 softmax、写回，全部在一个 kernel 的流水内完成，中间访存被彻底消除。

🔧 单 kernel 是怎么实现的？

实现分为 host 侧切分（tiling）与 device 侧流水线两部分。

Host 侧：按行切分，预算 UB

切分核心算法在 fused_bias_softmax_tiling_core.h，逻辑直观：

• 按行均分：把B*H*Q行均匀分给若干 AIV 核，不能整除的余数行交给"大核"多处理 1 行（见ComputeRowSplit）；
• UB 预算模型：每行要占多少 UB，按"x 载入 + y 写出 + 非 fp32 中转 + bias 载入/中转"逐项计价，入队张量按 double buffer（队列深度 2）计两份，再为高阶 SoftMax API 预留 32KB；
• 32B 对齐：行宽S对齐到 32 字节，padding 用 0 填充，保证向量单元高效搬运；
• dtypes 分发：tiling key 取0=fp32、1=fp16、2=bf16，kernel 侧据此实例化不同模板。

算子注册与约束校验见 FusedBiasSoftmax.cpp，要求输入必须是 4D，且S = width满足1 ≤ S ≤ 4096。

Device 侧：双缓冲三拍流水线

kernel 实现在 FusedBiasSoftmax.h（入口为 fused_bias_softmax.cpp），主循环是一组经典的三拍流水：

CopyIn ──→ Compute ──→ CopyOut （双缓冲预取下一拍） （fp32 域 Add + SoftMax） （写回 GM）

• CopyIn：DataCopyPad把 x 与 bias 行搬运入 UB，同时预取下一迭代的数据，用队列（FBS_QUE_DEPTH = 2）隐藏搬运延迟；广播形态的 bias1 只载入迭代覆盖到的 distinct 行，省 UB；
• Compute：非 fp32 输入先Cast到 fp32 域，bias 逐块Add，再调用高阶SoftMaxAPI 完成行归一化——精度上等价于 fp32 累加，写回时再Cast回原 dtype；
• CopyOut：结果按原布局写回 GM。

因为y与x共享同一块 device 内存（原地写回），连一份中间 buffer 都省掉了。

🚀 Python 调用：3 行上手

Python 封装在 fused_bias_softmax.py，支持 PyTorch autograd 与 CPU 回退：

import os os.environ["OIL_GAS_OPS_FUSED_BIAS_SOFTMAX_IMPL"] = "ascendc" from oil_gas_ops_prospect.fused_bias_softmax import fused_bias_softmax, attention_core x = torch.randn(2, 4, 64, 128, dtype=torch.bfloat16, device="npu") b1 = torch.randn(2, 4, 1, 128, dtype=torch.bfloat16, device="npu") # 广播 b2 = torch.randn(2, 4, 64, 128, dtype=torch.bfloat16, device="npu") # 全尺寸 y = fused_bias_softmax(x, b1, b2) # 原地写回 x

几个实用细节：

• 实现切换：环境变量OIL_GAS_OPS_FUSED_BIAS_SOFTMAX_IMPL取值torch（默认）/ascendc/auto，张量不在 NPU 上时自动回退到纯 PyTorch 参考实现，CPU 上也能跑通同一套代码；
• attention_core 组合封装：attention_core(q, k, v, bias1, bias2)等价 openfold 的AttentionCoreFunction.apply，前向走FusedBiasSoftmax，反向走配套融合算子FusedSoftmaxGrad；
• aclnn 直接调用：C++ 侧接口为aclnnFusedBiasSoftmax，样例见 test_aclnn_fused_bias_softmax.cpp。

✅ 精度与性能如何验证？

测试覆盖三层，均有现成脚本：

精度契约：fp32max_diff ≤ 1e-5；bf16 契约式atol=2e-2, rtol=4e-3。性能基准对 torch 参考实现逐 case 打印speedup（warmup 5、iter 20），覆盖[2,4,64,128]~[8,8,256,256]的典型注意力形状。运行入口：

bash build.sh -u --st --ops=fused_bias_softmax # 多 shape aclnn 系统测试

📚 延伸阅读

• 算子数学定义与约束全文：docs/zh/op_list.md
• 接口签名与环境变量：docs/zh/api_list.md
• 反向配套算子 FusedSoftmaxGrad 源码：ascendc/operators/fused_softmax_grad/
• 算子开发指南（想写自己的融合算子？）：docs/zh/develop/operator_development_guide.md

想动手跑起来？先克隆仓库（支持 CANN 9.0.0、Ascend 910B）：

git clone https://gitcode.com/cann/oil-gas-ops-prospect.git cd oil-gas-ops-prospect && bash build.sh --install

一句话总结：FusedBiasSoftmax 用"host 侧按行切分 + device 侧双缓冲三拍流水 + fp32 域中间计算"三板斧，把注意力前向的 3 次 kernel 发射压成 1 次，中间数据不落 HBM，是昇腾上做访存密集小算子融合的一个标准范式。

【免费下载链接】oil-gas-ops-prospect面向油气勘探（oil & gas exploration）领域的昇腾自定义算子库。仓名即 oil-gas-ops（油气算子）+ prospect（勘探目标），面向地震成像、全波形反演等勘探计算场景项目地址: https://gitcode.com/cann/oil-gas-ops-prospect