大模型推理优化实践：KV cache复用与投机采样_阿里开发者的专栏文章_微信文章

阿里妹导读

在本文中，我们将详细介绍两种在业务中实践的优化策略：多轮对话间的 KV cache 复用技术和投机采样方法。我们会细致探讨这些策略的应用场景、框架实现，并分享一些实现时的关键技巧。

背景

RTP-LLM 是阿里巴巴大模型预测团队开发的大模型推理加速引擎，作为一个高性能的大模型推理解决方案，它已被广泛应用于阿里内部。该引擎与当前广泛使用的多种主流模型兼容，并通过采用高性能的 CUDA 算子来实现了如 PagedAttention 和 Continuous Batching 等多项优化措施。RTP-LLM 还支持包括多模态、LoRA、P-Tuning、以及 WeightOnly 动态量化等先进功能。

随着大模型的广泛应用，如何降低推理延迟并优化成本已成为业界关注的焦点。我们不断地在这一领域内探索和挖掘新方法。在本文中，我们将详细介绍两种在业务中实践的优化策略： 多轮对话间的 KV cache 复用技术和投机采样方法 。我们会细致探讨这些策略的应用场景、框架实现，并分享一些实现时的关键技巧。

多轮对话复用KV cache

在淘宝问问中，有两类多轮对话的场景：一是问答类的场景，它每次请求模型时会拼接之前的问答；二是LangChain应用，它在模型生成结束后会调用外部插件，拼接插件返回的结果再次请求模型。这两类场景共同的问题是： 随着对话轮数的增加，请求长度变长，导致模型的First Token Time（下称FTT）不断变长。

模型的FTT变长，本质上是因为第一次进入模型时，越来越多的token需要生成KV cache。考虑到这两种多轮对话场景存在一个共同点：前一轮对话的输出构成后一轮对话输入的一部分，或者存在较长的公共前缀。且大部分自回归模型（除了chatglm-6b）的Attention Mask都是下三角矩阵：即某一位置token的注意力与后续token无关，因此两轮对话公共前缀部分的KV cache是一致的。进而能够想到的解决办法是： 保存上一轮对话产生的KV cache，供下一轮对话时复用，就能减少下一轮需要生成KV cache的token数，从而减少FTT。 根据这个思路改进前后的模型如下：

2.1 框架设计

用户请求对应的KV cache存放在机器显存中，因此不同轮次的对话需要请求同一台机器，才能复用KV cache。但是在生产环境中，模型部署在由多台机器组成的机器集群， 用户层的请求由统一域名服务转发到机器集群中某一台机器上 ，这样的架构设计导致不同轮对话命中同一台机器的概率微乎其微。

最直观的解决办法是让用户去记录首次请求的机器信息，并将后续请求同一台机器。这个方法可行但是不合理，用户不仅需要感知机器集群的具体信息，还需要对自己链路做大量改造；进而能想到的办法是 增加一层转发层，用户将多轮请求携带同样的标识id并发送给转发层，转发层感知集群信息并匹配标识id和下游机器。 这样不同轮对话就能打到同一台存有KV cache的下游机器。至于如何在转发机器间同步匹配信息，可以使用分布式数据库记录，我们采样的方法是使用统一的哈希算法，将相同id哈希到固定的机器。只要选择合适的哈希算法，就能在机器集群负载均衡的同时让多轮对话命中同一台机器。

在底层实现上，复用KV cache的逻辑和P-Tuning v2在实现上非常相似，通过复用参数，我们使用PTuning的算子支持了KV cache复用。

2.2 总结与反思

我们在Qwen13B/int8量化/A10机器的条件下，对不同输入和前缀长度的请求进行了测试：

可以看到在 复用KV cache功能极大程度的减少了FTT ，并且历史长度的变化对FTT的影响较小 ，FTT更多的取决于本次请求的输入长度。并且除了多轮对话场景外，KV cache复用功能也扩展到复用Ptuning前缀和长System Prompt的场景，降低FTT和显存占用。

虽然复用KV cache的功能能够显著减少多轮对话场景下的FTT，但是在服务压力过大时，存放历史KV cache的显存可能被新请求占用，导致后续请求出现cache miss请求时间变长，加剧服务压力最后导致雪崩。目前我们已实现的解决方案是 使用LRU算法优先移除较旧请求的KV cache。 未来进一步的策略是参照vllm的思路， 将过期的KV cache转移到内存，必要时重新加载至显存。 这种策略比重新计算快，有助于减轻极端情况下的请求延迟，防止服务雪崩。

投机采样

3.1 介绍

投机采样最早在2022年的Fast Inference from Transformers via Speculative Decoding提出，因为不久前的gpt4泄密而被更多人知道。投机采样的设计基于两点认知：在模型推理中，token生成的难度有差别， 有部分token生成难度低，用小参数草稿模型（下简称小模型）也能够比较好的生成； 在小批次情况下，原始模型（下简称大模型）在前向推理的 主要时间在加载模型权重而非计算 ，因此批次数量对推理时间的影响非常小。

基于以上两点认知，投机推理的每一轮的推理变成如下步骤: 1. 使用小模型自回归的生成N个token 2. 使用大模型并行验证N个token出现的概率，接受一部分或者全部token。由于小模型推理时间远小于大模型，因此投机采样在理想的情况下能够实现数倍的推理速度提升。同时，投机采样使用了特殊的采样方法， 来保证投机采样获得的token分布符合原模型的分布，即使用投机采样对效果是无损的。

上图是投机采样的运行过程，每一行的绿色token代表小模型生成并被大模型接受的部分，红色token是小模型生成但被大模型拒绝的部分，蓝色token是大模型根据最后接受token的logits重新采样出来的部分。由上可以看到使用投机采样，在合适的场景下能够大幅提高每轮生成的token数，降低平均单个token生成时间。

3.2 设计思路

我们在RTP-LLM中基于论文的思路，使用大小模型进行了投机采样的实践。在代码设计上我们一方面考虑系统的可维护性，希望这部分能够和原始流程解耦；其次投机采样优化需要与其他优化正交，使投机采样时两个模型都能够用上FT的其他优化。最后我们的设计是 为投机采样封装了一层编排层，对外提供统一的API，在内部组织参数顺序调用正常流程。

3.3 性能评估

在实现过程中，我们着重关注投机采样引入的额外负担。我们希望做到在系统每轮接受token数较少的情况下，也能有与原始模型相近的表现。在实践中，我们测得 额外时间消耗主要有两块：小模型顺序生成token引入时间和采样。

首先最直观的额外消耗，就是小模型推理所占用的时间。在小模型顺序生成N个token时，会有N*T_{small\_model}的时间，在N较大且接受token数少的情况下，这部分开销会非常大。值得一提的时最初我们假设模型消耗时间和参数规模成正比，而实际上这个猜测是错误的。我们测试得到 在参数规模减少的情况下, lm_head在模型调用的占比会显著增加。 以下是Qwen1.8B和Qwen13B在A10/half条件下，单个token在Transformer网络(transformer_layer)和输出层(lm_head)的时间对比:

造成上述比例不一致原因在于：模型参数规模从13B变成1.8B时, Transformer网络在层数(40 -> 24)和权重大小(5120 -> 2048)两个维度减少，而输出层的参数仅从[5120, 152064]变成[2048, 152064] 。同时因为模型词表通常很大，因此输出层的时间通常也比较长。除了输出层的影响以外， 在一些情况下小模型矩阵乘对硬件的利用率并没有大模型这么高 ，因此在选择小模型时需要对这部分开销进行更谨慎的估计。

其次重复多次的采样也引入了巨大的开销。从上面流程图可以看到， 在一轮投机采样流程中，需要进行N次小batch采样和1次大batch采样。 我们以a10/half/vocab_size=152064/top_k=0.5/top_p=0.95的情况下用huggingface采样逻辑进行了测试：

假设在原始请求batch为2, 投机采样每次出5个token的条件下，需要1.15*5+1.47=7.221.15∗5+1.47=7.22毫秒的时间，这接近上述1.8B小模型一次推理的总时长。

好在FT的采样流程针对存在top_k参数的情况，通过融合算子对原版（Huggingface流程）进行了优化，改进后的流程分成两步：对维度是[batch, vocab]的输入进行TopK采样后，使用输出维度是[batch, k]的tensor进行后续流程；省略TopP步骤，直接在采样过程中对TopP进行判断。改进前后的流程对比如下：

改进后的流程不影响结果分布，且大幅度减少了计算量和kernel数量，极大程度减少了采样需要的时间。我们测试优化后的采样流程需要的时间是原来的1/10。

3.4 总结

大模型推理优化实践：KV cache复用与投机采样

正文

请到「今天看啥」查看全文