论文

关注理论最小值:内存优化 Transformer 内核的数组数学框架

Attention at the Theoretical Minimum: A Mathematics of Arrays Framework for Memory-Optimal Transformer Kernels

模型推理推理加速

摘要

注意力机制是现代基于 Transformer 的人工智能的主要计算瓶颈。其标准实现会在序列长度 ~$n$ 中产生二次内存流量,并且 DRAM 访问比当代硬件上的算术运算多花费 100--1000$\times$ 的能量,因此任何仅关注 FLOP 计数的分析从根本上都错误地描述了瓶颈。我们提出了缩放点积注意力及其数值稳定的 softmax 的数组数学 (MoA) 重新表述,导出了指称范式 (DNF),通过代数构造而不是经验调整消除了所有中间数组 - 包括隐式转置密钥缓冲区和每个 softmax 临时数组。 DNF 实现了 $O(n_{dk} + n{_{dv}})$ 数据移动,而标准实现则为 $O(n^2 + n_{dk} + n_{dv})$,其中 $n$ 是序列长度,$dk$ 是键维度,$dv$ 是值维度,并且在具体输入上以全双精度浮点对 PyTorch 进行数值验证。与特定于硬件的加速器或 FlashAttention 等经验平铺方案不同,MoA 通过单个代数框架同时提供阵列融合、形状变换正确性和预测成本模型。内存最小性是在编写任何代码之前建立的定理。预测性能模型预计可实现 $2$--$100\times$ 的加速和 $2$--$50\times$ 的能耗,并且在百亿亿级时优势会扩大。该推导建立了一个从 Python 规范到 (ONF) 操作范式和维度提升的硬件映射的正式验证管道,提供与 DARPA 边缘部署和 DOE 百亿亿次优先级直接相关的性能可移植的 AI 内核。