生成式推荐(GR)系统正成为大规模个性化推荐的强大新方向。GR不再将推荐视为一系列孤立的检索、排序和预测阶段,而是将推荐重新定义为用户行为的序列建模问题。用户的交互、上下文、候选物品和行为都被转化为高基数事件流中的Token,模型学习从该序列中生成或预测下一个相关物品。

这种方法对现代推荐工作负载尤为适用,因为用户历史记录往往很长,物品目录持续变化,个性化质量依赖于对丰富序列行为的建模。但这也带来了服务层面的挑战:尽管存在长历史记录、大型嵌入表和序列密集型模型架构,GR模型仍需要实现低延迟推理。

NVIDIA Dynamo-Triton(原NVIDIA Triton Inference Server)现已通过NVIDIA recsys-examples代码库支持端到端的分层序列转导单元(HSTU)GR推理工作流。该工作流结合了HSTU、PyTorch预先编译(Ahead-of-Time Inductor)、基于FlexKV的KV缓存、原生C++验证、NV嵌入缓存以及Dynamo-Triton部署。

最终结果是一条具有出色延迟性能的HSTU排序模型服务路径。在NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU上,动态批次大小为8时,相比未使用KV缓存的同款AOTI配置,Dynamo-Triton搭配PyTorch AOTI在GPU KV缓存命中率达到100%的情况下,三层HSTU模型最高提速4.47倍,八层模型最高提速5.93倍。

本文将介绍如何借助NVIDIA Dynamo-Triton、PyTorch AOTI和FlexKV,将HSTU生成式推荐系统从PyTorch开发阶段推进至生产推理阶段。您将学习如何导出并预先编译模型,在Python和原生C++环境中验证最终的部署产物,并通过Dynamo-Triton进行服务部署,而无需为单独的运行时重写模型。

本文还探讨了基于GPU的KV缓存如何减少重复计算,并给出了基准测试结果,证明该部署工作流可将延迟降低最高达5.93倍,突显了该方案的实际性能优势。

为何使用HSTU进行生成式推荐

HSTU是为处理高基数、非平稳事件流的GR工作负载而提出的。在传统推荐系统中,检索和排序通常由一系列专门的模型和特征流水线构建而成。而GR将推荐建模为序列预测问题,使模型能够在一个具备序列感知能力的架构中,统一推理用户上下文、物品历史、行为历史和候选物品。

在NVIDIA HSTU排序示例中,模型输入由分类Token构建而成。上下文Token代表用户侧信息,物品Token代表具体物品,可选的行为Token则代表用户与这些物品的交互行为。

HSTU预处理路径会检索嵌入向量,在存在行为Token时将物品嵌入和行为嵌入交错排列,附加上下文信息,并应用位置编码。随后HSTU模块处理该序列,预测头输出多任务排序结果。

这种结构非常适合需要考虑时近性、顺序和重复交互模式的推荐系统。但这也意味着,当每个请求都要反复处理长历史序列时,推理成本会变得高昂。生产系统需要在保留HSTU建模优势的同时,减少服务过程中的冗余计算。

为何大型序列推荐系统的服务具有挑战性

服务大型序列推荐系统与服务小型稠密排序模型截然不同。服务架构必须处理不规则序列输入、庞大的分类嵌入状态、长历史记录,以及同一用户可能反复请求、每次只携带少量新信息的请求模式。每次请求都重新计算用户历史的完整键值状态会浪费计算资源并增加延迟。

这正是KV缓存变得重要的原因。KV缓存存储先前序列计算中可复用的键值数据,使模型避免重新计算用户历史中已缓存的部分。对于推荐系统推理而言,当用户的长期历史大体保持稳定、仅有新的候选物品或近期行为到来时,这种方式尤为有用。

NVIDIA HSTU推理工作流包含一个KVCacheManager,利用GPU内存和主机存储来缓存KV数据。GPU缓存以分页KV数据表的形式组织,支持查找、分配、追加和淘汰操作。当GPU缓存空间受限时,系统会按照类似LRU的策略淘汰较旧的用户数据。主机侧存储提供了另一层缓存KV数据的空间,该工作流还包含基于FlexKV的KV缓存运行时后端。

HSTU注意力内核可以从分页缓存中读取KV数据,导出的推理路径包含支持缓存感知的自定义操作,用于查找、分配、加载、追加和卸载。这使得服务路径能够在减少冗余计算的同时,保留模型的序列语义。

用于原生推理的PyTorch AOTI

PyTorch AOTI(预先编译Inductor)工作流从一个PyTorch模型开始,使用torch.export和PyTorch AOTI对其进行导出。AOTI会将模型提前编译为一个可由原生C++运行时加载的软件包。这降低了Python运行时开销,为Dynamo-Triton PyTorch AOTI后端提供了便于部署的产物。

导出的模型包包含AOTI模型存档以及元数据和嵌入表文件。在NVIDIA示例中,嵌入实现结合了DynamicEmb推理嵌入表和NV嵌入缓存,后者仅将常用嵌入存储在GPU内存中,同时将整个表保留在CPU内存中,从而降低GPU内存占用。导出路径会在编译后的.pt2存档旁写入层元数据和嵌入表数据,这样模型加载时就无需产生不必要的重复嵌入表副本。

该工作流通过多种方式验证同一导出产物。Python导出脚本生成软件包并重放张量数据。原生C++可执行文件加载并重放导出模型,以验证正确性和性能。Dynamo-Triton部署随后使用同一个AOTI包和重放路径,这有助于保持开发验证与生产服务的一致性。

Dynamo-Triton部署路径

Dynamo-Triton为导出的HSTU模型提供了生产级服务层。AOTI部署使用Dynamo-Triton PyTorch后端,平台设置为"torch_aoti"。这使得Dynamo-Triton能够加载并服务预先编译的PyTorch模型包。

完整工作流包含以下五个阶段:

构建所需的自定义算子和运行时库

使用PyTorch AOTI导出HSTU排序模型

启动基于FlexKV的KV缓存服务

通过原生C++重放验证导出的产物

使用Dynamo-Triton服务导出的KV缓存AOTI模型

这种方法之所以重要,是因为推荐系统的服务不仅仅需要一个快速的模型内核。Dynamo-Triton带来了模型仓库管理、请求处理、后端集成、指标监控和部署结构。AOTI带来了低开销的编译模型产物。NV嵌入缓存通过仅在GPU内存中保留嵌入表的热点部分,降低了GPU内存需求。FlexKV通过缓存注意力模块,减少了长用户历史的重复计算。这些组件共同构成了一个面向实际生产环境的生成式推荐推理服务架构。

HSTU服务延迟基准测试

本基准测试比较了不同Dynamo-Triton后端、模型规模、批次大小和KV缓存状态下的HSTU服务延迟。

测试目标是量化生产级HSTU服务架构的性能优势。测试比较了Dynamo-Triton PyTorch AOTI后端与Python后端,衡量了GPU KV缓存带来的额外延迟降低幅度,并评估了这些优势在不同模型深度和批次大小下的扩展表现。最终,这些结果向开发者展示了从未缓存的基于Python的推理,迁移到编译型、具备缓存感知能力的HSTU部署方案(基于Dynamo-Triton)后可获得的性能提升。

recsys-examples中的基准测试结果使用了单GPU上的KuaiRand-1K排序配置。模型结构包括三层和八层HSTU变体,隐藏层大小为512,4个注意力头,BF16模型权重,BF16 KV缓存,最大历史序列长度共8,192个Token(4,096组物品与行为对)的历史流,最大候选序列长度为100,以及6个上下文特征。对齐前的有效序列长度为8,298个Token,导出时的最大对齐序列长度为8,320。

基准测试协议按每个逻辑请求报告延迟。对于Dynamo-Triton AOTI基准测试,每次Dynamo-Triton调用包含一个逻辑批次,每个逻辑请求的延迟通过将端到端的处理时间除以Dynamo-Triton调用次数再除以逻辑批次大小计算得出。数据集加载、验证、重新分批、用户ID生成、服务器启动、预热以及预热后的休眠时间均未计入测量时间。

用于Dynamo-Triton后端对比和批次大小结果测试的硬件为NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU。

Dynamo-Triton后端对比

在Dynamo-Triton批次大小为2时,即便没有KV缓存命中,PyTorch AOTI相比Dynamo-Triton Python后端也能改善延迟表现。而在20GB GPU KV缓存命中的情况下,延迟改善幅度明显更大。

这些结果展示了两种不同的增益来源。首先,AOTI相比Python后端降低了服务开销。其次,KV缓存命中通过复用已缓存的序列状态减少了模型的计算量。在更深的八层模型上,这种缓存带来的收益尤为明显,因为避免重复计算所带来的影响更大。

AOTI与KV缓存下的批次大小扩展性

按批次大小划分的PyTorch AOTI后端结果显示,随着逻辑批次大小的增加,KV缓存的效果愈发显著。

在批次大小为8时,三层HSTU模型在GPU KV缓存命中的情况下,每个逻辑请求的延迟达到0.423毫秒。八层HSTU模型的延迟则达到0.678毫秒。对于长序列排序推理而言,这是相当出色的结果,展示了编译模型执行与缓存感知服务相结合所带来的价值。

加速HSTU GR推理有何好处

推荐系统在严格的延迟预算下运行。额外的排序延迟会影响页面加载时间、信息流响应速度和广告投放的时效要求。与此同时,日益具备序列感知能力和个性化能力的模型往往需要更多的推理算力,尤其是随着用户历史记录的增长。

HSTU服务工作流结合了多项互补技术来应对这一挑战。HSTU提供了生成式推荐架构,PyTorch AOTInductor生成预先编译的部署产物,基于FlexKV的KV缓存使之前计算的注意力状态得以复用,NVIDIA Dynamo-Triton则提供了生产级服务环境。

当连续请求共享用户交互历史中未发生变化的前缀部分时,这种方法尤为有价值。模型无需重新计算该部分序列的注意力,而是可以复用已缓存的键值状态,只需针对新追加的Token进行计算。随着序列变长、模型加深,这种节省效果会进一步放大,因为跨多个HSTU层的重复计算原本会带来显著的延迟增加。

开始加速HSTU GR推理

您可以通过NVIDIA/recsys-examples GitHub代码库复现并扩展此工作流。HSTU概述介绍了GR模型结构,包括上下文Token、物品Token、行为Token、嵌入表、HSTU模块和预测头。

AOTI推理指南详细介绍了构建所需镜像和库、准备KuaiRand-1K数据、训练检查点、导出KV缓存AOTI模型、使用C++重放进行验证、打包Dynamo-Triton运行时镜像,以及通过Dynamo-Triton服务器重放请求的完整流程。

要了解更多信息,请查阅以下相关资源:

HSTU生成式推荐系统概述

Dynamo-Triton、FlexKV与PyTorch AOTI集成方案

HSTU推理基准测试

Dynamo-Triton HSTU文档

NV嵌入缓存

致谢

本文是NVIDIA多个团队跨职能协作的成果。我们要感谢J、Runchu Zhao、Yulu Liu、Lin Hu、Zhuofan Li、Jacob Subag和Tomer Bar-On的贡献。

Q&A

Q1:什么是HSTU生成式推荐系统?

A:HSTU(分层序列转导单元)是一种用于生成式推荐的模型架构,它将推荐问题重新定义为序列建模任务,能够统一处理用户上下文、物品历史和行为历史,而不是依赖传统的检索、排序等孤立阶段。

Q2:使用Dynamo-Triton部署HSTU模型能带来多大的性能提升?

A:在NVIDIA RTX PRO 6000 Blackwell GPU上,动态批次大小为8时,相比未使用KV缓存的配置,三层HSTU模型最高提速4.47倍,八层HSTU模型最高提速5.93倍,GPU KV缓存命中率达到100%。

Q3:KV缓存在HSTU推理中起什么作用?

A:KV缓存存储先前序列计算中可复用的键值数据,使模型避免对用户历史中未变化的部分重新计算,仅需处理新追加的Token,从而显著降低长序列推理的延迟和计算成本。

NVIDIA