CANN-Ascend-C流水线编程-昇腾NPU上Cube和Vector怎么协作

CANN-Ascend-C流水线编程-昇腾NPU上Cube和Vector怎么协作

昇腾NPU的 AI Core 里有两种计算单元：Cube 做矩阵乘法，Vector 做逐元素运算。FlashAttention 这种融合算子需要 Cube 和 Vector 交替工作——先 Cube 算 Q·K^T，再 Vector 算 Softmax，再 Cube 算 Attn·V。流水线编程让两者不互相等待。

为什么需要流水线

没有流水线时，Cube 和 Vector 串行执行：

时间线： [Cube: Q·K^T] → [Vector: Softmax] → [Cube: Attn·V] → [Vector: RoPE] 120μs 30μs 120μs 10μs 总时间 = 280μs，Cube 利用率 = 240/280 = 86%

Cube 在 Vector 工作时空闲，反之亦然。流水线让多个 tile 的计算重叠：

Tile 0: [Cube: Q·K^T] → [Vector: Softmax] → [Cube: Attn·V] Tile 1: [Cube: Q·K^T] → [Vector: Softmax] → [Cube: Attn·V] ↑ 跟 Tile 0 的 Vector 重叠？不对——需要 Cube

实际上 Cube 和 Vector 不能真正同时工作——它们共享 L1 缓存，同时读写会冲突。昇腾NPU的流水线是通过指令队列实现的：Cube 指令和 Vector 指令分别排入队列，硬件自动调度。

TPipe 编程模型

Ascend C 用TPipe管理流水线。核心思想是把数据分成多个 stage，每个 stage 由不同的计算单元处理：

#include"kernel_operator.h"classFlashAttentionKernel{public:__aicore__inlinevoidInit(...){// 初始化三个 stage// Stage 1: Cube 计算 Q·K^T// Stage 2: Vector 计算 Softmax// Stage 3: Cube 计算 Attn·V}__aicore__inlinevoidProcess(){// 分块处理for(inttile=0;tile<num_tiles_;tile++){// Stage 1: Cube 计算 Q·K^TMatMul(qk_local_,q_local_,kt_local_);// Stage 2: Vector 计算 Softmax// Cube 释放 qk_local_，Vector 接管Softmax(attn_local_,qk_local_);// Stage 3: Cube 计算 Attn·VMatMul(out_local_,attn_local_,v_local_);}}};

真正的流水线需要 double/triple buffer——当 Stage 2 处理 Tile 0 的 Softmax 时，Stage 1 已经在准备 Tile 1 的 Q·K^T 数据。

Triple Buffer 实现

__aicore__inlinevoidProcessWithPipeline(){// 三个 buffer 对应三个 stageLocalTensor<half>qk_buf[3];// Stage 1 的输出 = Stage 2 的输入LocalTensor<half>attn_buf[3];// Stage 2 的输出 = Stage 3 的输入// 预热：先完成 Tile 0 的 Stage 1MatMul(qk_buf[0],q_local_,kt_local_);for(inttile=0;tile<num_tiles_;tile++){intcur=tile%3;intnext=(tile+1)%3;// Stage 1: 下一块的 Cube 计算（与当前块的 Vector 重叠）if(tile+1<num_tiles_){MatMul(qk_buf[next],q_local_,kt_local_[next]);}// Stage 2: 当前块的 Vector 计算Softmax(attn_buf[cur],qk_buf[cur]);// Stage 3: 上一块的 Cube 计算if(tile>0){MatMul(out_local_,attn_buf[(cur+2)%3],v_local_);}}}

三组 buffer 轮转：当 Stage 1 写 buffer[0] 时，Stage 2 读 buffer[2]，Stage 3 读 buffer[1]。三者互不冲突。

Cube-Vector 依赖关系

Cube 和 Vector 的数据传递通过 L1 缓存。关键约束：Cube 写完 L1 后，Vector 才能读。需要用pipe_barrier或SetFlag同步：

// Cube 写完MatMul(qk_buf,q,kt);SetFlag<PIPE_M>(PIPE_V);// 通知 Vector：数据准备好了// Vector 等待WaitFlag<PIPE_M>(PIPE_V);// 等待 Cube 通知Softmax(attn,qk);

如果漏了同步，Vector 可能读到 Cube 正在写的数据——结果就是随机错误，而且不每次都复现，极难 debug。

性能收益

FlashAttention 在昇腾NPU上的流水线效果：

Triple Buffer 比 Double Buffer 多一组 buffer 的 L1 空间（约 256KB），但 Cube 利用率提升 3 个点。在 L1 空间充裕时优先用 Triple Buffer。

流水线编程是 Ascend C 算子优化的核心技能。不搞流水线，Cube 和 Vector 总有一个在等，NPU 利用率上不去。记住三件事：分块、多 buffer、同步。仓库在这里：

https://atomgit.com/cann/opbase