Vetorização da Trie: Decodificação Restrita Eficiente para Recuperação Generativa Baseada em LLM em Aceleradores
Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators
February 26, 2026
Autores: 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
Resumo
A recuperação generativa emergiu como um paradigma poderoso para a recomendação baseada em LLM. No entanto, os sistemas de recomendação industriais frequentemente beneficiam da restrição do espaço de saída para um subconjunto limitado de itens com base na lógica de negócio (por exemplo, impondo frescura de conteúdo ou categoria de produto), o que a decodagem autoregressiva padrão não suporta nativamente. Além disso, os métodos de decodagem restrita existentes que utilizam árvores de prefixos (Tries) incorrem em severas penalidades de latência em aceleradores de hardware (TPUs/GPUs). Neste trabalho, introduzimos o STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding), uma técnica de decodagem restrita eficiente e escalável, projetada especificamente para recuperação generativa baseada em LLM de alta vazão em TPUs/GPUs. Ao achatar a árvore de prefixos numa matriz estática no formato Compressed Sparse Row (CSR), transformamos travessias irregulares da árvore em operações totalmente vetorizadas com matrizes esparsas, desbloqueando ganhos massivos de eficiência em aceleradores de hardware. Implantamos o STATIC numa plataforma industrial de recomendação de vídeos em larga escala que atende a bilhões de utilizadores. O STATIC produz um impacto significativo nas métricas do produto com sobrecarga de latência mínima (0,033 ms por passo e 0,25% do tempo de inferência), alcançando uma aceleração de 948x em relação a uma implementação de trie em CPU e uma aceleração de 47-1033x em relação a uma linha de base com busca binária acelerada por hardware. Adicionalmente, a sobrecarga de tempo de execução do STATIC permanece extremamente baixa numa ampla gama de configurações práticas. Até onde sabemos, o STATIC permite a primeira implantação em escala de produção de recuperação generativa estritamente restrita. Além disso, a avaliação em benchmarks académicos demonstra que o STATIC pode melhorar consideravelmente o desempenho em cenários de *cold-start* para recuperação generativa. O nosso código está disponível em 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.