汤不热吧1. 为什么大模型推理会受限于 HBM？
在大模型（LLM）推理过程中，Transformer 架构中存在大量的逐元素（Element-wise）操作，如 Add、LayerNorm、Mul 和 Silu 等。在标准的执行流中，每个算子都需要从高带宽内存（HBM）读取输入数据，计算完成后将结果写回 HBM。
这种频繁的“访存-计算-回写”模式会导致大量的 I/O 开销，使得 AI 芯片的算力利用率（TFLOPS）远低于理论峰值。华为昇腾（Ascend）算子融合技术的核心思想，就是通过 CANN（Compute Architecture for Neural Networks） 的图编译器，将多个连续的算子在图准备阶段合并为一个大算子，让数据在 AI Core 的内部缓冲区（Unified Buffer）中完成多次连续计算，从而大幅减少与 HBM 的交互次数。
2. 核心原理：Buffer 级融合
在昇腾 NPU 上，融合分为两个层面：
1. UB 级融合：将多个连续的操作合并，在计算单元内部的 Unified Buffer 中流转，无需写回 HBM。
2. L1/L2 缓存优化：优化张量在片上缓存的生命周期，减少对片外内存的依赖。
3. 实操演示：如何开启并验证算子融合
在使用 torch_npu 进行大模型迁移后，昇腾提供了多种方式来开启自动融合。
3.1 环境准备
确保你已安装 torch 和 torch_npu：
pip install torch-npu
3.2 启用图执行模式（Graph Mode）
默认情况下，PyTorch 是 Eager 模式。为了发挥昇腾算子融合的最大潜力，我们推荐使用华为提供的 jit_compile 或配合昇腾的 Graph Engine 执行模式。
import torch
import torch_npu
from torch_npu.contrib import transfer_to_npu
# 设置编译配置，开启算子并行与融合优化
option = {
\"ge.exec.enableOperatorFusion\": \"1\",
\"ge.common_graph_optimizer\": \"1\",
\"ge.fusion_optimizer_config_path\": \"./fusion_rule.cfg\" # 可选：自定义融合规则
}
torch_npu.npu.set_option(option)
# 定义一个简单的 Transformer Block 片段（Add + LayerNorm + Silu）
class SimpleBlock(torch.nn.Module):
def __init__(self, dim):
super().__init__()
self.norm = torch.nn.LayerNorm(dim)
def forward(self, x, residual):
# 在 Eager 模式下，这里是 3 次 HBM 读写
# 在 Graph 模式下，CANN 会将其融合为 1 个融合算子
return torch.nn.functional.silu(self.norm(x + residual))
model = SimpleBlock(4096).npu()
x = torch.randn(2, 2048, 4096).npu()
res = torch.randn(2, 2048, 4096).npu()
# 首次运行触发编译优化
with torch.npu.graph(torch_npu.npu.Stream()):
output = model(x, res)
4. 进阶：使用 Profiling 验证融合效果
如何确定你的算子真的融合了？我们可以使用昇腾的 Profiling 工具查看算子流水。
- 在脚本中加入 Profiling 代码：
with torch_npu.profiler.profile(
on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./log\"),
record_shapes=True,
with_stack=True
) as prof:
output = model(x, res)
prof.step()
- 打开生成的 trace.json，在 Chrome 浏览器中输入 chrome://tracing 加载文件。
- 对比结果：
- 优化前：你会看到 Add、LayerNorm、Silu 分布在不同的时间片，中间有明显的间隙。
- 优化后：这些算子被合并为一个类似 TbeFusedAddLayerNormSilu 的任务，执行时间显著缩短。
5. 常见问题：为什么有些算子没融合？
- Shape 变化：动态 Shape 可能会导致融合失效。尽量使用静态 Shape 或在 CANN 中开启 ge.dynamic_dims 优化。
- 算子不连续：如果在中间插入了 print 或不支持 NPU 加速的自定义 CPU 操作，会强制触发数据同步回写（Sync），打断融合链。
- 内存碎片：如果 HBM 剩余空间不足，编译器可能会限制融合以减小中间张量的生命周期压力。
总结
算子融合是昇腾大模型推理优化的“必修课”。通过 torch_npu 开启图模式优化，可以将原本被 I/O 阻塞的算力释放出来，让大模型在国产芯片上跑得更稳、更快。