TAG：BladeLLM 的纯异步推理架构

作者：张子鹏 PAI引擎团队

随着 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 推理引擎可以分为以下组件：

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

以下是一个完全同步的推理引擎的 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 在调度和分配 KV Cache Block 时，需要的信息只是每个请求 Token 的长度 $l$ ，与 Token ID 具体的值无关；
Model Runner 每执行一次，能够新生成 Token 数的上界（ $p$ ）都是确定的：对普通的 step，生成 Token 数为 1；如果使用了投机采样算法，每步生成 Token 数至多为 gamma + 1，其中 gamma 为 draft model 采样步数；
因此，第i步调度时，只需要在第 $i-1$ 步调度的 Token 数（ $l_{i-1}$ ）基础上，额外预留上一步最多生成的 Token 数，即可以满足 worker 运行逻辑，而不会导致 KV Cache Block 不足；
在 Model Runner 实际运行完成后，如果实际生成的 Token 数（ $g_i$ ）小于 $p$ ，则更新 scheduler 状态，下次调度只需要再预留 $n-m$ 个 token 的 KV Cache。这个状态的更新是异步的，所以调度第 $i$ 个 step 时，会比实际所需 token 多预留 $\sum_{j=i-t}^i({n_j-m_j})$ 个 token 的 KV Cache，其中 $t$ 为异步调度的最大步数。在我们的实际使用中， $t$ 取 2 即可达到完全异步的效果，因此，额外预留 token 数较小，不会影响整体 cache 容量。

TAG 架构

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

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

WebServer 接受用户请求，经过 Tokenizer 后，将请求添加到 Scheduler 的请求队列中；同时，并行从 Detokenizer 中接收 response，返回给用户；
Update Loop 主要负责处理 Model Runner 的返回消息队列，包括：

给 Scheduler 更新请求之前步执行实际生成 token 数，减少 token cache 浪费；
detokenize，并将生成 token 结果传给 WebServer ；
判停，及时删除结束的请求；
完成上述步骤后，释放信号量，让 Scheduler 进行下一次调度。

3. Scheduler 和 Update Loop 通过 Python 协程的方式进行并行，又通过信号量控制 Scheduler 最大调度步数；因此，Scheduler 的循环为：获取信号量 -> 预分配最大 token cache -> 调度结果发送给 Model Runner。

4. 由于 Scheduler 会同时有多次调度结果发送到 Model Runner，所以 Model Runner 能够不间断的获取请求，运行模型，而不会出现闲置。Model Runner 进程内部运行也是完全异步的，包括接收并反序列化请求->调用模型->序列化结果并返回三个部分，使得调用模型的过程不会中断。

TAG 的 Timeline 可以表示为：

跨进程通讯框架

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，结果如下：

注：

由于主要测试 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 推理引擎。