Векторизация префиксного дерева: эффективное ограниченное декодирование для генеративного поиска на основе больших языковых моделей с использованием акселераторов

Аннотация

Порождающее извлечение стало мощной парадигмой для рекомендательных систем на основе больших языковых моделей. Однако промышленные рекомендательные системы часто выигрывают от ограничения пространства выбора узким подмножеством элементов на основе бизнес-логики (например, обеспечение актуальности контента или категории продукта), что стандартное авторегрессионное декодирование изначально не поддерживает. Более того, существующие методы ограниченного декодирования, использующие префиксные деревья (боры), приводят к значительным задержкам на аппаратных акселераторах (TPU/GPU). В данной работе мы представляем STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding) — эффективную и масштабируемую технику ограниченного декодирования, специально разработанную для высокопроизводительного порождающего извлечения на основе LLM на TPU/GPU. Путем преобразования префиксного дерева в статическую разреженную матрицу в формате CSR мы трансформируем нерегулярные обходы дерева в полностью векторизованные операции с разреженными матрицами, достигая значительного прироста эффективности на аппаратных акселераторах. Мы развернули STATIC в крупномасштабной промышленной платформе рекомендаций видео, обслуживающей миллиарды пользователей. STATIC демонстрирует значительное влияние на продуктовые метрики при минимальных накладных расходах по задержке (0.033 мс на шаг и 0.25% времени вывода), достигая ускорения в 948 раз по сравнению с CPU-реализацией бора и в 47-1033 раза по сравнению с базовым методом на основе бинарного поиска на аппаратных ускорителях. Более того, runtime-накладные расходы STATIC остаются крайне низкими в широком диапазоне практических конфигураций. Насколько нам известно, STATIC позволяет осуществить первое промышленное внедрение строго ограниченного порождающего извлечения. Кроме того, оценка на академических бенчмарках показывает, что STATIC может значительно улучшить производительность порождающего извлечения в условиях холодного старта. Наш код доступен по адресу 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.

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

Аннотация

Support