跳到正文
arXiv cs.CL· 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·· 10 天前AI 评分24

STATIC:面向 TPU/GPU 上 LLM 生成式检索的高效约束解码

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

AI 导读

STATIC 将前缀树展平为静态 CSR 稀疏矩阵,把不规则树遍历转为全向量化稀疏矩阵运算,实现 TPU/GPU 上的高效约束解码。它已部署在服务数十亿用户的大规模工业视频推荐平台,每步仅 0.033 ms、占推理时间 0.25%,相比 CPU trie 提速 948 倍,相比硬件加速二分查找基线提速 47-1033 倍。据称这是严格约束生成式检索的首次生产级部署,代码已开源。

正文 · AI 翻译

作者: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

查看 PDF HTML(实验性)

摘要:生成式检索已成为基于 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 URL 获取。
评论: KDD 2026 定稿版
主题: 信息检索(cs.IR);计算与语言(cs.CL);机器学习(cs.LG)
引用为: arXiv:2602.22647 [cs.IR]
  (或使用 arXiv:2602.22647v3 [cs.IR] 查看此版本)
  https://doi.org/10.48550/arXiv.2602.22647

由 DataCite 颁发的 arXiv DOI

期刊 参考: 第 32 届 ACM SIGKDD 知识发现与数据挖掘会议论文集 V.2(KDD '26),2026 年,第 8009-8020 页
相关 DOI: https://doi.org/10.1145/3770855.3818506

链接到相关资源的 DOI

提交历史

来自:Zhengyang Su [查看电子邮件]
[v1] 2026 年 2 月 26 日星期四 06:00:56 UTC(266 KB)
[v2] 2026 年 7 月 20 日星期一 21:24:06 UTC(269 KB)
[v3] 2026 年 10 月 7 日星期三 03:10:09 UTC(269 KB)

来源:arXiv cs.CL · arxiv.org