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

Recommendation systems have evolved from embedding-based retrieval with ANN search (e.g., ScaNN) toward LLM-based generative retrieval. In TIGER, SEATER, and LIGER, items are represented as Semantic IDs, often produced via RQ-VAE, and Transformers autoregressively generate the target ID token by token. This removes the need for a separate nearest-neighbor index and allows pre-trained LLMs to be adapted to recommendation tasks. Industrial systems such as PLUM and OneRec further established the viability of this paradigm. However, the open-ended nature of generation creates a validity gap: the model may produce out-of-stock, stale, or policy-violating items unless decoding is constrained. Prior constrained decoding methods like NeuroLogic, Synchromesh, FSTs, and DISC-PPV address parts of the problem, but not the accelerator-efficiency bottleneck at scale.

The paper targets a precise systems problem: enforce strict output constraints during autoregressive retrieval without sacrificing accelerator throughput. The valid item set C may contain millions of Semantic IDs, and at each decoding step the model must mask invalid next tokens based on the current prefix. A naive trie requires pointer chasing and recursive branching, which produce random memory access, poor cache behavior, and CPU-offload overhead. On TPU/GPU systems, this is particularly harmful because the hardware prefers static graphs, contiguous reads, and compiler-friendly control flow. The challenge is therefore both algorithmic and architectural: the constraint check must be exact, fast, and fully on-device.

1) Trie-to-CSR vectorization: the prefix tree is converted offline into a static sparse transition matrix, so legality checking becomes sparse table lookup rather than dynamic traversal.

2) Branch-free decoding: dynamic slicing and mask arithmetic replace data-dependent branching, keeping the graph compatible with XLA/Inductor.

3) Accelerator-native execution: the method avoids host-device round trips and uses coalesced reads to exploit HBM/SRAM movement efficiently.

4) Strict rather than approximate control: unlike filtering after generation, STATIC prevents invalid tokens before they are sampled, ensuring the beam search explores only valid prefixes.

5) I/O-focused design: relative to PPV's O(log|C|) binary search, STATIC achieves effectively O(1) I/O with respect to the constraint set size.|

• �� Problem setup: Let V be the semantic-token vocabulary, L the fixed Semantic ID length, and C⊂V^L the allowed candidate set. The model generates y=(y_1,...,y_L) autoregressively.
• �� Validity rule: Define F_t(y_{<t},y_t)=I(∃c∈C s.t. (y_{<t},y_t)⊑c). If F_t=0, the next-token probability is forced to zero (implemented as −∞ log-probability).
• �� State mapping: Each unique trie prefix node is mapped to a discrete state s∈[S], where S is the number of trie nodes.
• �� Matrix construction: Build T∈Z^{S×|V|} with T_{s,v}=s_next if token v transitions from state s to the next state, and 0 otherwise (sink state).
• �� CSR encoding: Store row pointers P, column indices C, and values Vals. P marks row boundaries, column indices store valid token IDs, and values store the destination state IDs.
• �� Decoding loop: At each step, use the current beam state to slice the corresponding CSR row, obtain all legal next tokens, mask out illegal logits, and continue standard beam search.
• �� System property: Because the structure is static and array-based, the entire operation can be compiled and executed on-device, avoiding CPU pointer chasing and enabling high-throughput inference.

The paper evaluates STATIC in three settings. First, an industrial video recommendation deployment at YouTube-scale, where a generative retrieval model is constrained to a preset vocabulary of 20 million fresh items. Second, a runtime comparison against a CPU trie implementation and a hardware-accelerated binary-search baseline (DISC-PPV/PPV). Third, academic cold-start evaluation on Amazon Reviews. The reported metrics emphasize per-step latency, inference-time overhead, and speedup, with the production case serving as the main systems benchmark. The narrative suggests that the method was also tested across a range of practical configuration sizes to assess scalability.

The headline result is production viability: STATIC adds only 0.033 ms per step and accounts for just 0.25% of inference time, which is negligible for an online recommendation stack. Compared with a CPU trie, it reaches a 948× speedup, demonstrating that the main bottleneck was not the constraint itself but the pointer-based representation. Compared with the accelerator-aware PPV binary-search baseline, the 47–1033× speedup shows that even on-chip O(log|C|) verification becomes expensive at million-scale item sets. The paper also reports improved cold-start recommendation on Amazon Reviews, indicating that strict candidate restriction can improve generalization rather than merely enforce policy.

The most immediate use case is business-rule-constrained recommendation: freshness filters, region-specific catalogs, inventory-aware shopping, and category-limited feeds. A single large generative model can be reused across products by changing only the constraint set C. That reduces model proliferation and simplifies deployment. Beyond recommender systems, the approach may support any structured decoding workload on TPU/GPU where the valid output space can be represented as a trie, especially when exactness and low latency both matter.

STATIC is optimized for prefix-based constraints and fixed-length Semantic IDs, so it is not a universal solution for arbitrary logical predicates or long-range dependencies. Its offline CSR construction assumes the constraint set is sufficiently stable; if the item inventory changes frequently, the matrix must be rebuilt and redeployed, adding operational overhead. Finally, while the paper reports strong speedups and a successful production deployment, the publicly visible text leaves some benchmark details, hyperparameters, and ablation specifics unspecified, which makes independent reproduction less transparent.