返回 文章 build CMS 文章

使用 NVIDIA Dynamo-Triton 部署 HSTU 生成式推荐器

用 Dynamo-Triton + PyTorch AOTI + FlexKV 把 HSTU 生成式推荐器部署到生产,长序列推理延迟最高降低 5.93 倍。

生成式推荐HSTUDynamo-TritonPyTorch AOTI
成长分 / 100 78 综合收获、行动、留存与影响

使用 NVIDIA Dynamo-Triton 部署 HSTU 生成式推荐器
为什么值得读了解生成式推荐(GR)如何将推荐重构为序列建模,以及 HSTU 架构的输入与处理流程。

掌握从 PyTorch 模型导出、AOTI 提前编译、原生 C++ 验证到 Dynamo-Triton 服务的完整部署路径。

关键洞察
  1. HSTU 将用户上下文、物品、动作和候选物品编码为 token,通过序列预测完成推荐,适合高基数、非平稳事件流。
  2. 服务大型序列推荐器的核心挑战是长历史重复计算,KV 缓存可复用已计算的键值状态,避免冗余注意力计算。
  3. PyTorch AOTI 通过 torch.export 和提前编译生成原生 C++ 可加载的 .pt2 包,降低 Python 运行时开销。
转成行动

深入阅读

正文与原文对照

原文保真覆盖:全文原文字符:16524

生成式推荐(GR)系统正作为一种强大的新方法,用于大规模个性化。GR 不再将推荐视为一组孤立的检索、排序和预测阶段,而是将推荐重新表述为对用户行为的序列建模。用户的交互、上下文、候选物品和动作成为高基数事件流中的 token,模型学习从该序列中生成或评分下一个相关物品。

这种方法对于现代推荐工作负载尤其有吸引力,因为用户历史可能很长,物品目录不断变化,个性化质量取决于对丰富序列行为的建模。但它也带来了服务挑战:尽管历史很长、嵌入表很大且模型架构以序列为主,GR 模型仍需要低延迟推理。

NVIDIA Dynamo-Triton(原 NVIDIA Triton Inference Server)现在通过 NVIDIA recsys-examples 仓库支持端到端的分层序列转导单元(HSTU)GR 推理工作流。该工作流结合了

HTSU、PyTorch Ahead-of-Time Inductor 编译、FlexKV 支持的 KV 缓存、原生 C++ 验证、NV 嵌入缓存以及

Dynamo-Triton 部署。

其结果是提供了一条以强延迟性能服务 HSTU 排序模型的实用路径。在 NVIDIA RTX PRO 6000 Blackwell Workstation Edition GPU 上,动态批大小为 8 时,在 100% GPU KV 缓存命中率下,使用 PyTorch AOTI 的 Dynamo-Triton 相对于未使用 KV 缓存的相同 AOTI 配置,三层 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 表示用户与这些物品的交互。

传统深度学习推荐模型与生成式 RecSys HSTU 工作流程的对比图。

*图 1. 传统深度学习推荐模型(左)与生成式 HSTU 工作流程(右)*HSTU 预处理路径会检索嵌入,在存在行动 token 时交错排列物品嵌入与行动嵌入,追加上下文信息,并应用位置编码。随后 HSTU 块处理该序列,预测头生成多任务排序输出。

这种结构非常适合那些时效性、顺序和重复交互模式至关重要的推荐系统。然而,这也意味着当每个请求都反复处理长历史序列时,推理可能变得昂贵。生产系统需要在保留 HSTU 建模优势的同时,减少服务期间的冗余计算。

为什么服务大型序列推荐器具有挑战性?

服务大型序列推荐器不同于服务小型稠密排序模型。服务栈必须处理参差不齐的序列输入、大型分类嵌入状态、长历史记录,以及同一用户可能仅携带少量新信息反复返回的请求模式。在每个请求上重新计算用户历史的完整键值状态会浪费工作并增加延迟。

这正是 KV 缓存变得重要的地方。KV 缓存存储来自先前序列计算的可复用键值数据,使模型能够避免重新计算用户历史中已缓存的部分。对于推荐器推理,当用户的长期历史基本保持稳定,而新的候选物品或近期行动到达时,这一点尤其有用。

展示 HSTU 服务中序列输入历史与候选物品,以及掩码注意力机制的示意图。

图 2. HSTU 服务NVIDIA HSTU 推理工作流程包含一个 KVCacheManager,它使用 GPU 内存和主机存储来缓存 KV 数据。GPU 缓存被组织为分页 KV 数据表,并支持查找、分配、追加和驱逐。当 GPU 缓存空间受限时,较旧的用户可以根据 LRU 风格策略被驱逐。主机侧存储为缓存的 KV 数据提供了另一层,并且该工作流程包含一个由 FlexKV 支持的 KV 缓存运行时后端。

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

用于原生推理的 PyTorch AOTI

PyTorch AOTI(Ahead-of-Time Inductor)工作流程从 PyTorch 模型开始,并使用 torch.export

和 PyTorch AOTI 将其导出。AOTI 将模型提前编译为一个可由原生 C++ 运行时加载的包。这减少了 Python 运行时开销,并为 Dynamo-Triton PyTorch AOTI 后端提供了便于部署的产物。

PyTorch AOTI 工作流程图,从模型导出到原生 C++ 运行时加载。

图 3. PyTorch AOTI 工作流程导出的模型包包含 AOTI 模型归档以及元数据和嵌入表文件。在 NVIDIA 示例中,嵌入实现结合了 DynamicEmb 推理嵌入表和 NV Embedding Cache,后者通过仅在 GPU 内存中存储热门嵌入,同时将整个表保留在 CPU 内存中,从而减少 GPU 内存使用。导出路径将层元数据和嵌入表数据与编译后的 .pt2 归档一起写入,以便模型加载时无需不必要的重复嵌入表副本。

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

Dynamo-Triton 部署路径

Dynamo-Triton 为导出的 HSTU 模型提供生产服务层。AOTI 部署使用 Dynamo-Triton PyTorch 后端,配置为 platform: "torch_aoti"。这允许 Dynamo-Triton 加载并提供提前编译的 PyTorch 模型包。

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

  • 构建所需的自定义算子和运行时库
  • 使用 PyTorch AOTI 导出 HSTU 排序模型
  • 启动由 FlexKV 支持的 KV 缓存服务
  • 使用原生 C++ 重放验证导出的产物
  • 使用 Dynamo-Triton 提供导出的 KV 缓存 AOTI 模型服务

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

HSTU GenRec 推理栈的架构图,包括 Dynamo-Triton、NV Embedding Cache 和 KV Cache Manager。

图 4. 带有 Dynamo-Triton 的 HSTU GR 推理栈## 基准测试 HSTU 服务延迟

此基准测试比较了跨 Dynamo-Triton 后端、模型大小、批大小和 KV 缓存状态的 HSTU 服务延迟。

目标是量化生产 HSTU 服务栈的性能优势。它将 Dynamo-Triton PyTorch AOTI 后端与 Python 后端进行比较,测量 GPU KV 缓存带来的额外延迟减少,并评估这些优势如何在模型深度和批大小上扩展。最终,它向开发者展示了从无缓存的基于 Python 的推理转向使用 Dynamo-Triton 的编译、缓存感知 HSTU 部署时可以预期的性能。

recsys-examples 中的基准测试结果在单块 GPU 上使用 KuaiRand-1K 排序配置。模型结构包括三层和八层 HSTU 变体、隐藏维度 512、四个注意力头、BF16 模型权重、BF16 KV 缓存、总长度为 8,192 的最大历史序列长度(4,096 个 item 加 action 对)历史流、最大候选序列长度 100,以及六个上下文特征。对齐前的有效序列长度为 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 后端也能改善延迟。在来自 20 GB GPU KV 缓存的命中的情况下,延迟改善显著更大。

柱状图,对比三层和八层 HSTU 模型在 Dynamo-Triton Python 后端、PyTorch AOTI 以及带 KV 缓存命中的 PyTorch AOTI 下的延迟。

图 5. 三层和八层 HSTU 的 Dynamo-Triton 后端对比这些结果显示出两种不同的收益。首先,与 Python 后端相比,AOTI 降低了服务开销。其次,KV 缓存命中通过复用缓存的序列状态减少了模型计算量。缓存收益在更深的八层模型上尤为明显,因为避免重新计算在那里影响更大。

使用 AOTI 和 KV 缓存的批次大小扩展

PyTorch AOTI 后端按批次大小的结果显示,随着逻辑批次大小增长,KV 缓存变得越来越有效。

折线图,展示三层 HSTU 在无缓存与 GPU KV 缓存命中延迟下的批次大小扩展。

图 6. 三层 HSTU,使用 Dynamo-Triton 批次大小扩展折线图,展示八层 HSTU 在无缓存与 GPU KV 缓存命中延迟下的批次大小扩展。

图 7. 八层 HSTU,使用 Dynamo-Triton 批次大小扩展在批次大小为 8 时,三层 HSTU 模型在 GPU KV 缓存命中下达到每个逻辑请求 0.423 ms 的延迟。八层 HSTU 模型达到 0.678 ms。对于长序列排序推理而言,这些是强劲的结果,并展示了将编译模型执行与缓存感知服务相结合的价值。

加速 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 的贡献。