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

科技 2024-11-27 13:30 江苏

点击蓝字关注我们

关注并星标

从此不迷路

计算机视觉研究院

公众号ID｜计算机视觉研究院

学习群｜扫码在主页获取加入方式

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

计算机视觉研究院专栏

Column of Computer Vision Institute

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

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

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

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 在两种设置下都几乎实现了零开销，在支持未来高性能端侧智能体方面具有巨大潜力。

© THE END

转载请联系本公众号获得授权

计算机视觉研究院学习群等你加入！

ABOUT

计算机视觉研究院

计算机视觉研究院主要涉及深度学习领域，主要致力于目标检测、目标跟踪、图像分割、OCR、模型量化、模型部署等研究方向。研究院每日分享最新的论文算法新框架，提供论文一键下载，并分享实战项目。研究院主要着重”技术研究“和“实践落地”。研究院会针对不同领域分享实践过程，让大家真正体会摆脱理论的真实场景，培养爱动手编程爱动脑思考的习惯！

往期推荐

🔗

计算机视觉研究院

计算机视觉研究院主要涉及AI研究和落地实践，主要致力于目标检测、目标跟踪、图像分割、OCR、模型量化、模型部署等研究方向。研究院每日分享最新的论文算法新框架，提供论文一键下载，并分享实战项目。研究院主要着重”技术研究“和“实践落地”！

最新文章

从源头消除大模型“幻觉”，性价比吊打传统微调方法

Fast YOLO：用于实时嵌入式目标检测（附论文下载）

旋转角度目标检测的重要性！！！（附源论文下载）

YOLO-S：小目标检测的轻量级、精确的类YOLO网络

最后征稿+连续8届检索 | 第九届控制工程与人工智能国际会议(CCEAI 2025)会议地点已定！

智慧建筑：基于YOLOv7的建筑外墙缺陷检测

轻量级模型，重量级性能，TinyLlama、LiteLlama小模型火起来了

欢迎加入“计算机视觉研究院”学习群

年薪百万or惨遭裁员，AIGC开发者如何破局？

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

2024全球无人机市场洞察报告

粉丝福利！免费赠书中奖名单

凭什么YOLO是最强目标检测器，一文读懂！

大改Yolo框架 | 能源消耗极低的目标检测新框架（附论文下载）

智能体零样本解决未见过人类设计环境！全靠这个开放式物理RL环境空间

史上最通俗易懂的YOLO系列（v1-v10）模型解读！

更快、更灵活的Transformer图像去雾网络

“计算机视觉研究院”商务合作

利用先进技术保家卫国：深度学习进行小目标检测（适合初学者入门）

机场项目：解决飞行物空间大小/纵横比、速度、遮挡等问题引起的实时目标检测问题

2024新技术：远距离的小目标也可以准确检测

大改Yolo框架 | 能源消耗极低的目标检测新框架（附论文下载）

“计算机视觉研究院”商务合作

RestoreDet：低分辨率图像中目标检测

GPT理解的CV：基于Yolov5的半监督目标检测

机场项目：解决飞行物空间大小/纵横比、速度、遮挡等问题引起的实时目标检测问题

“计算机视觉研究院”商务合作

送书福利《计算机视觉之PyTorch数字图像处理》等你来拿！！！

YOLO-SLD：一种改进的注意机制YOLO车牌检测

多尺度特征融合：为检测学习更好的语义信息（附论文下载）

“计算机视觉研究院”商务合作

YOLO跌落神坛？新一代目标检测器又有新突破！

乌镇“网”事：看这道“光”！站在小镇望世界

欢迎加入“计算机视觉研究院”学习群

快速学会AI核心架构，Transformer！

大模型承重墙，去掉了就开始摆烂！苹果给出了「超级权重」

免费分享顶会的idea方法！！！

无卷积骨干网络：金字塔Transformer，提升目标检测/分割等任务精度（附源代码）

我用捡来的idea发了顶会！！！

揭示Transformer重要缺陷！北大提出傅里叶分析神经网络FAN，填补周期性特征建模缺陷

YoloV：视频中目标实时检测依然很棒（附源代码下载）

ICLR'25审稿质量历史最低？这么多全负正常么

NeurIPS 2024 | 无需训练，一个框架搞定开放式目标检测、实例分割

“计算机视觉研究院”商务合作

Yolo-Z：改进的YOLOv5用于小目标检测（附原论文下载）

多尺度特征融合：为检测学习更好的语义信息（附论文下载）

轻量级模型，重量级性能，TinyLlama、LiteLlama小模型火起来了

“计算机视觉研究院”商务合作

喜提 TPAMI 顶刊！！！

Token化一切，甚至网络！北大&谷歌&马普所提出TokenFormer，Transformer从来没有这么灵活过！

分类

时事

民生

政务

教育

文化

科技

财富

体娱

健康

情感

旅行

百科

职场

楼市

企业

乐活

学术

汽车

时尚

创业

美食

幽默

美体

文摘

原创标签

时事社会财经军事教育体育科技汽车科学房产搞笑综艺明星音乐动漫游戏时尚健康旅游美食生活摄影宠物职场育儿情感小说曲艺文化历史三农文学娱乐电影视频图片新闻宗教电视剧纪录片广告创意壁纸头像心灵鸡汤星座命理教育培训艺术文化金融财经健康医疗美妆时尚餐饮美食母婴育儿社会新闻工业农业时事政治星座占卜幽默笑话独立短篇连载作品文化历史科技互联网

发布位置

广东北京山东江苏河南浙江山西福建河北上海四川陕西湖南安徽湖北内蒙古江西云南广西甘肃辽宁黑龙江贵州新疆重庆吉林天津海南青海宁夏西藏香港澳门台湾美国加拿大澳大利亚日本新加坡英国西班牙新西兰韩国泰国法国德国意大利缅甸菲律宾马来西亚越南荷兰柬埔寨俄罗斯巴西智利卢森堡芬兰瑞典比利时瑞士土耳其斐济挪威朝鲜尼日利亚阿根廷匈牙利爱尔兰印度老挝葡萄牙乌克兰印度尼西亚哈萨克斯坦塔吉克斯坦希腊南非蒙古奥地利肯尼亚加纳丹麦津巴布韦埃及坦桑尼亚捷克阿联酋安哥拉