💬 NLP

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

本文提出了名为 STATIC 的高效约束解码技术，通过将前缀树扁平化为静态压缩稀疏行矩阵，在 GPU/TPU 上实现了向量化的大规模生成式检索，显著降低了延迟并首次在生产环境中成功部署了严格受限的生成式推荐系统。

原作者： 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

发布于 2026-02-27

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： 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

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

这篇论文讲述了一个关于如何让 AI 在“做选择题”时既快又准的故事。

想象一下，你正在使用一个超级聪明的 AI 推荐系统（比如 YouTube 的推荐算法）。这个 AI 就像一个才华横溢但有点“天马行空”的作家。它的任务是为你写下一个“推荐列表”（比如推荐几个视频）。

1. 遇到的问题：AI 的“幻觉”与“死胡同”

传统的 AI 推荐系统像是在一个巨大的图书馆里找书，它先算出哪些书可能好，再挑出来。但现在的新一代 AI（基于大语言模型）更像是一个直接“写”出书名的作家。

问题一：AI 会“瞎编”。
如果让这位作家自由发挥，它可能会写出一些根本不存在的视频 ID，或者推荐一些已经下架、过期的视频。这就好比你让厨师做菜，他可能会端出一盘“空气炒蛋”或者“过期的剩菜”。
问题二：业务规则太复杂。
公司要求：“只能推荐最近 7 天上传的视频”或者“只能推荐‘夏季服装’类别”。
以前的做法是：让 AI 随便写，写完后，再由一个人工检查员（CPU）拿着清单去核对。如果写错了，就扔掉重写。
- 后果：这太慢了！就像让一个赛车手跑完一圈，再停下来问裁判“你跑对了吗？”，如果错了，还得倒回去重跑。在每秒要处理成千上万个请求的互联网大平台上，这种“先跑后查”的方法会让系统卡死。

2. 旧方法的困境：在迷宫里乱撞

为了解决这个问题，工程师们通常用一种叫“前缀树”（Trie）的数据结构。你可以把它想象成一个巨大的迷宫地图。

每走一步（生成一个词），AI 就要看地图，确认“这条路通不通”。
痛点：在传统的电脑（CPU）上，这个检查过程就像是在迷宫里到处乱跑、东张西望（指针跳转）。
在加速器（TPU/GPU）上的灾难：现在的 AI 芯片（如 Google 的 TPU）是超级流水线工厂，它们喜欢整齐划一、批量处理的任务。让它们在迷宫里“乱跑”检查，就像让一群整齐划一的机器人突然开始玩“捉迷藏”，有的机器人跑得快，有的跑得慢，有的还在找路，导致整个工厂停工等待。这会让速度变慢几百倍。

3. 解决方案：STATIC（把迷宫变成“传送带”）

这篇论文提出了一种叫 STATIC 的新方法。它的核心思想是：别在迷宫里跑了，直接把迷宫变成一张“传送带”！

核心比喻：从“找路”变成“查表”

旧方法（迷宫跑图）：
AI 每写一个字，就要去翻一本厚厚的字典，问：“我现在走到第 3 步，下一个字能选 A 吗？能选 B 吗？”这需要不停地翻书、跳转，非常慢。
STATIC 方法（向量化的传送带）：
作者把整个“迷宫地图”提前拍成了一张巨大的、稀疏的表格（矩阵）。
- 想象一下：这不再是让 AI 在迷宫里跑，而是把迷宫的所有可能路径，预先印在一张超级宽大的传送带上。
- 如何工作：
  1. ** flattening（压平）**：把复杂的树状结构压扁成一张二维表格。
  2. 向量操作（Vectorizing）：AI 芯片不再是一个字一个字地“查”，而是像切香肠一样，一次性把这一整段可能路径的数据“切”下来。
  3. 掩码（Masking）：如果某个路径是死胡同（比如推荐了过期的视频），就在表格上给这个位置贴个“禁止通行”的标签（把概率设为负无穷）。
  4. 结果：AI 芯片只需要做一次超级快速的批量读取，就能知道哪些路能走，哪些不能走。

为什么这很酷？

速度极快：因为不需要在迷宫里乱跑，而是直接“切”数据。论文说，这比旧方法快了 47 到 1033 倍！
不卡顿：它完美适配了现代 AI 芯片（TPU/GPU）的“流水线”特性，让芯片一直满负荷运转，没有等待时间。
内存小：虽然表格很大，但因为大部分路径是空的（稀疏的），所以它只占用了很少的内存（就像压缩文件一样）。

4. 实际效果：YouTube 的实战

作者在 YouTube 上测试了这个方法：

场景：限制 AI 只能推荐“最近 7 天内上传”的视频。
结果：
- 速度：每生成一个词，只增加了 0.033 毫秒 的延迟（几乎感觉不到）。
- 质量：用户看到的视频更新鲜了，点击率提升了。
- 对比：以前的方法（在 CPU 上查）会让系统慢 30 多毫秒，这在互联网时代简直是“世纪慢”。

5. 总结：给 AI 戴上了“智能导航”

简单来说，这篇论文做了一件非常聪明的事：
它没有让 AI 在复杂的规则迷宫里笨拙地乱跑，而是把规则提前画成了一张超级清晰的“导航地图”，并让 AI 芯片用最擅长的“批量扫描”方式来读取这张地图。

这就好比：

以前：让一个快递员在复杂的巷子里找路，每走一步都要停下来问路，最后累得半死还送得慢。
现在（STATIC）：给快递员发了一张自动导航的传送带，他只需要把包裹放上去，传送带自动把他送到正确的门口，既快又准，还能同时送几千个包裹。

这项技术让未来的 AI 推荐系统不仅能“写得快”，还能“守规矩”，并且不会变慢，是工业界大规模应用 AI 推荐的关键一步。

论文标题

Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators
（向量化 Trie：面向加速器上基于 LLM 的生成式检索的高效约束解码）

1. 研究背景与问题 (Problem)

背景：

生成式检索 (Generative Retrieval, GR) 正在成为推荐系统的新范式。它不再使用传统的“检索 + 排序”双塔架构，而是利用大语言模型 (LLM) 直接自回归地生成代表物品的“语义 ID (Semantic IDs, SIDs)"。
工业界需求： 在实际工业场景（如 YouTube 推荐）中，业务逻辑通常要求限制输出空间。例如，只推荐“最近 7 天上传”的视频、特定类别的商品或“有库存”的物品。

核心痛点：

标准解码的局限性： 标准的自回归解码无法原生支持这种约束。如果不加干预，LLM 可能会生成无效或过期的物品 ID（即“幻觉”），导致推荐无效。
现有约束解码的瓶颈：
- 前缀树 (Trie) 的指针跳跃 (Pointer Chasing)： 传统的约束解码使用 Trie 树在每一步屏蔽无效 Token。这种基于指针的数据结构会导致非连续的随机内存访问，无法利用现代加速器（TPU/GPU）的高带宽内存 (HBM) 特性，且破坏了硬件的预取机制。
- 编译不兼容： 现代加速器（如 Google TPU 依赖的 XLA）需要静态计算图。Trie 的动态控制流和递归分支与静态编译范式不兼容。
- 性能代价： 现有的基于 CPU 卸载的 Trie 实现会导致推理延迟增加 2 倍以上；基于二进制搜索的方法（如 PPV）虽然是在设备端，但其 I/O 复杂度为 $O(\log |C|)$ ，在大规模词表（如 2000 万物品）下仍会造成显著延迟。

2. 方法论 (Methodology)

作者提出了 STATIC (Sparse Transition Matrix-Accelerated Trie Index for Constrained Decoding)，一种专为 TPU/GPU 设计的高效约束解码框架。

核心思想：
将 Trie 树的图遍历问题转化为向量化的稀疏矩阵操作，从而消除指针跳跃，实现完全的设备端（On-device）执行。

关键技术组件：

静态压缩稀疏行 (CSR) 矩阵表示：
- 将 Trie 树“扁平化”为一个静态的 CSR 矩阵。
- Row Pointers (P): 存储当前状态允许转换的起始索引。
- Column Indices (C): 存储触发转换的有效 Token ID。
- Values (V): 存储目标状态节点 ID。
- 优势： 将不规则的树遍历转化为连续的内存读取（Coalesced Reads），极大提升了 HBM 利用率。
无分支解码算法 (Branch-free Decoding)：
- 动态切片 (Dynamic Slicing)： 针对每个解码步 $t$ ，预先计算该层最大分支因子 $B_t$ 。无论当前节点实际有多少子节点，算法总是提取固定长度 $B_t$ 的块。
- 掩码算术 (Mask Arithmetic)： 使用 Range 和 Where 操作生成布尔掩码，区分有效子节点和填充的无效子节点。
- 结果： 整个解码过程变为静态计算图，消除了 GPU/TPU 上的线程发散 (Warp Divergence) 和动态控制流开销。
混合架构 (Hybrid Approach)：
- 浅层 (Dense Mask)： 对于前 $d$ 层（通常 $d=2$ ），由于前缀密集，使用稠密布尔张量进行 $O(1)$ 查找。
- 深层 (Sparse Matrix)： 对于深层，使用上述的 CSR 稀疏矩阵进行查找。
- 这种设计平衡了内存占用和计算效率。
I/O 复杂度优化：
- 相比现有方法的 $O(\log |C|)$ ，STATIC 的 I/O 复杂度为 $O(1)$ ，因为它通过单次合并读取 (Coalesced Read) 获取所有工作集。

3. 主要贡献 (Key Contributions)

算法创新： 提出将 Trie 约束扁平化为静态 CSR 矩阵的方法，实现了 $O(1)$ 的内存访问开销。
系统优化： 设计了基于动态切片和掩码算术的无分支解码算法，完全适配加速器原生架构，消除了主机 - 设备往返 (Host-Device Round-trips)。
大规模部署： 在 YouTube 视频推荐平台上进行了部署，支持 2000 万物品的约束集，证明了其在生产环境中的可行性。
性能突破： 实现了显著的加速比，同时保持了极低的延迟。
冷启动提升： 证明了在 Amazon Reviews 数据集上，利用约束解码可以显著提升生成式检索在冷启动物品上的表现。

4. 实验结果 (Results)

实验设置：

平台： YouTube 视频推荐系统（TPU v6e）。
模型： 基于 Gemini 的 30 亿参数生成式检索模型。
约束集： 2000 万“最近 7 天上传”的高质量视频。

性能对比 (每解码步延迟)：

方法	延迟增加 (ms)	占推理时间比例	相比 STATIC 的倍数
STATIC (本文)	+0.033	0.25%	1x
PPV Approximate (近似)	+1.56	11.9%	47x 慢
PPV Exact (精确)	+34.1	260%	1033x 慢
CPU Trie (CPU 卸载)	+31.3	239%	948x 慢
Hash Bitmap (布隆过滤器)	+12.3	94.0%	372x 慢 (且有误报)

加速比： STATIC 比 CPU Trie 快 948 倍，比硬件加速的二进制搜索基线 (PPV Exact) 快 1033 倍。
扩展性：
- 随着约束集大小 ( $|C|$ ) 从 $10^5$ 增加到 $10^8$ ，STATIC 的延迟几乎保持恒定（ $O(1)$ ），而其他方法呈对数或线性增长。
- 随着语义 ID 词表大小 ( $|V|$ ) 增加，STATIC 的延迟也保持极低且稳定。
在线 A/B 测试 (YouTube)：
- 在“最近 7 天”新鲜度约束下，7 天新鲜视频观看量提升 5.1%，3 天新鲜视频观看量提升 2.9%。
- 点击率 (CTR) 提升 0.15%，用户满意度提升 0.15%。
- 实现了 100% 的约束合规性（无无效推荐）。

冷启动实验 (Amazon Reviews)：

在 Amazon Reviews 数据集上，针对从未在训练集中出现过的物品（冷启动），STATIC 的 Recall@1 指标显著优于无约束模型和随机猜测基线（例如在 Beauty 类别的 2% 冷启动设置下，Recall 从 0% 提升至 4.29%）。

5. 意义与影响 (Significance)

填补了工业界空白： 首次在生产规模（数十亿用户、千万级物品）上实现了严格约束的生成式检索。解决了生成式推荐模型在工业落地中最大的痛点——无法控制输出空间。
硬件亲和性设计： 证明了通过向量化稀疏矩阵操作，可以将原本“反硬件”的 Trie 树遍历转化为适合 TPU/GPU 的高效计算，为未来的 LLM 推理优化提供了新的范式。
业务价值： 使得单一模型可以灵活适应多种业务场景（如 freshness, locality, inventory），无需为每个场景训练独立模型，同时保证了推荐质量。
冷启动解决方案： 提供了一种无需重新训练模型即可改善冷启动物品推荐效果的有效手段。

总结：
STATIC 通过巧妙的数据结构转换（Trie -> CSR）和算法优化（无分支向量操作），成功解决了生成式检索在加速器上的延迟瓶颈，使得在大规模工业系统中实时执行复杂的业务约束成为可能，是生成式推荐系统走向成熟的关键一步。