将字典树向量化：面向加速器上基于 LLM 的生成式检索的高效约束解码

arXiv:2602.22647v3 公告类型：替换版本 摘要：生成式检索已成为基于 LLM 的推荐的一种强大范式。然而，工业推荐系统通常需要根据业务逻辑将输出空间限制为物品的受约束子集（例如，确保内容的新鲜度或产品类别），而标准的自回归解码原生不支持这一点。此外，现有利用前缀树（Trie）的约束解码方法在硬件加速器（TPU/GPU）上会产生严重的延迟开销。在本工作中，我们提出了 STATIC（Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding），一种专为 TPU/GPU 上高吞吐量的基于 LLM 的生成式检索而设计的高效且可扩展的约束解码技术。通过将前缀树展平为静态的压缩稀疏行（Compressed Sparse Row，CSR）矩阵，我们将不规则的树遍历转换为完全向量化的稀疏矩阵运算，从而释放硬件加速器的巨大效率优势。我们在一个服务数十亿用户的大规模工业视频推荐平台上部署了 STATIC。STATIC 以极小的延迟开销带来了显著的产品指标提升（每步 0.033 毫秒，占推理时间的 0.25%），相较于 CPU 字典树实现实现了 948 倍加速，相较于硬件加速的二分搜索基线实现了 47—1033 倍加速。此外，在广泛的实际配置范围内，STATIC 的运行时开销始终极低。据我们所知，STATIC 实现了首个生产规模的严格约束生成式检索部署。此外，学术基准测试评估表明，STATIC 能够显著改善生成式检索的冷启动性能。我们的代码可在 https://github.com/youtube/static-constraint-decoding 获取。