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

2026-07-10

AI Summary

This work addresses the computational bottleneck in large-scale generative retrieval caused by CPU-bound Trie-based constrained decoding, which suffers from limited parallelism and significant latency at large beam widths. We propose the first GPU-accelerated constrained beam search framework, which keeps the entire Trie index in GPU high-bandwidth memory and introduces a unified CUDA kernel that jointly performs beam expansion, constraint validation, and pruning. By replacing irregular CPU-style memory accesses and heap operations with an integer-aware compact Trie layout and GPU-native parallel primitives, our design substantially improves warp utilization and mitigates branch divergence. Evaluated on a keyword vocabulary of 800 million tokens, our method achieves decoding latency under 3 ms at beam width 1000—up to 24× faster than an optimized CPU baseline—and demonstrates a statistically significant 0.71% revenue uplift in online A/B experiments.

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.

Problem

• constrained decoding
• generative retrieval
• beam search
• trie traversal
• GPU acceleration

Innovation

• GPU-accelerated constrained decoding
• succinct trie
• beam search
• CUDA kernel optimization
• generative retrieval