Vettorizzazione del Trie: Decodifica Vincolata Efficiente per il Recupero Generativo Basato su LLM su Acceleratori

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

February 26, 2026

Autori: 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

cs.AI

Abstract

La retrieval generativa è emersa come un paradigma potente per la raccomandazione basata su LLM. Tuttavia, i sistemi di raccomandazione industriali traggono spesso vantaggio dalla limitazione dello spazio di output a un sottoinsieme ristretto di elementi basato sulla logica di business (ad esempio, imponendo la freschezza dei contenuti o la categoria del prodotto), operazione che la decodifica autoregressiva standard non supporta nativamente. Inoltre, i metodi di decodifica vincolata esistenti che utilizzano alberi dei prefissi (Trie) comportano severe penalità di latenza sugli acceleratori hardware (TPU/GPU). In questo lavoro, introduciamo STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding), una tecnica di decodifica vincolata efficiente e scalabile, progettata specificamente per la retrieval generativa ad alto throughput basata su LLM su TPU/GPU. Appiattendo l'albero dei prefissi in una matrice statica Compressed Sparse Row (CSR), trasformiamo le traversate irregolari dell'albero in operazioni completamente vettorizzate su matrici sparse, sbloccando notevoli guadagni di efficienza sugli acceleratori hardware. Abbiamo implementato STATIC su una piattaforma di raccomandazione video industriale su larga scala che serve miliardi di utenti. STATIC produce un impatto significativo sulle metriche di prodotto con un overhead di latenza minimo (0,033 ms per passo e lo 0,25% del tempo di inferenza), ottenendo un'accelerazione di 948x rispetto a un'implementazione Trie su CPU e un'accelerazione di 47-1033x rispetto a una baseline con ricerca binaria accelerata via hardware. Inoltre, l'overhead runtime di STATIC rimane estremamente basso in un'ampia gamma di configurazioni pratiche. Per quanto a nostra conoscenza, STATIC abilita la prima implementazione su scala produttiva della retrieval generativa rigorosamente vincolata. In aggiunta, la valutazione su benchmark accademici dimostra che STATIC può migliorare considerevolmente le prestazioni in cold-start per la retrieval generativa. Il nostro codice è disponibile all'indirizzo https://github.com/youtube/static-constraint-decoding.

English

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.