其实,KV Cache 配置开启后,推理过程可以分为2个阶段:
- 预填充阶段:发生在计算第一个输出token过程中,这时Cache是空的,计算时需要为每个 transformer layer 计算并保存key cache操作,推理速度慢。
- 使用KV Cache阶段:发生在计算第二个输出token至最后一个token过程中,这时Cache是有值的,每轮推理只需读取Cache,同时将当前轮计算出的新的Key、Value追加写入至Cache;FLOPs降低,gemm变为gemv操作,推理速度相对第一阶段变快,这时属于Memory-bound类型计算。
这里用图可能更有助理解,下图是一个Decoder Block,含有Self-Attention和MLP,标红部分为KV Cache影响到的内容,即KV Cache开启后,标红的序列长度 sss 变为 1,当batch_size=1时,Self-Attention中的2个dense全都变为gemv操作,MLP中的dense也全都变为gemv操作。看懂这个图就可以答对上面的3个问题啦。
参考
- https://zhuanlan.zhihu.com/p/662498827
- Transformer Inference Arithmetic | kipply's blog
- https://www.zhihu.com/question/596900067/answer/3040011798
- 【2023 Meet TVM · 北京】冯思远-利用 MLC-LLM 把语言模型部署到任意设备_哔哩哔哩_bilibili
- https://zhuanlan.zhihu.com/p/624740065
- https://huggingface.co/docs/transformers/llm_tutorial_optimization