用GPU加速生成式检索中的约束解码,速度提升24倍。
FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval
- 设计了适合GPU的紧凑整数索引结构,减少内存占用并避免数据停滞。
- 在8亿关键词库上,1000宽束搜索延迟低于3毫秒,提速24倍。
- 适用于实时广告搜索等对延迟敏感场景,已在线上实验中带来显著收益。
生成式检索中,从查询直接生成的文档标识符必须精确匹配预定义的有效ID列表。大规模应用中,通常使用前缀树(trie)结合束搜索进行约束解码,但多数实现依赖CPU,受限于并行度,导致前缀树遍历和候选验证成为服务瓶颈,尤其当束宽增大时。本文提出FlashTrie,通过优化GPU上的约束束搜索解决此问题。它采用一种面向整数的紧凑前缀树布局,利用位压缩降低内存占用,并将完整索引保留在GPU高带宽内存中以减少内存停顿;同时设计协作式CUDA内核,在设备端完成束扩展、验证与剪枝,无需每步主机调度。此外,以GPU友好的并行原语替代传统的不规则查找与堆维护,提升了线程组利用率并减少了分支发散。这些设计显著降低解码延迟,提高吞吐量,同时保持检索质量。在包含8亿关键词的库上,束宽达1000时,前缀树搜索延迟低于3毫秒,相比高度优化的多线程CPU基线提升24倍。该改进使束大小在延迟敏感场景下可扩大5倍。在主流商业搜索引擎的大规模线上A/B实验中,该方法带来统计显著的+0.71%收入增长,实现了此前仅限离线处理的实时约束解码。代码将在评审结束后公开。
原文摘要 · Abstract (English)
Constrained decoding is essential in generative retrieval, where document identifiers generated directly from a query must exactly match a predefined library of valid IDs. At scale, decoding is often constrained using a trie with beam search but most implementations run on CPU. Limited parallelism then makes trie traversal and candidate validation a serving bottleneck as beam width grows. We present FlashTrie, which addresses this limitation by optimizing constrained beam search on GPUs. It introduces an integer-aware succinct trie layout that uses bit compression to reduce memory footprint while keeping the full index in GPU high-bandwidth memory reducing memory stalls, and a cooperative CUDA kernel that performs beam expansion, validation, and pruning entirely on-device without per-step host orchestration. It further replaces CPU-style irregular lookup and heap maintenance with GPU-aware parallel primitives, improving warp utilization and reducing divergence. Together, these designs significantly reduce decoding latency and increase throughput while preserving retrieval quality. On a library of 800M keywords with beam widths up to 1000, FlashTrie reduces trie-search latency to under 3 ms, achieving up to 24x speedup over a highly optimized multi-threaded CPU baseline. These improvements enable FlashTrie to scale beam sizes by up to 5x in latency-critical applications such as sponsored search. In a large-scale online A/B experiment on a popular commercial search engine, it delivers a statistically significant +0.71% revenue lift, enabling real-time constrained decoding at a scale previously feasible only offline. The FlashTrie code will be publicly released after the review process.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。