在 LLM（大语言模型）时代，Transformer 架构的计算效率直接决定了推理系统的吞吐量与首字延迟。ops-transformer 是专为达芬奇架构（Davinci Architecture）定制的高性能算子库，它突破了通用算子的性能瓶颈，通过对 Attention、FFN 等核心模块的深度定制，实现了计算与访存的完美平衡。

核心资源链接：

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

---

一、 算力与带宽的博弈：设计哲学与内存分级架构

Transformer 模型的推理过程，本质上是一场针对显存带宽的“战争”。特别是在解码（Decoding）阶段，矩阵乘法的算术强度（Arithmetic Intensity）较低，导致 AI Core 的计算单元经常处于等待数据的饥饿状态。

1.1 全局内存（GM）至片上内存（Local Memory）的下沉

ops-transformer 的核心设计理念是将计算尽可能地保留在片上。

• 减少搬运： 标准实现中，每一层网络的输出都需要写回 HBM（高带宽内存）。ops-transformer 通过算子融合，使得中间状态（如 $Q\times K^T$ 的结果）仅在片上 Unified Buffer (UB) 或 L1 Cache 中流转，极大减少了对 GM 的读写压力。
• 数据局部性： 利用 NPU 的多级缓存结构，算子库强制高频访问的数据（如 LayerNorm 的权重）常驻 L1，确保 Vector Unit 可以单周期获取操作数。

1.2 针对 NPU 架构的指令级适配

库中的算子并非简单的 C++ 实现，而是深度适配了 NPU 的 ISA（指令集架构）。例如，在处理 Softmax 时，直接调用硬件内置的高精度指数运算指令和向量归约指令，而非通过泰勒展开模拟，从而在保证精度的前提下提升了数倍性能。

二、 Self-Attention 的分块计算（Tiling）与显存优化

自注意力机制的计算复杂度随序列长度呈二次方增长（$O(N^2)$）。为了在有限的片上 SRAM 中处理无限增长的序列，ops-transformer 实施了精细的 Tiling 策略。

2.1 基于 Flash Attention 思想的块级调度

算子库采用了类似 Flash Attention 的切分逻辑，将巨大的 $Q,K,V$ 矩阵切分为适应 SRAM 大小的 Block。

• 动态计算 Block Size： 在 Host 侧的 Tiling 引擎会根据当前 NPU 的具体型号（如 L1/UB 容量差异）和输入 Shape，动态计算最优的 BLOCK_M 和 BLOCK_N。
• 分形对齐（Fractal Alignment）： NPU 的 Cube Unit 对数据排布有特殊要求（通常是 16x16 或 32x32 的分形格式）。Tiling 策略会自动处理边缘填充（Padding），确保每个 Block 都能完美填满 Cube 的指令流水线，避免硬件空转。

2.2 中间结果的片上累加

在计算 $Softmax(Q{K}^T)V$ 时，标准做法需要存储巨大的 Attention Score 矩阵。ops-transformer 通过在线 Softmax（Online Softmax）算法，在片上直接更新局部最大值和归一化因子，使得 $N\times N$ 的中间矩阵从未真正写入过显存，彻底消除了这一最大的显存瓶颈。

三、 FFN 层的高深度融合：从链式调用到单核执行

前馈神经网络（FFN）占据了模型参数量的三分之二。在 SwiGLU 等变体中，包含多次矩阵乘法和逐元素操作。

3.1 跨越 MatMul 与 Vector 的融合边界

传统的深度学习框架会将 FFN 拆解为 MatMul -> Add -> Swish -> Mul -> MatMul 等多个 Kernel。
ops-transformer 实现了一种全融合算子：

1. Cube 单元计算第一个 MatMul。
2. 结果不回写，直接流转至 Vector 单元进行激活函数运算（如 Silu/GELU）。
3. 运算结果直接作为下一个 Cube 单元 MatMul 的输入。
   这种融合将原本需要的 5-6 次 Kernel Launch 减少为 1 次，且中间数据完全在 L0/L1 寄存器间传递。

3.2 寄存器复用与指令并行

在融合算子内部，编译器能够更激进地进行指令调度。当 Cube Unit 正在忙于计算密集型的矩阵乘时，Vector Unit 可以并行处理上一块数据的激活操作，实现了异构计算单元的并行掩盖。

四、 KV Cache 的显存池化与非连续访存管理

在生成式任务中，Key 和 Value 状态的缓存（KV Cache）是显存占用的主要来源。随着对话轮数增加，KV Cache 会面临严重的显存碎片化问题。

4.1 逻辑块与物理页的解耦

ops-transformer 引入了类似操作系统分页管理的机制（Paged Attention）。

• Block Table 映射： 算子不再要求 KV Cache 在物理显存上连续。它接受一个 Block Table，通过逻辑索引查表得到物理地址。
• 零拷贝更新： 在解码阶段，新的 Token 生成的 K/V 向量被直接写入到预分配的显存页中，无需对旧数据进行任何搬运或 Resize 操作。

4.2 动态槽位管理（Slot Mapping）

为了支持高并发推理（Batching），算子库支持细粒度的槽位管理。
每个 Sequence 维护自己的读写指针。在计算 Attention 时，算子内核能够根据 cache_seqlens 数组，准确地从全局显存池中“收集”（Gather）出当前请求所需的历史 KV 数据，并在计算完成后“分散”（Scatter）回写。

五、 动态 Shape 的分档（Binning）与即时编译策略

大模型的输入 Token 数量是不确定的。如果为每个不同的长度都编译一个 Kernel，会导致严重的编译时间膨胀（Compilation Explosion）。

5.1 预置档位与模糊匹配

ops-transformer 采用分档策略来平衡性能与灵活性。

• 档位预设： 系统预定义一系列标准长度（如 128, 512, 1024, 2048）。
• 向上取整： 当实际输入长度为 300 时，系统会将其填充至 512 档位进行计算。虽然引入了少量的无效计算，但复用了已编译好的高性能二进制代码，消除了 Runtime 的编译延迟。

5.2 核心 Tiling 结构定义

为了支持运行时的动态调整，算子使用特定的结构体在 Host 和 Device 之间传递 Tiling 参数。

以下代码展示了典型的 Tiling 配置结构，它决定了 Kernel 如何拆解问题规模：

// Transformer 核心算子 Tiling 配置结构
struct TransformerTilingData {
    // 基础维度信息
    uint32_t batch_size;
    uint32_t head_num;
    uint32_t embedding_dim;
    uint32_t seq_len;

    // 针对 NPU Cube Unit 优化的分块参数
    uint32_t block_m;       // M 轴分块大小 (如 128)
    uint32_t block_n;       // N 轴分块大小 (如 128)
    uint32_t block_k;       // K 轴分块大小 (如 64)

    // 内存地址偏移量，用于支持非连续 KV Cache
    uint64_t kv_cache_offset;
  
    // 动态 Padding 修正参数，防止计算越界
    uint32_t padding_mask_len;
  
    // 硬件特定的 buffer 配置
    uint32_t ub_size_bytes; // Unified Buffer 可用大小
    uint32_t l1_size_bytes; // L1 Buffer 可用大小
};

六、 极致流水线：Double Buffering 与性能调优

在微观架构层面，ops-transformer 极度依赖流水线并行来掩盖数据搬运延迟（Latency Hiding）。

6.1 MTE 与 Cube 的乒乓操作

为了让计算核心（Cube）永不停止，算子实现了双缓冲（Double Buffering）机制。

• Ping-Pong Queue： 在片上内存划分出两块区域（Ping Buffer 和 Pong Buffer）。
• 并行执行： 当 MTE（Memory Transfer Engine，数据搬运引擎）正在将下一块数据从 GM 搬运到 Pong Buffer 时，Cube Unit 正在全速处理 Ping Buffer 中的数据。只要计算时间大于搬运时间，访存延迟就被完全掩盖了。

6.2 性能瓶颈分析与 Profiling

在实际调优中，开发者需关注 Profiling 数据中的“流水线气泡”。

• Vector Bound： 如果 Softmax 或 LayerNorm 耗时过长，会导致 Cube 单元空闲。此时需要优化 Vector 指令的 Tiling 大小。
• MTE Bound： 如果数据搬运是瓶颈，说明算术强度不足。此时可能需要增大 BLOCK_SIZE，或者在 Quantization（量化）层面减少数据位宽（如使用 INT8/FP16），以降低带宽需求。