探索vLLM分布式预填充与KV缓存：提升推理效率的前沿技术

最新推荐文章于 2025-04-09 20:00:00 发布

小天才学习机打游戏

最新推荐文章于 2025-04-09 20:00:00 发布

阅读量1.4k

点赞数 30

文章标签：分布式缓存知识图谱人工智能数据库 milvus oracle

本文链接：https://blog.csdn.net/m0_59164520/article/details/144409472

版权

本次分享来自Kuntai，一名芝加哥大学的研究生，主要内容大纲如下：

1.分布式推理与vLLM关于分布式预填充的PR。

2.KV 缓存共享的内容，包括更快的 KV 缓存传递以及多种 KV 缓存的组合。

首先，从分布式推理开始：在与许多行业人士交流后，我发现并非所有人都了解为何需要进行预填充与解码功能的解耦。简单来说，如果将预填充（prefill）和解码（decode）功能放在同一个节点中，会导致资源竞争。

预填充和解码功能的资源竞争的具体表现为：

1.当预填充任务与解码任务一起批处理时，单个 Token（ITL）的延迟会很高，因为预填充本身速度较慢。

2.如果将两者放在同一节点，它们必须共享相同的并行策略，无法单独优化首次 Token 生成时间（由预填充主导）或单个 Token 间延迟（由解码主导）。任何策略变更都会同时影响两者。

这种资源竞争问题还体现在其他场景中：1）推测解码，当工作负载较低时，GPU 会有空闲周期，可用于推测解码，提供额外的吞吐量。然而，当负载较高时，推测解码反而会占用 GPU 资源，导致推理速度变慢。因此，由于语言模型和推测解码器之间的资源竞争，我们无法在高负载下获得额外吞吐量。

这种资源竞争问题还体现在其他场景中：2）多模态处理，多模态中，图像编码器通常是计算密集型（高算力需求），而大语言模型则是内存密集型（高带宽需求）。为了获得最佳性能，理想情况下应将它们部署在不同设备上，以匹配各自的计算和内存需求。

推理架构可能需要从同一节点协作的集中式设计，转变为更分布式的架构。在分布式架构中，每个节点专注于单一任务，比如预填充、解码、推测解码或图像编码等，这些模块通过通信链路相互连接。这样，我们可以：

1.将计算密集型模块部署在计算性价比更高的硬件上；

2.将内存带宽密集型模块部署在内存性价比更高的硬件上；

3.从而整体提升性能。

核心挑战是构建中间层的张量传输（tensor transmission）模块。这需要先理解在这种架构中需要传输哪些张量。主要包括大语言模型请求状态的两部分：Token 和 KV 缓存。其中，KV 缓存远大于 Token，因此我们更关注 KV 缓存的高效传输。

为此，我提交了一个初始 PR，重点解决预填充与解码解耦的问题。之前有不少类似尝试，但大多失败或被搁置，主要原因在于：

1.这些尝试将预填充和解码工作线程直接集成到vLLM的现有推理引擎中，需要为其建立两套系统，包括内存分配器、KV 缓存引擎、调度器等。而这些模块本身已经非常复杂，进一步修改和同步两套系统会带来极高的开发和维护成本。

2.这些方法通常需要数千行核心代码改动，以及大量周边测试和新通信库的引入。

相比之下，我选择了一种更简洁的实现方法：

1. 将预填充和解码工作器分别放在不同的vLLM实例中。

2.使用NCCL建立实例间的通信。

结果是核心代码部分只需增加约数百行代码（不包括分布式通信部分）。性能表现令人满意，例如在传输约2000个token的情况下，延迟为60毫秒，并且还可以进一步优化。

在KV缓存共享方面，我还研究了跨实例传输KV缓存的方法。当前实现使用NVIDIA的 NCCL（用于集群通信），支持点对点（peer-to-peer）通信。这种方法与张量并行不同，更类似于流水线并行化，仅需要进行“发送”和“接收”操作。

针对分布式部署的一些常见问题，例如负载波动导致某些节点资源利用率低下，目前可以通过灵活节点的配置来解决。这些节点既可以执行预填充任务，也可以执行解码任务，从而动态调整资源分配。

当前的实现主要支持NVIDIA GPU平台。未来的计划是扩展到其他硬件平台，并进一步优化通信延迟和模块性能。

关于NVIDIA和AMD的基础设施支持与性能对比：NVIDIA的基础设施支持稍好一些。我尝试过AMD的通信系列，遇到了一些问题。目前我们计划在未来支持，但暂时还没有实现。

关于分离式预填充和块状预填充的性能比较：是的，有性能比较。分离式式预填充在某些特定工作负载下表现更好。例如，在一个节点上测试，分配4个GPU用于解码，4个GPU用于预填充，这种设置下，解码和预填充的运行时间需要大致平衡。这种情况下表现会更优，特别是在大约2000个输入token和10个输出token的理想化工作流中。然而，要实现真正的实际效益，需要多个预填充节点和解码节点，因为解码通常是主要的瓶颈。我目前尚未进行这类测试，因为缺乏硬件支持。

关于性能提升的估算：在最佳场景下，我们的改进可以提升约10%。但需要注意，分离式预填充并不能直接提高吞吐量，而是改善符合服务等级协议（SLO）的请求比例，比如更快的首个token响应时间。分离式预填充能够通过分配不同的并行策略显著加速预填充过程，这是块状预填充无法实现的。

处理不同长度输入的调度问题：对于输入长度差异较大的情况，例如短于100个token的提示和多达4000个token的提示，它们可能会被分配到不同批次，具体情况取决于调度器的实现。

大型语言模型需要使用长上下文（知识）来回答用户的短查询。这些长上下文可能包括书籍、新闻、商业文件、聊天记录等。例如，“今天的新闻是什么？”需要模型读取大量文档以准确回答。KV缓存（键值缓存）可以显著加速长上下文的处理。

我们的目标是通过共享KV缓存，仅对每种不同类型的文本进行一次预填充，避免在不同虚拟机实例间或同一实例的多次运行中重复预填充计算。这样可以提高首token响应时间和吞吐量，同时节省计算成本。

我们想要做的是共享我们称之为知识的缓存，其具体的好处是，它将具有更快的FFTFT (First-Time-To-Fill-Time) 和更高的吞吐量，因为它不需要重新运行预填充计算。而且，由于有二八原则，知识大多可以被重用，因此它不会太贵。实际上，与重新计算相比，即使命中率很低，将K缓存存储在磁盘上也会更便宜。

我们做了一个粗略的计算，如果您的命中率（可能是指缓存命中率）高于每小时一次，那么将其存储在磁盘上就是经济的。

我们设想，这种知识共享系统将有三个组件，并且会有几个挑战。

其中一个挑战是缓存很大，因此我们需要非常快地加载它。我们有一个关键技术可以做到这一点（KV编码，提速3到10倍）。

第二个挑战是，如果文本不再是前缀，那么缓存就不能直接重用。我们需要做一些技巧来确保我们将cation结合在一起，以便它仍然可以被重用。我们在这里还使用了另一个关键技术。

让我们先深入了解第一个关键技术。这将是我们最近的一项发表，其中我们对KV缓存进行了压缩，以更快地加载缓存。

在实际应用中，我们测试了一个示例查询：“用 10 个单词描述 FMAC”。在传统方法下，响应时间约为 2.34 秒，而使用优化后的引擎，响应时间缩短至 0.16 秒，实现了接近20 倍的速度提升。

尽管 KV Cache 可以加速预填充过程，但首先需要将其加载到 GPU 中，而加载时间可能非常漫长。为此，我们采用了以下流程：

1.压缩阶段：通过编码将 KV Cache 转换为二进制表示，并存储在磁盘中。

2.解压阶段：根据请求，检索二进制表示并进行 GPU 解码，将解压后的 KV Cache 加载到 VRAM 进行推理。

这种方法即使在远程磁盘加载 KV Cache 的情况下，也能够显著加快速度。

许多业内人士认为 KV Cache 是随机张量，难以压缩。但我们的研究表明，KV Cache 并非完全随机，且邻近 Token 的 KV Cache 具有相似性。这一性质使得 KV Cache 可以通过以下方式进行有效压缩：

1.利用相邻 Token 的相似性：通过 Delta 编码减少冗余，提高压缩率。

2.差异化量化：在不同层实施精细化的量化策略，兼顾压缩率与推理精度。

3.算术编码：优化数值表示，进一步减少存储和传输开销。

这些技术共同作用，使得 KV Cache 的加载与推理性能得到了显著提升，为大规模模型的高效部署提供了有力支持。

Token 维度的连续性：我们分析了 KV Cache 的一层，并将 Token 和通道的分布可视化。结果显示，沿 Token 维度的分布存在明显的水平纹理特征，这表明相邻 Token 的 KV 值在空间上具有较高的连续性。

压缩启示：由于相邻 Token 的 KV 值非常相似，我们可以通过计算相邻 Token KV Cache 的差值（Delta 值）来减少冗余。具体来说，对于 Token i和i+1的 Delta 值，其分布更接近零，相比于原始分布，Delta 分布更易于压缩。

量化敏感性差异：不同层的 KV Cache 对量化的敏感性各不相同。我们发现：

1.在前几层，量化过于激进会导致精度显著下降，因此需要采用较温和的量化策略。

2.在较深的层（如第 20 至第 23 层），即使进行更激进的量化，精度损失也非常小（仅约 2%-3%）。相比之下，前几层的量化可能导致精度下降高达 50%。

根据这一特性，我们对不同层采取了差异化量化策略：前几层使用温和量化以保留精度，而后几层采用激进量化以实现更高的压缩比。

算术编码的潜力：传统的浮点数表示方式在存储效率上存在不足，因为每个数值需要占用较大的字节数。通过采用算术编码，我们能够有效减少 KV Cache 的存储需求。

算术编码的优势：算术编码通过优化数值的二进制表示，显著降低了表示每个数值所需的字节数。这种方式不仅提升了压缩比，还能在推理过程中快速解码，从而保持高效的加载性能。

在优化KV Cache的过程中，我们进行了多维度的评估和实验，重点关注性能提升与压缩效率的权衡。以下是我们的一些关键发现和优化策略。

我们以Llama 70B模型为例，针对长文本对KV Cache进行了详细评估。在数据中心较低带宽环境（约3GB/s加载KV Cache速度）下，我们对不同精度的量化方法进行了对比分析

1.FP8量化（标记为蓝色圆圈）：在压缩率和准确率之间实现了较好的平衡。

2.FP6量化：进一步减少了缓存大小，但可能对精度产生较大影响。

实验结果表明，我们提出的绿色曲线（即优化策略）在准确率与缓存大小的权衡上表现更佳，同时在首次Token响应时间方面实现了显著提升。

接下来我们探讨第二个关键技术：灵活的合并多个KV Cache，在处理长上下文或多文档时，我们需要连接多个KV Cache。通过对比两种引擎：

1.标准引擎：直接使用vLLM的默认KV Cache策略。

2.改进引擎：支持KV Cache混合的VM引擎。

我们的工作负载包括从文档中检索五个数据块，总共大约有25,000个令牌。让我们通过一些查询来演示这一点。

例如，我们有一个关于FFmpeg的查询，要求用10个词描述FFmpeg。可以看到，使用我们的KV Cache混合特性，响应延迟大约是1.73秒，而标准vLLM的响应延迟大约是5.41秒。虽然输出略有变化，但这种优化并没有改变输出的语义含义，也不会影响准确性。

这种技术引入的原因是，如果一个数据块的缓存后面紧跟着一些新文本，那么这个缓存就不能被重用。这就是为什么我们使用前缀缓存。例如，如果文档一是一则简短的FIFA新闻，文档二是一则更全面的FIFA新闻，我们询问的问题是在FIFA世界杯上，梅西还是C罗进球最多。如果我们将这两个文本连接在一起，我们会得到正确的答案。但如果我们将缓存连接在一起，我们会得到错误的答案。因此，我们需要恢复那些跨注意力的部分，以获得更准确的结果。

使用文本作为输入进行查询时，我们可以在它们之间进行交叉注意力计算，计算它们之间的关系。但如果我们只是简单地将它们连接起来，它们之间就不会有交叉注意力，我们需要恢复这种交叉注意力以获得更准确的结果。不同方法之间存在一个粗略的权衡：如果我们不做K缓存的完整预处理，速度会很慢，但质量会更高；如果我们直接重用KV缓存，质量会大幅下降。我们希望在保持低推理时间的同时，仍然保持高答案质量。我们的直觉是，我们只需要选择性地预填充或F令牌。

为了获得准确的结果，我们在不同方法之间需要做出一个大致的权衡。如果我们不做KV缓存的完整预处理，虽然速度会很慢，但质量会更高。然而，如果我们直接重用KV缓存，质量会大幅下降。我们的目标是在保持低推理时间的同时，仍然保持高答案质量。我们的直觉是，我们只需要选择性地预填充。

基本上，我们仍然在做全部预填充，但在做预填充的同时，我们会检查哪些令牌的缓存在变化，哪些没有变化。我们发现只有一小部分令牌的缓存在变化，我们可以利用这个属性来节省计算资源。以下是示意图：完整的预填充会预填充所有层，但我们的选择性预填充只会根据需要预填充一部分令牌。第一层我们仍然会做完整的预填充，但我们发现一些令牌变化很大，而一些令牌没有变化，我们只会预填充那些变化很大的令牌，而不会预填充那些没有变化的令牌。

我们得到的代表性结果显示，与完整预处理和不预处理相比，我们选择性预处理策略具有更高的F1分数和更低的第一个令牌的时间。有时甚至比完整预处理更好。

数据集是Sansum数据集，这是手机的日常使用数据，模型是Lama 7tb，在两个A40 GPU上运行。你可以在这些演示中自己尝试。我们发现，在成本和延迟方面，所有其他方法都在改变这种权衡，但我们想做的是，我们想要有一个非常高效的知识共享，可以同时降低成本和延迟，那就是带有V的LM缓存。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

在这里插入图片描述