# Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators

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

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 22 of its 46 references that are in the map.
- First posted on arXiv 26 Feb 2026.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [In-Datacenter Performance Analysis of a Tensor Processing Unit](https://mlsanity.com/paper/1704.04760/) (2017)
- [Ups and Downs: Modeling the Visual Evolution of Fashion Trends with One-Class Collaborative Filtering](https://mlsanity.com/paper/1602.01585/) (2016)
- [Transformer Memory as a Differentiable Search Index](https://mlsanity.com/paper/2202.06991/) (2022)
- [Beam Search Strategies for Neural Machine Translation](https://mlsanity.com/paper/1702.01806/) (2017)
- [OneRec: Unifying Retrieve and Rank with Generative Recommender and Iterative Preference Alignment](https://mlsanity.com/paper/2502.18965/) (2025)
- [Synchromesh: Reliable code generation from pre-trained language models](https://mlsanity.com/paper/2201.11227/) (2022)
- [NeuroLogic A*esque Decoding: Constrained Text Generation with Lookahead Heuristics](https://mlsanity.com/paper/2112.08726/) (2021)
- [Towards Universal Sequence Representation Learning for Recommender Systems](https://mlsanity.com/paper/2206.05941/) (2022)
- [On the Theoretical Limitations of Embedding-Based Retrieval](https://mlsanity.com/paper/2508.21038/) (2025)
- [Automata-based constraints for language model decoding](https://mlsanity.com/paper/2407.08103/) (2024)
- [OneRec-Think: In-Text Reasoning for Generative Recommendation](https://mlsanity.com/paper/2510.11639/) (2025)
- [PLUM: Adapting Pre-trained Language Models for Industrial-scale Generative Recommendations](https://mlsanity.com/paper/2510.07784/) (2025)
- [OneRec Technical Report](https://mlsanity.com/paper/2506.13695/) (2025)
- [Unifying Generative and Dense Retrieval for Sequential Recommendation](https://mlsanity.com/paper/2411.18814/) (2024)
- [Generative Recommendation with Semantic IDs: A Practitioner's Handbook](https://mlsanity.com/paper/2507.22224/) (2025)
- [Mathematical Foundations of the GraphBLAS](https://mlsanity.com/paper/1606.05790/) (2016)
- [Eliminating Out-of-Domain Recommendations in LLM-based Recommender Systems: A Unified View](https://mlsanity.com/paper/2505.03336/) (2025)
- [Efficient and Asymptotically Unbiased Constrained Decoding for Large Language Models](https://mlsanity.com/paper/2504.09135/) (2025)

## Built on by (most cited first, 6 of 6)

- [LLMs Need Encoders for Semantic IDs Too](https://mlsanity.com/paper/2606.00324/) (2026)
- [UniPinRec: Unifying Generative Retrieval and Ranking at Pinterest Scale](https://mlsanity.com/paper/2606.00422/) (2026)
- [LatticeBridge: Rare-Event Sequential Inference for Faithful Structured Sequence Synthesis](https://mlsanity.com/paper/2606.11203/) (2026)
- [History-Conditioned Joint-Prefix Alignment for Generative Recommendation](https://mlsanity.com/paper/2608.29179/) (2026)
- [Trie Automata for Constrained Decoding over Large Finite Sets](https://mlsanity.com/paper/2608.12574/) (2026)
- [SIDScope: A Diagnostic Resource for Semantic-ID Interfaces in Generative Recommendation](https://mlsanity.com/paper/2608.18779/) (2026)

## Abstract (arXiv)

Generative retrieval has emerged as a powerful paradigm for LLM-based recommendation. However, industrial recommender systems often benefit from restricting the output space to a constrained subset of items based on business logic (e.g. enforcing content freshness or product category), which standard autoregressive decoding cannot natively support. Moreover, existing constrained decoding methods that make use of prefix trees (Tries) incur severe latency penalties on hardware accelerators (TPUs/GPUs). In this work, we introduce STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding), an efficient and scalable constrained decoding technique designed specifically for high-throughput LLM-based generative retrieval on TPUs/GPUs. By flattening the prefix tree into a static Compressed Sparse Row (CSR) matrix, we transform irregular tree traversals into fully vectorized sparse matrix operations, unlocking massive efficiency gains on hardware accelerators. We deploy STATIC on a large-scale industrial video recommendation platform serving billions of users. STATIC produces significant product metric impact with minimal latency overhead (0.033 ms per step and 0.25% of inference time), achieving a 948x speedup over a CPU trie implementation and a 47-1033x speedup over a hardware-accelerated binary-search baseline. Furthermore, the runtime overhead of STATIC remains extremely low across a wide range of practical configurations. To the best of our knowledge, STATIC enables the first production-scale deployment of strictly constrained generative retrieval. In addition, evaluation on academic benchmarks demonstrates that STATIC can considerably improve cold-start performance for generative retrieval. Our code is available at https://github.com/youtube/static-constraint-decoding.
