一种用于 Transformer 注意力分数计算中权重静态动态矩阵乘法的基于数字 SRAM 的存内计算宏单元

《Microelectronics Journal》：A Digital SRAM-Based Compute-In-Memory Macro for Weight-Stationary Dynamic Matrix Multiplication in Transformer Attention Score Computation

　　

摘要

存内计算 (CIM) 技术被广泛应用于能效型人工智能 (AI) 处理器中。它们缓解了由计算单元和存储单元之间大量数据移动引起的功耗和延迟瓶颈。为了将这些优势扩展到 Transformer，本文提出了一种用于注意力分数计算的数字 CIM 宏单元。为了消除动态矩阵乘法 (MM)，我们使用组合的 QK 权重矩阵将计算重构为静态 MM，以便输入可以直接馈送到单个 CIM 宏单元以获得分数结果。然而，这引入了双输入静态 MM 的新挑战。该计算进一步分解为四组位串行逻辑与加法运算。这允许双输入通过与门直接激活字线，从而以最小开销实现双输入静态 MM。引入了一种分层零值位跳过机制，以优先跳过双输入情况下的零值位。该机制有效利用了双输入的数据稀疏性，显著减少了冗余操作。在 65 纳米工艺中实现，0.35 平方毫米的宏单元在 1.24 毫瓦功耗下实现了 42.27 GOPS，对应的能效为 34.1 TOPS/W，面积效率为 120.77 GOPS/平方毫米。与评估的 CPU 和 GPU 基线相比，提出的宏单元能效分别提高了约 25 倍和 13 倍。当标准化到 28 纳米工艺时，预计其能效比先前的 Transformer-CIM 设计提高约 3 倍，面积效率提高约 1.6 倍，展示了其在能效型边缘智能方面的潜力。

章节摘要

引言

在算法创新和硬件能力双重突破的推动下，Transformer 已成为推动人工智能的核心引擎 [1], [2]。然而，它们惊人的计算需求和密集的内存访问导致了严重的性能瓶颈 [3]。这极大地加剧了传统冯·诺依曼架构中固有的著名“内存墙”困境 [4]。为了弥合这一差距，存内计算 (CIM) 得到了广泛探索。这种范式

W_K 和 W_Q 乘法的预计算

在传统 Transformer 中，注意力分数矩阵 S 的计算如公式 (1) 所示：其中 **S**∈?^N×N 是注意力图，**Q**∈?^N×d 是查询矩阵，**K**∈?^N×d 是键矩阵，**X**∈?^N×dmodel 是由 N 个令牌组成的输入矩阵，每个令牌的维度为 d_model。**W**_Q 和 **W**_K∈?^D×d 是对应的权重矩阵。然而，公式 (1) 中的 **Q** 和 **K** 是动态生成的。因此，所需的动态 MM 与常见的权重静止 CIM 范式不兼容。

在

硬件设计

注意力分数矩阵 **S** 可以被具有固定权重静止特性的单个 CIM 宏单元高效直接处理，从而避免与动态 MM 相关的 CIM 权重读/写以及额外的宏单元开销。如图 4(a) 所示，单个注意力分数 s_ij 的获取是通过循环 CIM 累加和近内存移位累加逐步积累每个计算项的结果来实现的。矩阵 S 的完整构建

评估

如图 10 所示，展示了为高效注意力计算设计的数字宏单元布局。SRAM-CIM 阵列使用 Push-Rule 位单元实现。64 个 64×8 位 SRAM 子阵列共享单个 WL 和输入控制块，并组织成 16 列 × 4 行的平面布局以缓解布线拥堵。该宏单元在 65 纳米工艺中实现，占用总面积为 0.35 平方毫米。它具有最大 64×64×8 位的权重重存容量。

结论

本工作提出了一种用于 Transformer 注意力机制的高效数字 SRAM 基 CIM 宏单元，旨在解决权重静止架构在处理动态 MM 时的瓶颈。注意力分数计算通过组合 QK 权重矩阵进行了重构，将动态 MM 转变为静态 MM，以完全消除频繁内存更新的开销。

未引用的参考文献

[30], [31]。

竞争利益声明

? 作者声明，他们没有任何已知的竞争性财务利益或个人关系，这些因素可能会影响本论文中报告的工作。

6. 致谢

本工作部分资金来自 重大研究计划 的 国家自然科学基金 在批准号 92464103 下的资助，部分来自 国家自然科学基金 在批准号 6233400 下的重点计划。

于建义 | 王晓肖 | 王宇轩 | 傅祥 | 乔飞 | 王颖 | 袁瑞 | 刘力元 | 史聪

中国，重庆，400044，重庆大学微电子与通信工程学院