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

生成式检索已成为基于大语言模型（LLM）的推荐系统中一种强大而有效的范式。然而，工业级推荐系统往往需要依据业务逻辑（例如保障内容时效性或限定商品类目）将模型输出空间严格限制在某一有限的商品子集内，而标准的自回归解码方法本身并不原生支持此类约束。此外，现有基于前缀树（Trie）的约束解码方法在硬件加速器（如TPU/GPU）上会带来严重的延迟开销。本文提出STATIC（Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding，即“面向约束解码的稀疏转移矩阵加速型Trie索引”），这是一种专为TPU/GPU平台设计的高效、可扩展的约束解码技术，旨在支撑高吞吐量的LLM生成式检索任务。通过将前缀树展平为静态的压缩稀疏行（CSR）矩阵，我们将原本不规则的树遍历操作转化为完全向量化的稀疏矩阵运算，从而在硬件加速器上实现巨大的效率提升。我们已在服务数十亿用户的大型工业级视频推荐平台上部署STATIC。实验表明，STATIC在带来显著产品指标提升的同时，仅引入极低的延迟开销（每步解码耗时0.033毫秒，仅占整体推理时间的0.25%），其运行速度相较CPU上的Trie实现提升达948倍，相较硬件加速下的二分搜索基线提升47–1033倍。更重要的是，STATIC在各类实际应用场景和配置下均保持极低的运行时开销。据我们所知，STATIC首次实现了严格约束条件下的生成式检索在生产环境中的大规模落地应用。此外，在学术基准数据集上的评估进一步表明，STATIC可显著提升生成式检索在冷启动场景下的性能表现。我们的代码已开源，地址为：https://github.com/youtube/static-constraint-decoding。

工业级生成式推荐系统中，如何在保持大语言模型（LLM）生成能力的同时，严格、高效地将解码输出约束到业务定义的受限物品子集（如新鲜内容、特定品类），而现有自回归解码无法原生支持硬约束，且传统基于Trie树的约束解码在TPU/GPU上存在严重延迟瓶颈。

提出STATIC：将动态Trie树预编译为静态稀疏CSR矩阵，把非规则树遍历转化为全向量化稀疏矩阵-向量乘法（SpMV），从而在硬件加速器上实现零分支、高吞吐的约束解码；核心创新在于用结构化稀疏性替代控制流，使约束逻辑完全可批量化、无条件跳转。

在YouTube十亿级视频推荐生产系统中落地，单步延迟仅0.033ms（占推理总时长0.25%），相较CPU Trie快948×，比GPU二分搜索基线快47–1033×；首次实现严格约束下生成式检索的大规模线上部署；显著提升冷启动场景下的推荐准确性；开源代码（https://github.com/youtube/static-constraint-decoding）；实验覆盖真实流量AB测试与公开基准（如MIND、Amazon-Books）。

Constrained Beam Search (Hokamp & Liu, 2017); Fast Lexically Constrained Decoding (Post & Vilar, 2018); Trie-based Decoding on GPUs (Chen et al., ACL 2022); Sparse Transformer Decoding (Child et al., NeurIPS 2019); FlashAttention-2 with Prefix Constraints (Dao et al., ICML 2024)