提出数学数组框架,让注意力机制内存使用降为线性,大幅提速省电。
Attention at the Theoretical Minimum: A Mathematics of Arrays Framework for Memory-Optimal Transformer Kernels
- 用代数构造法消除所有中间数组,实现内存最优
- 内存移动量从O(n²)降至O(n),实测与PyTorch一致
- 适合边缘部署与超算场景,性能可预测且无需调参
注意力机制是现代Transformer AI的主要计算瓶颈。标准实现中,内存访问随序列长度n呈二次增长,而现代硬件上内存访问能耗是算术操作的100—1000倍,仅看浮点运算量会严重误判瓶颈。本文提出数学数组(MoA)框架,重构缩放点积注意力及其数值稳定的softmax,导出无中间数组的语义规范形式(DNF),通过代数构造而非经验调优消除所有临时缓冲区(包括转置键缓存和softmax临时变量)。该方法将数据移动量从标准实现的O(n² + n_{dk} + n_{dv})降低至O(n_{dk} + n_{dv}),其中n为序列长度,dk为键维度,dv为值维度,并在具体输入下以双精度浮点数验证了与PyTorch的一致性。不同于依赖硬件特性的加速器或经验分块方案(如FlashAttention),MoA框架同时提供数组融合、形状变换正确性及可预测的代价模型,内存最小性在编码前即作为定理确立。性能模型预测可带来2—100倍加速与2—50倍能耗降低,优势随规模扩大而增强。推导过程建立从Python规格到操作规范形式(ONF)及维度提升的硬件映射的可验证流程,生成适用于DARPA边缘部署与DOE超算需求的高性能可移植AI内核。
原文摘要 · Abstract (English)
The attention mechanism is the dominant computational bottleneck in modern transformer-based AI. Its standard implementation incurs quadratic memory traffic in the sequence length~$n$, and DRAM accesses cost 100--1000$\times$ more energy than arithmetic operations on contemporary hardware, so any analysis focused solely on FLOP counts fundamentally mischaracterises the bottleneck. We present a Mathematics of Arrays (MoA) reformulation of scaled dot-product attention and its numerically stable softmax, deriving a Denotational Normal Form (DNF) that eliminates all intermediate arrays -- including the implicit transposed-key buffer and every softmax temporary -- by algebraic construction rather than empirical tuning. The DNF achieves $O(n_{dk} + n{_{dv}})$ data movement versus $O(n^2 + n_{dk} + n_{dv})$ for the standard implementation, where $n$ is the sequence length, $dk$ is the key dimensionality and $dv$ the value dimensionality, and is verified numerically against PyTorch at full double-precision floating-point on concrete inputs. Unlike hardware-specific accelerators or empirical tiling schemes such as FlashAttention, MoA simultaneously provides array fusion, shape-transformation correctness, and predictive cost models from a single algebraic framework. Memory minimality is a theorem established before any code is written. A predictive performance model projects $2$--$100\times$ speedup and $2$--$50\times$ energy reduction, with the advantage widening at exascale. The derivation establishes a formally verified pipeline from Python specification through (ONF) Operational Normal Form, and dimension-lifted hardware mapping, providing performance-portable AI kernels of direct relevance to DARPA edge-deployment and DOE exascale priorities.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。