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

---

在生成式 AI 的浪潮下，Transformer 架构已成为事实上的标准。然而，随着参数量突破千亿级以及上下文窗口（Context Window）的不断扩展，传统的算子实现已无法满足极致的推理性能需求。ops-transformer 仓库作为 CANN 平台针对 Transformer 架构的深度优化算子集合，通过对底层硬件指令的精细编排、显存管理的重构以及计算流的融合，解决了 LLM 推理中的“显存墙”与“算力利用率”两大核心矛盾。

1. 显存与算力的博弈：LLM 推理的性能瓶颈分析

大语言模型（LLM）推理过程具有显著的 Memory-Bound（访存受限） 特征，尤其是在 Decoding 阶段。每一次 Token 的生成都需要加载全量的模型权重和 KV Cache，而计算量却相对较小（矩阵乘向量）。

1.1 访存延迟与带宽利用率

在标准算子实现中，数据的搬运（HBM $\leftrightarrow$ Chip）往往占据了大部分时间。

• 低效的 IO 模式：如果算子粒度过细（如独立的 Add, LayerNorm），会导致中间结果频繁回写 HBM，浪费宝贵的显存带宽。
• Cache Miss：对于长序列输入，如果缺乏特定的 Tiling（分块）策略，片上缓存（L1/L0 Buffer）无法有效复用数据，导致流水线停顿。

1.2 动态 Shape 的挑战

LLM 推理的 Batch Size 和 Sequence Length 是动态变化的。静态编译的图往往难以适配所有形状，导致 Padding 带来的无效计算，或者频繁的重编译开销。ops-transformer 提供了动态 Shape 感知的算子实现，能够在运行时根据输入形状自适应调整 Tiling 参数。

2. 分组查询注意力（GQA）的内核级优化

Llama 3、Mistral 等现代模型广泛采用了 GQA（Grouped Query Attention）以在性能和显存之间取得平衡。相比于 MHA（多头注意力），GQA 多个 Query 头共享一组 Key/Value 头，这对算子的显存访问模式提出了新的要求。

2.1 共享 KV 的数据复用策略

在底层 Kernel 实现中，必须最大化 KV 数据的片上复用率。

• 广播式加载：算子利用 MTE（Memory Transfer Engine）将共享的 Key/Value 块一次性加载到 Unified Buffer。
• 多头并行计算：在计算 $Q{K}^T$ 时，Cube Unit 连续调度多个 Query 头与同一份驻留在 L0 Buffer 中的 KV 数据进行矩阵乘法。这种设计将 KV 数据的读取带宽需求降低了 $H_{query}/H_{kv}$ 倍。

2.2 向量化的 Score 处理

注意力分数的计算不仅仅是矩阵乘法，还包含缩放（Scale）、掩码（Mask）和 Softmax。

• 算子融合：ops-transformer 将这些操作完全融合在 Local Memory 中进行。Cube Unit 输出 Score 矩阵后，Vector Unit 立即介入执行 Scale 和 Exp 操作，中间数据不落盘。
• 高精度累加：为了防止 Softmax 溢出，算子内部维护了 FP32 格式的累加器，仅在最终输出时转换为 FP16/BF16。

3. 前馈网络重构：SwiGLU 的指令级融合

FFN（Feed-Forward Network）层占据了模型参数量的 2/3。SwiGLU 激活函数引入了门控机制：$\text{Swish}(X{W}_1)\odot (X{W}_2)$，这涉及两次大矩阵乘法和复杂的逐元素操作。

3.1 双流并行与计算掩盖

为了提升吞吐量，算子采用了双流设计。

• 并行矩阵乘：利用 Cube Unit 的高算力，并发执行 $X{W}_1$ 和 $X{W}_2$ 的计算。
• 流水线编排：当 $X{W}_1$ 的部分分块计算完成并流入 Vector Unit 进行 Swish 激活时，Cube Unit 已经开始计算 $X{W}_2$ 的后续分块。这种“乒乓”流水线彻底掩盖了 Vector 计算的延迟。

3.2 Ascend C 算子代码实战

下述代码展示了如何在 Kernel 层面通过 Ascend C 实现 SwiGLU 的核心逻辑，利用本地内存进行数据交换。

// 示例：SwiGLU 核心计算片段 (简化版)
// 假设 input_gate 和 input_value 已经是 MatMul 的结果，存储在本地 LocalTensor 中

template <typename T>
__aicore__ inline void SwiGLU_Compute(LocalTensor<T>& gate_tensor, LocalTensor<T>& value_tensor, 
                                      LocalTensor<T>& out_tensor, uint32_t data_len) {
    // 1. Swish 激活: Gate = Gate * Sigmoid(Gate)
    // 利用 Vector Unit 的 Sigmoid 指令
    Sigmoid(gate_tensor, gate_tensor, data_len);
    // 利用 Mul 指令完成 x * sigmoid(x)，注意这里可能需要临时 Buffer 或原位操作
    // 假设硬件支持 Silu 直接指令，或通过组合指令实现
    // Muls(gate_tensor, gate_tensor, (T)1.0, data_len); // 仅示意

    // 2. Element-wise Mul: Out = Gate * Value
    // Vector Unit 极速执行逐元素乘法
    Mul(out_tensor, gate_tensor, value_tensor, data_len);

    // 3. 结果写回 Global Memory
    // 这一步通常在 CopyOut 阶段通过 Queue 完成，此处省略
}

// 在主 Process 函数中，通过 TPipe 管理内存
__aicore__ void KernelProcess() {
    TPipe pipe;
    // 分配 Unified Buffer
    TBuf<TPosition::VECCALC> calc_buf; 
    pipe.InitBuffer(calc_buf, 1024 * sizeof(half)); // 举例

    // ... MatMul 计算逻辑，结果存入 gate_tensor 和 value_tensor ...
  
    // 调用计算函数，数据全程在片上流动
    SwiGLU_Compute(gate_tensor, value_tensor, out_tensor, len);
}

4. Paged Attention：非连续显存的物理映射

在长文本生成中，KV Cache 的显存占用随序列长度线性增长。传统的连续内存分配会导致严重的显存碎片和预分配浪费。ops-transformer 实现了 Paged Attention 机制，允许 KV Cache 存储在物理不连续的内存块（Block）中。

4.1 逻辑-物理地址转换表

算子维护了一张 Block Table，记录了每个请求（Sequence）的逻辑块 ID 到物理显存地址的映射。

• Gather 机制：在注意力计算加载 KV 时，MTE 单元不再执行连续读取，而是根据 Block Table 进行 Gather 操作。硬件直接根据索引从分散的物理地址抓取数据块拼接到片上的连续 Buffer 中。
• 零拷贝动态扩容：当上下文增长需要新显存时，系统只需在空闲池中申请一个物理块并更新 Block Table，无需迁移已有的数据。

5. 归一化算子的带宽突围：RMSNorm

RMSNorm（Root Mean Square Layer Normalization）是带宽敏感型算子。虽然计算量小，但需要读取全量输入。

5.1 单次读取与向量规约

ops-transformer 针对 RMSNorm 进行了极致的访存优化。

• One-Pass Algorithm：算子确保每个输入数据仅从 HBM 读取一次。在 Vector Unit 内部，数据被加载到寄存器后，立即进行平方、求和（ReduceSum）以及最终的归一化乘法。
• Welford 算法：为了在 FP16 输入下保持高精度，内部累加过程强制使用 FP32，并采用 Welford 算法减少浮点舍入误差，确保深层网络的数值稳定性。

6. 量化计算：INT8/FP8 混合精度加速

为了进一步提升吞吐量，ops-transformer 集成了对低比特量化的原生支持，特别是 W8A16（权重 INT8，激活 FP16）和 W8A8 模式。

6.1 动态量化参数补偿

在矩阵乘法层面，算子直接调用 Cube Unit 的 INT8 核心指令，算力密度提升 2 倍以上。

• Per-Token / Per-Channel 量化：为了减少精度损失，算子支持细粒度的量化参数（Scale）。
• 反量化融合：矩阵乘输出（通常为 INT32）在流出 Cube Unit 时，会立即在 Vector Unit 结合 Scale 参数转换回 FP16/BF16。这一过程被融合在 MatMul 的后处理流水线中，几乎不增加额外延迟。

7. 性能诊断与调优方法论

在使用 ops-transformer 进行模型开发时，配合 MSPROFILER 工具可以精准定位瓶颈。

7.1 关键指标解读

• MTE2/MTE3 利用率：如果 MTE2（DDR 到 AICore）利用率长期处于 100% 而 Cube 利用率低，说明瓶颈在显存带宽（如 KV Cache 读取）。此时应尝试增大 Batch Size 或优化 Paged Block Size。
• Vector/Cube Ratio：Transformer 类模型应由 Cube 计算主导。如果 Vector 占比过高，需检查是否由大量的 Cast（类型转换）或非融合的 Element-wise 算子导致。

7.2 调优建议

• 启用 Flash Attention：确保在长序列（SeqLen > 1024）场景下开启 Flash Attention 路径，以 $O(N)$ 复杂度替代 $O(N^2)$。
• 静态内存规划：在 GE 层面配置静态内存池，不仅能减少运行时的 malloc/free 开销，还能通过紧凑排布减少内存碎片，间接提升缓存命中率。