ELSA实现高精度视觉变换器的快速低内存注意力计算，兼容各类硬件。
ELSA: Exact Linear-Scan Attention for Fast and Memory-Light Vision Transformers
- 通过前缀扫描重构在线softmax，保证精确数学语义。
- 在A100上比内存高效SDPA快1.3–3.5倍，边缘设备提速1.5–1.6倍。
- 无需重训练、不依赖张量核心，可跨平台部署于云端与嵌入式设备。
现有注意力加速器常牺牲精确softmax语义，依赖融合张量核，或引入串行深度限制FP32吞吐率。本文提出ELSA，一种在线softmax注意力的算法重构方法：(i) 在实数运算中保持精确softmax语义，且有可证明的$µ\mathcal{O}(u\log n)$ FP32相对误差界；(ii) 将在线softmax更新建模为关于结合幺半群$(m,S,W)$的前缀扫描，实现$O(n)$额外内存和$O("log n)$并行深度；(iii) 独立于张量核心，使用Triton和CUDA C++实现，可作为即插即用替代品，无需重新训练或修改权重。与FlashAttention-2/3不同，ELSA不依赖HMMA/GMMA张量核指令，可在A100与资源受限的边缘设备（如Jetson TX2）上一致运行，是首个在全精度下将并行深度降至$O("log n)$的硬件无关精确注意力核。在A100 FP32基准测试（1K–16K tokens）中，相比内存高效SDPA提速1.3–3.5倍，BERT上提速1.97–2.27倍；在Jetson TX2上相较Math提速1.5–1.6倍，且在LLaMA-13B卸载场景下，≥32K序列时吞吐提升17.8–20.2%。在FP16下，长序列性能接近硬件融合基线，同时保留完整FP32能力，提供跨平台高精度推理统一内核。代码已开源于https://github.com/ming053l/ELSA。
原文摘要 · Abstract (English)
Existing attention accelerators often trade exact softmax semantics, depend on fused Tensor Core kernels, or incur sequential depth that limits FP32 throughput on long sequences. We present \textbf{ELSA}, an algorithmic reformulation of online softmax attention that (i)~preserves exact softmax semantics in real arithmetic with a \emph{provable} $\mathcal{O}(u\log n)$ FP32 relative error bound; (ii)~casts the online softmax update as a prefix scan over an associative monoid $(m,S,W)$, yielding $O(n)$ extra memory and $O(\log n)$ parallel depth; and (iii)~is Tensor-Core independent, implemented in Triton and CUDA C++, and deployable as a \emph{drop-in replacement} requiring no retraining or weight modification. Unlike FlashAttention-2/3, which rely on HMMA/GMMA Tensor Core instructions and provide no compatible FP32 path, ELSA operates identically on A100s and resource-constrained edge devices such as Jetson TX2 -- making it the only hardware-agnostic exact-attention kernel that reduces parallel depth to $O(\log n)$ at full precision. On A100 FP32 benchmarks (1K--16K tokens), ELSA delivers $1.3$--$3.5\times$ speedup over memory-efficient SDPA and $1.97$--$2.27\times$ on BERT; on Jetson TX2, ELSA achieves $1.5$--$1.6\times$ over Math (64--900 tokens), with $17.8$--$20.2\%$ throughput gains under LLaMA-13B offloading at $\ge$32K. In FP16, ELSA approaches hardware-fused baselines at long sequences while retaining full FP32 capability, offering a unified kernel for high-precision inference across platforms. Our code and implementation are available at https://github.com/ming053l/ELSA.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。