面向Transformer解码阶段的轻量化注意力机制：硬件感知的可扩展注意力机制
硬件体系结构
2025-01-15 v2 机器学习
摘要
Transformer-based模型已成为自然语言处理、自然语言生成和图像生成中最广泛使用的架构之一。该架构的规模不断增长，已达到数十亿参数。这些巨大的模型在内存上非常耗费，即使是在最先进的AI加速器(如GPU)上，也会造成显著的推理延迟。 Specifically，注意力操作的时间和空间复杂度随总上下文长度(即提示和输出标记)呈二次增长。因此，已提出各种优化技术，如键值张量缓存和FlashAttention计算，以满足此类大型模型的低延迟需求。然而，这些技术并不针对推理过程中不同阶段的计算特性进行设计。为此，我们提出LeanAttention，即一种针对解码器-only Transformer模型的标记生成阶段(解码阶段)计算自注意力的可扩展技术。LeanAttention通过重新设计解码阶段的执行流程，实现了对长上下文长度的可扩展注意力机制。我们确定，在线softmax的关联属性可被视为归约操作，从而允许我们在这些大型上下文长度上对注意力计算进行并行化。我们将“stream-K”风格的块状计算归约扩展到自注意力，从而实现并行计算，平均实现2.6倍的注意力执行加速，最高可达8.33倍512k上下文长度的加速。
关键词
引用
@article{arxiv.2405.10480,
title = {Lean Attention: Hardware-Aware Scalable Attention Mechanism for the Decode-Phase of Transformers},
author = {Rya Sanovar and Srikant Bharadwaj and Renee St. Amant and Victor Rühle and Saravan Rajmohan},
journal= {arXiv preprint arXiv:2405.10480},
year = {2025}
}
备注
13 pages, 10 figures