原文：https://zhuanlan.zhihu.com/p/717581669
大模型推理：一个简单的例子
为了便于理解大模型推理背后有哪些步骤，我们先假设大模型一次只处理一条文本（也就是先只考虑batch size为1的情形），例如：
输入是“What color is the sky”，在大模型推理中输入也称为prompt
输出是“The sky is blue.“，在大模型推理中输出也称为completion
使用基于transformer的大模型，从上述prompt生成completion的过程分为如下几个步骤：
第1步，将大模型的模型权重加载到GPU的显存中
第2步，在CPU上对prompt做tokenization（分词），并将token的张量表示（token_ids）从内存传输到到GPU的显存中：
欢迎加入自动驾驶实战群
第3步，让token_ids走一遍网络得到输出的第一个token：
这一步只需要走一遍网络，也被称为prefill（预填充）阶段。
第4步，将生成的token拼接到输入的token的结尾，并以此作为一个新的输入来生成下一个token。重复上述过程，直到生成一个终止符的token或者达到配置的最大生成长度。终止符token也被称为end-of-sequence (EOS) token。
这一步需要反复走很多遍网络，也被称为decoding（解码）阶段。
第5步，将生成的所有token从GPU的显存传输回CPU，然后由CPU将生成的id映射回文本，这就是最终生成的completion：
大模型推理背后的两阶段过程
总的来说，大模型推理是一个两阶段的过程：
预填充（prefill）阶段：在这个阶段，模型处理输入的全部prompt，并进行前向计算。这个阶段的目的是生成第一个输出token，即响应的起始点。
解码（decoding）阶段：一旦prefill阶段完成，模型进入decoding阶段，逐个生成剩余的响应token。
在上面的图示中，预填充阶段和解码阶段都需要经过神经网络的前向计算，仅从逻辑上看二者本没有本质区别。但是在接下来的文章中，我们可以看到，这两个阶段的分离是基于它们在资源需求和优化策略上的根本差异。预填充阶段由于其计算密集性，可以充分利用GPU的并行处理能力，而解码阶段则因为内存密集性，可能受到内存带宽的限制。此外，预填充阶段完成后，可以独立于解码阶段运行，这允许系统更灵活地分配资源，优化处理效率。
对于预填充阶段和解码阶段更详细的介绍，超出了本文的范围，我们可以先了解这两个阶段的性能评估指标，这些指标会在实验中用到。
预填充阶段的关键指标是TTFT（Time To First Token），即生成第一个token所需的时间。
解码阶段的关键指标是TPOT（Time Per Output Token），即生成每个响应token所需的平均时间。
这些指标对于评估和优化大模型推理性能至关重要。
引入KV Cache
KV Cache即为一种针对解码阶段的加速方法，KV Cache中的K和V指的是transformer的Self-Attention层中计算scaled dot-product attention时使用的Key和Value向量，如下图所示（图片来自https://lilianweng.github.io/posts/2018-06-24-attention/#full-architecturests/2018-06-24-attention/#full-architecture）：
在上图中Multi-Head Attention的输出等价于多个独立的单头注意力的拼接，假设输入的（行）向量序列为：
那么Multi-Head Attention的输出可以形式地表示为:
简单起见，这里省略了Attention矩阵的缩放因子。上面的公式中可以看出，在解码阶段，逐个生成token时，新预测出来的第t+1个token，并不会影响已经算好的
因此这部分结果我们可以缓存下来供后续生成调用，避免不必要的重复计算，这就是所谓的KV Cache。
使用KVCache后，Multi-Head Attention的计算过程表示如下：
KV Cache的图示
如果你觉得对KV Cache形式化的表达过于抽象，不妨参考下面这篇博客来理解KV Cache的实现原理：
https://medium.com/@joaolages/kv-caching-explained-276520203249
下面的Step1-4即为“Transformers KV Caching Explained”（https://medium.com/@joaolages/kv-caching-explained-276520203249）中提供的不使用KV Cache和使用KV Cache的对比：
测试KV Cache对大模型推理速度的提升
基于transfomers库测试KV Cache对于推理速度的提升，transformers库的generate方法通过use_cache参数来指定是否开启KV Cache：
https://discuss.huggingface.co/t/what-is-the-purpose-of-use-cache-in-decoder/958
测试代码如下：
import numpy as np
import time
from transformers import AutoModelForCausalLM, AutoTokenizer
checkpoint = "distilgpt2"
tokenizer = AutoTokenizer.from_pretrained(checkpoint)
model = AutoModelForCausalLM.from_pretrained(checkpoint)
for use_cache in (True, False):
times = []
for _ in range(10): # measuring 10 generations
start = time.time()
model.generate(**tokenizer("What is KV caching?", return_tensors="pt"),
use_cache=use_cache,
max_new_tokens=20)
times.append(time.time() - start)
print(f"TPOT {'with' if use_cache else 'without'} KV caching: "
f"{round(np.mean(times), 3)} +- {round(np.std(times), 3)} seconds")
for use_cache in (True, False):
times = []
for _ in range(10): # measuring 10 generations
start = time.time()
model.generate(**tokenizer("What is KV caching?", return_tensors="pt"),
use_cache=use_cache,
max_new_tokens=1)
times.append(time.time() - start)
print(
f"TTFT {'with' if use_cache else 'without'} KV caching: "
f"{round(np.mean(times), 3)} +- {round(np.std(times), 3)} seconds")以下是在我的笔记本电脑Mac M1上的运行结果：
TPOT with KV caching: 0.331 +- 0.011 seconds
TPOT without KV caching: 2.037 +- 0.056 seconds
TTFT with KV caching: 0.047 +- 0.002 seconds
TTFT without KV caching: 0.05 +- 0.005 seconds可以看到，KV Caching对于解码阶段的指标TPOT提升较大（Time Per Output Token），而对于预填充阶段的指标TTFT则没有影响（Time To First Token）。
KV Cache带来的问题
KV Cache实际上是一种用空间来换时间的做法，因此不可避免地带来一些新的问题：
第一，通常情况下，在gpu上做大模型的推理时，开启KV Cache会占用更多的GPU显存，更多的显存占用一方面限制了序列的总长度，影响了上下文的最大窗口长度；另一方面也限制了一次能处理的序列的总数量，也就是batch_size，进而影响到了系统的吞吐量。
第二，KV Cache减少了gpu在一次推理中的运算量，但是花费了更多的时间在读取缓存数据上。通常，加载1MB的数据到GPU核中的时间，比在GPU核上对1MB的数据做计算的时间更长，因此KV Cache开启后可能会导致GPU的使用率比较低，从成本上来讲这也是不划算的。
围绕以上两点，很多大模型推理框架都从减少KV Cache的大小/降低KV Cache的碎片化/提升KV Cache的访存效率等角度出发来做系统的优化。
最后别忘了，帮忙点“在看”。
您的点赞，在看，是我创作的动力。
AiFighing是全网第一且唯一以代码、项目的形式讲解自动驾驶感知方向的关键技术。
长按扫描下面二维码，加入知识星球。
