大模型推理阶段KV计算分析：prefill和decode

大模型推理阶段分为两个过程，prefill和decode。prefill是用户输入完prompt到生成首个token的过程，decode则为生成首个token到推理停止的过程。

在prefill（预填充）阶段，大模型一次性对prompt中所有token进行计算QKV，由于不同token的计算是独立的，因此该过程可以并行。在attention部分，由计算得到的QKV进一步计算出Output矩阵，再经过后续的FFN层和解码得到首字母token。

在decode（解码）阶段，计算原理和prefill完全相同，但计算方式和decode是不能一样的。原因有两个。一是随着seqlen增加，Attention计算复杂度平方级增长，直接计算代价很大，导致长序列的推理时间极慢，甚至不可行。二是对decode阶段的计算过程简单分析发现，该过程可以复用prefill阶段的KV结果，也可以复用docode阶段已经产生的KV结果。综上，可以把已产生的KV存起来，不必重新计算，这就是KV cache。对于Q矩阵，每次需要计算的只是Q的最后一行q，计算关于qKV的attention，而不是关于QKV的attention，这样复杂度降低了一个量级，实现以存换算。

转自：https://zhuanlan.zhihu.com/p/716688741