前言

想学 CANN 开发,光看文档不够,得跑代码。samples 仓是昇腾CANN 的官方示例代码库,位于第一层——应用与加速库(边缘计算场景)。这个仓里有 705 个示例,覆盖算子开发、模型适配、性能调优、应用部署四个方向。这篇文章拆开看怎么用这些示例代码。

samples 仓的目录结构

samples 仓的目录结构是按应用场景分的,不是按技术栈分的。根目录下有几个主要目录:

operator/:算子开发示例。里面有 100+ 个算子的完整实现,从简单的 Add、Mul 到复杂的 FlashAttention、RotaryEmbedding。每个算子都是一个独立的目录,包含:

  • op_kernel/:Ascend C 实现的算子 kernel
  • op_proto/:算子原型定义(输入输出的 dtype、shape 约束)
  • op_tiling/:Tiling 策略实现
  • testcases/:测试用例

modelzoo/:模型适配示例。里面有 50+ 个主流模型的端到端适配代码,包括 LLaMA、GPT、BERT、ResNet、YOLO 等。每个模型目录包含:

  • model/:模型定义(PyTorch 或 MindSpore)
  • scripts/:训练/推理脚本
  • args/:参数配置
  • README.md:复现指南

inference/:推理部署示例。里面有 30+ 个推理引擎的使用示例,包括 Acs、TensorRT、ONNX Runtime 等。每个示例展示怎么把一个训练好的模型部署成高性能推理服务。

tutorials/:教程示例。里面有 20+ 个手把手教程,从环境搭建到第一个算子开发,再到性能调优,每个教程都有配套的示例代码。

contrib/:社区贡献的示例。这里面是社区成员提交的示例代码,质量参差不齐,但很多是官方示例里没有的冷门场景。

跑通第一个算子示例

operator/add 这个最简单算子开始,跑通整个流程。这个示例展示了算子开发的完整流程:

# 1. 克隆 samples 仓
git clone https://atomgit.com/cann/samples.git
cd samples/operator/add

# 2. 查看目录结构
tree .
# .
# ├── op_kernel/
# │   └── add_kernel.cpp   # Ascend C 实现的 Add 算子
# ├── op_proto/
# │   └── add.h           # 算子原型定义
# ├── op_tiling/
# │   └── add_tiling.cpp  # Tiling 策略
# ├── testcases/
# │   └── test_add.py     # Python 测试用例
# └── README.md           # 复现指南

# 3. 编译算子
# 需要先把 CANN 的路径设好
export ASCEND_HOME=/usr/local/Ascend
export PATH=$ASCEND_HOME/bin:$PATH
export LD_LIBRARY_PATH=$ASCEND_HOME/lib64:$LD_LIBRARY_PATH

# 用 CANN 提供的编译脚本
python3 -m op_builder.build_op \
    --op_kernel=op_kernel/add_kernel.cpp \
    --op_proto=op_proto/add.h \
    --op_tiling=op_tiling/add_tiling.cpp \
    --output=build/

# 4. 跑测试用例
python3 testcases/test_add.py
# 输出应该是:
# [PASS] Add operator test passed!

op_kernel/add_kernel.cpp 里的核心代码:

// Add 算子的 Ascend C 实现
// 这个算子做 element-wise 加法:C = A + B
extern "C" __global__ __aicore__ void add_kernel(
    GM_ADDR a, GM_ADDR b, GM_ADDR c, int64_t numel)
{
    TPipe pipe;
    // 用 TQue 管理 L1 缓冲区
    TQue<QuePosition::VECIN, 1> a_q, b_q;
    TQue<QuePosition::VECOUT, 1> c_q;
    
    // 初始化缓冲区,大小是 L1 Buffer 的容量
    // 这里取 16KB,够放 8192 个 FP16 元素
    pipe.InitBuffer(a_q, 8192 * sizeof(half));
    pipe.InitBuffer(b_q, 8192 * sizeof(half));
    pipe.InitBuffer(c_q, 8192 * sizeof(half));
    
    // 分块处理,每次处理 8192 个元素
    int64_t tile_size = 8192;
    for (int64_t i = 0; i < numel; i += tile_size) {
        int64_t cur = min(tile_size, numel - i);
        
        // 从 HBM 加载 A 和 B
        LocalTensor<half> a_local = a_q.AllocTensor<half>();
        LocalTensor<half> b_local = b_q.AllocTensor<half>();
        DataCopy(a_local, a + i, cur * sizeof(half));
        DataCopy(b_local, b + i, cur * sizeof(half));
        
        // Vector 单元做加法
        LocalTensor<half> c_local = c_q.AllocTensor<half>();
        vec_add(c_local, a_local, b_local, cur, 1);
        
        // 写回 HBM
        DataCopy(c + i, c_local, cur * sizeof(half));
        
        // 释放缓冲区
        a_q.FreeTensor(a_local);
        b_q.FreeTensor(b_local);
        c_q.FreeTensor(c_local);
    }
}

这个代码展示了 Ascend C 编程的基本模式:分配缓冲区 → 从 HBM 加载 → Vector 单元计算 → 写回 HBM → 释放缓冲区。虽然 Add 算子很简单,但这个模式对所有 element-wise 算子都适用。

跑通第一个模型适配示例

modelzoo/LLaMA/7B 这个示例开始,看怎么把 LLaMA-7B 适配到昇腾NPU上跑推理:

# 1. 进入模型目录
cd samples/modelzoo/LLaMA/7B

# 2. 查看 README.md,按照里面的步骤准备环境
# 需要:CANN 8.0+, torch-npu, transformers 4.37+

# 3. 下载模型权重
# 从 HuggingFace 或 ModelScope 下载 LLaMA-7B 的权重
python3 scripts/download_weights.py --output_dir ./weights/

# 4. 转换权重格式(从 HuggingFace 格式转到昇腾格式)
python3 scripts/convert_weights.py \
    --input_dir ./weights/ \
    --output_dir ./weights_npu/ \
    --dtype fp16

# 5. 跑推理
python3 scripts/run_inference.py \
    --model_path ./weights_npu/ \
    --input_text "什么是人工智能?" \
    --max_new_tokens 100

# 输出:
# 人工智能(Artificial Intelligence,简称 AI)是计算机科学的一个分支,
# 致力于创建能够模拟人类智能的系统...

scripts/run_inference.py 里的核心代码:

# LLaMA-7B 推理脚本(简化版)
import torch
import torch_npu
from transformers import LlamaTokenizer
from atb_speed import AtbSpeed  # ATB 加速库

# 加载 tokenizer
tokenizer = LlamaTokenizer.from_pretrained("./weights/")

# 初始化 ATB 加速引擎
# ATB 会把 Transformer 层融合成一个大算子,大幅提升推理速度
atb = AtbSpeed(
    model_path="./weights_npu/",
    device="npu:0",
    precision="fp16",
    kv_cache_dtype="fp16",
    max_seq_len=2048,
)

# 编码输入
input_ids = tokenizer.encode("什么是人工智能?", return_tensors="pt").npu()

# 生成
generated_ids = input_ids.clone()
for i in range(100):  # 生成 100 个 token
    # forward 返回的是 logits (词表大小的概率分布)
    logits = atb.forward(generated_ids)
    
    # 取最后一个 token 的概率分布,采样下一个 token
    next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
    
    # 拼到生成序列里
    generated_ids = torch.cat([generated_ids, next_token], dim=1)
    
    # 如果生成了结束 token,就停止
    if next_token.item() == tokenizer.eos_token_id:
        break

# 解码生成的 token
generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
print(generated_text)

这段代码展示了用 ATB 加速库跑 LLaMA 推理的完整流程。核心是 AtbSpeed 这个对象,它封装了权重加载、KV-Cache 管理、算子融合等底层细节。

跑通第一个推理部署示例

inference/acs_llama 这个示例开始,看怎么用 Acs(昇腾推理引擎)部署 LLaMA 推理服务:

# 1. 进入推理部署目录
cd samples/inference/acs_llama

# 2. 查看 README.md,安装 Acs
# Acs 是昇腾的官方推理引擎,性能比直接用 PyTorch 好

# 3. 转换模型格式(从 PyTorch 转到 Acs 格式)
python3 scripts/convert_to_acs.py \
    --model_path ./weights_npu/ \
    --output_path ./acs_model/ \
    --precision fp16

# 4. 启动推理服务
python3 scripts/serve.py \
    --model_path ./acs_model/ \
    --port 8000 \
    --max_batch_size 8 \
    --max_seq_len 2048

# 服务启动后,可以用 curl 发请求
# curl -X POST http://localhost:8000/generate \
#   -H "Content-Type: application/json" \
#   -d '{"text": "什么是人工智能?", "max_new_tokens": 100}'

scripts/serve.py 里的核心代码:

# 用 Acs 部署 LLaMA 推理服务(简化版)
from flask import Flask, request, jsonify
import acs  # 昇腾推理引擎

app = Flask(__name__)

# 加载 Acs 模型
# Acs 的模型格式是 .om (Offline Model)
# 这种格式是编译后的,推理速度比 PyTorch 格式快很多
model = acs.Model("./acs_model/llama-7b.om")
model.init()

@app.route("/generate", methods=["POST"])
def generate():
    data = request.json
    input_text = data["text"]
    max_new_tokens = data.get("max_new_tokens", 100)
    
    # 编码输入
    input_ids = tokenizer.encode(input_text, return_tensors="pt").npu()
    
    # 推理
    generated_ids = input_ids.clone()
    for i in range(max_new_tokens):
        # Acs 的 forward 接口跟 PyTorch 不一样
        # 需要用 model.execute 来调用
        outputs = model.execute(
            input_ids=generated_ids,
            max_seq_len=2048,
        )
        logits = outputs["logits"]
        
        # 采样下一个 token
        next_token = torch.argmax(logits[:, -1, :], dim=-1, keepdim=True)
        generated_ids = torch.cat([generated_ids, next_token], dim=1)
        
        if next_token.item() == tokenizer.eos_token_id:
            break
    
    # 解码
    generated_text = tokenizer.decode(generated_ids[0], skip_special_tokens=True)
    return jsonify({"text": generated_text})

if __name__ == "__main__":
    app.run(host="0.0.0.0", port=8000)

这段代码展示了用 Acs 部署推理服务的完整流程。Acs 的模型格式是 .om,这种格式是编译后的,推理速度比 PyTorch 格式快很多。

samples 仓的学习路径

根据 samples 仓的示例,可以规划出一条完整的学习路径:

第一阶段:跑通 Hello World。把 operator/add 这个示例跑通,理解 Ascend C 编程的基本模式。这个阶段重点是理解缓冲区管理、数据搬运、Vector 单元计算这三个概念。

第二阶段:写一个进阶算子。参考 operator/matmuloperator/softmax,写一个稍微复杂一点的算子。这个阶段重点是理解 Cube 单元、Tiling 策略、双缓冲这三个概念。

第三阶段:适配一个模型。参考 modelzoo/LLaMA/7B,把一个 HuggingFace 上的模型适配到昇腾NPU上跑推理。这个阶段重点是理解 torch-npu 的桥接机制、ATB 加速库的用法。

第四阶段:部署推理服务。参考 inference/acs_llama,用 Acs 部署一个高性能推理服务。这个阶段重点是理解 Acs 的模型格式、推理 API、batch 策略。

第五阶段:性能调优。参考 tutorials/performance_tuning,用 msprof 工具分析性能瓶颈并做针对性优化。这个阶段重点是理解 Tiling 调优、算子融合、量化这三个优化手段。

注意事项

用 samples 仓的示例代码时,有几个坑要注意:

第一是版本对齐。samples 仓的示例代码是针对特定 CANN 版本写的,如果你用的 CANN 版本跟示例不一致,可能会编译失败或运行出错。建议在跑示例前先看一下 README.md 里的版本要求。

第二是环境依赖。有些示例需要额外的依赖(比如 Acs、ATB、torch-npu),这些依赖的安装比较复杂。建议先用 Docker 镜像(samples 仓提供了 Dockerfile)搭环境,避免把本地环境搞乱。

第三是性能数据仅供参考。samples 仓里的性能数据是在特定硬件上测的(比如 Ascend 910),如果你用的硬件不一样(比如 Ascend 310P),性能可能会有差异。

samples 仓是学习 CANN 开发的最佳入口。705 个示例覆盖了从入门到精通的各个环节,跟着示例一步步做,就能系统掌握昇腾NPU 的开发技能。

仓库地址:https://atomgit.com/cann/samples

Logo

作为“人工智能6S店”的官方数字引擎,为AI开发者与企业提供一个覆盖软硬件全栈、一站式门户。

更多推荐