大语言模型推理优化--键值缓存--Key-value Cache

文章目录

[一、生成式预训练语言模型 GPT 模型结构](#一、生成式预训练语言模型 GPT 模型结构)
[二、FastServe 框架](#二、FastServe 框架)
[三、Key-value Cache](#三、Key-value Cache)
- 1.大模型推理的冗余计算
- [2.Self Attention](#2.Self Attention)
- [3.KV Cache](#3.KV Cache)

一、生成式预训练语言模型 GPT 模型结构

目前，深度神经网络推理服务系统已经有一些工作针对生成式预训练语言模型 GPT 的独特架构和迭代生成模式进行优化。GPT 架构的主要部分是一堆叠的 Transformer 层，如图所示。在Transformer 层中，掩码自注意力模块是与其他深度神经网络架构不同的核心组件。对于输入中的每个词元，它派生出三个值，即查询（Query）、键（Key）和值（Value）。将查询与当前词元之前所有词元的键进行点积，从当前词元的角度衡量之前词元的相关性。由于 GPT 的训练目标是预测下一个词元，因此通过 Transformer 中的掩码矩阵实现每个词元不能利用其位置之后的信息。之后，对点积使用 Softmax 函数以获得权重，并根据权重对值进行加权求和以产生输出。

二、FastServe 框架

FastServe系统，是由北京大学研究人员开发，针对大语言模型的分布式推理服务进行了设计和优化。整体系统设计目标包含以下三个方面：（1）低作业完成时间：专注于交互式大语言模型应用，用户希望作业能够快速完成，系统应该在处理推理作业时实现低作业完成时间；（2）高效的 GPU 显存管理：大语言模型的参数和键值缓存占用了大量的 GPU 显存，系统应该有效地管理 GPU 显存，以存储模型和中间状态；（3）可扩展的分布式系统：大语言模型需要多个 GPU 以分布式方式进行推理，系统需要是可扩展的分布式系统，以处理大语言模型推理作业。

三、Key-value Cache

1.大模型推理的冗余计算

我们先简单观察一下基于Decoder架构的大模型的生成过程。假设模型只是一层Self Attention，用户输入"中国的首都"，模型续写得到的输出为"是北京"，模型的生成过程如下：

将"中国的首都"输入模型，得到每个token的注意力表示（绿色部分）。使用"首都"的注意力表示，预测得到下一个token为"是"（实际还需要将该注意力表示映射成概率分布logits，为了方便叙述，我们忽略该步骤。
将"是"拼接到原来的输入，得到"中国的首都是"，将其输入模型，得到注意力表示，使用"是"的注意力表示，预测得到下一个token为"北"。
将"北"拼接到原来的输入，依此类推，预测得到"京"，最终得到"中国的首都是北京"
在每一步生成中，仅使用输入序列中的最后一个token的注意力表示，即可预测出下一个token。但模型还是并行计算了所有token的注意力表示，其中产生了大量冗余的计算（包含qkv映射，attention计算等），并且输入的长度越长，产生的冗余计算量越大。例如：
在第一步中，我们仅需使用"首都"的注意力表示，即可预测得到"是"，但模型仍然会并行计算出"中国"，"的"这两个token的注意力表示。
在第二步中，我们仅需使用"是"的注意力表示，即可预测得到"北"，但模型仍然会并行计算"中国"，"的"，"首都"这三个token的注意力表示。

2.Self Attention

参考：https://mp.weixin.qq.com/s/7Fm8LbUN9jQ2HqxPbUU7UQ