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

科技 2024-11-26 12:18 北京

机器之心报道

编辑：Panda、杜伟

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

不管是编写和调试代码，还是通过函数调用来使用外部工具，又或是控制机器人，都免不了需要 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 在两种设置下都几乎实现了零开销，在支持未来高性能端侧智能体方面具有巨大潜力。

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

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

投稿或寻求报道：liyazhou@jiqizhixin.com

http://mp.weixin.qq.com/s?__biz=MzA3MzI4MjgzMw==&mid=2650944286&idx=3&sn=c5cf695b7aa64433e8b525a3c292fab4

机器之心

专业的人工智能媒体和产业服务平台

最新文章

向量数据库的中场战事：长期主义者Zilliz如何全球突围

世界首次！智源研究院实现数字孪生心脏电功能超实时仿真

rebuttal真的有用！这篇ICLR论文，所有审稿人都加了2分，直接跃升排名第9

12%计算量就能媲美原模型，Adobe、罗切斯特大学等提出YOPO剪枝技术

LLM破局泛化诊断难题，MSSP刊登北航PHM实验室健康管理大模型交叉研究

刚刚，Ilya的Seq2Seq、Ian的GAN获NeurIPS时间检验奖

Scaling Law 撞墙？复旦团队大模型推理新思路：Two-Player架构打破自我反思瓶颈

国产大模型首发中文逻辑推理，「天工大模型4.0」o1版来了

遗憾不？原来百度2017年就研究过Scaling Law，连Anthropic CEO灵感都来自百度

HuggingFace工程师亲授：如何在Transformer中实现最好的位置编码

跨模态大升级！少量数据高效微调，LLM教会CLIP玩转复杂文本

Sora就这么泄露了三小时，网友调侃Altman急拔网线，艺术家们也在抗议被「白嫖」

创业一年半，胖了30斤，AI大佬感叹：还是回谷歌好

吴恩达出手，开源最新Python包，一个接口调用OpenAI等模型

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

「毕昇一号」DNA活字存储喷墨打印机来了，低成本、高效率、全自动的DNA存储

撞墙还是新起点？自回归模型在图像领域展现出Scaling潜力

和梁朝伟同获港科荣誉博士，黄仁勋与沈向洋对谈Scaling Law、后训练、机器人和爱情

更新了！带Agent的Cursor太疯狂了

小学二年级数学水平，跟着这篇博客也能理解LLM运行原理

文本、图像、点云任意模态输入，AI能够一键生成高质量CAD模型了

智能体竟能自行组建通信网络，还能自创协议提升通信效率

AI版周扒皮！打字速度慢、鼠标超30秒未动，就被AI「警告」，Karpathy下场评论

RTX 4090可跑、完全开源，最快视频生成模型问世，实测一言难尽

RL「误人」？LeCun 在技术路线上又有何战略摇摆？

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

研究大模型门槛太高？不妨看看小模型SLM，知识点都在这

大半年过去，主流视频生成模型们超越Sora了吗？

这才是真・开源模型！公开「后训练」一切，性能超越Llama 3.1 Instruct

阿里国际版o1来了，Marco-o1：聚焦开放式问题推理

英伟达开源福利：视频生成、机器人都能用的SOTA tokenizer

NeurIPS 2024 Oral | 还原所见！揭秘从脑信号重建高保真流畅视频

如今的智能体，已经像人一样「浏览」视频了，国内就有

仅仅一天，Gemini就夺回了GPT-4o拿走的头名

上交大o1复现新突破：蒸馏超越原版，警示AI研发"捷径陷阱"

大模型不会推理，为什么也能有思路？有人把原理搞明白了

全球十亿级轨迹点驱动，首个轨迹基础大模型来了

扣子OpenAPI突进智能语音战场！点满低延时、定制化、随时打断和音色克隆技能（内测开启！）

推理性能直逼o1，DeepSeek再次出手，重点：即将开源

诺奖得主哈萨比斯新作登Nature，AlphaQubit解码出更可靠量子计算机

神级项目训练GPT-2仅需5分钟，Andrej Karpathy都点赞

NeurIPS 2024 | 水印与高效推理如何两全其美？最新理论：这做不到

大模型代肝，自动刷《崩铁》升级材料，Claude操纵计算机还能这么用！

实测昆仑万维对话AI「Skyo」，会读诗、知晓雷军摆拍

室温超导学术不端、多次Nature撤稿，这位印度裔学者被大学解雇

德国科学家激进观点：意识是虚拟的，存在于大脑构建的梦中

媲美OpenAI事实性基准，这个中文评测集让o1-preview刚刚及格

在「最难LLM评测榜单」上，阶跃万亿参数模型拿下中国第一

登上Nature的AI芯片设计屡遭质疑，谷歌发文反击，Jeff Dean：质疑者连预训练都没做

高通的自研架构芯片，正在整合生成式AI世界

分类

时事

民生

政务

教育

文化

科技

财富

体娱

健康

情感

旅行

百科

职场

楼市

企业

乐活

学术

汽车

时尚

创业

美食

幽默

美体

文摘

原创标签

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

发布位置

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