大模型推理优化实践：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 哈希到固定的机器。只要选择合适的哈希算法，就能在机器集群负载均衡的同时让多轮对话命中同一台机器。