大模型推理性能优化之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]， ………………………………