为特定应用定制大型语言模型(LLM)的方法之一是为它们提供详细的系统提示。随着 LLM 上下文长度的不断增加,你可以为它们提供更长的系统提示,以提高它们在应用程序中的性能。
然而,随着系统提示变得越来越长,它们可能会减慢语言模型的运行速度,并导致吞吐量瓶颈。香港城市大学和商汤科技的研究人员开发了一项名为RelayAttention的新技术,该技术提高了涉及长系统提示的LLM服务的效率。
RelayAttention通过消除计算因果注意力所需的冗余内存访问来提高推理速度,而且这样做不会降低模型响应的质量。
系统提示与LLM成本
LLM服务通常使用特定于应用的系统提示来指定任务的指令。系统提示可以包含各种组件,包括指令、领域特定知识和对话示例。对该服务的所有请求都共享系统提示。当用户发送请求时,它会与系统提示组合并发送到LLM。
长系统提示可能会严重降低服务的推理吞吐量和延迟。在推理过程中,转换器模型会预测序列中出现的下一个标记。
在每个时间步中,LLM会注意所有先前的标记来生成下一个标记。自然,随着序列变长,计算成本也会增加。
减少LLM推理计算成本的一种常见方法是计算“键值缓存”。LLM是自回归模型,这意味着每个新标记仅依赖于其前序标记。键值缓存计算并存储序列中较早标记的隐藏值,并在后续生成中重用它们。在生成新标记时,模型不再需要为所有标记计算注意力值,而只需计算新标记的值并将其添加到缓存中。
在使用应用级系统提示的LLM服务中,另一种优化推理的方法是预先计算系统提示的键值缓存,并在请求中重复使用它。各种技术利用推理计算算法来减少LLM的计算和内存占用。
RelayAttention
RelayAttention提出了一种新的方法来提高使用长提示的LLM服务的效率。
研究人员写道:"我们观察到的关键问题是,在因果注意计算过程中,不仅存在与系统提示相对应的冗余内存占用和计算,而且还存在不必要的内存访问。”
尽管应用程序中的所有请求都共享系统提示,但当前的注意力算法(如PagedAttention和FlashAttention)在批处理中的每个请求中都会从内存中读取其隐藏状态(即KV值)。
“为了消除这种冗余的内存访问,我们提出了RelayAttention,这是一种基于其数学重新表述来精确计算因果注意力的算法。”研究人员写道。
RelayAttention的关键思想是将与系统提示相对应的矩阵向量乘法组合成矩阵矩阵乘法。这允许LLM服务从DRAM中一次性加载系统提示的隐藏状态,以供批处理中的所有请求标记使用。
这个过程将因果注意力层的计算分为三个步骤:系统注意力步骤、上下文注意力步骤和中继融合步骤。
“在系统注意力和上下文注意力步骤中,我们计算两个中间注意力输出,就好像LLM仅由共享的系统提示/特定于请求的上下文触发一样。”研究人员写道。“在中继融合步骤中,我们将最终输出计算为两个中间输出的凸组合。”
研究人员实施了两项主要的架构创新,以在现有的推理系统中改进RelayAttention。首先,他们没有使用单个KV缓存来存储系统提示和特定于请求的上下文,而是使用了一个单独的缓存来存储系统KVs,并在为模型提供服务之前离线填充它。
“这可以看作是PagedAttention中的前缀共享和PromptCache的组合,前者消除了系统KVs的冗余内存占用,后者消除了提示阶段中的冗余计算。”他们写道。
其次,由于系统KVs已经离线计算完成,他们通过系统提示的长度偏移特定于请求的上下文标记的位置,以确保位置嵌入的正确性。
RelayAttention在实际应用中的表现
研究人员在ShareGPT和MMLU数据集以及三种不同的GPU上测试了RelayAttention。他们使用了三种不同配置的vLLM,vLLM是一个流行的开源库,专为高吞吐量LLM服务设计。
两种配置使用PagedAttention和PromptCaching,而第三种配置使用RelayAttention。研究人员发布了使用RelayAttention与vLLM的代码。
将RelayAttention集成到vLLM中,对于聊天机器人工作负载的Llama2-7B模型,可持续请求率提高了高达2.2倍,吞吐量提高了2倍。研究人员观察到,对于其他几个流行的LLM,也实现了类似的效率提升。随着系统提示的变长,效率增益继续增长。有趣的是,随着提示大小的增加,RelayAttention能够保持其吞吐量,这对于能够支持跨越数十万个标记的上下文的新模型尤其有用。
值得注意的是,RelayAttention特别适合批量推理。批次大小越大,RelayAttention的效率就越高。另一方面,当只有一个请求时,例如在设备上运行LLM时,RelayAttention并没有帮助。“因此,RelayAttention适用于云服务场景。”研究人员写道。
随着LLM在不同环境和设备上的部署,研究人员正在寻找新的方法来使它们运行得更快且占用更少的内存。RelayAttention是优化LLM推理的几种技术之一。最近,Apple推出了“快速LLM”技术,该技术用于减少LLM在边缘设备(如笔记本电脑和智能手机)上的内存占用。另一篇来自苏黎世联邦理工学院的研究人员的论文提出,通过重新排列Transformer架构来消除不必要的计算,从而将推理速度提高高达300倍。