AI大模型推理性能优化之KV Cache

前言

本篇介绍KV Cache。

KV Cache（键-值缓存）是一种在大模型推理中广泛应用的优化技术，其核心思想是利用缓存 key 和 value 来避免重复计算，从而提高推理效率。代价是显存占用会增加。

核心思想

在自注意力层的计算中，对于给定的输入序列，模型会计算每个token的key和value向量。这些向量的值在序列生成过程中是不变的。因此，通过缓存这些向量，可以避免在每次生成新token时重复计算，只需计算新token的query向量，并使用缓存的key/value向量进行自注意力计算 。

具体来说，decoder一次推理只输出一个token，输出token会与输入tokens 拼接在一起，然后作为下一次推理的输入，这样不断反复直到遇到终止符。

在上面的推理过程中，每 step 内，输入一个 token序列，经过Embedding层将输入token序列变为一个三维张量[b, s, h]，经过一通计算，最后经logits层将计算结果映射至词表空间，输出张量维度为[b, s, vocab_size]。

当前轮输出token与输入tokens拼接，并作为下一轮的输入tokens，反复多次。可以看出第