FlashTrie：一种用于生成式检索的GPU加速约束束搜索方法

# FlashTrie：面向生成式检索的GPU加速约束束搜索

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

Dakshitha Anandakumar¹, Anurag Mukkara², Wenxiang Hu¹, Jiusheng Chen¹, M Akash Kumar¹, Ting Ye¹, Qiang Lou¹, Jian Jiao¹

¹微软，美国华盛顿州雷德蒙德  
²英伟达，美国加利福尼亚州圣克拉拉  

\{danandakumar, tiy\}@microsoft.com

###### 摘要

约束解码在生成式检索中至关重要，因为直接从查询生成的文档标识符必须精确匹配预定义的合法ID库。在大规模场景下，解码通常借助字典树与束搜索进行约束，但大多数实现运行在CPU上。随着束宽增大，有限的并行性使得字典树遍历和候选验证成为服务瓶颈。我们提出**FlashTrie**，通过在GPU上优化约束束搜索来解决这一限制。它引入了一种整数感知的紧凑字典树布局，利用位压缩减少内存占用，同时将完整索引保留在GPU高带宽内存中，从而减少内存停顿；以及一个协作式CUDA内核，完全在设备上执行束扩展、验证和剪枝，无需每步主机编排。它进一步用GPU感知的并行原语替换了CPU风格的不规则查找和堆维护，提高了线程束利用率并减少了发散。这些设计共同显著降低了解码延迟并提高了吞吐量，同时保持了检索质量。在包含8亿关键词的库上，当束宽达到1000时，FlashTrie将字典树搜索延迟降至3毫秒以下，相较于高度优化的多线程CPU基线实现了高达24倍的加速。这些改进使得FlashTrie能够在延迟关键型应用（如赞助搜索）中将束大小扩展高达5倍。在商业搜索引擎的大规模在线A/B实验中，它带来了统计显著的+0.71%收入提升，使得此前仅能离线完成的实时约束解码成为可能。FlashTrie代码将在评审过程后公开发布。

# FlashTrie：面向生成式检索的GPU加速约束束搜索

Dakshitha Anandakumar¹, Anurag Mukkara², Wenxiang Hu¹, Jiusheng Chen¹, M Akash Kumar¹, Ting Ye¹, Qiang Lou¹, Jian Jiao¹

¹微软，美国华盛顿州雷德蒙德  
²英伟达，美国加利福尼亚州圣克拉拉  

\{danandakumar, tiy\}@microsoft.com

## 1 引言

生成式检索（Tay et al., 2022 (https://arxiv.org/html/2607.10044#bib.bib1); Metzler et al., 2021 (https://arxiv.org/html/2607.10044#bib.bib54)）将文档检索重新定义为序列到序列任务，直接将查询映射到标识符（docID），取代了双编码器索引（Karpukhin et al., 2020 (https://arxiv.org/html/2607.10044#bib.bib36)）。这一范式对大规模搜索和推荐系统具有吸引力，因为原则上它可以避免昂贵的检索流水线，同时实现紧凑的端到端建模。然而，在实践中，在线服务受严格的延迟预算限制，解码过程成为主要瓶颈。自回归（AR）解码器（Sutskever et al., 2014 (https://arxiv.org/html/2607.10044#bib.bib55)）是顺序且昂贵的，而非自回归（NAR）解码器（Gu et al., 2018 (https://arxiv.org/html/2607.10044#bib.bib12); Sun and Yang, 2020 (https://arxiv.org/html/2607.10044#bib.bib37)）恢复了并行性，但其每个位置独立性经常产生无效标识符，必须通过约束搜索进行过滤（Ziem et al., 2023 (https://arxiv.org/html/2607.10044#bib.bib7); Pradeep et al., 2023 (https://arxiv.org/html/2607.10044#bib.bib39)）。

#### 先前工作与挑战。

几种策略将解码限制为有效输出，包括对数几率掩码（Tay et al., 2022 (https://arxiv.org/html/2607.10044#bib.bib1)）、有限状态编译（Willard and Louf, 2023 (https://arxiv.org/html/2607.10044#bib.bib11)）和谓词逻辑框架（Lu et al., 2021 (https://arxiv.org/html/2607.10044#bib.bib10); Anderson et al., 2017 (https://arxiv.org/html/2607.10044#bib.bib40)）。当标识符库扩展到数百万或数十亿条目时，常用方法是字典树约束束搜索（Hokamp and Liu, 2017 (https://arxiv.org/html/2607.10044#bib.bib25)），用于扁平令牌方案（DSI（Tay et al., 2022 (https://arxiv.org/html/2607.10044#bib.bib1)）、GENRE（Cao et al., 2021 (https://arxiv.org/html/2607.10044#bib.bib2)））和结构化语义标识符（NCI（Wang et al., 2022 (https://arxiv.org/html/2607.10044#bib.bib6)）、SEAL（Bevilacqua et al., 2022 (https://arxiv.org/html/2607.10044#bib.bib5)）以及最近的工作（Penha et al., 2025 (https://arxiv.org/html/2607.10044#bib.bib43)））。然而，基于指针的字典树表示（Morrison, 1968 (https://arxiv.org/html/2607.10044#bib.bib17); Aoe, 1989 (https://arxiv.org/html/2607.10044#bib.bib16)）存在不规则内存访问和硬件利用率低的问题，导致生产系统依赖优化的CPU实现（Cao et al., 2021 (https://arxiv.org/html/2607.10044#bib.bib2)）。紧凑设计如MARISA（Yata, 2011 (https://arxiv.org/html/2607.10044#bib.bib14)），基于LOUDS（Jacobson, 1989 (https://arxiv.org/html/2607.10044#bib.bib15)）和最小无环有限状态自动机（Daciuk et al., 2000 (https://arxiv.org/html/2607.10044#bib.bib18)），提高了空间效率，但仅支持字符级键，限制了其在具有大令牌词汇表和批量束感知遍历的生成式检索中的适用性。这些结构也难以映射到GPU，导致线程束发散和未合并访问（Merrill et al., 2012 (https://arxiv.org/html/2607.10044#bib.bib19)）。虽然现代推理运行时（Wang et al., 2021 (https://arxiv.org/html/2607.10044#bib.bib21); Dao et al., 2022 (https://arxiv.org/html/2607.10044#bib.bib26); Kwon et al., 2023 (https://arxiv.org/html/2607.10044#bib.bib22)）在GPU上优化了模型执行，但约束实施通常位于核心GPU解码路径之外，这可能引入额外的编排开销。这激发了三个关键问题：*(i)* 紧凑字典树能否在不牺牲空间效率的情况下适配令牌级词汇表；*(ii)* 约束束搜索能否完全在GPU上执行；*(iii)* 此类系统能否在不降低检索质量的情况下满足生产延迟约束？

#### 我们的贡献。

我们提出**FlashTrie**，一个面向生成式检索的GPU原生约束解码框架，联合优化了约束表示和解码计算。基于紧凑字典树，我们重新设计了MARISA以适配整数令牌词汇表，使用位压缩布局减少索引大小，并将约束结构驻留在GPU高带宽内存（HBM）中。我们还引入了GPU友好的搜索和排序原语、并行字典树子节点匹配和无堆束选择，以替换指针密集型遍历并减少发散。约束束搜索通过一个协作式多步CUDA内核完全在设备上执行，该内核执行扩展、验证和剪枝，无需在解码步骤间重复启动内核的开销。两级并行性（束并行和Top-K并行执行）利用层级GPU并发，饱和数百个流式多处理器（SM），并在整个设备上保持高占用率。FlashTrie将约束解码扩展到GPU上的十亿级约束库，这一规模此前仅离线可行，同时实现了比优化CPU基线高达24倍的加速，并将字典树搜索延迟降至3毫秒以下。在大规模赞助搜索系统的生产A/B测试中，这些延迟增益使得在严格服务预算内进行更宽的束搜索成为可能，转化为+0.71%的收入提升。

## 2 设计与实现

我们考虑使用词汇表 $\mathcal{V}$ 上的字典树 $\mathcal{T}$ 进行约束解码。FlashTrie适用于AR和NAR解码；我们在实验中专注于NAR，因为更大的分支因子使得约束检查更加昂贵。在解码步骤 $t$，基础模型发出Top-K候选 $(x_{t,i}, \ell_{t,i})$，其中 $\ell_{t,i} = \log p(x_{t,i} \mid \theta_{\text{sent}}, \text{previous tokens})$。约束束搜索确保每个beam中的轨迹对应于 $\mathcal{T}$ 中的有效前缀。我们在GPU上紧凑地表示 $\mathcal{T}$，并设计CUDA内核以最大化并行性。

```
算法1 约束束搜索（单查询，基于字典树 $\mathcal{T}$，词汇表 $\mathcal{V}$）
输入：基础模型top-K对数几率；字典树 $\mathcal{T}$；束宽 $B$
输出：$B$ 个最佳有效标识符
1: 初始化：$\mathit{cur} \leftarrow \{ (\epsilon, \epsilon, 0) \}$  
2: for $t = 1$ to $T$ do  
3:   $\mathit{next} \leftarrow \emptyset$  
4:   for each beam $(s, u, \log p)$ in $\mathit{cur}$ do  
5:     从基础模型获取top-K $(x_i, \ell_i)$  
6:     for each $(x_i, \ell_i)$ do  
7:       if $\exists \rho_k$ s.t. $\mathcal{T}.\text{child}(s, x_i) = \rho_k$ then  
8:         $\rho_k \leftarrow \text{Step2FlashTrie}\{\text{cond}\}$  
9:         附加 $(s, u, \rho_k)$ 到 $\mathit{next}$  
10:       end if  
11:     end for  
12:   end for  
13:   障碍；$\mathit{cur} \leftarrow \text{Top-}B(\mathit{next})$  \{在设备上\}  
14: end for  
15: 返回 $\text{Backtrace}(\mathit{cur})$
```

## 3 实验设置

我们使用来自生产NAR模型的13,000个检索请求进行评估，词汇表大小为220万。每个请求在 $T=8$ 个解码步骤上生成Top-K提议（$\in \{100, \ldots, 1000\}$，束宽 $BW = K$）。约束字典树由8亿个关键词序列构建，驻留在单个A100 80GB GPU上，配备AMD EPYC 7V13主机。对于MARISA-Opt，我们使用8个CPU工作线程以饱和8核分配；完整的硬件隔离细节见附录E (https://arxiv.org/html/2607.10044#A5)。我们测量每请求延迟（均值、p50/p90/p95/p99）、批量服务下的吞吐量（$b \in \{1,4,8,16,32\}$ 查询）、磁盘索引大小和构建时间，以及由Transformer教师评分的Precision@100/200（第4.5节 (https://arxiv.org/html/2607.10044#S4.SS5)）。延迟计时覆盖完整请求路径（C++入口到结果返回，包括GPU流同步）。所有计时平均10次遍历完整的13,000请求集，丢弃一次预热。详细信息见附录E (https://arxiv.org/html/2607.10044#A5)。

## 4 结果

### 4.1 字典树构建时间与索引大小

我们首先评估FlashTrie的存储重新设计如何影响两个直接的构建结果：字典树构建时间和磁盘索引大小。随着约束集增长，FlashTrie两者均减少（图2 (https://arxiv.org/html/2607.10044#S4.F2)）。在1000万到8亿序列范围内，不断扩大的差距反映了主导成本的转移：在小规模下，固定的LOUDS和尾部构建开销占主导（附录A (https://arxiv.org/html/2607.10044#A1)），而在大规模下，每边链接机制和遍历记账成为瓶颈（附录C (https://arxiv.org/html/2607.10044#A3)）。MARISA-Int和MARISA-Opt保持接近，因为它们保留了相同的全宽节点存储布局，而FlashTrie通过缩小链接偏移量并将每个32位标签拆分为紧凑基字段和打包的高位来减少每节点工作（第2.2节 (https://arxiv.org/html/2607.10044#S2.SS2)）。随着语料库增长，这些节省累积。即使在8亿这样的大约束索引下，FlashTrie中的字典树大小也仅约3 GB，使得字典树完全适合GPU HBM，并可在较小的GPU（如T4或A100 MIG）上部署。此外，FlashTrie的构建时间比MARISA-Opt快3.4倍，索引空间减少22%。总之，这些构建时的增益确立了存储重新设计的实际价值，这是后续延迟结果的关键促成因素之一。

参考图注：图2：字典树构建在1000万到8亿关键词序列上的扩展，（左）构建时间，（右）磁盘上的约束索引大小；运行时，整个字典树驻留在GPU内存中。在1000万键时，FlashTrie构建需0.15分钟，使用0.061 GB，而MARISA-Opt为0.48分钟/0.060 GB，MARISA-Int为0.58分钟/0.075 GB。在8亿键时，FlashTrie达到13.4分钟和3.1 GB，而MARISA-Opt为45.3分钟/4 GB，MARISA-Int为51.1分钟/4 GB。

### 4.2 延迟结果

接下来，我们评估GPU占用率敏感的内核优化，以及约束索引完全驻留设备对字典树搜索延迟的影响，随着束宽增加。在GPU上，每个请求运行在单个持久协作内核中，该内核在设备上执行每个解码步骤，并通过网格级障碍跨步骤同步（第2.2节 (https://arxiv.org/html/2607.10044#S2.SS2)）。主机每查询支付一次启动开销，而非每步一次，使得延迟在 $BW$ 增长时保持窄带（图3 (https://arxiv.org/html/2607.10044#S4.F3)）。GPU平均延迟从 $K=100$ 时的0.56毫秒上升到 $K=1000$ 时的1.91毫秒，p95低于2.79毫秒，p99低于3.31毫秒。相反，MARISA-Opt在固定大小的工作线程池上顺序扩展束（第2.1节 (https://arxiv.org/html/2607.10044#S2.SS1)），其平均延迟从9.03毫秒上升到46.30毫秒，p99从15.08毫秒上升到76.71毫秒。由此产生的平均加速比随着 $BW/K$ 增加从16.3倍增长到24.2倍（图3 (https://arxiv.org/html/2607.10044#S4.F3)）。我们在 $K \approx 500$ 和800附近观察到轻微下降，这是因为设备上的Top-B排序器改变了线程容量（附录D (https://arxiv.org/html/2607.10044#A4)）。MARISA-Int（第2.1节 (https://arxiv.org/html/2607.10044#S2.SS1)）在十亿关键词规模下不可行（10查询子样本的单个 $K=1000$ 查询 >3小时），因此我们采用MARISA-Opt作为主要CPU比较对象；完整的MARISA-Int测量见附录H (https://arxiv.org/html/2607.10044#A8)。保持低延迟随着 $BW$ 增加表明FlashTrie的优化在大规模下仍然有效；接下来，我们按阶段分解内核运行时，以确定随着 $BW$ 增长哪些组件占主导。

参考图注：图3：每请求字典树搜索延迟和GPU加速比 vs. 束宽 $BW$，$b=1$（$|\mathcal{V}|=2.2M$，8亿键字典树）。（左）均值（实线）和p95（虚线）延迟。CPU均值从9.0毫秒（$BW=100$）上升到46.3毫秒（$BW=1000$），在最高端p95为65.5毫秒，p99为76.7毫秒。GPU均值在整个扫描范围内低于2毫秒（0.55–1.91毫秒），p95 ≤ 2.79毫秒，p99 ≤ 3.31毫秒。（右）平均GPU加速比从16倍攀升至24倍。完整百分位表：附录F (https://arxiv.org/html/2607.10044#A6)。

### 4.3 运行时分解

为了分析延迟趋势，我们检测FlashTrie的协作内核，并测量GPU跨四个阶段的运行时（图1 (https://arxiv.org/html/2607.10044#S2.F1)）：束扩展、验证、剪枝和网格同步。图4 (https://arxiv.org/html/2607.10044#S4.F4)（左）显示主导成本随束宽变化：在 $K=100$ 时，扩展占主导（≈66%），因为字典树遍历是瓶颈。随着 $BW$ 增长，扩展仍然便宜，因为并行束共享上层字典树节点。相反，存活候选的数量随 $B \times K$ 扩展，增加了验证和剪枝成本；在 $BW=1000$ 时，验证成为主导（约占总运行时的43%）。剪枝保持在15–19%之间稳定，而同步开销很小（≤6%，降至3.5%），证实持久内核避免了每步主机编排且屏障开销可忽略。测量细节见附录D.4 (https://arxiv.org/html/2607.10044#A4.SS4)。总体而言，瓶颈在大BW时从字典树遍历转移到候选处理。接下来我们使用批量吞吐量测量来量化系统级影响。

### 4.4 批量大小下的吞吐量扩展

延迟改进如果在实际批处理下转化为更高的服务容量，则具有操作意义。在图4 (https://arxiv.org/html/2607.10044#S4.F4)（右）中，吞吐量曲线反映了不同的批处理机制。FlashTrie的吞吐量从 $b=1$ 到 $b=8$ 急剧上升，因为批处理填充了更多GPU执行资源，在 $BW=1000$ 时达到峰值1,621 q/s。这比MARISA-Opt在其自身峰值批处理大小下的吞吐量高70倍。