1. 自回归生成的特点
大模型(如 GPT 等)在推理时通常采用自回归生成的方式:
- 模型逐个生成 token,每次生成一个新 token 时,需要重新计算注意力。
- 在生成第
t
个 token 时,模型需要基于前t-1
个已生成的 token 来计算注意力权重。
由于自回归生成是逐步进行的,k
和 v
的性质决定了它们可以被重复利用,而 q
则需要每次都重新计算。
2. k
和 v
可以缓存的原因
(1) k
和 v
是基于历史 token 计算的
k
和v
是从输入序列的 token 中生成的,且它们只依赖于每个 token 的嵌入表示。- 在自回归生成中,前
t-1
个 token 的k
和v
已经计算过,并且不会因为后续生成新 token 而改变。 - 因此,这些
k
和v
可以直接缓存下来,在生成新的 token 时重复使用。
(2) k
和 v
是全局共享的
- 在注意力机制中,所有 token 的
k
和v
都会被集中起来形成全局的K
和V
矩阵,供当前 token 的q
查询。 - 这意味着
k
和v
的计算结果是可以复用的,无需每次重新生成。
(3) 减少重复计算
- 如果不缓存
k
和v
,每次生成新 token 时都需要重新计算前t-1
个 token 的k
和v
,这会导致大量的冗余计算。 - 缓存
k
和v
后,只需在生成新 token 时计算该 token 对应的k
和v
,并将其追加到缓存中即可。
3. q
不需要缓存的原因
(1) q
是针对当前 token 的
q
是由当前正在生成的 token 的嵌入表示计算得到的,因此它只与当前 token 相关。- 每次生成新 token 时,
q
都会发生变化,无法复用之前的q
。 - 因此,没有必要缓存
q
。
(2) q
的计算成本较低
q
的计算只需要对当前 token 的嵌入表示进行一次线性变换即可完成,计算量相对较小。- 即使每次都重新计算
q
,也不会显著增加推理时间。
4. 缓存 k
和 v
的实际操作
在实际实现中,缓存 k
和 v
的流程如下:
- 初始化缓存 :在生成第一个 token 时,计算该 token 的
k
和v
,并将它们存储到缓存中。 - 追加缓存 :在生成后续 token 时,计算新 token 的
k
和v
,并将其追加到现有的缓存中。 - 复用缓存 :在计算注意力时,直接从缓存中读取
k
和v
,而不需要重新计算。
这种方式可以显著减少计算和内存访问的开销,尤其是在生成长序列时。
5. 为什么 k
和 v
的缓存对性能至关重要?
(1) 加速推理
- 缓存
k
和v
后,每次生成新 token 时只需要计算该 token 的q
、k
和v
,而不需要重新计算整个序列的k
和v
。 - 这使得推理速度大幅提高,尤其是在生成长序列时。
(2) 降低内存带宽压力
- 如果不缓存
k
和v
,每次生成新 token 时都需要重新计算并加载前t-1
个 token 的k
和v
,这会对内存带宽造成巨大压力。 - 缓存
k
和v
后,可以直接从高速缓存中读取,减少了内存访问次数。
(3) 支持高效的硬件优化
- 现代硬件(如 GPU 或 TPU)对矩阵运算有专门的优化,缓存
k
和v
可以让注意力计算更加高效。 - 例如,通过批处理技术,可以一次性处理多个 token 的
k
和v
,从而充分利用硬件资源。