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

Zhengyang Su; Isay Katsman; Yueqi Wang; Ruining He; Lukasz Heldt; Raghunandan Keshavan; Shao-Chuan Wang; Xinyang Yi; Mingyan Gao; Onkar Dalal; Lichan Hong; Ed Chi; Ningren Han

Source-linked AI summary

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

Zhengyang Su, Isay Katsman, Yueqi Wang, Ruining He, Lukasz Heldt, Raghunandan Keshavan, Shao-Chuan Wang, Xinyang Yi, Mingyan Gao, Onkar Dalal, Lichan Hong, Ed Chi, Ningren Han

arXiv:2602.22647v2Feb 26, 2026cs.IRcs.CLcs.LG

TL;DR

Generative retrieval needs inference-time control over its output space, but trie-based constraints incur accelerator-hostile latency. STATIC flattens trie constraints into CSR sparse operations and achieves large speedups, production deployment, and improved cold-start performance.

Problem
Generative retrieval lacks native output-space control for business rules, while trie constraints create irregular memory access and compilation problems on accelerators.
Method
STATIC flattens prefix-tree constraints into static CSR matrices and applies vectorized, accelerator-native decoding constraints during inference.
Results
47–1033× speedup over alternative on-device methods accompanies production-scale strict constraints and considerable cold-start improvements.
Takeaways & Limitations
STATIC demonstrates that strict constrained generative retrieval can operate at scale without compromising serving latency.
Takeaways & Limitations
Sparse transition-matrix construction is offline, so real-time inventory changes require future dynamic sparse updates rather than full model recompilation.

Abstract

from arXiv · show

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.

1 Introduction

Generative retrieval needs inference-time output control to enforce business rules, but trie-based constraints are poorly suited to accelerator hardware. STATIC reformulates trie traversal as vectorized sparse operations for scalable constrained decoding, with industrial and cold-start benefits.

Motivation: Generative retrieval lacks native control over outputs needed for freshness, locality, category, and inventory constraints.Without intervention, models can generate stale, unavailable, or legally restricted items, while post-generation filtering may leave no valid recommendations.
Motivation: Trie constraints mask invalid tokens during decoding, but pointer chasing causes irregular memory access and compilation incompatibility on TPUs and GPUs.Random accesses prevent coalescing and prefetching, while data-dependent branching conflicts with static computation graphs.
STATIC: STATIC flattens prefix-tree constraints into static CSR matrices and uses branch-free decoding with dynamic slicing and mask arithmetic.The design enables coalesced reads, vectorized sparse operations, and accelerator-native execution without host-device round-trips.
STATIC: O(1) I/O complexity with respect to constraint-set size contrasts with logarithmic scaling for existing binary-search methods.The stated I/O measure counts costly data transfers between off-chip HBM and on-chip SRAM on TPUs.
Deployment and evaluation: STATIC was deployed on YouTube for a 20-million-item fresh-vocabulary setting, improving key online metrics with minimal latency overhead.The system’s latency remains extremely low across varied constraint-set and Semantic ID vocabulary sizes.
Deployment and evaluation: Constrained decoding with STATIC improves cold-start recommendation performance on Amazon Reviews and enables accelerator-compatible constrained generation at scale.The paper positions this work at the intersection of generative recommendation, constrained token generation, and hardware-aware optimization.

3 Background

Generative retrieval predicts Semantic IDs token by token, while constrained decoding restricts these sequences to valid item prefixes. STATIC illustrates how restricted Semantic-ID vocabularies form tries that can be used to mask invalid beam-search continuations.

Generative retrieval: Generative retrieval directly predicts candidate identifiers token by token, with each Semantic ID capturing an item’s semantic properties.This replaces nearest-neighbor search in an embedding space with sequence generation.
Semantic IDs: RQ-VAE constructs Semantic IDs by iteratively quantizing item-feature residuals across levels, producing a tuple of codebook indices.The residual update is r_d+1 := r_d − e_y_d, and semantically similar items can share prefix tokens.
Decoding: During inference, beam search autoregressively decodes fixed-length Semantic IDs while tracking token prefixes and cumulative log-probability scores.At each step, the highest-scoring M candidates are retained and the rest are pruned.
Constrained decoding: Constrained decoding restricts generation to a vocabulary of desirable Semantic IDs and masks invalid tokens by assigning them log-probability −∞.This ensures beam search selects only valid sequences from the restricted set.
Constrained decoding: A restricted Semantic-ID vocabulary forms a prefix tree whose valid prefixes determine which tokens may be selected at each decoding step.Figure 1 connects prefix-tree construction with its transition matrix, CSR representation, and decoding-time vocabulary constraint.

4 Methodology

STATIC replaces dynamic trie traversal with a static CSR transition matrix and accelerator-compatible sparse operations for constrained Semantic ID decoding. Its vectorized state transitions combine sparse CSR lookups with limited dense prefix lookups for efficient inference.

STATIC replaces CPU-based pointer-chasing prefix-tree traversal with a static sparse matrix lookup compatible with TPU/GPU execution and XLA/Inductor compilation.The approach recasts constrained decoding as vectorized sparse matrix operations.
Sparse Transition Matrix Conversion: Each unique trie prefix becomes an integer state in a sparse transition matrix T whose entries encode valid token-triggered transitions.The matrix has one row per prefix state and one column per semantic token.
Sparse Transition Matrix Conversion: CSR stores row pointers, valid token IDs, and target state IDs, enabling hardware-optimized sparse operations instead of dynamic control flow.The construction is performed offline and is memory efficient for highly sparse constraint sets.
Accelerator-Native Decoding: A beam-wise transition-state vector tracks each beam’s current trie node while vectorized lookups identify valid semantic extensions.The state vector has dimensions B × M for batch size B and beam count M.
Accelerator-Native Decoding: Dense prefix masks are used for early layers, but exponential |V|^d growth makes d≤2 practical in most real-world cases; deeper layers use the sparse CSR matrix.The dense tensor is a one-time construction cost, whereas CSR handles deeper levels without dense enumeration.

1 Phase 1: Log-Space Projection

STATIC’s transition kernel converts variable-width trie expansion into fixed-shape, branch-free accelerator operations. It uses dynamic slicing, masking, projection, and beam-state gathering to preserve valid transitions within a static computation graph.

Vectorized Kernel Operations: STATIC maintains next-node states for each beam and applies the resulting dense boolean mask directly to model log-probabilities.This enforces the prefix constraint during the main decoding step.
Vectorized Kernel Operations: The decoding kernel uses vectorized primitives including DenseLookup, VNTK, Scatter, LogSoftmax, and Gather to process constraints and update beams.VNTK performs sparse transition lookup, while Gather selects the states and scores of top-performing beams.
Hardware Alignment: Branch-free execution addresses accelerator bottlenecks caused by dynamic branching on TPUs and warp divergence from mismatched child counts on GPUs.The design targets portability across both TPU and GPU hardware.
Hardware Alignment: The kernel slices B_t entries at each level regardless of actual child count, then masks and sanitizes invalid entries with Range and Where operations.Fixed-length slicing keeps the computation graph static and arithmetic units saturated.

5 Large-scale Deployment on YouTube

The YouTube-scale evaluation measures STATIC’s industrial deployment efficiency using TPU latency and throughput comparisons. The evaluation includes its vectorized transition kernel and assesses system efficiency against baselines.

System Efficiency: The YouTube-scale evaluation tests STATIC along industrial deployment dimensions including system efficiency, latency, and throughput.The stated system-efficiency comparison uses TPU accelerators and baselines.
Vectorized Node Transition Kernel: The evaluation includes the Vectorized Node Transition Kernel as the algorithmic component used for accelerator-native constrained decoding.The kernel takes token cardinality, maximum branch factors, CSR column indices, and values as inputs and produces next nodes and a logit mask.

4 Phase 2: Speculative Slicing

STATIC provides low-latency constrained decoding for large Semantic ID vocabularies while maintaining manageable memory use and favorable scaling. In production, it enforces freshness constraints and improves fresh-content and user-experience metrics.

System Efficiency Analysis: +0.033 ms per decoding step is STATIC’s latency overhead, representing 0.25% of inference time.The comparison uses means over 100 trials on a 3 billion parameter model.
System Efficiency Analysis: 948× speedup over CPU Trie and 1033× over PPV Exact demonstrate STATIC’s latency advantage.STATIC also outperforms PPV Approximate by 47×; PPV uses binary search with logarithmic I/O scaling in constraint-set size.
Scalability: Approximately 90 MB of memory per 1 million restricted vocabulary items supports a 20-million-item deployment with about 1.8 GB maximal HBM usage.The practical memory requirement is usually at most 75% of this upper limit.
Scalability: STATIC maintains low latency across constraint-set sizes and exhibits almost constant latency across tested SID vocabulary sizes.Figure 2 varies |C| from 10^5 to 10^8 with |V| = 2048, while Figure 3 varies |V| from 256 to 32k with |C| = 10^7.
Online A/B Testing: 100% compliance with the Last 7 Days constraint enabled a 5.1% increase in 7-day fresh video views and a 2.9% increase in 3-day fresh video views.The constrained model strictly produced valid 7-day fresh videos, unlike the unconstrained model’s frequent older-video generations.
Online A/B Testing: +0.15% CTR and +0.15% user satisfaction were observed in the online experiment.The user-satisfaction result concerns a strategic user segment in the Home Feed setting.

6 Cold-Start Retrieval on Amazon Reviews Datasets

The Amazon Reviews experiments test whether constraining generative retrieval to cold-start item sets can improve retrieval performance. STATIC substantially outperforms unconstrained decoding, random guessing, and dense retrieval in both tested cold-start settings.

Evaluation setup: Cold-start items are defined using review age, with the newest 2% and 5% isolated into separate evaluation sets.Training sequences containing these items are removed, and test sequences target cold-start items.
Baselines and setup: The study compares unconstrained beam search, constrained random guessing, dense retrieval, and STATIC.The experiments use separate Amazon subdatasets with Semantic IDs generated by RQ-VAE models.
Evaluation setup: STATIC restricts the pretrained generative retrieval model to the cold-start item set during every decoding step.The approach uses the Transformer trained without cold-start training sequences.
Results: STATIC improves Recall@1 considerably over all three comparison methods in both the 2% and 5% cold-start settings.Table 3 reports Recall@1 percentages for cold-start test sequences.

7 Conclusion

The paper concludes that STATIC makes strictly constrained generative retrieval practical on accelerator hardware and at production scale. It also reports improved cold-start performance from constrained decoding alone, while identifying dynamic sparse updates as future work.

Conclusion: STATIC transforms pointer-chasing trie lookups into vectorized sparse matrix operations for accelerator-compatible constrained decoding.The framework bridges prefix-tree constraints and vector-based TPU/GPU hardware.
Conclusion: 47–1033× speedups over alternative on-device methods support strict constraints at scale without compromising serving latency.The system was demonstrated in production environments serving billions of users.
Conclusion: Constrained decoding alone improves cold-start performance on the Amazon Reviews datasets.The conclusion presents this as evidence for the viability of constrained decoding in cold-start recommendation.
Future work: Dynamic sparse updates remain future work because sparse transition matrix construction is currently offline.Real-time inventory changes would otherwise require avoiding full model recompilation.

Appendix

The appendix details STATIC’s accelerator implementation, memory representation, and portability. It replaces dynamic trie control flow with fixed-shape gathers and masks, while discussing hybrid dense/sparse processing and compatibility with variable-length settings.

Hardware implementation: XLA’s static-shape requirement makes ordinary dynamic trie traversal incompatible with compiled accelerator execution.Variable child counts and pointer chasing can trigger compilation errors.
Hardware implementation: STATIC processes the maximum branch factor B_ℓ at each level using fixed-size vectorized gathers executed for every beam in parallel.This avoids variable-length child processing in the decoding kernel.
Hardware implementation: A validity mask converts dynamic loops over children into static data flow, enabling one XLA graph with loop unrolling and pipelining.Nodes with fewer than B_ℓ children receive masked invalid entries.
Memory layout: STATIC stacks CSR column indices and values so each transition fetches its token ID and next-node pointer in one coalesced memory transaction.The layout reduces the number of random memory accesses compared with separate arrays.
Hybrid representation: The implementation uses dense boolean masks for the first d levels and sparse traversal thereafter to handle high early branching factors.This trades limited static memory for higher throughput in expensive initial layers.
Scope and portability: The CSR transition matrix and VNTK support variable-length sequences and variable-size codebooks, and the gather strategy ports to PyTorch/CUDA.The reported evaluation nevertheless focuses on fixed-length SIDs with a fixed-size codebook.

B.2 Calculation for YouTube

The YouTube memory analysis estimates STATIC’s HBM requirements from dense early-level masks and CSR storage for later trie levels. For 20 million constraints, the stated upper bound is approximately 1.5 GB, while practical usage is lower because Semantic IDs collide in prefix space.

YouTube configuration: The YouTube configuration fixes |V| = 2048, L = 8, K_1 = 12, K_2 = 4, d = 2, and |C| = 20 · 10^6.These values determine the dense-mask and constraint-phase memory calculations.
Memory calculation: The dense mask for levels 1 and 2 contributes approximately 17.3 MB.This phase stores the mask and state IDs for |V|^d prefixes.
Memory calculation: The later constraint phase contributes 1.44 GB for six levels at 20 million constraints and 12 bytes per item-level contribution.The calculation covers levels 3 through 8.
Memory calculation: The resulting maximum per-chip usage is approximately 1.46 GB, validating the approximately 1.5 GB HBM estimate for 20 million constraints.The actual requirement is often lower because clustered Semantic IDs create prefix collisions.
Practical usage: In practice, STATIC uses at most 75% of the upper-bound memory estimate for most restricted vocabularies.The stated reduction is attributed to non-uniform item distributions and prefix-space collisions.
Capacity planning: The production rule of thumb is approximately 90 MB of additional HBM per million constraints, with an observed average of approximately 73 MB per million in the 20-million-item corpus.The 90 MB figure is presented as a reliable linear upper-bound approximation.
Serving architecture: Replicating the approximately 1.5 GB transition matrix on every device avoids cross-chip communication during constraint checks.Each chip computes validity masks for its local beam batch, preserving linear serving scalability.
Boundary: Scaling item corpora toward billions may require hierarchical storage or sharding to prevent device-memory saturation.This is identified as future research rather than a demonstrated production capability.

D Detailed Latency Analysis

This section defines how per-step constraint overhead is measured and evaluates STATIC’s hardware scaling as branching factor grows. It also specifies the cold-start evaluation protocol and dense-retrieval comparison baseline.

Latency measurement: Per-step latency is measured as additional constraint-enforcement time relative to unconstrained decoding across 100 trials and 8 decoding steps.The analysis reports means and standard deviations for the full decoding process.
Hardware scaling: STATIC’s masking kernel exhibits asymptotically linear O(B) scaling with maximum branching factor.The benchmark varies the vocabulary and branching factor while fixing the constraint set at |C| = 106.
Cold-start evaluation: Cold-start experiments report peak Recall@1 over 70 training epochs for autoregressive Semantic ID decoding and dense retrieval.The protocol uses the Amazon cold-start split and selects the peak result across training.
Dense-retrieval baseline: The dense-retrieval baseline uses Gemma 1B and an MLP to map user sequences and T5 item embeddings into 128-dimensional vectors.Constrained recommendations are generated with MIPS over cold-start item vectors.

G STATIC Code: JAX Implementation

The JAX implementation represents constraints with a CSR transition matrix and performs each decoding step through JIT-compiled, vectorized sparse operations. The implementation supports both TPUs and GPUs with nearly identical latency overhead.

Portability: The example implementation uses JAX and Flax and is designed to run on TPU and GPU with almost identical latency overhead.The code defines CSR arrays for row pointers and transition data, with optional dense optimizations.
Decoding step: The hardware-accelerated decoding step converts logits to log probabilities and uses a CSR transition matrix to extract valid sparse candidates.The implementation is exposed as a JAX-jitted decoding function with static vocabulary, branching-factor, and step arguments.
Sparse transition kernel: The Vectorized Node Transition Kernel gathers contiguous transition entries using CSR row pointers, offsets, and validity masks.Valid candidate log probabilities are selected from the vocabulary logits, while invalid positions receive NEG_INF.
Masking and state updates: The implementation handles initial decoding steps with optional packed start and dense level-one masks before applying sparse candidate extraction.Sparse candidates are scattered into a vocabulary-sized masked-logit array, with next-node states returned alongside logits.