- 作者:老汪软件技巧
- 发表时间:2024-09-21 15:01
- 浏览量:
背景
随着 GQA/MLA/MoE 等模型结构不断发展,大语言模型的推理逐步解除了显存限制,逐渐向着高并发、高吞吐的方向发展。推理引擎的运行时开销也变得不可忽视。主流 LLM 推理框架的运行时开销大致来自:
Python 性能:考虑用户易用性和开发效率,业界主流框架都采用 Python 为主要开发语言、C++实现模型和算子的方式。Python 一直存在让人诟病的 GIL 问题,框架中很多的代码无法有效并行。这导致在高并发场景,Python 执行速度成为了显著的瓶颈。通信开销:为了规避Python GIL的限制来提升框架的并行度,即使是在单机内,主流框架通常都会引入多进程的设计。此外大模型推理也逐渐扩展至分布式推理,系统模块可能存在多个 Prefill、Decode 实例。系统在运行过程中,不同的模块、实例间需要频繁进行交互。而进程间的通信、消息的序列化和反序列化也不可避免的成为了瓶颈。同步执行逻辑:主流的推理框架都会采用分页的 KV Cache 管理,执行流程是严格按照 KV Cache Block分配、模型计算、KV Cache Block更新的迭代式的过程。目前业界主流框架都是基本按照同步的设计,各个组件的开销不能被掩盖,当并发变大后累计的开销可能会达到模型计算的量级。
近期,大模型推理社区(vLLM,SGLang 等)普遍开始关注框架运行时开销,提出了多步调度、异步输出处理、独立 API Server 进程等工作,来分摊或掩盖部分开销。
在我们的实际业务场景中,也观察到高额的框架开销严重限制了系统吞吐,特别是在高并发(>1k)场景下,运行时开销已经接近或高于 GPU 运行时间,导致资源严重浪费和性能下降。为此,BladeLLM 设计并实现了基于 Python 的纯异步 LLM 推理架构 -- TAG (Totally Asynchronous Generator) ,以最大程度提高 GPU 利用率,提升引擎性能。
LLM 推理引擎简介
一个典型的 LLM 推理引擎可以分为以下组件:
以下是一个完全同步的推理引擎的 timeline 示意图:
社区现状
针对运行时开销逐渐变大的问题,社区也在逐渐优化中,提出了包括异步输出处理、多步调度等方案来尽可能掩盖或分摊运行时开销。
以目前最新的 vLLM 0.6.0 为例,开启 MultiStepScheduling 和异步输出处理优化后,整体的 timeline 大致可以表示为:
可以发现,开启优化后, Model Runner 在连续 N 个 step 之间,可以不和 Scheduler 进行同步和交互,连续进行多次 forward,从而保证了该段时间内 GPU 利用率变高。SGLang 也提供了类似功能。
然而我们仍然可以观察到, 这样的调度方式存在以下问题:
连续 N 个 step 结束之后,Model Runner 还是会处于等待状态,直到 Scheduler 调度新的一轮请求。所以MultiStepScheduling 只能平摊、但不能完全消除 Scheduler 的开销。而随着并发增大,Scheduler 和 Decoder 等环节开销变大,Model Runner 的等待时间仍然会增加。Model Runner 需要多个 step 后再进行返回,而这中间,用户是不能接收到这些新生成的 token。所以从用户视角,TTFT 和 ITL 会显著上升,影响用户体验。纯异步 LLM 推理引擎 TAG异步Scheduler
目前看来,社区推理引擎关键问题在于调度和 Model Runner 执行存在同步关系。以 vLLM 为例,这个同步关系主要来源于: Scheduler 进行调度时,会为每个请求分配 KV Cache Block,而分配决策依赖于每个请求的 Token;而当前 step 调度所需要的 Token需要等待上一步Model Runner 运行结果返回给 Scheduler ,从而形成了依赖,导致 Scheduler 和 Model Runner 执行需要同步串行进行。
BladeLLM 团队针对这一场景,重新分析了该依赖关系,发现:
Scheduler在调度和分配KVCacheBlock时,需要的信息只是每个请求Token的长度
,与TokenID具体的值无关;ModelRunner每执行一次,能够新生成Token数的上界(
)都是确定的:对普通的step,生成Token数为1;如果使用了投机采样算法,每步生成Token数至多为gamma+1,其中gamma为draftmodel采样步数;因此,第
步调度时,只需要在第
步调度的Token数li−1l_{i-1}li−1基础上,额外预留上一步最多生成的Token数,即可以满足worker运行逻辑,而不会导致KVCacheBlock不足;在ModelRunner实际运行完成后,如果实际生成的Token数gig_igi小于
,则更新scheduler状态,下次调度只需要再预留
个token的KVCache。这个状态的更新是异步的,所以调度第i个step时,会比实际所需token多预留∑j=i−ti(nj−mj)\sum_{j=i-t}^i({n_j-m_j})∑j=i−ti(nj−mj)个token的KVCache,其中
为异步调度的最大步数。在我们的实际使用中,
取2即可达到完全异步的效果,因此,额外预留token数较小,不会影响整体cache容量。
TAG 架构
在消除了 Scheduler 和 ModelRunner 之间的同步点后,我们发现,整个推理系统可以完全异步起来,因此我们提出 TAG (Totally Asynchronous Generator) 的架构设计:
图中,所有组件之间的消息传递都是异步完成的。整体流程分为:
WebServer 接受用户请求,经过 Tokenizer 后,将请求添加到 Scheduler 的请求队列中;同时,并行从 Detokenizer 中接收 response,返回给用户;Update Loop 主要负责处理 Model Runner 的返回消息队列,包括:Scheduler 和 Update Loop 通过 Python 协程的方式进行并行,又通过信号量控制 Scheduler 最大调度步数;因此,Scheduler 的循环为:获取信号量 -> 预分配最大 token cache -> 调度结果发送给 Model Runner。由于 Scheduler 会同时有多次调度结果发送到 Model Runner,所以 Model Runner 能够不间断的获取请求,运行模型,而不会出现闲置。Model Runner 进程内部运行也是完全异步的,包括接收并反序列化请求->调用模型->序列化结果并返回三个部分,使得调用模型的过程不会中断。
TAG 的 Timeline 可以表示为:
跨进程通讯框架
Scheduler 和 ModelRunner 采用的是 CPU/GPU 多进程分离的设计,这种多进程的架构不可避免的会引入通信带来的包括序列化和反序列化、消息传输一系列开销,引擎在早期的通信选型调研过包括:RPC、消息队列、Socket、管道等多种方式。
因此结合上述考虑后,各组件间的通信采用了 Unix Domain Socket 配合共享内存的方式,在线上常见的50~100 并发以内的时候,进程间通信的开销在 1ms 以内,避免了引擎引入过大额外的开销。
性能测试
我们对比了 TAG 和 vLLM(0.6.0)和 SGLang(0.3.0)在不同并发(64/256/512)下的 QPS 和 TBT,结果如下:
注:
由于主要测试 runtime 开销,所以数据集输入长度固定位 10,输出长度为 20-500 的正态分布,并设置 ignore_eos=True;vLLM 开启 multi step scheduling 后,TBT 会显著高于 SGLang 和 BladeLLM-TAG;
由这两个图可以发现,TAG 由于完全掩盖了 Worker 进程外的所有开销,所以相比社区方案能做到更低延迟、更高吞吐,支持更高并发。
总结展望
我们设计并实现了一个纯异步的 LLM 推理引擎 TAG,解决了推理引擎各组件的同步依赖,最大程度的提高 GPU 利用率,提升服务吞吐,降低请求延迟。内部测试中,在满足 SLO 的情况下,Decode 实例能够同时支持上千并发,而 runtime overhead 不会成为性能瓶颈。后续,我们将聚焦于 Model Runner 内的运行开销,致力于打造更极致优化的 LLM 推理引擎。
招聘
阿里云人工智能平台PAI长期开放推理优化方向的研究型实习生、校招和社招岗位。团队致力于从模型和系统两方面对大语言模型推理进行协同优化,工作内容覆盖模型压缩、高性能算子、推理框架和运行时、分布式等工作。
欢迎投递简历: