STATIC:面向 TPU/GPU 上 LLM 生成式检索的高效约束解码
Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators
STATIC 将前缀树展平为静态 CSR 稀疏矩阵,把不规则树遍历转为全向量化稀疏矩阵运算,实现 TPU/GPU 上的高效约束解码。它已部署在服务数十亿用户的大规模工业视频推荐平台,每步仅 0.033 ms、占推理时间 0.25%,相比 CPU trie 提速 948 倍,相比硬件加速二分查找基线提速 47-1033 倍。据称这是严格约束生成式检索的首次生产级部署,代码已开源。
Authors: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
Abstract: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 this https URL.
| Comments: | KDD 2026 camera-ready |
| Subjects: | Information Retrieval (cs.IR); Computation and Language (cs.CL); Machine Learning (cs.LG) |
| Cite as: | arXiv:2602.22647 [cs.IR] |
| (or arXiv:2602.22647v3 [cs.IR] for this version) | |
| https://doi.org/10.48550/arXiv.2602.22647 arXiv-issued DOI via DataCite |
|
| Journal reference: | Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V.2 (KDD '26), 2026, pp. 8009-8020 |
| Related DOI: | https://doi.org/10.1145/3770855.3818506
DOI(s) linking to related resources |
Submission history
From: Zhengyang Su [view email]
[v1]
Thu, 26 Feb 2026 06:00:56 UTC (266 KB)
[v2]
Mon, 20 Jul 2026 21:24:06 UTC (269 KB)
[v3]
Wed, 7 Oct 2026 03:10:09 UTC (269 KB)
来源:arXiv cs.CL · arxiv.org