LLM自回归推理文本生成原理

看到一篇最清晰的讲解，把两阶段（Prefill和Decode）的计算过程和维度变化写得很明白。读完这段文字应该就对LLM的生成过程有一个把握了。
原文：LLM Inference Unveiled: Survey and Roofline Model Insights [arXiv 2024.3]

哪里存在KV Cache？

KV cache发生在多个token生成的步骤中，并且只发生在decoder中（例如，decoder-only模型，如 GPT，或在encoder-decoder模型，如T5的decoder部分），BERT这样的encoder模型不是生成式模型（而是判别性模型），因此没有KV cache。

以下动图是GPT-2以自回归形式生成文本的动态图：

下图是Attention的标准计算方式：

什么是KV Cache？

通过缓存以前的键（Key）和值（Value），我们可以只关注计算新token的注意力。
如下图，每当来一个新的token $q_{new}$时，计算得到新的$k_{new}$和$v_{new}$，并将其拼接（concat）到缓存的$K_{prev}$和$V_{prev}$中。

下面这个动图非常清晰：

假设$T$是序列长度，$D$是维度（也就是上图的emb_size）。
在没有cache的情况下：

• $Q:(T,D)$
• $K^T:(D,T)$
• $V:(T,D)$
• $Q{K}^T:(T,T)$
• $AttentionOutput:(T,D)$

可以看到，每来一个新的query token后，都需要重新计算一遍注意力，复杂度是$O(T^2)$，这也就是原始的Transformer平方复杂度。 不过在Prefill阶段确实是这个计算过程。

在有cache的情况下：

• $Q:(1,D)$
• $K^T:(D,T)$
• $V:(T,D)$
• $Q{K}^T:(1,T)$
• $AttentionOutput:(1,D)$

可以看到，每来一个新的query token后，不需要重新计算一遍注意力，而是只关注计算新token的注意力，复杂度是$O(T)$，降低到了线性。

为什么这个优化很重要？
如上图所示，通过KV cache获得的矩阵要比没有KV cache小得多，这导致了更快的矩阵乘法。

Memory Usage分析

Transformer中注意力层中KV的存储开销计算公式：

下面是一个具体的例子，可以看到KV cache的大小竟然是模型的3倍，这在推理场景非常常见。在内存使用中，KV cache往往是主导因素。

几个非常清楚的学习文档

• Transformers KV Caching Explained
• The KV Cache: Memory Usage in Transformers
• Mastering LLM Techniques: Inference Optimization