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

arXiv: https://arxiv.org/abs/2607.10044 · HTML: https://mlsanity.com/paper/2607.10044/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 16 of its 30 references that are in the map.
- First posted on arXiv 10 Jul 2026.

## Builds on (most cited first, 16 of 16)

- [Efficient Memory Management for Large Language Model Serving with PagedAttention](https://mlsanity.com/paper/2309.06180/) (2023)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [Dense Passage Retrieval for Open-Domain Question Answering](https://mlsanity.com/paper/2004.04906/) (2020)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Latent Retrieval for Weakly Supervised Open Domain Question Answering](https://mlsanity.com/paper/1906.00300/) (2019)
- [KILT: a Benchmark for Knowledge Intensive Language Tasks](https://mlsanity.com/paper/2009.02252/) (2020)
- [Non-Autoregressive Neural Machine Translation](https://mlsanity.com/paper/1711.02281/) (2017)
- [Autoregressive Entity Retrieval](https://mlsanity.com/paper/2010.00904/) (2020)
- [Transformer Memory as a Differentiable Search Index](https://mlsanity.com/paper/2202.06991/) (2022)
- [NeuroLogic Decoding: (Un)supervised Neural Text Generation with Predicate Logic Constraints](https://mlsanity.com/paper/2010.12884/) (2020)
- [Autoregressive Search Engines: Generating Substrings as Document Identifiers](https://mlsanity.com/paper/2204.10628/) (2022)
- [A Neural Corpus Indexer for Document Retrieval](https://mlsanity.com/paper/2206.02743/) (2022)
- [Guided Open Vocabulary Image Captioning with Constrained Beam Search](https://mlsanity.com/paper/1612.00576/) (2016)
- [An EM Approach to Non-autoregressive Conditional Sequence Generation](https://mlsanity.com/paper/2006.16378/) (2020)
- [Semantic IDs for Joint Generative Search and Recommendation](https://mlsanity.com/paper/2508.10478/) (2025)

## Abstract (arXiv)

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.