[LLM性能优化]聊聊长文本推理性能优化方向

[LLM性能优化]聊聊长文本推理性能优化方向

原创 阿杰 大模型新视界 2024年07月09日 08:00 四川

原文：https://zhuanlan.zhihu.com/p/698308542

近期，LLM 的长文本能力越来越受到关注。LLM 处理长文本的能力可以应用在多个应用场景中，例如 LLM Agent 场景：假设 Agent 会调用不同的工具解决用户给出的任务，所以当用户对 Agent 提出一个任务时，Agent 会先调用一次 LLM，对给定的任务生成一系列的 Funtion Call，然后依次调用不同的 Funtion，Agent 将 Funtion 的所有输出结果作为输入，再调用一次 LLM，生成最终呈现给用户的自然语言。其中，通过 Function 返回的结果可能很长，多个Function结果拼起来可能是一个很长的输入，这样 Agent 模型就需要具备长文本处理能力。除了 Agent 以外，RAG、文本摘要都需要 LLM 模型具备长文本处理能力，这些应用在落地时需要 LLM 推理服务具备很高的长文本推理效率。笔者之前在介绍 vLLM 的文章中 介绍过与 LLM 推理服务性能最关键的因素：KVCache 显存大小。KVCache 显存占用的计算公式为：

其中:

通过公式（1）我们可以看到， LLM 推理服务为每个 Token 分配了显存空间，其显存空间大小与模型层数、头维度、头数量以及KVCache存储的数据类型大小四个维度相关。通过公式（2）我们可以看到，KVCache 的总显存开销与 LLM 推理服务处理的总 Token 数以及每个 Token 所占的显存开销相关。在长文本推理服务场景中，我们的目标是在给定显存空间，推理服务要处理的 Token 数尽可能的多，即 T 尽可能的大。通过公式 2 可知，提升 T 有两种策略：

下面将分别介绍这两种策略相关的工作。

1. 压缩 KVCache 长度

LLM 最核心的模块是 Attention：通过 Q 、 K 两个矩阵进行矩阵乘，计算O(T2)大小的注意力得分矩阵，并使用注意力得分矩阵对 V 矩阵求加权平均值，得到注意力层的输出矩阵。其中注意力得分矩阵每一行表示 Q 中每个 Token 和 K 中每个 T