[{"data":1,"prerenderedAt":-1},["ShallowReactive",2],{"doc-detail-85797-en":3,"doc-seo-85797-105":30,"detail-sidebar-cat-0-en-105":91},{"code":4,"msg":5,"data":6},0,"success",{"doc_id":7,"user_id":8,"nickname":9,"user_avatar":10,"doc_module":4,"category_id":11,"category_name":12,"doc_title":13,"doc_description":14,"doc_content":15,"file_id":16,"file_url":17,"file_type":18,"file_size":19,"view_count":20,"is_deleted":4,"is_public":21,"is_downloadable":21,"audit_status":21,"page_count":22,"language":23,"language_code":24,"site_id":25,"html_lang":24,"table_of_contents":26,"faqs":27,"seo_title":13,"seo_description":14,"update_tm":28,"read_time":29},85797,8796095462418,"Noah","https://ap-avatar.wpscdn.com/avatar/80000253c1241d02b47?x-image-process=image/resize,m_fixed,w_180,h_180&k=1778826106357471780",8,"Research & Report","FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative Retrieval","Constrained decoding is central to generative retrieval, where generated document identifiers must exactly match a predefined ID library. At large scale, trie-based constrained beam search often runs on CPU, making trie traversal and candidate validation a serving bottleneck as beam width increases. FlashTrie accelerates this workflow on GPUs by using an integer-aware succinct trie with bit compression to reduce memory footprint while keeping the index in GPU high-bandwidth memory. A cooperative CUDA kernel executes beam expansion, validation, and pruning entirely on-device, replacing irregular CPU-style lookups with GPU-parallel primitives to reduce stalls, divergence, and latency. Experiments on an 800M-keyword library with beam widths up to 1000 show trie-search latency under 3 ms, up to 24× speedup versus a tuned CPU baseline, and a +0.71% statistically significant revenue lift in online A/B testing. ","FlashTrie: A GPU-Accelerated Constrained Beam Search for Generative  \nRetrieval  \nDakshitha Anandakumar1 Anurag Mukkara2 Wenxiang Hu1 Jiusheng Chen1  \nM Akash Kumar 1 Ting Ye 1 Qiang Lou 1 Jian Jiao 1  \n1Microsoft, Redmond, WA, USA 2Nvidia, Santa Clara, CA, USA  \n{danandakumar, [tiy](tiy}@microsoft.com)[}](tiy}@microsoft.com)[@microsoft.com](tiy}@microsoft.com)  \narXiv :2607 . 10044v 1 [ cs .LG] 10 Jul 2026  \nAbstract  \nConstrained 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.  \nLimited parallelism then makes trie traversaland candidate validation a serving bottleneck as beam width grows.  \nWe 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 mainte  \nnance with GPU-aware parallel primitives, improving warp utilization and reducing divergence.  \nTogether, 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 24× speedup over a highly optimized multi-threaded CPU baseline. These improvements enable FlashTrie to scale beam sizes by up to 5 × in latency-critical applications such as sponsored search. In a large-scale online A/B experiment on a commercial search engine, it delivers a statistically significant +0 .71% revenue lift, enabling realtime constrained decoding at a scale previously feasible only offline. The FlashTrie code will be publicly released after the review process.  \n1 Introduction  \nGenerative retrieval (Tay et al., 2022 ; Metzler et al., 2021) reframes document retrieval as a sequence-  \nto-sequence task, mapping a query directly to an identifier (docID), replacing dual-encoder indexing (Karpukhin et al., 2020) . This paradigm is attractive for large-scale search and recommendation systems because it can, in principle, avoid expensive retrieval pipelines while enabling compact end-to-end modeling. In practice, however, online serving is constrained by strict latency budgets, and the decoding procedure becomes the dominant bottleneck. Autoregressive (AR) (Sutskever et al., 2014) decoders are sequential and costly, while non-autoregressive (NAR) (Gu et al., 2018 ; Sun and Yang, 2020) decoders regain parallelism but their per-position independence frequently produce invalid identifiers that must be filtered by constrained search (Ziems et al., 2023 ; Pradeep et al., 2023) .  \nPrior Work and Challenges. Several strategies restrict decoding to valid outputs, including logit masking (Tay et al., 2022), finite-state compilation (Willard and Louf, 2023), and predicatelogic frameworks (Lu et al., 2021 ; Anderson et al., 2017) . As the identifier library scales to millions or billions of entries, a common approach is trie-constrained beam search (Hokamp and Liu, 2017), used across flat-token schemes (DSI (Tayet al., 2022), GENRE (Cao et al., 2021)), and structured semantic-identifiers (NCI (Wang et al., 2022), SEAL (Bevilacqua et al., 2022), and recent work (Penha et al., 2025)) .  \nHowever, pointer-based trie representations (Morrison, 1968 ; ichi Aoe, 1989) suffer from irregular memory access and poor hardware utilization, leading production systems to rely on optimized CPU implementations (Cao et al., 2021) . Succinct designs such as MARISA (Yata, 2011), built on LOUDS (Jacobson, 1989) and minima","cbCaijC4RzDxLtl6","https://ap.wps.com/l/cbCaijC4RzDxLtl6","pdf",680559,4,1,19,"English","en",105,"# Abstract\n# Introduction\n## Prior Work and Challenges\n## Our Contribution","[{\"question\":\"What problem does FlashTrie address in generative retrieval?\",\"answer\":\"FlashTrie targets the latency bottleneck caused by running trie-constrained beam search on CPU, where trie traversal and candidate validation become slow as beam width grows.\"},{\"question\":\"How does FlashTrie represent the constraint trie on GPU?\",\"answer\":\"It redesigns a succinct trie for integer token vocabularies using bit-compressed layout so the full constraint index can remain resident in GPU high-bandwidth memory.\"},{\"question\":\"What GPU execution strategy does FlashTrie use for constrained decoding?\",\"answer\":\"It performs beam expansion, validation, and pruning fully on-device using a cooperative multi-step CUDA kernel, avoiding per-step host orchestration and reducing irregular access patterns.\"}]",1784206330,48,{"code":4,"msg":31,"data":32},"ok",{"site_id":25,"language":24,"slug":33,"title":13,"keywords":34,"description":14,"schema_data":35,"social_meta":86,"head_meta":88,"extra_data":90,"updated_unix":28},"flashtrie-a-gpu-accelerated-constrained-beam-search-for-generative-retrieval","",{"@graph":36,"@context":85},[37,53,68],{"@type":38,"itemListElement":39},"BreadcrumbList",[40,44,48,51],{"item":41,"name":42,"@type":43,"position":21},"https://docshare.wps.com","Home","ListItem",{"item":45,"name":46,"@type":43,"position":47},"https://docshare.wps.com/document/","Document",2,{"item":49,"name":12,"@type":43,"position":50},"https://docshare.wps.com/document/research-report/",3,{"item":52,"name":13,"@type":43,"position":20},"https://docshare.wps.com/document/flashtrie-a-gpu-accelerated-constrained-beam-search-for-generative-retrieval/85797/",{"url":52,"name":13,"@type":54,"author":55,"headline":13,"publisher":57,"fileFormat":60,"inLanguage":24,"description":14,"dateModified":61,"datePublished":62,"encodingFormat":60,"isAccessibleForFree":63,"interactionStatistic":64},"DigitalDocument",{"name":9,"@type":56},"Person",{"url":41,"name":58,"@type":59},"DocShare","Organization","application/pdf","2026-07-26","2026-07-16",true,{"@type":65,"interactionType":66,"userInteractionCount":20},"InteractionCounter",{"@type":67},"ViewAction",{"@type":69,"mainEntity":70},"FAQPage",[71,77,81],{"name":72,"@type":73,"acceptedAnswer":74},"What problem does FlashTrie address in generative retrieval?","Question",{"text":75,"@type":76},"FlashTrie targets the latency bottleneck caused by running trie-constrained beam search on CPU, where trie traversal and candidate validation become slow as beam width grows.","Answer",{"name":78,"@type":73,"acceptedAnswer":79},"How does FlashTrie represent the constraint trie on GPU?",{"text":80,"@type":76},"It redesigns a succinct trie for integer token vocabularies using bit-compressed layout so the full constraint index can remain resident in GPU high-bandwidth memory.",{"name":82,"@type":73,"acceptedAnswer":83},"What GPU execution strategy does FlashTrie use for constrained decoding?",{"text":84,"@type":76},"It performs beam expansion, validation, and pruning fully on-device using a cooperative multi-step CUDA kernel, avoiding per-step host orchestration and reducing irregular access patterns.","https://schema.org",{"og:url":52,"og:type":87,"og:title":13,"og:site_name":58,"og:description":14},"article",{"robots":89,"canonical":52},"index,follow",{"doc_id":7,"site_id":25},{"code":4,"msg":5,"data":92},[93,97,101,105,110,115,120,123,128,131,135],{"id":21,"doc_module":4,"doc_module_name":46,"category_name":94,"show_sort_weight":95,"slug":96},"Story & Novel",90,"story-novel",{"id":47,"doc_module":4,"doc_module_name":46,"category_name":98,"show_sort_weight":99,"slug":100},"Literature",80,"literature",{"id":20,"doc_module":4,"doc_module_name":46,"category_name":102,"show_sort_weight":103,"slug":104},"Exam",70,"exam",{"id":106,"doc_module":4,"doc_module_name":46,"category_name":107,"show_sort_weight":108,"slug":109},5,"Comic",60,"comic",{"id":111,"doc_module":4,"doc_module_name":46,"category_name":112,"show_sort_weight":113,"slug":114},6,"Technology",50,"technology",{"id":116,"doc_module":4,"doc_module_name":46,"category_name":117,"show_sort_weight":118,"slug":119},7,"Healthcare",40,"healthcare",{"id":11,"doc_module":4,"doc_module_name":46,"category_name":12,"show_sort_weight":121,"slug":122},30,"research-report",{"id":124,"doc_module":4,"doc_module_name":46,"category_name":125,"show_sort_weight":126,"slug":127},9,"Religion & Spirituality",20,"religion-spirituality",{"id":126,"doc_module":4,"doc_module_name":46,"category_name":129,"show_sort_weight":126,"slug":130},"World Cup","world-cup",{"id":132,"doc_module":4,"doc_module_name":46,"category_name":133,"show_sort_weight":132,"slug":134},10,"Lifestyle","lifestyle",{"id":22,"doc_module":4,"doc_module_name":46,"category_name":136,"show_sort_weight":106,"slug":137},"General","general"]