当前位置：首页 > news >正文

昇腾NPU 的“后厨五人组“：CANN 架构原理一把抓

news 2026/5/23 20:04:09

昨天晚上和做推理引擎的朋友吃饭，他问我：“昇腾NPU 的软件栈为什么这么复杂？我写CUDA 的时候，直接调用cuBLAS 就行，你们为什么要分五层？”

我想了一下，跟他说：“你把NPU 想象成一个超大的后厨。做一道菜（跑一个模型）需要五组人配合——每组人只做自己擅长的事，但缺了哪组都不行。”

他眼睛亮了：“继续说。”

第一组人：控制面板（Framework Adaptor）

后厨门口站着一个服务员，负责把客人点的菜（PyTorch 模型）翻译成后厨能看懂的菜单（NPU 能执行的计算图）。这个人就是控制面板。

你在PyTorch 里写nn.Linear，控制面板把它翻译成NPU 能看懂的matmul调用。你在PyTorch 里写scaled_dot_product_attention，控制面板把它路由到ops-transformer 的FlashAttention 实现。

# 控制面板的工作（简化版）importtorchimporttorch_npu# 你在PyTorch 里写的代码Q=torch.randn(4,32,2048,64,dtype=torch.float16).npu()K=torch.randn(4,32,2048,64,dtype=torch.float16).npu()V=torch.randn(4,32,2048,64,dtype=torch.float16).npu()# 控制面板把它路由到NPU 的算子# 如果装了torch-npu，这里会调用NPU 的scaled_dot_product_attentionoutput=torch.nn.functional.scaled_dot_product_attention(Q,K,V,is_causal=True)# 验证：看看控制面板是不是真的路由到了NPU 算子print(f"Q 的设备:{Q.device}")# 应该是 npu:0print(f"output 的设备:{output.device}")# 应该是 npu:0# 如果都是 npu:0，说明控制面板工作正常

控制面板不做事，只做路由。但它的路由规则决定了你的模型能不能在NPU 上跑起来。

第二组人：酱料包仓库（AOL 算子库，ops-transformer 在这）

后厨里有一个大仓库，里面放满了配好的酱料包。每个酱料包都是一个优化好的算子实现——你不用自己调调料，拆包就能用。

这个仓库就是AOL（Ascend Operator Library）算子库。ops-transformer 是其中的一个子仓库，专门放Transformer 架构需要的算子（FlashAttention、LayerNorm、RoPE 等）。

# 酱料包仓库的工作（ops-transformer 示例）fromflash_attention_opsimportflash_attention_npu# 你不用自己实现FlashAttention，直接从仓库拿Q=torch.randn(4,32,2048,64,dtype=torch.float16).npu()K=torch.randn(4,32,2048,64,dtype=torch.float16).npu()V=torch.randn(4,32,2048,64,dtype=torch.float16).npu()output=flash_attention_npu(Q,K,V,causal=True)# 这个flash_attention_npu 就是从酱料包仓库拿的# 它已经在昇腾NPU 上高度优化过了# 验证：和PyTorch 原生实现对比误差output_native=torch.nn.functional.scaled_dot_product_attention(Q,K,V,is_causal=True)max_err=(output.cpu().float()-output_native.cpu().float()).abs().max().item()print(f"酱料包 vs 原生实现最大误差:{max_err:.6f}")print("误差 < 1e-3，酱料包质量合格！"ifmax_err<1e-3else"误差过大，检查酱料包！")

酱料包仓库的价值：你不用懂算子怎么优化，只要会拆包（调用接口），就能拿到接近最优的性能。

第三组人：拼菜师傅（GE 图引擎）

后厨里有一个拼菜师傅，他的工作是把多个步骤合并成一个步骤执行。比如菜单上写"炒A→炒B→炒C"，拼菜师傅发现A、B、C 可以合并成一个步骤，就自动把它们拼在一起。

这个拼菜师傅就是GE（Graph Engine）图引擎。它在编译期分析你的计算图，发现可以融合的算子序列，就把它们融合成一个算子执行。

# 拼菜师傅的工作（GE 融合）# 你写的PyTorch 代码importtorch Q=torch.randn(4,32,2048,64,dtype=torch.float16).npu()K=torch.randn(4,32,2048,64,dtype=torch.float16).npu()V=torch.randn(4,32,2048,64,dtype=torch.float16).npu()# 这三个算子（MatMul → Softmax → MatMul）会被GE 识别到# 如果符合条件（dtype= float16, seq_len=2的幂次方），GE 会把它们融合成一个FlashAttentionKerneloutput=torch.nn.functional.scaled_dot_product_attention(Q,K,V,is_causal=True)# 验证：用Profiler 看GE 有没有拼菜fromtorch_npu.profilerimportprofile,ProfilerActivitywithprofile(activities=[ProfilerActivity.NPU],export_name="ge_fusion_check.json"):output=torch.nn.functional.scaled_dot_product_attention(Q,K,V,is_causal=True)torch.npu.synchronize()print("Profiler trace 已保存到 ge_fusion_check.json")print("打开这个文件，搜索 FlashAttentionKernel——如果搜得到，说明拼菜师傅工作了")

拼菜师傅的价值：减少算子之间的数据搬运（不用把中间结果写回HBM），从而加速训练。

第四组人：催菜员（Runtime 运行时）

后厨里有一个催菜员，他的工作是协调每个菜的出餐顺序：哪些菜可以同时做、哪些菜必须按顺序做、哪些菜要先做因为后面还有几道在等。

这个催菜员就是Runtime。它负责把GE 融合后的计算图，拆成一个个任务，调度到NPU 的计算单元上。它还负责数据搬运：一边让当前任务的计算进行，一边把下一个任务的数据提前搬到计算单元旁边（overlap）。

# 催菜员的工作（Runtime 调度）# 你不用直接调用Runtime，但你可以通过npu-smi 看它的工作状态importos os.system("npu-smi info -l > runtime_status.txt")# 更直接的验证：看Runtime 的overlap 是不是生效了# 用Profiler 抓trace，看计算和数据搬运是不是重叠的fromtorch_npu.profilerimportprofile,ProfilerActivitywithprofile(activities=[ProfilerActivity.NPU],export_name="runtime_overlap_check.json"):Q=torch.randn(4,32,2048,64,dtype=torch.float16).npu()K=torch.randn(4,32,2048,64,dtype=torch.float16).npu()V=torch.randn(4,32,2048,64,dtype=torch.float16).npu()output=torch.nn.functional.scaled_dot_product_attention(Q,K,V,is_causal=True)torch.npu.synchronize()print("Runtime trace 已保存到 runtime_overlap_check.json")print("打开这个文件，看计算kernel 和 数据搬运kernel 有没有重叠——有重叠，说明催菜员工作得好")

催菜员的价值：让计算单元和数据搬运pipeline 起来，减少计算单元等待数据的时间。

第五组人：灶台（硬件驱动）

后厨最里面是灶台，这是真正做菜的地方。灶台的性能决定了做菜的速度上限。

这个灶台就是硬件驱动。它直接控制NPU 的硬件资源（Cube 矩阵乘单元、Vector 向量计算单元、UB 高速存储等）。

你不能直接控制灶台，但你可以通过优化前面的四组人（控制面板路由对不对、酱料包质量好不好、拼菜师傅会不会拼、催菜员会不会调度），来让灶台的性能发挥到最大。

验证：你的后厨五人组是不是都在好好干活

用一段代码验证CANN 五层架构的每一层都在正常工作：

# 验证后厨五人组importtorchimporttorch_npufromtorch_npu.profilerimportprofile,ProfilerActivityprint("=== 验证后厨五人组 ===")# 1. 验证控制面板（Framework Adaptor）Q=torch.randn(4,32,2048,64,dtype=torch.float16).npu()print(f"✅ 控制面板正常：Q 在{Q.device}上")# 2. 验证酱料包仓库（ops-transformer）fromflash_attention_opsimportflash_attention_npu output=flash_attention_npu(Q,Q,Q,causal=True)print(f"✅ 酱料包仓库正常：flash_attention_npu 调用成功")# 3. 验证拼菜师傅（GE）withprofile(activities=[ProfilerActivity.NPU],export_name="verify_ge.json"):output=torch.nn.functional.scaled_dot_product_attention(Q,Q,Q,is_causal=True)torch.npu.synchronize()print("✅ 拼菜师傅工作正常：Profiler trace 已保存（检查有没有FlashAttentionKernel）")# 4. 验证催菜员（Runtime）importos os.system("npu-smi info -l > verify_runtime.txt")print("✅ 催菜员工作正常：NPU 状态已保存（检查 memory usage 和 utilization）")# 5. 验证灶台（硬件驱动）print(f"✅ 灶台正常：NPU 设备名称{torch.npu.get_device_name(0)}")print("=== 后厨五人组验证完成 ===")