用于大型有限集合约束解码的Trie自动机

# 面向大规模有限集合约束解码的 Trie 自动机

来源：https://arxiv.org/html/2608.12574

###### 摘要

大型语言模型越来越需要生成符合预定义模式的结构化输出，其中一个常见约束是从有限的有效字符串集合中进行选择。当前的约束解码系统通过通用语法编译来处理这一问题，但随着有效值数量增长到数千，这变得极其缓慢，形成所谓的*基数墙*。我们引入了*Trie 自动机*，一种专门机制，利用有限集结构（共享前缀、有限深度、已知基数），通过 Aho-Corasick 多模式匹配预计算每个节点的 token 掩码。与 XGrammar（vLLM 和 SGLang 中的主要后端之一）相比，Trie 的每步有效 token 计算速度快 7 倍（0.65μs 对 5.8μs），并且在 K≥300K≥300 时编译速度快 2–6.5 倍。由于预计算掩码实现了无状态的服务路径，绕过了引导解码流程，这种优势在批量服务中进一步放大：在批大小 256 时，端到端 vLLM 吞吐量达到 219 req/s，而 XGrammar 为 7.5 req/s（29 倍）。这 29 倍结合了算法加速与仅预计算掩码才能实现的集成路径节省。在七种 tokenizer 家族（32K–262K 词表）中，Trie 在 K=10,000K=10{,}000 时保持低于 100ms 的编译时间，且无论集合大小如何，每步成本都持平，同时保证 100% 的输出有效性。

## 1 引言

约束解码已成为保证 LLM 输出符合模式的标准机制（34 (https://arxiv.org/html/2608.12574#bib.bib1)；15 (https://arxiv.org/html/2608.12574#bib.bib6)）。通过在每一步生成时屏蔽无效 token，它消除了格式错误的 JSON、幻觉字段名和无效值。各大 LLM 提供商现已提供此功能，Outlines、XGrammar（10 (https://arxiv.org/html/2608.12574#bib.bib4)）和 SGLang（35 (https://arxiv.org/html/2608.12574#bib.bib2)）等开源引擎也使其可供任何应用使用。这些系统将 JSON 模式（或语法）编译为通用自动机（有限状态机、下推自动机或 Earley 解析器），并在每个解码步骤使用它来屏蔽 token。这种架构处理任意模式，包括嵌套对象、递归结构和复杂正则表达式。然而，它对所有约束应用相同的通用编译流程，而不管其实际复杂度如何。深度嵌套的递归 JSON 模式和包含 1,000 个工具名称的平面列表都经历相同的编译流程，这是通用引擎与简单约束之间的根本性错配。这种统一性成为生产环境中最常见约束之一的瓶颈：*从已知有限集合中选择一个字符串*。OpenAI 的结构化输出设定了 1,000 个枚举限制（26 (https://arxiv.org/html/2608.12574#bib.bib25)），Google Gemini 在约 120 个枚举值时开始失效（17 (https://arxiv.org/html/2608.12574#bib.bib26)），Anthropic 的 180 秒编译超时（3 (https://arxiv.org/html/2608.12574#bib.bib27)）在几百个值时也暗示了类似的墙。随着 LLM 应用从开放式生成转向结构化工具使用，这些限制越来越重要。在智能体工作流中，LLM 必须从可能包含 500–5,000+ API 的注册表中选择要调用的工具（28 (https://arxiv.org/html/2608.12574#bib.bib19)；11 (https://arxiv.org/html/2608.12574#bib.bib20)；14 (https://arxiv.org/html/2608.12574#bib.bib21)）；随着 Model Context Protocol (MCP) 生态系统的增长以及组织将内部服务暴露为工具，这些注册表迅速扩展，通常在部署后几个月内就超过提供商的枚举限制。同样的模式也出现在零样本分类中，如产品分类法（1,500+ 类别）、ICD-10-CM 医疗代码（2026 CMS 版本中有 74,719 个代码（6 (https://arxiv.org/html/2608.12574#bib.bib10)））或法律案件类型（10,000+）；在针对知识库的实体链接中（9 (https://arxiv.org/html/2608.12574#bib.bib18)）有数万个条目；以及在检索增强系统中的动态每次查询约束中，有效集合随查询而变化，无法摊销编译成本。在所有这些情况下，约束都是字符串的有限并集 s1|s2|⋯|sKs_{1}|s_{2}|\cdots|s_{K}。虽然这是一种没有 Kleene 星号、递归或嵌套结构的正则语言，但当前系统通过用于任意语法的正则到 NFA 再到 DFA 的相同流程来编译它，成本随字符串数量和字母表大小增长。这造成我们所谓的*基数墙*：超过某个最大 KK 后，约束解码变得不切实际地缓慢。

核心洞见在于，不同类型的约束需要不同的执行机制。有限字符串集具有可利用的结构：共享前缀、有限深度和已知基数。我们引入了 *Trie 自动机*，作为现有约束解码流程中 FSM 层的即插即用替代品，专门针对有限集约束。它（1）直接从集合构建字符级 Trie（13 (https://arxiv.org/html/2608.12574#bib.bib16)），（2）使用 Aho-Corasick 多模式匹配（1 (https://arxiv.org/html/2608.12574#bib.bib22)）在每个 Trie 节点预计算词表感知的 token 掩码，以将 BPE token 与字符级 Trie 路径对齐，以及（3）在解码时通过 O(1)\mathcal{O}(1) 缓存查找提供服务掩码。核心算法挑战是 BPE-Trie 对齐：单个 BPE token 可能跨越多个 Trie 节点，而 32K32\text{K}–262K262\text{K} 规模的词表必须与每个节点进行匹配。据我们所知，这个对齐问题在约束解码文献中尚未得到解决；先前基于 Trie 的工作（9 (https://arxiv.org/html/2608.12574#bib.bib18)）通过以 token 粒度操作来回避它。我们证明它可以归结为多模式字符串匹配，求解时间与 Trie 大小成线性关系，而不是与词表大小成二次关系，从而在 K=10,000K=10{,}000 时实现亚 100ms 编译。图 1 (https://arxiv.org/html/2608.12574#S1.F1) 展示了基数墙以及 Trie 自动机如何克服它，将实际限制从约∼\sim1,000 扩展到约∼\sim100,000 个值，同时保持 100% 约束合规性。

这项工作有两个主要贡献：（1）我们引入了 *Trie 自动机*，一种针对有限集约束的专用约束解码后端，结合字符级 Trie、Aho-Corasick 多模式匹配和预计算 token 掩码，实现 O(|valid[st]|)\mathcal{O}(|\text{valid}[s_{t}]|) 的每步掩码（经验上在 3-4 个字符前缀后为 10-100 个 token，相对于通用 FSM 方法的 O(V⋅l)\mathcal{O}(V\cdot\ell) 成本，产生实际上恒定的成本）。BPE-Trie 对齐问题归结为多模式字符串匹配，产生 2–6.5 倍的编译速度提升，以及批量服务中高达 29 倍的端到端吞吐量提升：7 倍来自更快的每步掩码，加上仅预计算掩码才能使用的更简单服务路径。（2）我们通过七个 tokenizer 家族（32K–262K 词表）经验性地刻画了基数墙，表明将执行机制与约束结构匹配（包括集成路径）可以克服扩展瓶颈。预计算掩码使 Trie 完全绕过引导解码流程；FSM 方法无法做到这一点。

枚举基数 (KK)
编译时间
10
100
1K
10K
100K
1ms
10ms
100ms
1s
Gemini
Anthropic
OpenAI
40×(a) 编译时间
批大小
吞吐量 (req/s)
1
2
4
8
16
32
64
128
256
1
10
100
29×(b) vLLM 吞吐量 (K=1,000K{=}1{,}000)
Trie (ours)
FSM (XGrammar)

图 1：(a) 编译时间与枚举基数的关系（Qwen3-8B，对数-对数刻度）。虚线标记文档中记载的提供商限制（Gemini∼\sim120，Anthropic∼\sim200，OpenAI 1,000）。Trie 在 30–67ms 处保持平坦；XGrammar 在 K≈300K{\approx}300 处超过 Trie，并在 K=100KK{=}100\text{K} 时达到 2.7s。(b) 端到端 vLLM 吞吐量（对数刻度）。在 B=256B{=}256 处的 29 倍差距结合了 7 倍的每步算法优势与预计算掩码带来的集成路径节省（第 5 节 (https://arxiv.org/html/2608.12574#S5)）。

## 2 背景与问题表述

约束解码在每一步 tt 将模型的输出分布限制为能够导致有效完成的 token。给定大小为 VV 的词表 V\mathcal{V} 和由模式定义的正则语言 L\mathcal{L}，约束分布为：

pc(yt∣y,)。对于任何枚举 E\mathcal{E} 和任何前缀 y

等等（这里原文似乎有缺失，但我们就按原文翻译，保留 LaTeX 片段）。

...

（注意：原文第 2 节后面的内容似乎有省略号，我们按原文翻译。）

...

对于 K>50,000K\>50{,}000 的情况，我们在附录 I (https://arxiv.org/html/2608.12574#A9) 和 J (https://arxiv.org/html/2608.12574#A10) 中概述了两个初步扩展，作为未来方向：分层模式重写（O(K)\mathcal{O}(\sqrt{K}) 有效每步基数）和投机式短路。

## 5 实验

我们在 NVIDIA A100 GPU 上评估了 Trie 自动机在各种开源模型系列上的延迟、编译时间和准确性/有效性。主要比较对象是 xgrammar（10 (https://arxiv.org/html/2608.12574#bib.bib4)），即 vLLM 和 SGLang 使用的后端。Trie 自动机用 Rust 实现，通过 PyO3 提供 Python 绑定；XGrammar 用 C++ 实现，提供 Python 绑定。完整的实验细节见附录 C (https://arxiv.org/html/2608.12574#A3)。

### 5.1 延迟与可扩展性

表 4 (https://arxiv.org/html/2608.12574#S5.T4) 显示了 K∈{10,100,1,000,10,000}K\in\{10,100,1{,}000,10{,}000\} 下 XGrammar（10 (https://arxiv.org/html/2608.12574#bib.bib4)）、LLGuidance（15 (https://arxiv.org/html/2608.12574#bib.bib6)）和我们的 Trie 的性能细分。所有基准测试均在 NVIDIA A100 GPU（80GB）上运行，搭配 AMD EPYC 7R32 CPU（96 核）。三种方法在编译与掩码的权衡中占据不同位置（表 18 (https://arxiv.org/html/2608.12574#A13.T18)）。LLGuidance 实现接近零编译（0.6–24ms），但每步掩码成本为 73–141μs。XGrammar 在两者之间取得平衡（3–239ms 编译，5–10μs 掩码）。Trie 通过预计算最小化每步成本（0.65μs），并且编译时间几乎平坦（30–40ms）。对于编译占主导的一次性动态模式（例如，K<500K<500 时每次查询都有新的枚举集的检索增强设置），LLGuidance 的 1–3ms 编译可能更可取，尽管每步成本更高；当每步成本占主导时，Trie 的优势具有决定性，特别是在批量服务中（附录 L (https://arxiv.org/html/2608.12574#A12)）。

Trie 在批量服务中的优势具有决定性，因为 GPU 前向传播由 BB 个请求共享，但掩码在每个请求的 CPU 上运行。在 B=128B=128 时，LLGuidance 掩码（3.7ms）消耗了 GPU 前向传播时间的 37%，成为吞吐量瓶颈；XGrammar 的 783μs 占 7.8%；Trie 的 10μs 可以忽略不计（0.1%）（附录 M.1 (https://arxiv.org/html/2608.12574#A13.SS1) 中表 19 (https://arxiv.org/html/2608.12574#A13.T19) 给出了 BB 上的每批步掩码成本）。我们通过 K=1,000K=1{,}000 的端到端 vLLM 吞吐量验证了这一点（表 3 (https://arxiv.org/html/2608.12574#S5.T3)）：在 B=256B=256 时，Trie 达到 219 req/s，而 XGrammar 为 7.5 req/s（29 倍）。Trie 还超过了无约束吞吐量（219 对 104 req/s），因为约束解码在 Trie 叶节点终止（3.2 tokens/请求 对 8.7 无约束），减少了 GPU 前向传播次数。两种约束方法生成相同数量的 token，因此 Trie 与 XGrammar 的比率隔离了掩码和集成差异。

29 倍的差距复合了两个效应。第一，每步算法优势：预计算掩码查找成本为 0.65μs，而 XGrammar 的动态计算为 5.9μs（约∼\sim7 倍；表 4 (https://arxiv.org/html/2608.12574#S5.T4)）。第二，集成路径：由于 Trie 的掩码是预计算的，它作为无状态 `LogitsProcessor` 集成，每步返回一个缓存的位掩码。XGrammar 目前不支持该路径；其架构需要通过 vLLM 的引导解码流程进行动态掩码计算，包括按请求的语法编译、顺序 FSM 状态管理和调度开销。原则上，XGrammar *可以*为枚举约束预计算并缓存每状态掩码，但这样做实际上会重建 Trie：有限集的最小 DFA 与 Trie 同构（命题 1 (https://arxiv.org/html/2608.12574#Thmproposition1)），因此缓存其每状态掩码会产生相同的数据结构。因此，Trie 是优化基于 FSM 的有限集解码的自然终点。

#### 部署影响。批量服务结果对生产中的 GPU 利用率有直接影响。在 B=128B=128 时，每个解码步骤包含一次 GPU 前向传播（约∼\sim10ms），随后是 CPU 掩码计算。使用 XGrammar，掩码需要 783μs，其中 7.8% 的步骤时间 GPU 处于空闲状态等待掩码。使用 LLGuidance（3.7 ms），这一比例上升到 27%，使掩码成为主要瓶颈。Trie 将空闲时间减少到步骤时间的 0.1%，保持 GPU 饱和。对于动态枚举约束（例如，每次查询有效集都变化的检索增强工具选择），Trie 的 33–40 ms 编译足够快，可以即时运行而不影响服务延迟，而 XGrammar 在 K≥1,000K\geq 1{,}000 时的 75–239 ms 编译会增加可感知的延迟。AC 自动机可以按 tokenizer 缓存，并在所有枚举模式之间共享，因此每次模式变化只需运行 Trie 遍历（O(Nchars⋅l)\mathcal{O}(N_{\text{chars}}\cdot\ell)，通常<5{<\}\,5ms）。

表 3：K=1,000K=1{,}000 下在合成工具基准（附录 C.2 (https://arxiv.org/html/2608.12574#A3.SS2)）上的端到端 vLLM 吞吐量（req/s）。Trie 作为无状态 `LogitsProcessor`（预计算掩码）集成；XGrammar 需要 vLLM 的引导解码流程（动态掩码计算）。Trie 超过无约束吞吐量，因为它在 Trie 叶节点提前终止（3.2 对 8.7 tokens/请求）；两种约束方法都生成 3.2 tokens/请求（已验证），因此 Trie 与 XGrammar 的比率隔离了掩码和集成差异。Qwen3-8B，A100，贪心解码，3 次运行的中位数。

表 4：按 KK 划分的编译时间和每步掩码成本（Qwen3-8B，151K 词表，合成工具基准；附录 C.2 (https://arxiv.org/html/2608.12574#A3.SS2)）。每步测量仅包括掩码*计算*（确定有效 token），不包括应用于 logits 的 O(V)\mathcal{O}(V) 位掩码操作（约∼\sim31μs，通过张量运算，所有方法相同）。XGrammar 的每步成本非单调（K=10K=10 时为 9.5μs，K=100K=100 时为 5.4μs），反映了其词表分区：在较小的 KK 时，落入“自适应”分区的 token 较少，导致更多缓存未命中；分区在 K≥100K\geq 100 时趋于稳定。平均值±\pm标准差，10 次运行。

#### 内存与部署。Trie 的内存占用适度增长（附录 D (https://arxiv.org/html/2608.12574#A4)，表 10 (https://arxiv.org/html/2608.12574#A4.T10)）：K=10,000K=10{,}000 时为 0.9 MB，K=100,000K=100{,}000 时为 8 MB，而 FSM 约为∼\sim2 GB（附录 H (https://arxiv.org/html/2608.12574#A8)）。AC 自动机对于最大的词表（Gemma3 262K）需要约∼\sim150 MB，可在使用该 tokenizer 的所有枚举模式之间摊销。预计算掩码在构建后是不可变的，因此并发读取 a...