机器之心原创

2024/11/26 13:22

陈天奇团队LLM结构化生成新引擎XGrammar：百倍加速、近零开销

现在，大语言模型的结构化生成有了一个更加高效、灵活的引擎。

不管是编写和调试代码，还是通过函数调用来使用外部工具，又或是控制机器人，都免不了需要 LLM 生成结构化数据，也就是遵循某个特定格式（如 JSON、SQL 等）的数据。

但使用上下文无关语法（CFG）来进行约束解码的方法并不高效。针对这个困难，陈天奇团队提出了一种新的解决方案：XGrammar。

XGrammar 是一个开源软件库，可实现高效、灵活且可移植的结构化生成。该团队在博客中表示：「我们毫不妥协地实现了这三个目标，并致力于一个核心使命：将灵活、零开销的结构化生成带到任何地方。」

论文标题：XGrammar: Flexible and Efficient Structured Generation Engine for Large Language Models
论文地址：https://arxiv.org/pdf/2411.15100
代码地址：https://github.com/mlc-ai/xgrammar

对于结构化生成，一种常用方法是约束解码。在每个解码步骤中，约束解码都会检查词表，并通过将无效 token 的概率设置为零来过滤掉违反指定结构的 token。为了支持多种多样的结构格式，需要一种灵活的机制来指定和检查这些约束。

^{使用 JSON 方案实现约束解码}

上下文无关语法（CFG）就能提供一种通用方法，即通过一组规则来定义结构。其中每条规则都包含一个字符序列或其他规则，并允许递归组合来表示复杂的结构。相比于正则表达式等其它格式，CFG 由于支持递归结构，因而能提供更大的灵活性，使其适合描述 JSON、SQL 和领域特定语言（DSL）等常见语言。

下图展示了一个用于数组和字符串的 CFG，可以清楚地看到其中的递归结构。

但是，也正因为 CFG 很灵活，所以直接将其应用于约束解码的效率并不高。首先，每个解码步骤都需要对词表中每个可能的 token 解释 CFG，在 Llama 3.1 中，这个词表的大小可能高达 128k。此外，CFG 解释需要一个堆栈状态来跟踪之前匹配的递归规则，因此无法提前计算和缓存堆栈模式的所有组合。最后，LLM 生成结果中的每个 token 都包含多个字符，这些字符可能会跨越语法元素的边界，并在运行时执行期间导致进一步的递归或堆栈弹出。这种未对齐的边界问题很棘手，需要在语法执行期间小心处理它们。

XGrammar 便是为解决上述难题而生的，并且效果卓越：相比于之前的 SOTA 方法，XGrammar 可以将上下文无关语法的每 token 延迟减少多达 100 倍！此外，他们还基于 Llama3.1 模型实验了集成了 XGrammar 的 LLM serving 引擎；在 H100 GPU 上，这能将通过结构化输出实现端到端 LLM serving 的速度提升 80 倍！

该团队表示：「我们正在开源 XGrammar 并将其集成到主要的开源 LLM 框架中。」

XGrammar 概览

如图 1 所示，Grammar 利用了字节级下推自动机（byte-level pushdown automaton）来解释上下文无关语法。

这种字节级设计允许每个字符边缘包含一个或多个字节，处理不规则的 token 边界并支持包含 sub-UTF8 字符的 token 。该自动机的结构经过优化以加快匹配速度。

在预处理阶段，会生成一个自适应 token 掩码缓存，它会通过预先计算与上下文无关的 token 来加快运行时的掩码生成。上下文扩展（context extension）能进一步提升这种缓存的有效性。

在运行时，token 掩码缓存会快速生成大部分掩码，而持续性执行堆栈会高效处理其余的上下文相关 token。

此外，掩码生成和 LLM 推理是互相重叠的，以最大限度地减少约束解码的开销。一旦 LLM 在掩码约束下生成新 token，就会使用此 token 来更新下推自动机的堆栈状态，以进行下一次掩码生成。

具体来说，陈天奇团队首先得到了一个见解：虽然无法预先计算下推自动机（PDA）无限多个状态的完整掩码，但可以预先计算掩码中相当一部分（通常超过 99%）的 token。因此，可将这些 token 分成两类：

上下文无关 token：仅通过查看 PDA 中的当前位置而不是堆栈即可确定其有效性的 token。
上下文相关 token：必须使用整个堆栈来确定其有效性的 token。

下图展示了一组上下文相关和无关 token 的示例。大多数情况下，上下文无关 token 占大多数。我们可以预先计算 PDA 中每个位置的上下文无关 token 的有效性，并将它们存储在自适应 token 掩码缓存中。此过程称为语法编译（grammar compilation）。

下图则展示了自适应存储格式。

在运行时，首先检索来自缓存的上下文无关 token 的有效性。然后，高效地执行 PDA 来检查其余的上下文相关 token。通过跳过运行时检查大多数 token，便可以显著加快掩码生成速度。XGrammar 执行时间的整体工作流程见图 1。

此外，他们还设计了一组额外的算法和系统优化方法，以进一步提高掩码生成速度并减少预处理时间，包括上下文扩展、持续性执行椎栈、下推自动机结构优化、并行式语法编译。

上下文扩展

该团队提出的方法是检测语法中每个规则的额外上下文信息，并将其用于减少上下文相关 token 的数量，并进一步加快运行时检查速度。

持续性执行堆栈

为了加快由于多种可能的扩展路径而导致的拆分和合并期间多个并行堆栈的维护速度，他们设计了一个基于树的数据结构，可以有效地同时管理多个堆栈。

它还可以存储以前的状态并实现高效的状态回滚，从而加快上下文相关 token 的运行时检查速度。

下推自动机结构优化

研究者进行了额外的优化，以改进下推自动机的结构，加快最终执行的效率。这些优化借鉴了传统的编译器优化概念，它们对于高效约束解码特别有用。

一是规则内联。在指定的上下文无关语法中，可能有许多片段规则，即只有少数元素的规则，然后在下推自动机中将其转换为小的 FSA（有限状态自动机）。

为了解决这个问题，研究者为片段规则引入了一种自动内联策略。他们迭代地选择不引用其他规则的规则并将它们内联到父规则中。为了避免自动机大小的爆炸式增长，研究者将内联规则和内联结果的大小限制为常量。该内联过程几乎消除了片段规则，从而提高了 token 检查的效率并增强了上下文扩展的有效性。

二是下推自动机节点合并。对于下推自动机，在许多情况下，歧义来自具有相同标签的节点的多个外向边。在匹配 token 时，如果到达此节点，并且下一个字符恰好与标签匹配，则匹配堆栈将被拆分为多个堆栈，每个外向边一个。堆栈数量增多会增加计算量，这是因为需要检查每个堆栈的上下文相关 token 并合并 token 掩码。

为了减少这种歧义，节点合并算法会合并满足以下两个条件的后续节点，a）它们由来自同一点的具有相同标签的边指向，b）它们没有被其他边指向。

以上两种优化保留了自动机的等效性，但减少了节点和边的数量。运行时，减少了堆栈的数量和 token 检查所需的计算量，从而加快了掩码的生成过程。

重叠掩码生成和 LLM 推理

通过上述优化，token 掩码生成过程显著加快，但仍需要 CPU 计算。为了进一步消除约束解码的开销，研究者将 mask 生成计算与 LLM 推理过程重叠，如下图 8 所示。

研究者观察到，mask 生成过程和 LLM 推理过程可以重叠，原因在于 mask 生成只需要 CPU，并且只依赖于之前生成的 token。LLM 推理过程（除采样阶段外）只需要 GPU，并且也只依赖于之前生成的 token。因此可以将 CPU 上的 mask 生成过程与 GPU 上的 LLM 推理过程并行化。

评估结果

研究者利用 12,000 行核心 C++ 代码来实现 XGrammar，并提供了 Python 捆绑包以方便与 LLM 推理框架无缝集成。他们在评估 XGrammar 过程中回答以下几个问题：

XGrammar 能否高效支持约束解码的每个步骤？
XGrammar 能否在 LLM serving 中实现端到端结构化生成的最小开销？
XGrammar 能否部署在更广泛的平台上？

语法引擎效率

本节中评估了语法引擎的性能。研究者在 Llama-3.1-8B Instruct 上评估了他们的方法和基线，该模型能够遵循人类的指令。

结果如下图 9 所示，在 JSON 模式设置中，XGrammar 可以实现高达 3 倍的加速；在 JSON 语法用例下，可以实现超过 100 倍的加速。与 JSON 模式（更受限制）相比，JSON 的上下文无关语法包含更复杂的规则，因为它可以包含递归列表和字典，导致语法引擎更难有效地执行它。

在这两种情况下，XGrammar 都可以在不到 40 微秒的时间内生成每个 token 的掩码，使其成为低延迟 LLM 推理的理想选择。

端到端 LLM 引擎评估

本节在 LLM serving 设置下来评估 XGrammar。研究者将 XGrammar 集成到端到端 LLM 推理框架中，并与其他 LLM serving 框架进行效率比较。同时，他们还与其他支持结构化生成的 LLM 引擎进行效率比较，包括集成 Outlines 的 vLLM (v0.6.3) 和内置语法引擎的 llama.cpp。

实验结果如下图 10 所示，XGrammar 在 CFG 和 JSON 模式的所有基线中实现了最佳的 TTFT 和 TPOT。vLLM 和 llama.cpp 的计算受到其语法引擎更长预处理和每个 token 处理时长的阻碍。

在批量较大的情况下，vLLM 中 TPOT 速度的下降尤为明显。与现有解决方案相比，XGrammar 引擎总体上可以将输出 token 的速度提高 80 倍。这种加速来自 XGrammar 带来的性能优化。

研究者还在下表 1 中研究了语法处理的开销问题。由于 token 掩码生成效率和语法 GPU 重叠，语法过程在 TPOT 中几乎不产生任何开销。

跨平台部署

本节探讨如何将 XGrammar 引入各种平台。研究者利用 Emscripten 将 XGrammar 编译成 WebAssembly 并构建 JavaScript 捆绑包。他们进一步将 web-binding 与浏览器内 LLM 推理框架 WebLLM 集成，以实现结构化生成。

研究者使用 JSON-mode-eval 数据集评估端到端性能，在装有 Google Chrome 的 MacBook Pro M3 Max（macOS 14.5）上使用 4 位量化模型 Llama-3.1-8B-Instruct，并在装有 Safari 的 iPhone 14 Pro Max（iOS 18）上使用 Qwen2.5-0.5B-Instruct。

结果如下图 11 所示，研究者比较了使用 XGrammar 进行结构化生成和非结构化生成时的第一个 token 时间 (TTFT) 和每个输出 token 时间 (TPOT)，同时确保生成的 token 数量相同。结果表明，XGrammar 在两种设置下都几乎实现了零开销，在支持未来高性能端侧智能体方面具有巨大潜力。

更多技术细节请参阅原论文。

工程XGrammar