RelayAttention：让大型语言模型更高效地处理长提示符

一、前言

虽然大型语言模型 (LLM) 近年来取得了非常显著的进展，也在各种自然语言处理任务中展现出强大的能力。然而，LLM 的在实际的应用落地层面也面临着一些实际挑战，其中之一就是效率和成本问题，导致了在垂直行业实际落地的应用非常少。因为现有的 LLM 通常需要大量的计算资源和内存，这限制了它们在实际应用中的部署。

特别是在处理长文本输入时，LLM 的效率问题尤为突出。这是因为 LLM 需要存储和处理大量的中间结果，而长文本输入会导致中间结果的数量急剧增加。为了解决这个问题，研究人员提出了各种方法来提高 LLM 的效率，例如模型压缩、剪枝和量化等。

最近，由香港城市大学朱磊和商汤研究院联合发布的研究论文《RelayAttention for Efficient Large Language Model Serving with Long System Prompts》中提到 RelayAttention 是一种新的注意力机制，可以有效地提高 LLM 处理长文本输入的效率。RelayAttention 通过减少内存访问冗余来实现这一点，它允许 LLM 在处理一批输入词时仅从内存中读取一次隐藏状态，而不是多次读取。这可以显著减少 LLM 的计算量和内存占用，从而提高其效率。

从上图中的测试结果对比，可以看到 RelayAttention 将 LLM 的吞吐量提高了 75%。这意味着使用 RelayAttention 的 LLM 可以在相同的时间内处理比使用传统注意力机制的 LLM 多 75% 的请求。

1、延迟与系统提示长度： 左侧图表显示了不同注意力机制的延迟随系统提示长度的增加而变化的情况。RelayAttention（蓝色条形图）在所有提示长度中都具有最低的延迟，对于最长的提示长度，它比 PagedAttention（橙色条形图）快 7 倍。这是因为 RelayAttention 可以从 DRAM 中读取隐藏状态一次，用于一批输入词，而 PagedAttention 需要多次读取它们。

2、吞吐量： 右侧图表显示了不同 LLM 配置的吞吐量随请求速率的增加而变化的情况。vLLM-RA（ours）（蓝线）在所有请求速率中都具有最高的吞吐量，对于最高的请求速率，它比 vLLM-PS（橙线）高 75%。这是因为 RelayAttention 每秒可以处理比 PagedAttention 更多的词，因为它的延迟更低。

二、RelayAttention 解决了什么问题？

在实际应用中，大型语言模型 (LLM) 需要高效地处理大量请求，尤其是当这些请求包含用于引导模型响应的冗长系统提示符时。然而，长提示符会增加计算量，并随着序列长度的增加而导致生成新词的成本不断提高，最终造成模型吞吐量和延迟的瓶颈。RelayAttention 旨在就是为了解决这一问题，在不影响生成质量的前提下，提高 LLM 在处理长系统提示符时的效率。

RelayAttention 的主要优势在于它能够减少内存访问冗余。传统的注意力机制需要将隐藏状态从片外 DRAM 传输到片上 SRAM 多次，才能处理每个请求。而 RelayAttention 只需要将隐藏状态从 DRAM 中读取一次，就可以处理一批输入词。这可以显著减少 LLM 的计算量和内存占用，从而提高其效率。

RelayAttention 的另一个优势是它不需要对模型进行重新训练。这意味着它可以很容易地集成到现有的 LLM 中，而无需更改模型本身。

RelayAttention 可以大幅提高 LLM 的吞吐量，因为它可以降低注意力计算的延迟。这意味着 LLM 可以更快地处理每个请求，从而在相同的时间内处理更多的请求。

三、RelayAttention 如何解决问题？

RelayAttention 是一种创新的注意力算法，它针对现有因果注意力算法在处理系统提示符时内存使用效率低下的问题进行了优化。现有的算法需要反复访问内存，将缓存的隐藏状态从片外 DRAM 传输到片上 SRAM，而 RelayAttention 允许这些隐藏状态仅从 DRAM 中读取一次，并用于处理一批输入词。这种对内存访问冗余的减少不仅提高了效率，同时也保证了生成质量。RelayAttention 的另一个优点是简单易用，它基于对因果注意力机制的数学公式的重新构造，无需对模型进行重新训练。

四、RelayAttention 的实际应用情况

研究人员在 ShareGPT 和 MMLU 数据集以及三种不同的 GPU 上测试了 RelayAttention 的性能。他们使用了 vLLM 的三种不同版本，vLLM 是一个专门为高吞吐量 LLM 服务设计的流行开源库。

其中两种配置使用了 PagedAttention 和 PromptCaching 技术，而第三种配置使用了 RelayAttention。研究人员还发布了将 RelayAttention 与 vLLM 一起使用的代码，方便其他研究人员进行测试和应用。

测试结果表明，将 RelayAttention 集成到 vLLM 中，可以将 Llama2-7B 模型在聊天机器人工作负载中的可持续请求率提高 2.2 倍，吞吐量提高 2 倍。研究人员还观察到，其他几种流行的 LLM 也取得了类似的效率提升。

值得注意的是，随着系统提示符长度的增加，RelayAttention 带来的效率提升也更加明显。更重要的是，即使提示符的长度持续增加，RelayAttention 也能够保持其吞吐量，这对于能够支持跨越数十万个词的上下文的新模型来说特别有用。

RelayAttention 特别适用于批量推理，因为批量大小越大，RelayAttention 的效率就越高。但是，当只有一个请求时，例如在设备上运行 LLM 时，RelayAttention 就没有明显的优势。因此，RelayAttention 更适合云服务场景。

随着 LLM 被部署在不同的环境和设备中，研究人员正在寻找新的方法来提高它们的运行速度和降低内存占用。RelayAttention 是优化 LLM 推理的几种技术之一。最近，苹果公司推出了“LLM in a flash”技术，可以减少 LLM 在笔记本电脑和智能手机等边缘设备上的内存占用量。苏黎世联邦理工学院的研究人员也提出了一种新的方法，可以通过重新排列 Transformer 架构来消除不必要的计算，并将推理速度提高高达 300 倍。

五、RelayAttention 的未来发展？

RelayAttention 的出现为 LLM 的运行效率带来了显著提升，特别是在处理系统密集型任务方面。由于它可以无需修改模型本身即可集成，因此有望被 LLM 服务快速采用并得到广泛应用。此外，RelayAttention 的代码已经公开[2]，LLM 社区可以对其进行进一步测试和优化，推动其发展和应用。

六、References

[1]. Lei Zhu, Xinjiang Wang, Wayne Zhang, Rynson W.H. Lau, R. (2024). RelayAttention for Efficient Large Language Model Serving with Long System Prompts: https://github.com/rayleizhu/vllm-ra

[2]. vLLM with RelayAttention integration, A vLLM fork with RelayAttention implemented: https://github.com/rayleizhu/vllm-ra