向量化Trie：面向加速器的高效约束解码生成式检索 Vectorizing the Trie: Efficient Constrained Decoding for LLM-based Generative Retrieval on Accelerators
STATIC方法将前缀树转换为CSR矩阵，实现硬件加速的约束解码，延迟极低。
前置知识
生成式检索（Generative Retrieval）
生成式检索是一种推荐系统范式，与传统基于双编码器的最近邻搜索不同，它使用Transformer模型自回归地解码目标物品的语义ID。模型直接预测候选物品的标识符序列，而不是在嵌入空间中搜索最近邻。这种方法将物品语料库内化于模型参数中，绕过了大规模嵌入和近似最近邻搜索的复杂性，能够捕捉更深层的语义关系。
本文针对生成式检索的约束解码问题提出解决方案，理解生成式检索的基本原理是理解本文技术挑战的基础。
语义ID（Semantic ID）
语义ID是物品的离散、基于token的标识符，通过RQ-VAE（残差量化变分自编码器）等方法生成。RQ-VAE首先将物品特征编码为潜在向量z，然后通过L级迭代量化生成离散token序列。每一级使用特定的码本，残差定义为$r_1 := z$，在第d级将残差$r_d$与最近的嵌入$e_{y_d}$匹配，索引$y_d$成为语义ID的第d个码字。这种递归方法使语义相似的物品共享前缀token。
语义ID是生成式检索的核心表示，约束解码需要在语义ID空间上进行，理解其结构对理解约束解码机制至关重要。
约束解码（Constrained Decoding）
约束解码是在解码过程中通过掩码无效token来限制输出空间的技术。在标准beam search中，模型可能生成不存在或不符合业务逻辑的物品标识符。约束解码通过在每个解码步骤中，将无效token的对数概率设置为$-\infty$，确保beam search只选择有效序列。这通常通过前缀树（trie）实现，其中约束函数$F_t(y_{<t}, y_t)$判断将当前token$y_t$添加到前缀$y_{<t}$是否形成有效路径。
本文的核心贡献是改进约束解码的实现方式，使其能在硬件加速器上高效运行，理解约束解码的基本概念对理解本文技术方案至关重要。
前缀树（Trie）
前缀树是一种树形数据结构，用于存储和检索字符串集合。在约束解码中，前缀树表示所有有效的语义ID序列，每个节点代表一个token，从根节点到叶子节点的路径对应一个完整的语义ID。标准实现依赖指针追踪，这在CPU上高效，但在硬件加速器上会导致随机内存访问模式，无法利用高带宽内存（HBM）的突发传输能力。
本文的核心创新是将前缀树转换为稀疏矩阵表示，理解前缀树的传统实现及其局限性是理解本文技术贡献的前提。
硬件加速器（TPU/GPU）
现代硬件加速器如TPU和GPU针对高吞吐量并行计算优化，但要求静态计算图以进行编译优化。XLA（加速线性代数）编译器要求张量在编译时具有精确的静态形状。这些加速器依赖高带宽内存（HBM）和突发传输能力，随机内存访问模式会导致严重的内存延迟和缓存抖动。此外，数据依赖的控制流与静态计算图范式不兼容，需要重新设计算法以适应硬件特性。
本文的核心目标是实现硬件原生的约束解码，理解硬件加速器的特性和约束对理解本文的设计选择至关重要。
研究动机
在工业推荐系统中，基于LLM的生成式检索模型需要将输出空间限制为基于业务逻辑的约束子集，例如强制内容新鲜度（如“最近7天内上传”）、地域性（如地区推荐）、产品类别（如“推荐夏季服装”）或库存可用性（如“仅库存商品”）。然而，现有的生成式检索模型缺乏对输出空间的控制。标准自回归解码无法原生支持这些约束，模型可能会自信地生成不存在、过期或法律限制的物品标识符。依赖后生成过滤在计算上是浪费的，模型可能将整个推理预算用于生成无效物品，导致零有效推荐。前缀树（trie）是实现约束解码的标准算法解决方案，但指针追踪的trie实现在硬件加速器（TPU/GPU）上存在严重延迟问题。在初步实验中，CPU卸载的trie实现使推理时间增加2倍，无法满足每解码步骤≤10ms的目标延迟。更复杂的方法如二分搜索也存在显著延迟惩罚。有限状态转换器（FST）在应用于生产规模的语义ID词汇（如数百万物品）时面临状态爆炸问题，且缺乏原生TPU支持。
本文的目标是本文旨在实现一种高效且可扩展的约束解码技术，专门为基于LLM的生成式检索在TPU/GPU上的高吞吐量设计。具体目标包括：将约束解码的每步延迟开销控制在极低水平（实际达到0.033ms，仅占推理时间的0.25%）；实现O(1)的I/O复杂度（相对于约束集大小）；支持大规模生产部署（如YouTube平台服务数十亿用户）；在约束解码的同时改善冷启动推荐性能。这些目标旨在解决现有约束解码方法在硬件加速器上的效率瓶颈，使严格的输出空间约束能够在不牺牲服务延迟的情况下大规模实施。
与已有工作不同的是，本文的独特切入角度在于将约束解码从图遍历问题重新表述为一系列向量化稀疏矩阵操作。现有工作如DISC-PPV虽然将约束解码适配到硬件加速器，但依赖二分搜索引入了O(log|C|)的I/O复杂度。本文观察到，在超大规模词汇（如2000万物品）中，即使对数级扩展也成为I/O瓶颈。关键洞察是：前缀树遍历本质上是稀疏矩阵查找，通过将前缀树展平为静态压缩稀疏行（CSR）矩阵，可以将不规则树遍历转换为完全向量化的稀疏矩阵操作，从而在硬件加速器上实现巨大的效率提升。这种方法将I/O复杂度降低到O(1)，通过单次向量化VNTK内核实现，并消除了CPU-GPU通信开销。此外，本文首次展示了如何通过约束解码来改善冷启动推荐性能，这是现有生成式检索工作忽略的重要场景。
核心方法
STATIC方法的核心思想可以用一个简单的类比来理解：将动态的树遍历转换为静态的矩阵查找。想象一下，传统前缀树遍历就像在一个巨大的迷宫中逐个房间探索，需要不断回头检查路标；而STATIC方法则是预先绘制出整个迷宫的地图（CSR矩阵），然后直接查找从当前位置可以到达的房间。技术路线分为三个阶段：首先离线构建阶段，将约束集（如2000万新鲜视频）转换为前缀树，然后展平为CSR稀疏矩阵表示；其次在线解码阶段，在每个解码步骤中，通过向量化节点转移内核（VNTK）执行单次内存访问，获取所有可能的下一个token及其对应的节点转移；最后约束应用阶段，使用动态切片和掩码算术，将无效token的对数概率设置为$-\infty$，确保beam search只选择有效序列。整个设计完全在设备上执行，消除了主机-设备往返，支持XLA/Inductor编译，实现真正的硬件原生约束解码。
STATIC方法的核心创新在于将前缀树遍历的指针追踪问题转换为稀疏矩阵的向量化查找问题。与已有方法的本质区别在于：1) I/O复杂度：现有PPV方法具有O(log|C|)的I/O复杂度（二分搜索），而STATIC通过CSR矩阵实现O(1)的I/O复杂度，单次内存访问获取所有工作集；2) 内存访问模式：前缀树的指针追踪导致随机内存访问，无法利用HBM突发传输能力，而CSR矩阵的行指针和列索引是连续存储的，支持内存合并访问；3) 编译兼容性：动态控制流与XLA静态计算图不兼容，而STATIC的向量化操作可以完全展开为静态计算图，支持端到端ML编译；4) 执行模式：GPU的warp发散问题被消除，所有beam同时处理固定数量的候选（最大分支因子），保持SIMT并行性。这种转换使得约束解码从“CPU友好”转变为“加速器原生”，实现了数量级的速度提升。
方法步骤详情
STATIC方法包含以下步骤：1) 离线CSR矩阵构建：给定约束集C（如2000万物品），每个物品的语义ID长度为L（如8个token），词汇大小为|V|（如2048）。构建前缀树并映射每个唯一前缀节点到状态整数$s \in [S]$。定义稀疏矩阵$T \in Z^{S \times |V|}$，其中$T_{s,v} = s_{next}$如果存在转移$s \xrightarrow{v} s_{next}$，否则为0。CSR表示包括行指针P、列索引C和值数组V。2) 加速器原生解码算法：维护转移状态向量$n_t \in Z^{B \times M}$跟踪每个beam的当前节点索引。对于前$d$层（通常$d=2$）使用密集张量掩码D进行快速查找，对于更深层使用CSR稀疏矩阵。3) 向量化节点转移内核（VNTK）：对于每个解码步骤$t$，执行边界查找获取当前节点的子节点数量$N_{child}$，通过动态切片获取固定数量$B_t$的候选（无论实际子节点数量如何），使用Range生成的掩码进行清理，最后通过Scatter操作生成密集布尔掩码$m$。4) 约束应用：使用Where操作将掩码$m$应用到对数概率上，无效token的概率设置为$-\infty$。5) Beam search优化：选择前M个最可能的有效序列，更新token缓冲区、累积分数和节点状态。
技术新颖性
STATIC方法的技术新颖性体现在多个方面：首先，它是第一个将前缀树遍历重新表述为稀疏矩阵操作用于约束解码的工作，桥接了经典数据结构（Double-Array Trie）和现代深度学习编译器（XLA/Inductor）之间的差距。其次，CSR矩阵表示实现了O(1)的I/O复杂度，通过单次内存访问获取所有工作集，这与现有二分搜索方法的O(log|C|)复杂度形成鲜明对比。第三，分支自由的VNTK内核设计确保整个解码步骤保持单一静态计算图，通过动态切片和掩码算术处理可变数量的子节点，避免了动态控制流。第四，密集层优化策略在浅层使用密集张量掩码（$d \leq 2$），在深层使用稀疏CSR矩阵，平衡了内存使用和计算效率。第五，堆叠CSR内存布局将列索引和数据值交错存储在单个张量中，减少随机内存访问次数。这些创新共同实现了在TPU/GPU上几乎无开销的约束解码。
实验结果
STATIC方法在YouTube大规模生产部署和学术基准测试中均取得了显著成果。在YouTube部署中，约束词汇包含约2000万最近7天上传的高质量视频，模型为基于Gemini的30亿参数生成式检索模型，批量大小为2（每芯片），beam大小为70。系统效率分析显示，STATIC的每步延迟开销仅为+0.033ms，仅占推理时间的0.25%。相比之下，CPU Trie方法的延迟为31.3ms（比STATIC高948倍），PPV Exact方法为34.1ms（高1033倍），PPV Approximate方法为1.56ms（高47倍）。内存使用方面，STATIC在2000万约束词汇下仅需约1.5GB HBM，约为每100万约束项90MB。在线A/B测试中，STATIC实现了100%约束合规性，7天新鲜视频观看量提升5.1%（95%置信区间[5.0%, 5.2%]），3天新鲜视频观看量提升2.9%（[2.8%, 3.0%]），点击率（CTR）提升0.15%（[0.01%, 0.29%]），战略用户群体满意度提升0.15%（[0.03%, 0.27%]）。在Amazon Reviews数据集的冷启动实验中，STATIC在Beauty子集上，2%冷启动设置下Recall@1为4.29%（对比无约束为0.00%，随机约束为0.42%），5%冷启动设置下为1.60%（对比无约束为0.00%，随机约束为0.17%）。在Sports & Outdoors子集上，2%冷启动为1.24%，5%冷启动为1.17%。在Toys & Games子集上，2%冷启动为4.39%，5%冷启动为2.25%。扩展性分析显示，STATIC的延迟随约束集大小|C|从10^5到10^8几乎保持恒定（约0.023-0.039ms），随语义ID词汇大小|V|从256到32768也几乎恒定（约0.034-0.044ms），而其他方法呈现显著增长。最大分支因子的扩展性测试显示，STATIC呈现渐近线性O(B)扩展。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 约束解码延迟开销 | 每步延迟（ms） | 0.033 | CPU Trie: 31.3, PPV Exact: 34.1, PPV Approximate: 1.56, Hash Bitmap: 12.3 | 比CPU Trie快948倍，比PPV Exact快1033倍，比PPV Approximate快47倍 |
| 在线A/B测试 | 7天新鲜视频观看量提升 | +5.1% | 无约束生成式检索 | 95%置信区间[5.0%, 5.2%] |
| 在线A/B测试 | 点击率提升 | +0.15% | 无约束生成式检索 | 95%置信区间[0.01%, 0.29%] |
| 冷启动推荐（Beauty子集） | Recall@1（2%冷启动） | 4.29% | 无约束: 0.00%, 随机约束: 0.42% | 比无约束提升4.29个百分点，比随机约束提升3.87个百分点 |
| 约束集扩展性 | 延迟随|C|变化 | 0.023-0.039ms（|C|从10^5到10^8） | CPU Trie: 27-31ms, PPV Exact: 6-38ms, PPV Approximate: 0.8-1.7ms | 几乎恒定延迟，其他方法随|C|显著增长 |
局限与改进
STATIC方法存在几个局限性：首先，稀疏矩阵构建是离线过程，无法支持约束集的动态更新。当库存变化频繁时（如商品上下架），需要完全重新构建CSR矩阵，这在实时环境中可能不切实际。其次，内存占用随约束集规模线性增长，虽然每100万约束项仅需约90MB，但对于超大规模约束集（如10亿物品）可能达到GB级别。第三，密集层优化参数$d$需要仔细选择：$d$过大会导致内存爆炸（$|V|^d$指数增长），$d$过小则浅层稀疏查找效率降低。第四，虽然STATIC在延迟上表现优异，但CSR矩阵本身的构建和存储仍需要额外的离线计算资源。第五，方法假设语义ID长度固定且为$L$，对于变长标识符需要额外适配。最后，论文主要评估在YouTube和Amazon数据集上，对于其他领域（如NLP、语音识别）的适用性需要进一步验证。
独立分析的弱点
STATIC方法存在几个可改进的弱点：首先，离线构建限制了实时更新能力。在库存频繁变化的场景中（如电商实时上下架），约束集可能需要每分钟更新，而当前方法需要完全重新构建CSR矩阵。改进方向包括开发增量更新机制，只修改受影响的矩阵行，或者使用动态稀疏数据结构。其次，内存占用虽然可控，但线性增长限制了超大规模应用。对于10亿物品的约束集，内存需求可能达到数十GB。改进方向包括分层存储策略，将热点约束放在快速内存，冷约束放在慢速内存，或者使用压缩稀疏表示。第三，当前方法假设所有约束具有相同的权重，但在实际业务中，不同约束可能有不同的重要性（如库存约束比新鲜度约束更严格）。改进方向包括引入加权约束解码，允许不同的掩码强度。第四，虽然VNTK内核实现了向量化，但最大分支因子$B_t$的选择影响性能：过小可能截断有效候选，过大则增加计算开销。改进方向包括自适应分支因子选择，根据实际数据分布动态调整。第五，密集层优化参数$d$需要手动调优，不同数据集可能需要不同的$d$值。改进方向包括自动选择最优$d$，基于内存预算和计算约束进行优化。
未来方向
基于STATIC框架，未来研究可以朝多个方向发展：首先，开发动态稀疏更新机制，支持约束集的实时变化而无需完全重新构建矩阵。这可以通过增量更新CSR矩阵的特定行实现，或者使用可逆的稀疏表示。其次，探索分层存储策略，当约束集规模达到数十亿时，可以将约束分布在多个设备上，使用All-Gather操作收集约束掩码，但这需要重新设计VNTK内核以支持分布式执行。第三，将约束解码扩展到其他生成任务，如代码生成、对话系统、文档摘要等，其中输出空间存在语法或语义约束。第四，结合强化学习，将约束解码作为动作空间约束，优化推荐策略的同时满足业务逻辑。第五，研究约束解码与模型训练的交互，探索是否可以在训练时注入约束知识，减少推理时的约束解码需求。最后，开发自动化的约束管理工具，允许业务人员以声明式方式定义约束（如“库存>0且上传时间<7天”），自动转换为CSR矩阵表示。
复现评估
STATIC方法的复现性评估如下：代码已开源在GitHub（https://github.com/youtube/static-constraint-decoding），使用JAX和Flax框架实现，支持TPU和GPU。数据方面，YouTube部署使用专有数据集，但Amazon Reviews数据集是公开的，可用于复现冷启动实验。算力需求方面，实验在Google TPU v6e上进行，但代码设计具有跨平台可移植性，可在PyTorch/Inductor上运行。复现难度中等：需要理解稀疏矩阵表示和XLA编译约束，但论文提供了详细的算法描述和代码示例。关键配置包括语义ID长度$L=8$，词汇大小$|V|=2048$，密集层参数$d=2$，beam大小$M=70$。对于学术研究，可以在较小的约束集（如1万物品）上验证方法有效性，然后扩展到更大规模。需要的主要资源包括：支持XLA的硬件（TPU或GPU），JAX环境，以及足够的内存存储CSR矩阵（每100万约束项约90MB）。论文还提供了详细的硬件优化讨论（附录A）和内存分析（附录B），有助于理解和复现关键实现细节。
论文图表
