FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval
2026-07-10 • Machine Learning
Machine Learning
AI summaryⓘ
The authors developed FlashTrie, a new way to speed up a step called constrained decoding used in search engines. Normally, this step is slow because it uses CPUs that can't handle many tasks at once, but FlashTrie uses GPUs to do this faster by organizing data efficiently and doing many tasks in parallel. This change lets the decoding happen much quicker without losing accuracy, handling very large sets of document IDs in just a few milliseconds. When tested in a big search engine, it also helped increase revenue slightly by making search results faster and more reliable.
Constrained decodingGenerative retrievalTrie data structureBeam searchGPU parallelismCUDA kernelLatencyThroughputSponsored searchOnline A/B testing
Authors
Dakshitha Anandakumar, Anurag Mukkara, Wenxiang Hu, Jiusheng Chen, M Akash Kumar, Ting Ye, Qiang Lou, Jian Jiao
Abstract
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.