Transformer 推理的瓶颈不在矩阵乘法，在内存带宽。LayerNorm → MatMul → GELU → MatMul → LayerNorm，每两个算子之间要写一次 HBM、再读一次 HBM，带宽瓶颈比计算瓶颈更致命。ATB（Ascend Transformer Boost）做的是把这堆小算子融合成一个大算子，中间结果塞在片上 SRAM，不落 HBM。

昇腾NPU 的达芬奇架构有 Cube 核（矩阵运算）和 Vector 核（逐元素运算），两种核之间的数据搬运走 SRAM，延迟比 HBM 低两个数量级。融合策略就是把 Cube 和 Vector 的计算捏成一段，不让中间结果回 HBM。

---

融合策略的物理意义

不融合的推理路径

输入 hidden_states (HBM)
    ↓ 读 HBM（2000 GB/s 带宽）
    ↓ LayerNorm（Vector 核）
    ↓ 写 HBM（中间结果）
    ↓ 读 HBM
    ↓ MatMul（Cube 核）
    ↓ 写 HBM（中间结果）
    ↓ 读 HBM
    ↓ GELU（Vector 核）
    ↓ 写 HBM（中间结果）
    ↓ 读 HBM
    ↓ MatMul（Cube 核）
    ↓ 写 HBM（输出）

HBM 访问次数：读 4 次 + 写 4 次 = 8 次 HBM 往返。

每层的延迟里，HBM 访问占 60-70%，Cube/Vector 计算只占 30-40%。

融合后的推理路径

输入 hidden_states (HBM)
    ↓ 读 HBM（1 次）
    ↓ LayerNorm + MatMul + GELU + MatMul（融合成一个 Kernel）
    ↓ 中间结果存在 SRAM（不落 HBM）
    ↓ 写 HBM（1 次，输出）

HBM 访问次数：读 1 次 + 写 1 次 = 2 次 HBM 往返。

HBM 访问减少 75%，延迟对应下降 30-50%（取决于模型宽度和序列长度）。

---

ATB 支持的典型融合模式

1. MLP 融合（最常用）

Transformer 的 FFN 层是 LayerNorm → Dense → GELU → Dense → Add（残差连接）。

ATB 把整个 FFN 融合成一个 Kernel：

# 不融合（PyTorch 原生，每个算子独立 Kernel）
def ffn_nofusion(hidden_states):
    # 1. LayerNorm（独立 Kernel，写 HBM）
    x = torch.nn.functional.layer_norm(hidden_states, ...)
    # 2. 第一个 Dense（独立 Kernel，读 HBM + 写 HBM）
    x = torch.matmul(x, w1) + b1
    # 3. GELU（独立 Kernel，读 HBM + 写 HBM）
    x = torch.nn.functional.gelu(x)
    # 4. 第二个 Dense（独立 Kernel，读 HBM + 写 HBM）
    x = torch.matmul(x, w2) + b2
    # 5. 残差 Add（独立 Kernel，读 HBM + 写 HBM）
    return x + hidden_states

# ATB 融合（一个 Kernel，中间结果在 SRAM）
import torch_npu
from atb import MlpLayer  # ATB 的 MLP 层

def ffn_fused(hidden_states):
    mlp = MlpLayer(
        in_features=4096,
        intermediate_size=11008,  # LLaMA-2 7B 的 FFN 中间维度
        activation='gelu',
        fuse_layernorm=True,      # ← 关键：融合 LayerNorm
        fuse_residual=True         # ← 关键：融合残差连接
    ).npu()
    
    # 一次 Kernel 调用，完成整个 FFN
    return mlp(hidden_states)

Fusion 的触发条件（ATB 自动判断）：

条件	说明
fuse_layernorm=True	LayerNorm 和第一个 Dense 融合
fuse_residual=True	残差 Add 和最后一个 Dense 融合
activation='gelu' 或 'silu'	GELU/SiLU 和第二个 Dense 融合
intermediate_size <= 16384	中间维度太大，SRAM 放不下，自动取消融合

2. Attention 融合

Transformer 的 Attention 层是 QKV 生成 → Softmax → Dropout → 输出投影。

ATB 把整个 Attention 融合成一个 Kernel（包括 FlashAttention 的优化）：

# ATB 的 Attention 融合配置
from atb import AttentionLayer

attn = AttentionLayer(
    hidden_size=4096,
    num_heads=32,
    max_seq_len=2048,
    fuse_qkv=True,          # ← QKV 三个矩阵乘融合成一个 Kernel
    fuse_softmax=True,       # ← Softmax 和 Dropout 融合
    use_flash_attention=True,# ← 用 FlashAttention（减少 HBM 访问）
    sparse_head_skipping=True # ← 稀疏注意力（可选，LLaMA-2 70B 用）
).npu()

# 一次 Kernel 调用，完成整个 Attention
output = attn(hidden_states, attention_mask)

FlashAttention 和融合的关系：

FlashAttention 本身就是一种融合策略（把 Softmax 和 Matrix Multiply 融合，不存完整的 Attention Matrix 到 HBM）。ATB 在 FlashAttention 的基础上再融 QKV 生成和输出投影。

3. LayerNorm 融合（独立开关）

LayerNorm 是 Transformer 里调用次数最多的算子（每个 Layer 2 次）。ATB 支持把 LayerNorm 融到前一个算子和后一个算子：

# LayerNorm 的融合方向
# 前向：... → LayerNorm → MatMul → ...
#         ↑融合        ↑ 不融（跨 Kernel 边界）
# 后向：... → MatMul → LayerNorm → ...
#                  ↑融合   ↑ 不融

# ATB 配置
mlp = MlpLayer(
    fuse_layernorm='both',  # ← 'pre'（融前面）/ 'post'（融后面）/ 'both'
    ...
)

为什么不融所有的 LayerNorm：LayerNorm 的计算需要全局统计量（均值和方差），如果前一个算子输出太大（超过 SRAM 容量），LayerNorm 只能单独跑。

4. Bias 融合（可选）

Dense 层的 Bias 加法可以融到前面的激活函数里（GELU 或 SiLU）：

# 不融合：GELU 输出 → 读 HBM → BiasAdd → 写 HBM → MatMul
# 融合：GELU + BiasAdd 在一个 Kernel 里完成

dense = DenseLayer(
    in_features=4096,
    out_features=11008,
    bias=True,
    fuse_bias_with_activation=True  # ← Bias 融到 GELU/SiLU
)

收益：小（Bias Add 的计算量很小），但积少成多（每层省 1 次 HBM 读写）。

---

融合策略的配置方法

ATB 的融合策略是自动的（默认开），但可以手动调。

全局配置（环境变量）

# 开所有融合（默认）
export ATB_OP_FUSION=1

# 关所有融合（Debug 用）
export ATB_OP_FUSION=0

# 只开 MLP 融合，关 Attention 融合
export ATB_FUSE_MLP=1
export ATB_FUSE_ATTENTION=0

逐层配置（Python API）

from atb import MlpLayer, AttentionLayer

# MLP 层：手动配融合策略
mlp = MlpLayer(
    in_features=4096,
    intermediate_size=11008,
    activation='silu',          # LLaMA 用 SiLU，BERT 用 GELU
    fuse_layernorm=True,        # 融 LayerNorm
    fuse_residual=True,          # 融残差连接
    fuse_bias_with_activation=True  # 融 Bias
).npu()

# Attention 层：手动配融合策略
attn = AttentionLayer(
    hidden_size=4096,
    num_heads=32,
    fuse_qkv=True,              # 融 QKV 生成
    fuse_softmax=True,          # 融 Softmax
    use_flash_attention=True    # 用 FlashAttention
).npu()

融合策略的冲突处理

有些融合策略冲突（不能同时开）：

冲突组合	行为
fuse_layernorm=True + fuse_residual=True	自动合并（融整个 FFN）
fuse_qkv=True + fuse_softmax=False	警告（QKV 融了但 Softmax 没融，收益减半）
use_flash_attention=True + sparse_head_skipping=False	不冲突，FlashAttention 照常跑

ATB 会自动检测冲突，打印警告日志（不崩溃）。

---

融合策略的性能收益

延迟对比（LLaMA-2 7B，Batch=1，SeqLen=128）

融合策略	Prefill 延迟 (ms)	Decode 延迟 (ms/tok)
无融合（PyTorch eager）	~120	~35
只融 MLP	~95 (-21%)	~28 (-20%)
只融 Attention	~105 (-13%)	~25 (-29%)
MLP + Attention 都融	~75 (-38%)	~20 (-43%)

关键：Decode 阶段的收益比 Prefill 大（Decode 是 Memory-Bound，融合减 HBM 访问的效果更明显）。

融合策略的适用场景

场景	推荐融合策略	原因
短序列（≤ 512）推理	MLP + Attention 都融	SRAM 够放中间结果
长序列（≥ 2048）推理	只融 Attention	MLP 的中间结果太大，SRAM 放不下
训练	不融（或只融 LayerNorm）	训练需要存中间结果算梯度
低延迟要求（< 100ms）	MLP + Attention 都融	减 Kernel 调用次数
高吞吐要求（> 100 QPS）	只融 Attention	MLP 融合增加 Kernel 长度，可能降低 Occupancy

---

融合策略的调试方法

1. 看融合是否生效（日志）

# 开 ATB 的融合日志
export ATB_LOG_LEVEL=INFO

# 运行推理，看日志
python infer.py
# 输出示例：
# [INFO] MLP fusion enabled: LayerNorm + Dense + GELU + Dense + Residual
# [INFO] Attention fusion enabled: QKV + Softmax + OutputProjection
# [WARN] MLP fusion skipped: intermediate_size=32768 > SRAM capacity

2. 逐层关闭融合（定位问题）

# 如果某一层融合后结果不对，逐层关融合定位
mlp = MlpLayer(
    fuse_layernorm=False,   # 先关 LayerNorm 融合
    fuse_residual=False,     # 再关残差融合
    fuse_bias_with_activation=False  # 最后关 Bias 融合
)

# 哪一层关了融合后结果对了，就是那一层的融合有 Bug

3. 性能 Profiling（Timeline）

from atb.utils import profile

# Profiling 融合前后的 Kernel 调用次数
with profile() as prof:
    output = model(input_ids)

# 输出 Timeline（看每个 Kernel 的耗时）
prof.export_timeline("timeline.json")

正常情况：融合后 Kernel 调用次数减半（从 ~50 降到 ~25）。

---

融合策略的硬件限制

SRAM 容量限制

昇腾NPU 的片上 SRAM 是 16MB（Ascend 910）。融合后的中间结果必须能塞进 16MB。

计算 SRAM 占用：
  MLP 融合：intermediate_size * hidden_size * 2 (FP16) ≤ 16MB
  → intermediate_size ≤ 16384（LLaMA-2 7B 是 11008，够）

  Attention 融合：seq_len * num_heads * head_dim * 4 (FP16) ≤ 16MB
  → seq_len ≤ 2048（LLaMA-2 7B 是 2048，够）

超了会怎么样：ATB 自动取消融合（回退到不融合），并打印警告日志。

Cube/Vector 核的流水线冲突

融合 Kernel 里同时有 Cube 指令（矩阵乘）和 Vector 指令（逐元素）。如果 Cube 和 Vector 的流水没排好，会互相抢 SRAM 带宽。

ATB 的融合 Kernel 用了双缓冲（Double Buffering）——Cube 算当前 Tile，Vector 算上一个 Tile，不冲突。

---

如果你的模型推理延迟太高（> 50ms/tok），先开 ATB 的融合策略——export ATB_OP_FUSION=1，不用改模型代码。

如果融合后结果不对，逐层关融合定位问题（先关 LayerNorm 融合，再关残差融合，最后关 Bias 融合）。

融合策略的配置示例在 cann/ascend-transformer-boost 仓库的 examples/fusion/ 目录下。atb/python/atb/layers/ 目录下是 MlpLayer 和 AttentionLayer 的完整实现。

ATB 的代码在 AtomGit 的 cann/ascend-transformer-boost 仓库：

https://atomgit.com/cann/ascend-transformer-boost

examples/fusion/ 目录下是 6 个融合策略的参考配置。docs/fusion_strategy.md 是融合策略的完整文档。