CANN 组织链接： https://atomgit.com/cann
ops-nn 仓库链接： https://atomgit.com/cann/ops-nn

---

在当前的大模型（LLM）时代，Transformer 架构中的自注意力机制（Self-Attention）占据了推理与训练计算量的核心地位。尽管 CANN 的 ops-nn 仓库提供了诸如 MatMul、Softmax 等高度优化的原子算子，但在追求极致性能的场景下，单纯的算子堆叠往往会遭遇“内存墙（Memory Wall）”的瓶颈。本文将深入探讨如何利用 CANN 的底层编程能力，将离散的 ops-nn 算子逻辑融合成一个高效的自定义核函数，实现计算与访存的深度流水线并行。

1. 传统算子级联的性能瓶颈与内存墙危机

在标准的深度学习框架实现中，多头注意力（MHA）的计算流程通常被拆解为一系列独立的算子调用：两个矩阵乘法（BMM）、缩放（Scale）、掩码（Mask）以及归一化（Softmax）。

1.1 全局内存（HBM）的带宽损耗

每一个独立的算子在执行时，都遵循“从 HBM 读取数据 $\to$ 片上计算 $\to$ 结果写回 HBM”的模式。

• 冗余读写：$Q\times K^T$ 的中间结果矩阵（Score Matrix）通常非常巨大（$N\times N$），将其写回 HBM 再读出进行 Softmax 计算，消耗了大量的显存带宽。
• 访存延迟：HBM 的访问延迟远高于片上 SRAM（Unified Buffer/L1），频繁的换入换出导致计算单元（Cube/Vector Core）经常处于等待数据的空闲状态。

1.2 算子调度与启动开销

对于 Batch Size 较小或序列长度较短的推理场景，算子的执行时间极短（微秒级）。

• Kernel Launch 开销：Runtime 下发每一个算子都需要经过 CPU 到 NPU 的控制流传输，这种固有的启动开销在小算子场景下占比显著。
• 同步阻隔：算子间的依赖关系导致必须插入同步屏障，破坏了硬件流水线的连续性。

1.3 融合优化的理论上限

通过算子融合（Operator Fusion），我们将整个 Attention 计算图通过 Tiling（分块）策略映射到一个核函数中。中间结果完全在片上高速缓存（Unified Buffer/L0 Buffer）中流转，彻底消除了 Score Matrix 对 HBM 的读写需求，理论上可将访存量减少至原来的 $1/N$（$N$ 为分块数量）。

2. ops-nn 核心算子原语解析

在构建融合算子之前，我们需要深入理解 ops-nn 仓库中提供的基础算子原语，它们是构建复杂逻辑的基石。

2.1 矩阵运算单元（Cube Unit）的抽象

MatMul 是 ops-nn 中最核心的原语，映射到 NPU 的 Cube 单元。

• 分形存储格式：为了适配 Cube 单元的 $16\times 16$ 脉动阵列，数据在 L1/L0 缓存中通常采用分形（Fractal）或块状格式存储。
• 高维度并行：Cube 单元擅长处理 $M\times K\times N$ 的稠密矩阵乘法，单指令周期可完成数千次乘加运算。

2.2 向量运算单元（Vector Unit）的适配

Softmax、Add、Mul 等操作映射到 NPU 的 Vector 单元。

• 单指令多数据（SIMD）：Vector 单元一次可处理 128 或 256 个元素，适合执行 Element-wise 或 Reduction 类操作。
• 精度转换：在 Attention 中，矩阵乘通常使用 FP16/BF16，而 Softmax 为保证数值稳定性通常需要转换为 FP32 计算，Vector 单元提供了高效的精度转换指令。

2.3 数据搬运与同步指令

高效的计算离不开高效的数据搬运。

• DataCopy：支持在 Global Memory、Unified Buffer、L1 Buffer 之间进行复杂的步长（Stride）搬运和填充（Padding）。
• Queues：用于不同流水线阶段（如搬运与计算）之间的同步与通信。

3. 融合架构设计：从 Global Memory 到 Local Memory

为了实现 Fused Attention，我们需要设计精细的内存管理方案，利用 Ascend C 编程语言提供的显式内存控制能力。

3.1 核心数据流设计

我们的目标是：$Q,K,V$ 从 HBM 加载 $\to$ 片上计算 $\to$ $O$ 写回 HBM。

1. 加载阶段：将 $Q$ 的一个分块（Query Block）和 $K,V$ 的对应分块加载到 Local Memory。
2. 计算阶段 I：计算 $S=Q_{block}\times K_{block}^T$。结果保留在 L0C 或 Unified Buffer，不写回 HBM。
3. 向量处理：在 Unified Buffer 中对 $S$ 进行 Scale 和 Softmax 操作。
4. 计算阶段 II：计算 $O=S_{softmax}\times V_{block}$。
5. 累加与输出：利用 Online Softmax 算法更新最终结果，并最终写回 HBM。

3.2 Tiling 策略与分块管理

由于片上内存有限，必须采用分块策略。

• 外层循环：遍历 Query 的分块。
• 内层循环：遍历 Key/Value 的分块。
• 动态切分：根据硬件版本（如 buffer 大小）和输入 shape 动态计算 block size，确保占满 buffer 但不溢出。

3.3 内存复用机制

为了节省宝贵的 Unified Buffer 空间，我们需要复用内存块。例如，计算 $Q{K}^T$ 的临时 buffer 在完成 Softmax 和下一次矩阵乘之后，可以立即释放给下一个迭代使用。

4. Ascend C 编程实战：算子类的构建

使用 Ascend C 开发融合算子时，通常采用面向对象的封装方式。

4.1 Kernel 类结构定义

我们需要定义一个 C++ 类来管理内存张量（Tensor）和通信管道（Pipe）。

#include "kernel_operator.h"

using namespace AscendC;

template<typename T>
class FusedAttentionKernel {
public:
    __aicore__ inline void Init(GM_ADDR q, GM_ADDR k, GM_ADDR v, GM_ADDR out, 
                                uint32_t total_len, uint32_t tile_len) {
        // 初始化全局内存地址
        q_gm.SetGlobalBuffer((__gm__ T*)q);
        k_gm.SetGlobalBuffer((__gm__ T*)k);
        v_gm.SetGlobalBuffer((__gm__ T*)v);
        out_gm.SetGlobalBuffer((__gm__ T*)out);

        // 初始化流水线与队列
        pipe.InitBuffer(q_queue, 2, tile_len * sizeof(T)); // 双缓冲
        pipe.InitBuffer(k_queue, 2, tile_len * sizeof(T));
        pipe.InitBuffer(out_queue, 1, tile_len * sizeof(T));
      
        // 核心：初始化用于 Cube 和 Vector 计算的临时 Buffer
        pipe.InitBuffer(matmul_res_buf, 1, tile_len * tile_len * sizeof(float)); 
        this->tile_len = tile_len;
    }

    __aicore__ inline void Process();

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, 2> q_queue, k_queue;
    TQue<QuePosition::VECOUT, 1> out_queue;
    TBuf<QuePosition::VECCALC> matmul_res_buf; // 中间结果存储

    GlobalTensor<T> q_gm, k_gm, v_gm, out_gm;
    uint32_t tile_len;
};

4.2 流水线编排 (Pipeline Orchestration)

在 Process 方法中，我们不再是顺序执行，而是通过 CopyIn -> Compute -> CopyOut 的逻辑链条，结合 TPipe 实现多级流水线并行。

4.3 算子融合的关键代码逻辑

下面的代码展示了如何在 Local Memory 中完成矩阵乘后直接接续 Softmax，中间无 HBM 交互。

__aicore__ inline void FusedAttentionKernel::Compute(LocalTensor<T>& q_local, 
                                                     LocalTensor<T>& k_local, 
                                                     LocalTensor<T>& v_local) {
    // 1. 获取中间结果 Buffer
    LocalTensor<float> score_matrix = matmul_res_buf.Get<float>();

    // 2. 矩阵乘 Q * K^T -> score_matrix (Cube Unit)
    // 这里的 mm 是 MatMul 对象的封装实例
    mm.SetTensorA(q_local);
    mm.SetTensorB(k_local);
    mm.IterateAll(score_matrix); 

    // 3. 向量计算：Scale (Vector Unit)
    // 直接在 Local Buffer 上操作，无需搬运
    Muls(score_matrix, score_matrix, scalar_val, this->tile_len * this->tile_len);

    // 4. 向量计算：Softmax (Vector Unit)
    // 对每一行进行 Softmax，CANN 提供了高性能的 Softmax 指令
    // 注意：这里需要处理 ReduceSum 和 Exp
    for (int i = 0; i < this->tile_len; i++) {
         Softmax(score_matrix[i * this->tile_len], score_matrix[i * this->tile_len], this->tile_len);
    }

    // 5. 矩阵乘 Score * V -> Output (Cube Unit)
    // 最终结果写入输出队列
    LocalTensor<T> out_local = out_queue.AllocTensor<T>();
    mm_output.SetTensorA(score_matrix);
    mm_output.SetTensorB(v_local);
    mm_output.IterateAll(out_local);

    out_queue.EnQue(out_local);
}

5. 极致性能优化：双缓冲与指令并行

仅实现逻辑融合是不够的，我们需要利用硬件特性榨干每一滴算力。

5.1 双缓冲技术（Double Buffering）

在 Init 函数中，我们为输入队列分配了 2 个块的深度。这意味着当 Cube 单元正在计算第 $i$ 个块时，DMA 搬运单元（MTE）可以并行地从 HBM 加载第 $i+1$ 个块的数据。

• Ping-Pong 机制：通过 AllocTensor 和 EnQue 的配合，硬件自动管理两个缓冲区的切换。
• 掩盖延迟：只要计算时间大于数据搬运时间，理论上可以将通信延迟完全掩盖（Hide Latency）。

5.2 向量与矩阵单元的并发

NPU 的 Cube 单元和 Vector 单元是独立工作的。在计算 Attention 时，我们可以精心编排指令序列。

• 并行发射：在等待 Cube 完成 $Q\times K$ 计算的同时，Vector 单元可以处理上一轮结果的后处理（如类型转换），或者预处理下一轮的 Bias 数据。
• 同步原语优化：使用 SetFlag 和 WaitFlag 进行细粒度的资源依赖控制，避免粗暴的全局阻塞。

5.3 内存对齐与 Padding

为了发挥 Burst 传输的性能，所有涉及 Global Memory 的读写操作都应满足 32 字节或 64 字节对齐。

• ND2NZ 格式：在 Cube 计算前，通常需要将数据从 ND（普通连续格式）转换为 NZ（分形格式）。CANN 提供了专门的指令完成此转换，利用片上转换单元，效率极高。

6. 精度控制与验证

融合算子虽然快，但数值稳定性至关重要，尤其是在 FP16 精度下。

6.1 Online Softmax 技巧

在标准 Softmax 中，需要先求出全局最大值来防止指数溢出。在分块计算（Tiling）场景下，我们无法一次性看到整行数据。

• 算法调整：采用 Online Softmax 算法，维护局部的 max 和 sum，在处理新的分块时，动态更新之前的累加结果。这需要保存额外的中间状态，但避免了多次扫描 HBM。

6.2 精度模式选择

Ascend C 允许开发者指定计算精度模式。

• 高精度模式：中间累加器使用 FP32，仅在最终输出时转回 FP16/BF16。这对于 Attention 中的指数运算至关重要，能有效防止梯度消失或爆炸。

6.3 算子验证

开发完成后，必须使用 ops-nn 提供的标准算子作为 Benchmark 进行精度对齐。

• 单元测试：针对 $Q{K}^T$、Softmax、Output 三个阶段分别注入桩数据（Stub Data）进行中间值比对。
• 边界测试：重点测试 SeqLen 不能被 BlockSize 整除的边缘情况，确保 Padding 逻辑正确，没有读取越界内存。

---

通过上述深度的融合优化，我们不仅仅是在写代码，而是在编排硬件的交响乐。利用 CANN 和 ops-nn 提供的强大底座，自定义融合算子能够突破传统框架的性能天花板，为大模型的推理加速提供坚实的动力。