TAG:BladeLLM 的纯异步推理架构

本文涉及的产品
模型训练 PAI-DLC,100CU*H 3个月
模型在线服务 PAI-EAS,A10/V100等 500元 1个月
交互式建模 PAI-DSW,每月250计算时 3个月
简介: 近期,大模型推理社区(vLLM,SGLang 等)普遍开始关注框架运行时开销,提出了多步调度、异步输出处理、独立 API Server 进程等工作,来分摊或掩盖部分开销。在我们的实际业务场景中,也观察到高额的框架开销严重限制了系统吞吐,特别是在高并发(>1k)场景下,运行时开销已经接近或高于 GPU 运行时间,导致资源严重浪费和性能下降。为此,BladeLLM 设计并实现了基于 Python 的纯异步 LLM 推理架构 -- TAG (Totally Asynchronous Generator) ,以最大程度提高 GPU 利用率,提升引擎性能。

随着 GQA/MLA/MoE 等模型结构不断发展,大语言模型的推理逐步解除了显存限制,逐渐向着高并发、高吞吐的方向发展。推理引擎的运行时开销也变得不可忽视。主流 LLM 推理框架的运行时开销大致来自:

  1. Python 性能:考虑用户易用性和开发效率,业界主流框架都采用 Python 为主要开发语言、C++实现模型和算子的方式。Python 一直存在让人诟病的GIL问题,框架中很多的代码无法有效并行。这导致在高并发场景,Python 执行速度成为了显著的瓶颈。
  2. 通信开销:为了规避 Python GIL 的限制来提升框架的并行度,即使是在单机内,主流框架通常都会引入多进程的设计。此外大模型推理也逐渐扩展至分布式推理,系统模块可能存在多个 Prefill、Decode 实例。系统在运行过程中,不同的模块、实例间需要频繁进行交互。而进程间的通信、消息的序列化和反序列化也不可避免的成为了瓶颈。
  3. 同步执行逻辑:主流的推理框架都会采用分页的 KV Cache 管理,执行流程是严格按照 KV Cache Block 分配、模型计算、KV Cache Block 更新的迭代式的过程。目前业界主流框架都是基本按照同步的设计,各个组件的开销不能被掩盖,当并发变大后累计的开销可能会达到模型计算的量级。

近期,大模型推理社区(vLLM,SGLang 等)普遍开始关注框架运行时开销,提出了多步调度、异步输出处理、独立 API  Server 进程等工作,来分摊或掩盖部分开销。在我们的实际业务场景中,也观察到高额的框架开销严重限制了系统吞吐,特别是在高并发(>1k)场景下,运行时开销已经接近或高于 GPU 运行时间,导致资源严重浪费和性能下降。为此,BladeLLM 设计并实现了基于 Python 的纯异步 LLM 推理架构 -- TAG (Totally Asynchronous Generator) ,以最大程度提高 GPU 利用率,提升引擎性能。


LLM 推理引擎简介

一个典型的 LLM 推理引擎可以分为以下组件:

  • API Server:负责请求接收和响应;
  • Scheduler:负责请求调度、KV Cache Block 分配等;
  • Model Runner:主要负责模型计算和采样;
  • Decoder:负责将采样得到的 Token ID 转化为字符形式输出。

以下是一个完全同步的推理引擎的 timeline 示意图:

image.png


社区现状

针对运行时开销逐渐变大的问题,社区也在逐渐优化中,提出了包括异步输出处理、多步调度等方案来尽可能掩盖或分摊运行时开销。以目前最新的 vLLM 0.6.0 为例,开启 MultiStepScheduling 和异步输出处理优化后,整体的 timeline 大致可以表示为:

image.png

可以发现,开启优化后, Model Runner 在连续 N 个 step 之间,可以不和 Scheduler 进行同步和交互,连续进行多次 forward,从而保证了该段时间内 GPU 利用率变高。SGLang 也提供了类似功能。然而我们仍然可以观察到, 这样的调度方式存在以下问题:

  1. 连续 N 个 step 结束之后,Model Runner 还是会处于等待状态,直到 Scheduler 调度新的一轮请求。所以 MultiStepScheduling 只能平摊、但不能完全消除 Scheduler 的开销。而随着并发增大,Scheduler 和 Decoder 等环节开销变大,Model Runner 的等待时间仍然会增加。
  2. 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 团队针对这一场景,重新分析了该依赖关系,发现:

  1. Scheduler 在调度和分配 KV Cache Block 时,需要的信息只是每个请求 Token 的长度  ,与 Token ID 具体的值无关;
  2. Model Runner 每执行一次,能够新生成 Token 数的上界(  )都是确定的:对普通的 step,生成 Token 数为 1;如果使用了投机采样算法,每步生成 Token 数至多为 gamma + 1,其中 gamma 为 draft model 采样步数;
  3. 因此,第  步调度时,只需要在第  步 调度的 Token 数(  )基础上,额外预留上一步最多生成的 Token 数,即可以满足 worker 运行逻辑,而不会导致 KV Cache Block 不足;
  4. 在 Model Runner 实际运行完成后,如果实际生成的 Token 数(  )小于p ,则更新 scheduler 状态,下次调度只需要再预留  个 token 的 KV Cache。这个状态的更新是异步的,所以调度第 i 个 step 时,会比实际所需 token 多预留   个 token 的 KV Cache,其中  为异步调度的最大步数。在我们的实际使用中,  取 2 即可达到完全异步的效果,因此,额外预留 token 数较小,不会影响整体 cache 容量。

image.png

TAG 架构

在消除了 Scheduler 和 ModelRunner 之间的同步点后,我们发现,整个推理系统可以完全异步起来,因此我们提出 TAG (Totally Asynchronous Generator) 的架构设计:

image.png

图中,所有组件之间的消息传递都是异步完成的。整体流程分为:

  1. WebServer 接受用户请求,经过 Tokenizer 后,将请求添加到 Scheduler 的请求队列中;同时,并行从 Detokenizer 中接收 response,返回给用户;
  2. Update Loop 主要负责处理 Model Runner 的返回消息队列,包括:
  1. 给 Scheduler 更新请求之前步执行实际生成 token 数,减少 token cache 浪费;
  2. detokenize,并将生成 token 结果传给 WebServer ;
  3. 判停,及时删除结束的请求;
  4. 完成上述步骤后,释放信号量,让 Scheduler 进行下一次调度。
  1. Scheduler 和 Update Loop 通过 Python 协程的方式进行并行,又通过信号量控制 Scheduler 最大调度步数;因此,Scheduler 的循环为:获取信号量 -> 预分配最大 token cache -> 调度结果发送给 Model Runner。
  2. 由于 Scheduler 会同时有多次调度结果发送到 Model Runner,所以 Model Runner 能够不间断的获取请求,运行模型,而不会出现闲置。Model Runner 进程内部运行也是完全异步的,包括接收并反序列化请求->调用模型->序列化结果并返回三个部分,使得调用模型的过程不会中断。

TAG 的 Timeline 可以表示为:

image.png

跨进程通讯框架

Scheduler 和 ModelRunner 采用的是 CPU/GPU 多进程分离的设计,这种多进程的架构不可避免的会引入通信带来的包括序列化和反序列化、消息传输一系列开销,引擎在早期的通信选型调研过包括:RPC、消息队列、Socket、管道等多种方式。

  • 以 gRPC 为典型代表的RPC应用场景最为广泛,支持多种协议,跨机通信,缺点是 gRPC 框架的性能开销不可忽略,对于频繁的小包通信的效率不是最优的。
  • 消息队列的通信方式在消息较小的时候性能尚可,但随着消息变大,性能下降比较严重。
  • 共享内存通过虚拟地址的方式,将同一块物理内存映射到多个进程中,使得共享内存传输大块的数据是效率最高。

因此结合上述考虑后,各组件间的通信采用了 Unix Domain Socket配合共享内存的方式,在线上常见的50~100 并发以内的时候,进程间通信的开销在 1ms 以内,避免了引擎引入过大额外的开销。


性能测试

我们对比了 TAG 和 vLLM(0.6.0) 和 SGLang(0.3.0) 在不同并发(64/256/512)下的 QPS 和 TBT,结果如下:

image.png

image.png


注:

  1. 由于主要测试 runtime 开销,所以数据集输入长度固定位 10,输出长度为 20-500 的正态分布,并设置 ignore_eos=True;
  2. vLLM 开启 multi step scheduling 后,TBT 会显著高于 SGLang 和 BladeLLM-TAG;

由这两个图可以发现,TAG 由于完全掩盖了 Worker 进程外的所有开销,所以相比社区方案能做到更低延迟、更高吞吐,支持更高并发。


总结展望

我们设计并实现了一个纯异步的 LLM 推理引擎 TAG,解决了推理引擎各组件的同步依赖,最大程度的提高 GPU 利用率,提升服务吞吐,降低请求延迟。内部测试中,在满足 SLO 的情况下,Decode 实例能够同时支持上千并发,而 runtime overhead 不会成为性能瓶颈。后续,我们将聚焦于 Model Runner 内的运行开销,致力于打造更极致优化的 LLM 推理引擎。

招聘

阿里云人工智能平台PAI长期开放推理优化方向的研究型实习生、校招和社招岗位。团队致力于从模型和系统两方面对大语言模型推理进行协同优化,工作内容覆盖模型压缩、高性能算子、推理框架和运行时、分布式等工作。


欢迎投递简历:xiafei.qiuxf@alibaba-inc.com

相关实践学习
使用PAI-EAS一键部署ChatGLM及LangChain应用
本场景中主要介绍如何使用模型在线服务(PAI-EAS)部署ChatGLM的AI-Web应用以及启动WebUI进行模型推理,并通过LangChain集成自己的业务数据。
机器学习概览及常见算法
机器学习(Machine Learning, ML)是人工智能的核心,专门研究计算机怎样模拟或实现人类的学习行为,以获取新的知识或技能,重新组织已有的知识结构使之不断改善自身的性能,它是使计算机具有智能的根本途径,其应用遍及人工智能的各个领域。 本课程将带你入门机器学习,掌握机器学习的概念和常用的算法。
相关文章
|
5月前
|
人工智能 测试技术 数据处理
首个Mamba+Transformer混合架构多模态大模型来了,实现单卡千图推理
【10月更文挑战第18天】《LongLLaVA: Scaling Multi-modal LLMs to 1000 Images Efficiently via Hybrid Architecture》提出了一种新型多模态大模型LongLLaVA,结合了Mamba和Transformer架构,通过系统优化实现在单张A100 80GB GPU上处理近千张图像的突破。该模型在视频理解、高分辨率图像分析和多模态智能体任务中表现出色,显著提升了计算效率。
219 65
|
12天前
|
自然语言处理
Scaling Law 撞墙?复旦团队大模型推理新思路:Two-Player架构打破自我反思瓶颈
复旦大学研究团队提出Two-Player架构,通过分离推理和批评模型的角色,突破大语言模型(LLM)在复杂推理任务中的自我反思瓶颈。该架构利用批评模型提供逐步反馈,监督推理模型,提升其性能。研究开发了AutoMathCritique框架,收集76,321个响应数据,实验表明批评模型显著提高演员模型的探索效率和解决方案多样性。论文地址:http://arxiv.org/abs/2411.16579
30 2
|
13天前
|
人工智能 自然语言处理 并行计算
MeteoRA:多任务AI框架革新!动态切换+MoE架构,推理效率提升200%
MeteoRA 是南京大学推出的多任务嵌入框架,基于 LoRA 和 MoE 架构,支持动态任务切换与高效推理。
47 3
|
3月前
|
人工智能 自然语言处理 PyTorch
Bamba-9B:基于 Mamba2 架构的仅解码语言模型,旨在提高大型语言模型在推理时的效率
Bamba-9B 是由 IBM、普林斯顿大学、卡内基梅隆大学和伊利诺伊大学香槟分校联合推出的基于 Mamba2 架构的仅解码语言模型。该模型在开放数据集上训练,旨在提高大型语言模型的推理效率,特别是在处理长文本时的内存带宽瓶颈。Bamba-9B 在推理时相较于标准变换器模型展现出 2.5 倍的吞吐量提升和 2 倍的延迟加速。
106 12
Bamba-9B:基于 Mamba2 架构的仅解码语言模型,旨在提高大型语言模型在推理时的效率
|
2月前
|
机器学习/深度学习 人工智能 自然语言处理
一文彻底讲透GPT架构及推理原理
本篇是作者从开发人员的视角,围绕着大模型正向推理过程,对大模型的原理的系统性总结,希望对初学者有所帮助。
|
3月前
|
存储 人工智能 缓存
Mooncake:月之暗面Kimi联合清华等机构推出的大模型推理架构
Mooncake是由月之暗面Kimi联合清华大学等机构共同开源的大模型推理架构。该架构以KVCache为中心,通过分布式设计和资源优化,显著提升了大模型推理的吞吐量和效率,同时降低了算力开销。Mooncake在长上下文场景中表现出色,支持多种应用场景,如自然语言处理、内容推荐系统和搜索引擎等。
284 6
Mooncake:月之暗面Kimi联合清华等机构推出的大模型推理架构
|
3月前
|
存储 人工智能 监控
【AI系统】推理系统架构
本文深入探讨了AI推理系统架构,特别是以NVIDIA Triton Inference Server为核心,涵盖推理、部署、服务化三大环节。Triton通过高性能、可扩展、多框架支持等特点,提供了一站式的模型服务解决方案。文章还介绍了模型预编排、推理引擎、返回与监控等功能,以及自定义Backend开发和模型生命周期管理的最佳实践,如金丝雀发布和回滚策略,旨在帮助构建高效、可靠的AI应用。
293 15
|
3月前
|
机器学习/深度学习 人工智能 调度
【AI系统】推理引擎架构
本文详细介绍了推理引擎的基本概念、特点、技术挑战及架构设计。推理引擎作为 AI 系统中的关键组件,负责将训练好的模型部署到实际应用中,实现智能决策和自动化处理。文章首先概述了推理引擎的四大特点:轻量、通用、易用和高效,接着探讨了其面临的三大技术挑战:需求复杂性与程序大小的权衡、算力需求与资源碎片化的矛盾、执行效率与模型精度的双重要求。随后,文章深入分析了推理引擎的整体架构,包括优化阶段的模型转换工具、模型压缩、端侧学习等关键技术,以及运行阶段的调度层、执行层等核心组件。最后,通过具体的开发流程示例,展示了如何使用推理引擎进行模型的加载、配置、数据预处理、推理执行及结果后处理。
240 0
|
8月前
|
开发框架 人工智能
AI问题之Chains架构如何固化推理过程
AI问题之Chains架构如何固化推理过程
|
机器学习/深度学习 TensorFlow Go
美团视觉GPU推理服务部署架构优化实战
美团视觉GPU推理服务部署架构优化实战
227 0