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
n
e
w
q_{new}
qnew时,计算得到新的
k
n
e
w
k_{new}
knew和
v
n
e
w
v_{new}
vnew,并将其拼接(concat)到缓存的
K
p
r
e
v
K_{prev}
Kprev和
V
p
r
e
v
V_{prev}
Vprev中。
下面这个动图非常清晰:
假设
T
T
T是序列长度,
D
D
D是维度(也就是上图的emb_size)。
在没有cache的情况下:
- Q : ( T , D ) Q: (T, D) Q:(T,D)
- K T : ( D , T ) K^T: (D, T) KT:(D,T)
- V : ( T , D ) V: (T, D) V:(T,D)
- Q K T : ( T , T ) QK^T: (T, T) QKT:(T,T)
- A t t e n t i o n O u t p u t : ( T , D ) Attention Output: (T, D) AttentionOutput:(T,D)
可以看到,每来一个新的query token后,都需要重新计算一遍注意力,复杂度是 O ( T 2 ) O(T^2) O(T2),这也就是原始的Transformer平方复杂度。 不过在Prefill阶段确实是这个计算过程。
在有cache的情况下:
- Q : ( 1 , D ) Q: (1, D) Q:(1,D)
- K T : ( D , T ) K^T: (D, T) KT:(D,T)
- V : ( T , D ) V: (T, D) V:(T,D)
- Q K T : ( 1 , T ) QK^T: (1, T) QKT:(1,T)
- A t t e n t i o n O u t p u t : ( 1 , D ) Attention Output: (1, D) AttentionOutput:(1,D)
可以看到,每来一个新的query token后,不需要重新计算一遍注意力,而是只关注计算新token的注意力,复杂度是 O ( T ) O(T) O(T),降低到了线性。
为什么这个优化很重要?
如上图所示,通过KV cache获得的矩阵要比没有KV cache小得多,这导致了更快的矩阵乘法。
Memory Usage分析
Transformer中注意力层中KV的存储开销计算公式:
下面是一个具体的例子,可以看到KV cache的大小竟然是模型的3倍,这在推理场景非常常见。在内存使用中,KV cache往往是主导因素。