前言
前一段时间sglang-v0.3.0和vllm-v0.6.0前后脚发布之后,就一直想总结梳理一下现在主流的大模型推理引擎。因为我觉得这也算是一个有意义的节点吧,从此开源大模型推理引擎总算是由"非常粗糙,但是能用"的阶段迈入到了"好用,稍微有那么点粗糙"的阶段。
大模型的推理引擎实际也就是近一两年才开始飞速发展,从最开始的tgi和vllm并驾齐驱到如今sglang、lmdeply的异军突起,整个开源社区都是非常有活力的。但是正如之前所说,从长远的一个视角看如今的开源引擎实际上都还是比较粗糙的,大家都是在摸索中前进。另一方面也是因为现在全世界的目光都聚焦在llm这里,新技术的更新换代太快了,做好一个大模型的推理引擎要做的事情实在是太太太太多了。除了要支持日新月异的新模型和新硬件,还要不断关心学术界最新的paper并且想方设法落地实现。而这些新的想法可能涉及到模型结构、计算策略、调度策略、存储策略、cuda内核、硬件加速等各个层级,这就需要开发者有非常广泛的知识范围和过硬的工程能力。我一直认为大模型推理引擎最难的地方就在于:对模型和硬件的广泛支持以及如何将各种角度的不同优化方法兼容实现。因为写paper的人可以只关心他自己的idea,在transformer库的基础上写个简单demo就行,但是在推理引擎里落地的时候往往就会与其它模块有冲突,需要想办法去做各种兼容。退一步说,即使没有冲突的情况,你也需要对其他基础的优化比较熟悉,你才能在这些的基础上完成新功能的开发。
本文主要分为两部分:
1.总结了一些现在主流推理引擎的现状,给大家在实际应用的时候一些参考吧算是。
2.总结了比较重要的大模型推理优化方法,会持续更新。
主流推理引擎现状
tgi
先从tgi说起吧,作为原本和vllm并驾齐驱的开拓者,有着huggingface的背景可以说是天时地利,但问题真的就是开发人员太不行了,目前是2个月一个版本看起来是属于放弃状态了。这让我想起一模一样的情况就是百度,可以说都是起了个大早赶了个晚集。那么tgi主要有哪些问题呢:
第一,tgi宣称使用了pagedattn技术,但实际上只是在decode时候用了名字叫做"pagedattn"的kernel。我不知道他们的开发人员是真的不明白pagedattn概念还是说开发难度太大,总之结果就是tgi的吞吐会非常差。因为没有pagedattn,预分配显存的时候是按"max_new_tokens"来分配的,导致显存的浪费会非常严重,decode的batch_size上不去。
第二,tgi的cpu和gpu调度是串行模型(vllm都已经改成优化了),导致cpu调度计算时候gpu是闲置的,这也会让吞吐变差。
第三,tgi使用了rust来实现调度逻辑。实际上我觉得没有必要,纯python的开发效果真的比rust高太多了。
最后也就是根本原因,还是开发人员投入不够,版本更新太慢了,各种新的功能都没有,提issue也没什么反馈。
vllm
vllm原本只是作为pagedattn的一个开源实现,但发展到今天已经成为llm推理引擎的标杆了。其优势在
第一,背靠ucb,有着大量且稳定的开发者,作者基本都是在读博士生,github上Contributors已经快600人了,相比于sglang的113人、tgi的134人、lmdeploy的75人、TensorRT的74人,vllm的开发人员投入是最高的。因此vllm对模型的支持和对硬件的支持都是最完善的,以及各种功能也往往是最齐全的。
第二,vllm的社区活跃度是最高的,github上issue和pr都很多。且大量paper都是以vllm作为baseline来开发demo,因此各种新技术的引入vllm是具有更大优势的。基础的各种优化以及进阶的权重量化、kv压缩、speculate decode、chunked prefill、promot cache、constrainted decoding等功能都是完备的。
同时vllm也有一些缺点:
第一,在v0.6.0版本之前,vllm中cpu和gpu的调度是串行的,导致cpu计算时候gpu是闲置的,这也会让吞吐变差。又因为vllm里合入了太多功能,这个情况又更加严重了,vllm高负载下gpu利用率甚至会降低到50%左右。而推动vllm去优化这一块也是因为看到sglang在吞吐的性能上远超于他,才会在v0.6.0里着重优化了这一部分。
第二,vllm里被合入的功能实在是太多了,代码非常臃肿复杂,导致二次开发变得非常困难。月之暗面的mooncake引擎就是基于vllm早期版本的二次开发,不确定他们是固定在那个版本不再同步更新,还是说花了大量工作来保持vllm的同步升级。总之如果你要对开源引擎大改,sglang会是比vllm更好的一个选择。
sglang
sglang也是ucb的团队,但是跟vllm是不同的一拨人,核心团队基本都是交大的。有借鉴了一款叫做lightllm的推理引擎,也import很多vllm代码,后续会完全去掉对vllm的依赖。其优势在于:
第一,sglang的性能是目前最优的。这里说的性能主要是吞吐。sglang通过多进程zmp传输中间数据来cover掉cpu的开销,高负载下gpu利用率可以到80%以上。还有一些美中不足的就是,实现prompt cache的逻辑也被放在了forward()主进程里,不然的话GPU利用率应该可以跟lmdeploy一样保持在95%左右。我看最近好像也已经在优化这部分了,到时候吞吐还能提10%。除此之外,sglang在推理的各个环节都有做各种细致的工程优化。
第二,sglang的代码可拓展性很高,主流功能都有支持的情况下,代码比vllm清晰简单很多,这对于二次开发来说是很重要的。
第三,sglang的社区活跃度虽然比不上vllm,但是作者都很积极地回复issue。而且现在开发节奏基本上是一周一个版本,有些功能还不太完善的地方下一个版本基本马上就更新了。
其缺点在于:
sglang的开发人员属实还是太少了,有些功能确实还不太完善。正如之前所说,一款完备好用的推理引擎要做的事情可太多了。期待能够有更多的稳定开发者能投入进来吧。
lmdeploy
lmdeploy是上海人工智能实验室团队开发的,当时第一次尝试也是让我眼前一亮的。相比如vllm和sglang的python实现,lmdeploy的调度甚至是部分模型代码都是用了C++实现。其优势在于:
第一,cpu调度策略是最优的,高负载下gpu利用率可以稳定在95%左右。也重写了很多cuda kernel,可以看出来在性能这块也是做了很多细致的研究,在sglang新版本出来之前lmdeploy就是当时最快的。
第二,对多模态模型的支持很好,支持了大量的多模态模型。
第三,对国内GPU厂商的硬件支持较好。
其缺点在于:
也是开发人员太少了,功能相比vllm和sglang来说还是少了挺多重要的功能。其次纯python的实现在推理引擎真的够用了,我自己是rust的支持者,但推理引擎里引入C++和rust是否利大于弊就不好说了。在AI场景python的开发效率比其他语言优势大太多了。
tensorrt-llm
tensorrt-llm是英伟达官方的推理引擎,其优势在于:
底层kernel有nvidia黑魔法加持,引擎性能总体来说不弱。
其缺点在于:
不完全开源,代码不可控,不适合用于生产环境。
总结
说在前面:
上述评价都是基于我个人的见解,很难对每个引擎都了如指掌,大家可以作为参考即可。
benchmark是个技术活,不同模型不同硬件不同场景下benchmark的结果差异很大,大家也不要过于纠结。
客观评价:
如果团队不具备二次开发能力,想要快速部署开箱即用,那么vllm是最好的选择。
如果是国内GPU厂商的硬件或者是部署多模态模型,那么lmdeploy是不错的选择。
如果对性能有极致追求,团队具备很强的二次开发能力,sglang是最好的选择。
出于个人情感:
建议大家多支持sglang和lmdeploy。写这篇文章的目的之一也是看到这两个开源项目真的非常优秀,团队都有建微信群可以跟作者直接沟通,大家可以github上搜一下。
大模型推理的主要优化技术
这里主要是列一个目录吧,给大家分享一下近几年收录的优秀的技术博客,顺便简单说几句我的理解。建议大家先看博客,再看paper原文,最后研究代码实现。
pagedattn
猛猿:图解大模型计算加速系列之:vLLM核心技术PagedAttention原理[1]
避免显存的浪费,完整版本实现还是挺复杂的,看这一篇就可以了。
kv cache
Young:大模型推理性能优化之KV Cache解读[2]
为什么q不cache,只有k和v的cache?因为sefl-attention的计算过程里上一轮的k和v会被复用,每一轮的q都是新的
Continous Batching
幻方AI:Continuous Batching:一种提升 LLM 部署吞吐量的利器[3]
方佳瑞:大模型推理核心技术之Continuous Batching和我的WXG往事[4]
推理引擎调度的核心都是围绕这里展开的。
Flash Attention
猛猿:图解大模型计算加速系列:FlashAttention V1,从硬件到计算逻辑[5]
猛猿:图解大模型计算加速系列:Flash Attention V2,从原理到并行计算[6]
这里很重要的是理解计算强度这个概念,这是做kernel层性能优化的基础。
Prefix Caching
DefTruth:[Prefill优化][万字]🔥原理&图解vLLM Automatic Prefix Cache(RadixAttention): 首Token时延优化[7]
首token优化的重要方法,sglang和vllm都有实现。
Speculative Decoding
方佳瑞:大模型推理妙招—投机采样(Speculative Decoding)[8]
小冬瓜AIGC:【手撕LLM-Medusa】并行解码范式: 美杜莎驾到, 通通闪开!![9]
chatGPT的泄密文件里有提到这种技术,是降延迟的重要方法,核心思想是某些token是"显而易见"的,不需要大模型来推理,用一个小的模型来推就行,从而实现加速的效果。
比如:问题"中国的首都是哪里",当decode到"中国的首都是北"的时候,下一个token"京"用小模型也是可以正确推理的。当然了这里一般都是一次性猜测后续多个token,再用大模型验证,命中率越高加速越大。
Speculative Decoding也有很多种,主流的方法是medusa,目前只有vllm和tgi有实现。需要注意的是,Speculative Decoding是有N倍的算力开销的,高水位情况下要关掉。低水位情况下可以用冗余闲置的算力来换时间。
量化
紫气东来:LLM(十一):大语言模型的模型量化(INT8/INT4)技术[10]
现在主流都是训练后量化,目的是为了节约显存,理论上也可以加速很多,但是要看配套kernel的实现。很多kernel没有开发int4/fp8版本的,所以有时候量化后是会比fp16慢很多的。这一块也是vllm支持的比较好。
kv cache 压缩
kvcache压缩有两个方向,都是对准确率有损的:
方佳瑞:LLM推理技术之StreamingLLM:如何拥有无限长生成能力[11]
紫气东来:LLM(二十):漫谈 KV Cache 优化方法,深度理解 StreamingLLM[12]
第一种是面向超长prompt的,对kv cache里slots数量的压缩,比如长度128Ktoken的prompt,可以挑选1024个tokens的kvcache来存储就行。主流有StreamingLLM snapKV等方法。
https://lmdeploy.readthedocs.io/zh-cn/v0.3.0/quantization/kv_int8.html)[13]
第二种是kv cache里slots数量不变,对kv cache从fp16压缩成int8/int4。vllm、sglang、lmdeploy都有实现。
zhang:大模型百倍推理加速之KV cache篇[14]
两种方法理论上时候正交的,主要目的都是省显存,也可以带来一些加速效果。
prefill-decode 分离架构
ZHANG Mingxing:Mooncake (1): 在月之暗面做月饼,Kimi 以 KVCache 为中心的分离式推理架构[15]
目前唯一生产环境应用的只有mooncake了,核心是解决Continous Batching在decode中会被插入prefill从而导致decode卡顿以及decode阶段MFU低下这两个问题。可以用相同数量的GPU服务更多的用户,vllm尝试合入过但是失败了。sglang也已经将prefill-decode分离列入了roadmap。也期待mooncake有一天能全部开源吧。
chunked Prefill
猛猿:图解大模型计算加速系列:分离式推理架构2,模糊分离与合并边界的chunked-prefills[16]
算是prefill-decode的简化版吧,主要也是为了解决decode阶段MFU低的问题,也能解决一些decode卡顿的问题。paper原文写得非常好,建议大家仔细读一遍,会让你对llm推理有更深刻的理解。
constrained decoding
手抓饼熊:SGLang技术分析[17]
openai的api也已经支持了这个方法,上文直接跳到第三章,这个功能有两层作用:
第一个是能约束模型以固定格式输出,对精确度有正向的提升。
第二个是能加速decode,因为指定了输出的格式后,很多token(比如json里的固定的key和{},:)就可以自动生成而不需要真正去推理计算。
sglang里已经完整实现,但是因为jump出来的token是没有kvcache的,这里需要一些trick来解决。目前sglang的这一版是jump后再次进行prefill,所以这里其实还是有点问题的,高负载下服务会不断循环的prefill,服务的吞吐会变低。但是低水位下单条请求的延迟会降低非常多,这取决于被跳过的token有多少。
CUDA Graph
改名机会要过期:CUDA效率优化之:CUDA Graph[18]
主要目的是降低GPU指令(kernel)反复提交的开销,基本各个引擎都是必备了现在,但是也容易跟其它优化冲突。
FlashInfer
yzh119:用FlashInfer加速大语言模型推理中的自注意力操作[19]
Flashattn的升级版,最近各引擎都集成了FlashInfer,大家也可以关注下。
多模态推理
文生图:
小小将:文生图模型之Stable Diffusion[20]
图生文:
HelloWorld:LLaVA(四)图解 LLaVA 推理流程[21]
文生图模型一般都不会集成到大模型推理引擎里,我们讨论的多模态一般指的是图生文系列。实际上就是文生文LLM多了一个将图片转成embeding的小模型,后续流程都是一致的。
大模型RAG
产品经理大群:一文读懂:大模型RAG(检索增强生成)含高级方法[22]
RAG是部署私有知识领域Agent的重要技术,跟SFT相比就是不用重新训练,RAG的数据库可以不断更新。
function call
战士金:大模型工具调用(function call)原理及实现[23]
这也是openai的api已经支持的功能,可以用增强LLM的能力。核心思路就是定义一堆工具,模型推理时候决定是否要用哪个工具,比如大模型不会做数学题,但是可以调用计算器api来直接计算得到准确的结果。
多卡推理
Tim在路上:一文讲清 NCCL 集合通信原理与优化[24]
多卡的优势是能省显存,但是通信有额外开销。LLM训练可以有很多种并行的方案,但是LLM推理一般都只会用TP。在多卡的场景里,除了计算瓶颈和内存瓶颈又多了一个通信瓶颈。PCIE、NCLINK、NCCL这些基础也都是需要去了解的。
至于为什么推理引擎都是做成多卡TP而不是PP,主要是因为从服务器视角看PP的吞吐上限更高,但是从单个请求视角看TP的延迟会更低,在线服务往往不需要那么高的吞吐,延迟更加重要。
nvidia GPU工作原理
GPU的工作原理[25]
无需多言,这个视频务必看5遍以上,所有kernel的优化万变不离其宗。
全文总结
正如之前所言,LLM是近几年才真正火起来的,LLM的推理引擎实际上也都处于摸索起步的阶段。其难点在于要面向各种五花八门的模型和硬件平台,且各种不同角度的优化方法错综复杂需要去想办法兼容。
比如当你阅读到一篇推理优化的paper,你应该有敏锐的直觉:
1.性能提升有多大,可信度怎么样?对比的参照物是什么?能否拓展到所有模型所有场景?
2.是无损的方法吗?是否会对模型准确性造成影响?
3.论文是否给出实现demo?开发难度有多大,开发投入是否值得
3.是完整实现paper还是某些地方可以trick简化?
5.是否跟引擎里其它优化方案有冲突?如何优雅兼容其它模块?
未来展望
在sglang横空出世之前,我真的不认为vllm的开发人员清晰意识到了他们的问题。大家感兴趣可以去看下vllm的V2引擎的设想,作者也总结了很多之前的经验和教训。这里面也是很多我想说的,vllm现在形势确实一片大好,但不彻底重构的话迭代维护会越来越困难。如果V2版本能彻底解决CPU侧调度和各种功能之间的冲突兼容问题,那长远看vllm还是比sglang要有优势的(毕竟人多)。
引用链接
[1]
猛猿:图解大模型计算加速系列之:vLLM核心技术PagedAttention原理: https://zhuanlan.zhihu.com/p/691038809
[2]
Young:大模型推理性能优化之KV Cache解读: https://zhuanlan.zhihu.com/p/630832593
[3]
幻方AI:Continuous Batching:一种提升 LLM 部署吞吐量的利器: https://zhuanlan.zhihu.com/p/657586838
[4]
方佳瑞:大模型推理核心技术之Continuous Batching和我的WXG往事: https://zhuanlan.zhihu.com/p/676109470
[5]
猛猿:图解大模型计算加速系列:FlashAttention V1,从硬件到计算逻辑: https://zhuanlan.zhihu.com/p/669926191
[6]
猛猿:图解大模型计算加速系列:Flash Attention V2,从原理到并行计算: https://zhuanlan.zhihu.com/p/691067658
[7]
DefTruth:[Prefill优化][万字]🔥原理&图解vLLM Automatic Prefix Cache(RadixAttention): 首Token时延优化: https://zhuanlan.zhihu.com/p/693556044
[8]
方佳瑞:大模型推理妙招—投机采样(Speculative Decoding): https://zhuanlan.zhihu.com/p/651359908[9]
小冬瓜AIGC:【手撕LLM-Medusa】并行解码范式: 美杜莎驾到, 通通闪开!!: https://zhuanlan.zhihu.com/p/686000524
[10]
紫气东来:LLM(十一):大语言模型的模型量化(INT8/INT4)技术: https://zhuanlan.zhihu.com/p/627436535
[11]
方佳瑞:LLM推理技术之StreamingLLM:如何拥有无限长生成能力: https://zhuanlan.zhihu.com/p/659875511
[12]
紫气东来:LLM(二十):漫谈 KV Cache 优化方法,深度理解 StreamingLLM: https://zhuanlan.zhihu.com/p/659770503
[13]
https://lmdeploy.readthedocs.io/zh-cn/v0.3.0/quantization/kv_int8.html): https://lmdeploy.readthedocs.io/zh-cn/v0.3.0/quantization/kv_int8.html
[14]
zhang:大模型百倍推理加速之KV cache篇: https://zhuanlan.zhihu.com/p/685853516
[15]
ZHANG Mingxing:Mooncake (1): 在月之暗面做月饼,Kimi 以 KVCache 为中心的分离式推理架构: https://zhuanlan.zhihu.com/p/705754254
[16]
猛猿:图解大模型计算加速系列:分离式推理架构2,模糊分离与合并边界的chunked-prefills: https://zhuanlan.zhihu.com/p/710165390
[17]
手抓饼熊:SGLang技术分析: https://zhuanlan.zhihu.com/p/711167552
[18]
改名机会要过期:CUDA效率优化之:CUDA Graph: https://zhuanlan.zhihu.com/p/467466998[19]
yzh119:用FlashInfer加速大语言模型推理中的自注意力操作: https://zhuanlan.zhihu.com/p/681506469[20]
小小将:文生图模型之Stable Diffusion: https://zhuanlan.zhihu.com/p/617134893[21]
HelloWorld:LLaVA(四)图解 LLaVA 推理流程: https://zhuanlan.zhihu.com/p/696654492
[22]
产品经理大群:一文读懂:大模型RAG(检索增强生成)含高级方法: https://zhuanlan.zhihu.com/p/675509396
[23]
战士金:大模型工具调用(function call)原理及实现: https://zhuanlan.zhihu.com/p/663770472
[24]
Tim在路上:一文讲清 NCCL 集合通信原理与优化: https://zhuanlan.zhihu.com/p/720502061
[25]
GPU的工作原理: https://www.zhihu.com/zvideo/1421421171497304064