Transformer-based models have emerged as one of the most widely used
architectures for natural language processing, natural language generation, and
image generation. The size of the state-of-the-art models has increased
steadily reaching billions of parameters. These huge models are memory hungry
and incur significant inference latency even on cutting edge AI-accelerators,
such as GPUs. Specifically, the time and memory complexity of the attention
operation is quadratic in terms of the total context length, i.e., prompt and
output tokens. Thus, several optimizations such as key-value tensor caching and
FlashAttention computation have been proposed to deliver the low latency
demands of applications relying on such large models. However, these techniques
do not cater to the computationally distinct nature of different phases during
inference.
To that end, we propose LeanAttention, a scalable technique of computing
self-attention for the token-generation phase (decode-phase) of decoder-only
transformer models. LeanAttention enables scaling the attention mechanism
implementation for the challenging case of long context lengths by re-designing
the execution flow for the decode-phase. We identify that the associative
property of online softmax can be treated as a reduction operation thus
allowing us to parallelize the attention computation over these large context
lengths. We extend the "stream-K" style reduction of tiled calculation to
self-attention to enable parallel computation resulting in an average of 2.6x
attention execution speedup over FlashAttention-2 and up to 8.33x speedup for
512k context lengths.

LeanAttention 是一种可扩展的自注意力计算技术，通过重新设计解码阶段的执行流程，将自注意力机制的实现扩展到具有挑战性的长上下文长度情况，以并行计算的方式提供 2.6 倍的平均注意力执行加速和最多 8.33 倍的速度提升。