Tokasaurus：为高吞吐量优化的 LLM 推理引擎

Jordan Juravsky Stanford

Ayush Chakravarthy Stanford

Ryan Ehrlich Stanford

Sabri Eyuboglu Stanford

Bradley Brown Stanford

Joseph Shetaye Stanford

Christopher Ré Stanford

Azalia Mirhoseini Stanford

我们正在发布 Tokasaurus，一个针对吞吐量密集型工作负载优化的新型 LLM 推理引擎。对于小型模型，Tokasaurus 受益于极低的 CPU 开销和动态 Hydragen 分组来利用共享前缀。对于较大的模型，Tokasaurus 支持为具有 NVLink 的 GPU 提供异步张量并行处理，以及为没有 NVLink 的 GPU 提供快速的流水线并行处理实现。在吞吐量焦点的基准测试中，Tokasaurus 的性能可以比 vLLM 和 SGLang 高出 3 倍以上。

优化小型模型

优化大型模型

随着 LLM 变得更加智能、更快、更便宜，社区不断找到新的使用方式。我们最近的工作探索了使用模型扫描代码库中的每个文件、为数学和代码问题采样 10,000 次尝试，以及与其他模型协作以最小化云成本。推理现在也是训练过程的重要组成部分，我们使用模型生成合成数据或作为生成和训练模型补全的 RL 管道的一部分。

最关键的是，这些新的推理工作负载与原始的 LLM 用例（提供聊天机器人服务）看起来相当不同。在这里，我们主要关心的是完成大量序列所需的总时间和成本，而对单个生成的个体延迟的关注程度要少得多（如果有的话）。换句话说，我们需要高吞吐量！

开源推理引擎（即专门用于运行高效 LLM 推理的系统）如 FlexGen、vLLM 和 SGLang 对社区非常有价值。受这些项目的启发并从中学习，我们构建了一个新引擎 Tokasaurus，从一开始就设计用于处理吞吐量焦点的工作负载。我们已经针对小型和大型模型进行了优化，使其能够在吞吐量基准测试中超越现有引擎。在本博客的其余部分，我们将详细介绍其中一些优化，并展示 Tokasaurus 真正大放光彩的几个场景。

优化小型模型

为了使用小型模型对 Tokasaurus 进行基准测试，我们使用两个工作负载：

从 ShareGPT 数据集完成聊天机器人提示（这是测试推理引擎的常见基准）。

重现《Large Language Monkeys》中的一个实验，我们从 GSM8K 数学数据集中取 128 个问题，并对每个问题采样 1024 个答案。这个工作负载的特点是序列之间有大量的前缀共享。

Tokasaurus 在这两个基准测试上的表现都优于 vLLM 和 SGLang，特别是在《Large Language Monkeys》工作负载上实现了超过其他引擎 2 倍的吞吐量。两个主要特性对小型模型的这些优势有贡献：

最小化 CPU 开销

LLM 引擎在 CPU 上执行许多不同的任务，如处理网络请求、标记化输入/去标记化输出、管理 KV 缓存分配，以及为模型准备输入。如果这些 CPU 端的任务导致 GPU 端的模型出现停滞，吞吐量会大幅下降。为了避免这些停滞，推理引擎通常会使许多 CPU 端的任务异步运行：当 GPU 对批次 N 运行前向传播时，引擎的 CPU 端会后处理来自批次 N-1 的结果，并为批次 N+1 准备输入。

Tokasaurus 更进一步，使引擎的 CPU 端（我们称之为管理器）既异步又自适应。管理器的目标是维持一个深队列的输入供模型运行前向传播。管理器监控这个队列的大小，并可以检测模型是否接近耗尽它（从而导致 GPU 停滞）。在这些情况下，管理器会自动开始跳过可选步骤（如检查停止字符串和加入新序列），直到模型的输入队列再次足够深。异步和自适应的组合使 Tokasaurus 能够以更少的 CPU 开销提供小型模型。

动态前缀识别和探索

前缀共享在 LLM 推理中非常常见——不仅在像《Large Language Monkeys》基准测试中重复采样时出现，也在对长文档提出许多问题或跨许多聊天机器人对话重用系统提示时出现。

共享前缀允许更高效地计算注意力。我们去年通过 Hydragen（也称为级联注意力和分支注意力）首次探索了这个想法，但当时我们没有解决在序列不断开始和结束的引擎中检测这些共享前缀的问题。通过 Tokasaurus，我们通过在每次模型前向传播前运行贪心深度优先搜索算法来解决这个检测问题，该算法迭代地找到尽可能长的共享前缀。Hydragen 对小型模型最有影响，这些模型在总 FLOP 中花费相对较大的部分在注意力上。

优化大型模型

Tokasaurus 也可以高效地跨多个 GPU 提供大型模型的服务！这里最重要的优化是我们的流水线并行处理（PP）和张量并行处理（TP）实现，这允许我们最大化具有或不具有 NVLink 的 GPU 上的吞吐量。

GPU 匮乏情况下的流水线并行处理

我们最初使用 Tokasaurus 的目标之一是在我们实验室的 L40S GPU 上高效运行多 GPU 推理，这些 GPU 没有快速的 GPU 间 NVLink 连接。没有 NVLink，在 8 个 GPU 的节点上运行 TP 的通信成本是巨大的。因此，高效支持 PP（需要少得多的 GPU 间通信）是一个高优先级。PP 需要大批次才能高效运行，因为来自管理器的批次被细分为分布在流水线阶段上的微批次。当针对吞吐量进行优化时，我们通常已经在使用适合 GPU 内存的最大批次大小，所以 PP 通常是吞吐量焦点工作负载的自然选择。在使用 Llama-3.1-70B 在八个 L40S GPU 上对 vLLM 和 SGLang 的流水线并行实现进行基准测试时，Tokasaurus 将吞吐量提高了超过 3 倍：

GPU 充裕情况下的异步张量并行处理

如果你拥有带有 NVLink 的 GPU（如 B200 和某些型号的 H100 和 A100），Tokasaurus 也为你准备了一些东西！Tokasaurus 中的模型可以进行端到端的 torch 编译，允许我们利用异步张量并行处理（Async-TP）。这是 PyTorch 中的一个相对较新的特性，可以重叠 GPU 间通信与其他计算，部分隐藏通信的成本。在我们的基准测试中，我们发现 Async-TP 为模型前向传播增加了大量 CPU 开销，仅在非常大的批次大小（如 6k+ 个 token）时才开始改善吞吐量。Tokasaurus 维护我们模型的 torch 编译版本，支持和不支持启用 Async-TP，允许我们在批次大小足够大时自动切换开启 Async-TP：

Tokasaurus 最初是一个内部实验室工作，旨在更快地运行我们的推理实验，我们很高兴能更广泛地分享它！你可以在 GitHub 上查看 Tokasaurus 代码，并使用以下命令从 PyPI 安装包：

pip install tokasaurus

目前，我们支持 Llama-3 和 Qwen-2 系列中的模型，并支持单个节点内数据、张量和流水线并行的任何组合。

Tokasaurus 是用纯 Python 编写的（尽管我们确实使用来自优秀的 FlashInfer 包的注意力和采样操作）。我们希望这使引擎更易于分叉和修改，就像 GPT-fast 一样。

重现我们基准测试的命令可在此处获取。对于每个基准测试，我们为所有引擎配置相同的 KV 缓存大小和最大运行请求数。我们已尽最大努力调整每个引擎的其余参数。我们报告完成预热运行后的平均吞吐量。对于每个基准测试，所有引擎都在同一台机器上运行。

我们为 ShareGPT 基准测试使用了 SGLang 的这个脚本，为《Large Language Monkeys》基准测试使用了这个自定义脚本。为了标准化我们的基准测试脚本和接口，所有实验都通过 OpenAI API 发送请求。我们还在 Llama-1B 上的《Large Language Monkeys》基准测试中使用 vLLM 的 Python API（即 LLM.generate()）进行了实验，并测量到大约 5% 的吞吐量增加（感谢 vLLM 团队的建议！）。

非常感谢 Prime Intellect 和 Together AI 为这个项目提供的计算资源。

另外，我们对 Dan Biderman、Simon Guo、Manat Kaur 和 Avanika Narayan 对引擎的测试版反馈表示感谢！

如果你发现 Tokasaurus 有用，请使用以下引文：

@misc{juravsky2025tokasaurus,
  author       = {Jordan Juravsky and Ayush Chakravarthy and Ryan Ehrlich and Sabri Eyuboglu and Bradley Brown and Joseph Shetaye and Christopher R{\'e} and Azalia Mirhoseini},
  title        = {Tokasaurus: An LLM Inference Engine for High-Throughput Workloads},
  year         = {2025},
  howpublished = {\url{https://scalingintelligence.stanford.edu/blogs/tokasaurus/}}
}