服务器 频道

大模型推理优化实践: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 总结

  我们在店铺起名和文案生成两类任务,对原模型和投机采样模型进行了性能对比。其中原模型是int8量化的Qwen13B模型,投机采样使用量化后的Qwen13B和Qwen1.8B模型,在A10机器测试结果如下:  

  在两类任务下投机采样对模型均有加速,效果随接受token数和输入token长度变化。我们测得在使用上述条件每轮生成5个token的情况下,短序列跑一轮需要60ms,长序列跑一轮需要70ms。而原模型跑一轮需要30ms,因此长/短序列只有在拒绝全部token的情况下会劣于原模型,其他情况则是与原模型平均token时间相近或者优于原模型。由于测试条件限制,相比原论文的大小模型比例(70B:6B),我们大模型和小模型的规模(7B:1.8B)更接近,并且Qwen中文模型的词表大小是152064,相比其他模型(如Llama词表大小是32000)大了数倍,这也额外开销时间变大,因此在其他测试场景下应该还能有更好的表现。

  存在的问题

  上文我们着重介绍了两个优化对推理速度的影响,但除了推理速度外,并行度也是影响大模型吞吐的重要因素。影响并行度的主要因素是显存,大模型的显存占用分三块:模型权重占用显存、运行时显存和KV cache显存,KV cache显存越多,模型能够同时承载的请求数越多,并行度越大。

  序列长度较长时,对运行时显存运行最大的是Softmax Buffer,它的大小和序列的平方成正比,知名优化FlashAttention除了能降低模型第一次运行时间外,更重要的作用是消除了Softmax Buffer对显存的占用。但比较遗憾的是,FlashAttention优化的开源实现要求Attention计算的QKV维度一致,而KV csache复用和投机采样都未满足这个条件,导致对显存有额外的占用。除此之外,投机采样因为要额外加载小模型的权重,且运行时需要多保存一份小模型的KV cache,还需要额外的显存。

  总结与致谢

  以上是我们在大模型推理上做的一些优化尝试,有根据业务场景和实际问题的,也有参考论文实现的,并且都取得了一定的加速效果。但是从极致性能的角度,我们做的还远远不算完美,这些功能在算子层和框架层都还有优化空间,这些是我们后续需要改进的。

  除了上述介绍的功能外,RTP-LLM还支持了非常多的功能,和上文相关的对System Prompt进行缓存的Multi Task Prompt复用Medusa投机采样,以及动态LoRA和不规则剪枝模型支持。未来我们也会持续的添加新功能,优化底层算子性能,打造更好的大模型推理框架。

  我们的项目主要基于FasterTransformer,并在此基础上集成了TensorRT-LLM的部分kernel实现。FasterTransformer和TensorRT-LLM为我们提供了可靠的性能保障。Flash-Attention2和cutlass也在我们持续的性能优化过程中提供了大量帮助。我们的continuous batching和increment decoding参考了vllm的实现;采样参考了hf transformers,投机采样部分集成了Medusa的实现,多模态部分集成了llava和qwen-vl的实现。感谢这些项目对我们的启发和帮助。

0
相关文章