0 Token 间间隔 100% GPU 利用率，百度百舸 AIAK 大模型推理引擎极限优化 TPS

1.   什么是大模型推理引擎

大模型推理引擎是生成式语言模型运转的发动机，是接受客户输入 prompt 和生成返回 response 的枢纽，也是拉起异构硬件，将物理电能转换为人类知识的变形金刚。

大模型推理引擎的基本工作模式可以概括为，接收包括输入 prompt 和采样参数的并发请求，分词并且组装成 batch 输入给引擎，调度 GPU 执行前向推理，处理计算结果并转为词元返回给用户。

• 和人类大脑处理语言的机制类似，大模型首先会把输入的 prompt 进行统一理解，形成具有记忆能力的上下文。这个阶段通常称为 Prefill 阶段。
• 在结束 Prefill 阶段之后，大模型引擎会根据生成的上下文不停地推断下一个可能出现的词语，如此往复循环，直到遇到停止符或者满足采样参数中的停止条件。这是一个自回归过程，通常称为 Decoder 阶段。

由于 Prefill 阶段和 Decoder 阶段所完成的任务不同，通常来讲，会从用户视角出发使用 SLO（Service Level Object）: TTFT（Time To First Token）和TPOT（Time Per Output Token）去评测引擎。

• TTFT 就是首 token 延迟，用于衡量 Prefill 阶段的性能。也就是用户发出请求之后，收到第一个词元返回的间隔，也就是系统的反应时间。对于客户来说，这个指标越低越好。
• TPOT 就是出字间隔，用于衡量 Decoder 阶段的性能。也就是每生成两个词元之间的间隔。通常需要比人眼阅读文字的速度要快，这个指标同样也是越低越好。

当然，只用这些 SLO 并不能完全评测推理引擎对资源的使用状态，所以，和其他使用异构资源的系统一样，会使用吞吐来评测引擎对资源的使用效率，常用的指标就是极限出字率。

极限出字率 TPS（Tokens Per Second ）就是系统在满载的情况下，使用所有可用的资源在 1s 内可以生成的词元的最大数量。这个指标越高，代表硬件的效率越高，可以支持的用户规模就越多。

目前市面上流行的推理引擎有很多，比如说 vLLM、SGLang、LMDeploy、TRT-LLM 等。其中 vLLM 是业界第一个完美解决了大模型不定长特性来各种问题的推理引擎，也是市面上使用最多，社区最活跃的推理引擎。

vLLM 的原创性高效显存管理、高吞吐、极易用、易拓展、模式众多、新特性支持快，社区活跃等特性是其受欢迎的原因。但是，vLLM 对复杂调度逻辑的处理没有做到极致，引入了大量的 CPU 操作，拉长了 TPOT。TPOT 的拉长会降低用户体验，降低了出字率，造成了 GPU 资源浪费。

2.   影响 TPOT 的罪魁祸首 —— Token 间间隔

区别于小模型推理以 batch 为最小推理单位，大模型推理的最小单位是 step。这也是由大模型推理中自回归的特点所决定的。

每一次 step 会给 batch 内部的每个请求生成一个词元，如果有请求生成了结束符，那么这个请求将会提前结束，并且从下个 step 的 batch 中剔除，空余出来的资源将会被引擎动态的分配给其余正在排队的请求。用户可以感知到的观测指标 TPOT，就是每次 step 的执行时间。

每个 step 的执行逻辑可以简单的概括为一下两部分：前向推理和 Token 间间隔。

• 前向推理是调用 GPU 计算资源对 Transfomer 结构进行运算的过程，是一个典型的 GPU 密集计算型任务。
• Token 间间隔，则负责做词元拼接、结束检测、用户响应、请求调度、输入准备等工作，是典型的 CPU 逻辑密集型任务。

优化推理引擎的终极目标其实就是，极限提升前向推理的吞吐，同时极限压缩 Token 间间隔，最终提高极限出字率。

然而，vLLM 的实现中，这两者天然存在着矛盾。极限提升前向推理的吞吐，（即充分发挥 GPU 算力）要求在适当范围内尽可能增加 batch 内的请求数。然而更多的请求数却拉长了 Token 间间隔，这样不仅会使 TPOT 拉长，还会导致 GPU 断流，出现空闲。在最差的情况下（比如 batch 为 256），Token 间间隔和前向推理时间几乎相同，GPU 的利用率只有 50%-60%。

为了提升极限出字率，同时确保高 GPU 利用率，优化 Token 间间隔成为了提升推理速度的关键。

3.   百度百舸 AIAK 优化 Token 间间隔的方案

百度百舸的 AI 加速套件 AIAK 完美继承了 vLLM 的优势，在优化 TPOT 持续发力，并且始终保持着对社区在同周期的技术领先。

3.1.    解决方案 1：多进程架构

这个方案的目标是尽可能缩短 Token 间间隔，将 detokenizer 所耗费的时间从 TPOT 中拿去。

我们发现在处理输入请求和生成返回的过程中，tokenize/detokenize 过程（token id 和字符串的转换）是完全可以独立于 GPU 推理运算的逻辑操作。

所以，我们借助 NVIDIA Triton 框架，将 tokenize/detokenize 的过程从推理流程中抽象出来作为单独的 Triton 模型部署，借助 Triton 的 ensemble 机制，把串行过程转变为 3 阶段（ 3 进程）流水，实现了 tokenize/detokenize 和 GPU 推理 overlap，有效缩短了 Token 间隔时间。尽管这个优化只把 Token 间间隔中一部分 CPU 操作消除了，但是依然有将近 10% 的收益。

3.2.    解决方案 2：静态 Slot 方案

这个方案主要改造了 vLLM 的调度逻辑，全方位优化了词元拼接、结束检测、用户响应、请求调度、输入准备，提高了各个模块的并行效率，实现了对上一个方案中的「剩余部分」耗时的压缩。

我们发现 vLLM 的调度逻辑是面向全局视角的。也就是说每个 step 的调度都会从全局中进行重新筛选，相当于当前 step 结束之后，调度器会把当前 batch 中的句子「放回」全局请求池子中，然后在下一个 step 开始前，从这个全局池子中「取回」适当请求进行运算，这一放一取引入了额外的 overhead。

为了实现全局调度，vLLM 在词元拼接等其他环节引入了大量的 for 循环去串行的处理每个请求，由于这些操作都在发生在 CPU 上，导致在输入打包过程中，必须要引入耗时较长的 host to device 操作。

事实上，step 之间的很多信息是可以复用的（每次放回去的请求和取回来的请求很大一部分是重复的）。也正是基于这个洞见，百度百舸的 AIAK 把 GPU 每次可以迭代的 batch 当成一批固定的 slot，一旦某个请求被调度到某个 slot 后，在完成请求所有推理迭代之前，都不会被唤出，也正是有了这些固定 slot 抽象，AIAK 实现了：

• 将全局调度改造为局部调度。也就是在下一个 step 调度时，最大程度复用上一个 step 的信息，避免全局搜索，只做增量调度。
• 串行转并行。也正是有了 slot 的引入，词元拼接、结束检测等这些原本串行的操作可以用 CUDA Kernel 做并发处理，耗时从 ms 级别降低到 us 级别。
• 避开 host to device 操作。输入打包的工作得以复用前序的显存，有效避开了 host to device 操作。

3.3.    方案 3：异步化执行

多进程架构将逻辑上容易独立的部分解耦到其他进程做流水并行，静态 Slot 方案则直面 token 间耗时问题，优化调度模式压榨各个环节的耗时。有了这两个方案，Token 间间隔已经从 35ms 降低到 14ms，GPU 的利用率已经从 50% 提升到了 75%，但是距离 100% 的 GPU 利用率和零耗时 Token 间间隔的目标还有不少距离。

百度百舸 AIAK 通过异步调度模式，将前一个方案中的「剩余部分」全部取出，最终实现了上述极限目标。

简单来讲，就是将 CPU 操作密集的 Token 间间隔和 GPU 计算密集的前向推理完全分开到两条流水线上做二级流水并行。

• 从逻辑上来讲，核心调度逻辑摆脱了对前向推理的同步依赖，实现异步化调度。
• 从效果上来说，GPU 避免 token 间同步导致的断流问题，处于一直繁忙状态，实现了推理过程中 100% 利用率和 0 Token 间间隔。

为了简化实现，我们将操作相对简单的前向推理当做一个任务放在后台线程中进行运行，主线程则运行核心的复杂的调度逻辑。两个线程通过一个队列进行交互，分别冲当生产者和消费者，通过线程信号量和 GPU 流上的事件进行信号同步，实现二级流互相 overlap。

和其他任何使用 GPU 类似的硬件作为加速器的系统一样，追求 100% 的利用率一直是所有工程师的终极目标。百度百舸的 AI 加速套件 AIAK 在优化 TPOT，同时打满 GPU 利用率这一目标上经历漫长而又艰辛的探索，最终才彻底实现了 0 Token 间间隔和 100% 利用率这一目标。

当然，除去在这个过程中使用的诸多巧妙的优化手段外，百度百舸的 AIAK 还在量化、投机式、服务化、分离式、多芯适配等领域做了大量工作，致力于实现一个适用于全场景、多芯片、高性能的推理引擎，助力用户在「降低推理成本，优化用户体验上」更上一层楼。