🤖 machine learning

FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval

FlashTrie 是一个通过采用位压缩字典树布局和协作式 CUDA 内核来消除 CPU 瓶颈，从而优化生成式检索中受限束搜索的 GPU 加速系统，在大型商业搜索应用中实现了高达 24 倍的加速和 0.71% 的营收提升。

原作者： Dakshitha Anandakumar, Anurag Mukkara, Wenxiang Hu, Jiusheng Chen, M Akash Kumar, Ting Ye, Qiang Lou, Jian Jiao

发布于 2026-07-14

📖 1 分钟阅读☕ 轻松阅读

原作者： Dakshitha Anandakumar, Anurag Mukkara, Wenxiang Hu, Jiusheng Chen, M Akash Kumar, Ting Ye, Qiang Lou, Jian Jiao

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你是一个超级聪明的机器人，正试图根据刚刚听到的一个问题，写出一份秘密代码列表（比如“DocID: 4592”）。但这里有一个限制：你只能编写那些确实存在于一本拥有 8 亿条有效条目的巨大、预先批准的电话簿中的代码。如果你猜了一个不在书中的代码，就算失败。

很长一段时间以来，机器人都是通过向一位非常快速、非常有条理的图书管理员（运行在标准计算机芯片或 CPU 上）询问来完成这项工作的，让它检查每一个猜测。但随着猜测列表的增长，图书管理员不堪重负。检查电话簿变成了一场交通拥堵，导致速度变慢。机器人必须一步一步地排队等待，看它的猜测是否被允许。

FlashTrie 出现了。微软和英伟达的研究人员决定解雇这位图书管理员，并将整本 8 亿条条目的电话簿直接移入机器人的超高速、高带宽内存（GPU）中。但他们不仅仅是移动了这本书；他们重新构建了它。

“位填充”电话簿的魔力

把旧的电话簿想象成一个巨大的图书馆，其中的每一本书都存放在一个巨大的、浪费了很多空间的空房间里。FlashTire 缩小了这些书的体积。它使用了一种被称为“位压缩”（bit compression）的聪明技巧，将信息紧紧挤压在一起，就像高效地打包行李一样，能将 8 亿个关键词压缩到仅 3.1 GB 的空间内。这个体积足够小，可以完全放入机器人的高速内存中，因此它永远不需要等待缓慢的外部硬盘来获取页面。

协作之舞

在旧系统中，机器人会做一个猜测，询问图书管理员进行检查，等待答案，然后做出另一个猜测，如此循环。这是一个孤独的、顺序的过程。

FlashTrie 改变了游戏规则。它使用了一个“协作式 CUDA 核函数”（cooperative CUDA kernel），这就像一个拥有 512 名舞者（线程）在舞池中完美同步协作的巨大舞池。

扩张（The Expansion）： 它不再是一个人检查一个猜测，而是数百名舞者同时检查数千个猜测。
验证（The Validation）： 它们使用“并行二分查找”（一种超快速的查找方式）来查看猜测是否与电话簿匹配。
修剪（The Pruning）： 如果一个猜测不好，它们会立即将其丢弃。如果是好的，它们就会保留。

因为一切都发生在舞池（GPU）上，而不需要机器人在每一步之后都停下来与主计算机（CPU）交谈，整个过程变得异常迅速。

结果：速度与智慧

团队在 8 亿个关键词 的库上测试了该系统。

速度： 当他们增加猜测的数量（即“束宽/beam width”）到 1,000 时，旧的 CPU 系统大约需要 46 毫秒，并且随着列表增长而变慢。FlashTrie 则将时间保持在 3 毫秒 以内（具体而言，平均为 1.91 毫秒，且最慢的 1% 也低于 3.31 毫秒）。
提升： 这意味着 FlashTrie 比高度优化的 CPU 版本快了高达 24 倍。
质量： 至关重要的是，变快并不意味着牺牲准确性。FlashTrie 找到的正确代码与慢速系统一样多。事实上，由于 FlashTrie 非常快，机器人可以在不突破时间限制的前提下，检查 600 个猜测 而不仅仅是 200 个。

现实世界的影响：金钱测试

研究人员不仅在计算机实验室进行了测试。他们在一个真实的、实时的商业搜索引擎（即你可能用来寻找信息的搜索工具）中测试了 FlashTrie。他们在不同国家进行了为期 16 天 的实验。

通过使用 FlashTrie 来检查更多猜测，搜索引擎展示了更好的广告。
这带来了 0.71% 的收入增长（广告带来的收益）。
同时，对于英文查询，点击率增加了 0.17%；对于非英文查询，点击率增加了 0.20%。
重要的是，广告质量并未下降；“缺陷率”（显示的错误广告）保持不变。

FlashTrie 不是什么

值得注意的是，这篇论文明确指出了哪些方法在此处不起作用或是不需要的。研究人员明确排除在 GPU 上使用旧式的“基于指针的”（pointer-based）库，因为它们会导致过多的混乱并减慢舞者的速度。他们还表明，如果只是简单地将旧系统移动到 GPU 而不重新设计数据结构（例如使用“线性探测/Linear-probe”方法），其速度会比他们的新方法慢 71 到 209 倍。这种提速来自于电话簿的具体设计和这种“舞蹈”方式，而不仅仅是使用了更快的硬件。

底线

FlashTrie 证明了你不必在速度和准确性之间做选择。通过重新设计如何存储“电话簿”以及如何进行“检查”，他们将一个缓慢的、顺序的瓶颈变成了一个快速的、并行的派对。这使得机器人能够在大脑中进行更大规模的思考（检查更多选项）并保持极高的速度，同时仍能满足实时互联网搜索所需的严格时间限制。该系统的代码将在评审过程结束后向公众发布，以便他人尝试这种全新的搜索方式。

技术摘要：FlashTrie

问题陈述

生成式检索将文档检索重新定义为一个序列到序列（sequence-to-sequence）的任务，即将查询直接映射到文档标识符（docID）。虽然这种范式提供了紧凑的端到端建模，但在线服务受到严格延迟预算的限制。解码过程，特别是对于非自回归（NAR）解码器，经常会生成无效的标识符，需要通过约束搜索进行过滤。

目前的先进方法利用**前缀树约束束搜索（trie-constrained beam search）**来确保生成的标识符与预定义的库相匹配。然而，现有的实现依赖于运行在 CPU 上的基于指针的前缀树表示（例如标准的 MARISA）。这些设计存在以下问题：

不规则内存访问： 导致硬件利用率低下和高延迟。
并行性受限： 随着束宽（beam width）的增加，基于 CPU 的遍历和候选验证成为主要的在线服务瓶颈。
在 GPU 上效率低下： 传统的紧凑型前缀树（如 MARISA）是为字符级键优化的，难以很好地适配 GPU 架构，从而导致线程束分歧（warp divergence）和非合并内存访问（uncoalesced memory access）。
编排开销： 约束强制执行通常在核心 GPU 解码路径之外实现，导致每个解码步骤都需要进行重复的主机-设备同步。

本文探讨了三个关键问题：能否在不牺牲空间效率的情况下，将紧凑型前缀树适配于 Token 级的词表？能否在 GPU 上完全执行约束束搜索？此类系统能否在不降低检索质量的前提下满足生产环境的延迟约束？

方法论：FlashTrie

FlashTrie 是一个 GPU 原生的约束解码框架，它对约束表示和解码计算进行了联合优化。它通过全设备协作式 CUDA 内核取代了基于 CPU 的编排。

1. 数据结构：整数感知型紧凑前缀树（Integer-Aware Succinct Trie）

FlashTrie 重新设计了 MARISA 前缀树的布局，以支持大规模整数 Token 词表（超过 100 万个 token），同时保持适合 GPU 高带宽内存（HBM）的小内存占用。

窄化 LOUDS（Narrow LOUDS）： 作者将 LOUDS（层序单进制度序列）表示扩展到了 32 位整数 Token。
位压缩： 为了减少内存使用，该布局将节点标签和后缀索引拆分为 2 字节的低位部分（存储在基础数组中）和位填充的高位部分（存储在需要它的二级数组中）。
优势： 这种“窄化”处理使主数组的大小减半，缩小了约束索引（将 8 亿个关键词压缩至约 3 GB），并通过减少每次访问时获取的比特数提高了内存局部性。

2. 算法：设备端协作式束搜索（Cooperative On-Device Beam Search）

FlashTrie 在单个持久化 CUDA 内核内执行整个约束束搜索，消除了每一步的主机编排。

两级并行性：
- 束并行（Beam-Parallel）： 一个查询中的活跃束被映射到不同的线程块（CTA）。
- 提议并行（Proposal-Parallel）： 在每个束内部，线程执行并行验证。每个线程通过对排序后的子节点标签进行并行二分查找，针对语言模型的 Top- $K$ 个提议进行验证。
内核工作流：
1. 扩展（Expansion）： 使用 Top- $K$ 提议扩展活跃束。
2. 验证（Validation）： 线程在前后缀树的子标签上执行二分查找以验证前缀。
3. 剪枝（Pruning）： 一个并行 Top- $B$ 选择内核取代了传统的 CPU 堆（heap）。它使用追加式候选缓冲区和并行排序来将束剪枝到目标宽度，从而避免锁竞争和指针追踪。
同步： 全局网格屏障（grid-wide barrier）同步各个步骤，允许内核在单次启动中运行 $T$ 个解码步骤。

核心贡献

GPU 常驻约束索引： 一种新型位压缩前缀树布局，可将十亿级规模的约束库完整保留在 GPU HBM 中，减少内存停顿。
协作式 CUDA 内核： 单内核执行模型，在设备端完成束扩展、验证和剪枝，消除了重复内核启动和主机编排的开销。
GPU 感知原语： 用并行二分查找和无堆 Top- $K$ 选择取代了不规则的指针遍历，显著提高了线程束（warp）利用率并减少了分歧。
可扩展性： 该系统可在单个 GPU 上将约束解码扩展到 8 亿个关键词，这在以前只能在离线设置中实现。

实验结果

系统在生产环境的 NAR 模型上进行了评估，该模型具有 220 万 Token 词表和包含 8 亿个关键词的约束前缀树，运行在 NVIDIA A100 80GB GPU 上。

性能指标

延迟： 即使在束宽（$BW$）为 1000 时，FlashTrie 也将前缀树搜索延迟降低至 3 ms 以下（p99）。
- 在 $BW=1000$ 时，FlashTrie 的平均延迟为 1.91 ms，而优化的 CPU 基准测试（MARISA-Opt）为 46.30 ms。
- 这比高度优化的多线程 CPU 基准测试快了 24 倍。
吞吐量： 在批处理服务（ $b=8$ ）下，FlashTrie 在 $BW=1000$ 时峰值达到 1,621 查询/秒，吞比 CPU 基准测试高出 32 倍至 71 倍。
构建： FlashTrie 的索引构建速度比 MARISA-Opt 快 3.4 倍，且磁盘空间占用减少了 22%。

检索质量

与 CPU 持平： 在所有束宽下，FlashTrie 的检索质量与 CPU 基准测试（MARISA-Opt）的精度差异在 0.001 以内。
通过束宽获得的质量提升： 由于 FlashTrie 保持了低延迟，它能够使用 CPU 无法实现的更大束宽（高达 1000）。将束宽从 100 增加到 1000，使 Precision@100 从 0.53 提高到 0.78。

在线 A/B 测试

在商业赞助搜索引擎的大规模在线实验中：

配置： 处理组使用 $BW=600$ 的 FlashTrie，而控制组使用 $BW=200$ 的 MARISA-Opt（这是 CPU 在延迟约束内能达到的最大值）。
结果： FlashTrie 带来了统计显著的 +0.71% 营收增长 以及 0.17–0.20% 的点击率（CTR）提升，且未降低广告质量（缺陷率保持不显著）。
延迟： 处理组维持了严格的延迟预算，与控制组的限制相比，其 p95 延迟降低了 32%。

重要性与主张

论文指出，约束解码不必始终是 CPU 的瓶颈。通过为并行执行重新设计数据结构并将整个约束强制逻辑移至 GPU，FlashTrie 实现了：

实时扩展： 它允许生成式检索系统在延迟敏感型应用（如赞助搜索）中，在不违反服务预算的前提下，将束大小扩展多达 5 倍。
解决质量与性能的权衡： 它将检索质量与延迟约束解耦，允许系统利用更大的束宽来进行更好的候选选择，同时保持低延迟。
广泛适用性： 该方法适用于任何具有大规模静态词表的约束生成任务，例如实体链接或产品目录检索。

作者总结道，FlashTrie 展示了一个更广泛的模式：当其访问模式被重新设计为适用于批处理和并行执行时，传统上在 CPU 上实现的数据结构可以在 GPU 上变得非常高效。