@PyTorch: PyTorch AOTI 后端在 NVIDIA 的 HSTU 推理测试中,相较于 Python 后端,实现了 1.14 倍至 1.28 倍的加速,当……

X AI KOLs Following 新闻

摘要

NVIDIA 报告了在 HSTU 推理测试中使用 PyTorch AOTI 后端的显著加速,在理想情况下最高可达 2.38 倍,并强调了用于大规模生成式推荐系统的工具如 recsys-examples 和 nv-embedding-cache。

当使用 Triton Inference Server 部署时,PyTorch AOTI 后端在 NVIDIA 的 HSTU 推理测试中,相较于 Python 后端,实现了 1.14 倍至 1.28 倍的加速。 使用 PyTorch AOTI 后端和 KV 缓存,@nvidia 报告在理想的全 GPU 缓存命中场景中,加速了 2.20 倍至 2.38 倍。 这些结果来自 NVIDIA 的 recsys-examples 仓库,这是一个展示使用 PyTorch 在 NVIDIA GPU 上训练和部署生成式推荐器最佳实践的示例集合。它包括优化的 HSTU 和 Semantic ID 实现,涵盖训练和推理工作流。 对于 HSTU 推理,recsys-examples 支持 PyTorch AOTInductor 在 Torch C++ 运行时中执行模型。NVIDIA 开发者博客还介绍了 nv-embedding-cache,它提供 PyTorch 兼容的模块以加速大规模嵌入查找。 🔗 阅读完整文章:
查看原文
查看缓存全文

缓存时间: 2026/09/04 02:14

PyTorch AOTI 后端在 NVIDIA 的 HSTU 推理测试中,使用 Triton Inference Server 部署时,相比 Python 后端实现了 1.14 倍至 1.28 倍的加速。

结合 PyTorch AOTI 后端与 KV 缓存,@nvidia 报告在理想的全 GPU 缓存命中场景下实现了 2.20 倍至 2.38 倍的加速。

这些结果来自 NVIDIA 的 recsys-examples 仓库,该仓库包含一系列示例,展示了如何在 NVIDIA GPU 上使用 PyTorch 训练和部署生成式推荐器的最佳实践。它包括优化的 HSTU 和 Semantic ID 实现,涵盖了训练和推理工作流。

对于 HSTU 推理,recsys-examples 支持 PyTorch AOTInductor 在 Torch C++ 运行时中执行模型。NVIDIA 开发者博客还介绍了 nv-embedding-cache,它提供了与 PyTorch 兼容的模块,用于加速大规模嵌入查找。

🔗 阅读完整文章:


生成式推荐系统如何大规模重塑 RecSys

来源:https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/ 推荐系统(RecSys)是消费互联网行业中最普遍的机器学习问题之一,但众所周知,其在大规模训练和服务方面难度很高。大语言模型(LLM)的兴起催生了从传统的基于嵌入相似度的目标向生成式目标的转变,即目标是根据用户历史序列,在大型目录中预测下一个动作或项目。

本文介绍了向生成式推荐器(GR)的架构转变、由此带来的生产挑战,以及 NVIDIA recsys-examplesnv-embedding-cache 如何解决这些挑战。

为何传统 RecSys 在大规模下会失效https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#why_traditional_recsys_breaks_at_scale

数据类型和规模https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#data_type_and_volume

用户历史是主要的 RecSys 数据类型,它记录了用户如何与目录中的项目交互。与文本或图像等模态不同,用户历史涉及混合了分类和连续特征的数据,且随时间频繁变化。在工业规模下,这些数据每天可达太字节(TB)甚至拍字节(PB)级别。即使在最高端的硬件加速器上,这种规模的数据也无法完全装入 GPU 高带宽内存(HBM),从而在训练和推理过程中引入诸多瓶颈。

稀疏性和长尾问题https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#sparsity_and_the_long-tail_problem

推荐系统中的长尾问题描述了目录中少数热门项目获得了大部分交互的现象。这个问题的产生是因为项目目录可能远超用户数量,导致用户-项目交互数据非常稀疏。由于与项目交互的概率分布严重偏向该项目的受欢迎程度,训练数据为大多数代表真实用户偏好的利基项目提供的信号非常有限。

冷启动问题https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#cold_start_problem

当新用户或新项目加入 RecSys 平台时,没有交互历史来立即生成高质量的嵌入。相反,必须从更少的特征集推断,这可能对推荐的轨迹产生负面影响。通过关联相似项目和语义描述可以部分解决这个问题,但初始推荐质量较差并降低用户体验的风险依然存在。

严格的延迟要求https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#strict_latency_requirements

在生产环境中,RecSys 模型通常在线为数百万用户提供服务,并受制于严格的服务水平协议(SLA),延迟的轻微增加都可能影响用户体验。与可能容忍自回归解码延迟的 LLM 工作负载不同,RecSys 模型必须在几毫秒内频繁检索和排序数千个候选项目。

生成式推荐器https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#generative_recommenders

与使用几何相似度建模用户-项目偏好的传统基于嵌入的推荐器不同,GR 将推荐重新构建为一个类似于 LLM 的序列建模问题。

目标是基于用户历史序列,建模下一个动作或项目的概率分布:

P(next_item | user_history)

向更同质化、类似 Transformer 架构的转变可以更好地利用缩放定律,可能将检索和排序统一在单个模型中,并更自然地与快速发展的 LLM 生态系统集成。

在推荐器中实现这一目标的两种主流方法是层次化序列转换单元(HSTU)和语义ID:

HSTUhttps://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#hstu

HSTU 由 Meta 于 2024 年提出,是一种基础性的 GR 模型架构,它在生成式目标下重构了 RecSys,并引入了关键创新以实现生产规模下的高效训练和服务。

HSTU 将输入数据表示为按时间戳排序的、由项目和动作(如点赞、点击等)交错组成的每个用户序列。它还摒弃了对显式特征工程的依赖(这是传统 RecSys 模型的常见做法),转而采用通过用户-项目交互的注意力学习到的序列表示。这种表述使用户历史类似于 LLM 中的下一个 token 预测,在训练中你可以在序列内和跨批次获得有意义的学习信号。

与标准 Transformer 注意力不同,HSTU 通过使用基于 SiLU 的加权替换 softmax 归一化、引入相对注意力偏置以及在输出投影前应用逐元素门控,修改了注意力聚合机制。这些修改在长序列中保留了更强的幅度信息,同时实现了更高效的内核融合和更低延迟的推理。

语义IDhttps://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#semantic_ids

在一个大型项目语料库上,使用稀疏的用户-项目交互集来建模下一个项目预测,会带来许多挑战:完整 softmax 计算的瓶颈、来自长尾项目的弱训练信号,以及对语义相似项目泛化能力差。

语义ID(SIDs)由 Google 提出,通过基于项目嵌入的层次化聚类生成更小的新的词表token集来缓解这些问题。诸如 TIGER、PLUM、OneRec v1/v2 以及许多现代 GR 架构都使用语义ID作为可扩展自回归推荐的基础。

传统的端到端语义ID GR 架构显示了一个冻结权重的自编码器模型生成语义ID,然后输入到 GR 中进行下一个 token 预测。图 1. 语义ID GR 端到端架构 与传统 RecSys 不同,语义ID的自回归解码直接生成推荐,而不是在嵌入空间中搜索,并通过输出 logits 自然地提供排序。这允许像集束搜索(beam search)这样的搜索方法在一次前向传递中产生多个语义ID,提高了吞吐量,并允许选择聚类中的利基项目。

recsys-examples 仓库https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#recsys-examples_repository

recsys-examples (https://github.com/NVIDIA/recsys-examples/tree/main) 仓库是一系列示例的集合,旨在展示如何在 NVIDIA GPU 上使用 PyTorch 训练和部署生成式推荐器的最佳实践。它包括 HSTU 和语义ID模型的优化实现,涵盖了训练和推理工作流。该仓库还整合了三个模块化组件:用于嵌入层的 DynamicEmb、为推荐系统量身定制的 KV 缓存和存储管理器,以及用于 HSTU 和语义ID集束搜索解码的高效 CUDA 操作。

动态嵌入https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#dynamic_embedding

RecSys 中的传统嵌入表假定一个固定、静态的词汇表。然而,在生产环境中,新的用户和项目不断出现,项目目录的长尾增长速度远超任何单个 GPU HBM 的容量。过度配置静态表会浪费内存在永远不会被访问的行上,而配置不足则会导致昂贵的复制操作,可能降低模型性能和质量。

DynamicEmb 用一个针对 GPU 优化的带评分的哈希表替换了静态表,该表按需将任意特征ID映射到嵌入行。仅为模型实际看到的ID分配行,并且该表驻留在 HBM 和固定的主机内存中,因此其容量可以远超单个 GPU 的容量。该实现基于 HierarchicalKV (https://arxiv.org/pdf/2603.17168) 哈希表设计中的算法构建。准入控制和基于分数的驱逐相结合,使得容量可以用于对模型训练重要的ID,并使长尾问题在大规模下变得可行。

它作为 TorchRec 后端提供,表按行跨 rank 分片,使用 EmbeddingBagCollectionEmbeddingCollection API。融合的 CUDA 内核处理 SUMMEAN 和序列池化模式的查找和梯度归约。预取技术将频繁访问的嵌入保留在 HBM 中以实现高效访问。

HSTU 支持https://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#hstu_support

Recsys-examples 为 HSTU 提供了生产级的训练和推理栈。项目、用户、动作和上下文嵌入表通过 TorchRec 管理,DynamicEmb 为高基数表提供动态容量和缓存。对于密集层,HSTU 骨干网络使用 Megatron-Core,因此单次训练运行可以利用数据、张量、序列和流水线并行。该库设计在整个栈中保持模块化,因此组件可以与自定义架构即插即用。

在训练期间,TorchRec/DynamicEmbMegatron-Core 无缝集成,以协调嵌入模块和密集模块之间的分片和并行。训练流水线包含动态打乱以平衡 rank 间的工作负载,将嵌入通信和预取与密集计算重叠,并包含带有融合操作和针对 NVIDIA Ampere、Hopper 和 Blackwell GPU 优化的 FBGEMM 注意力内核的 HSTU 层。这些优化共同将两台 DGX H100 节点上的端到端模型 FLOP 利用率(MFU)从 7.65% 提高到 31.40%,展示了训练 (https://github.com/NVIDIA/recsys-examples/blob/main/examples/hstu/training/benchmark/E2E_BENCHMARK.md) 效率的巨大提升。

对于推理,recsys-examples 旨在满足严格的低延迟要求,支持 PyTorch AOTInductor 在 Torch C++ 运行时中执行模型,同时与 NVIDIA Triton Inference Server 兼容。

频繁访问的嵌入使用 nv-embedding-cache 保持在 GPU 附近,计算通过定制的、启用 FlexKV 的 KV 缓存进一步减少,该缓存将缓存条目分布到多个内存层中。

使用 Triton Inference Server 部署时,使用 Pytorch AOTI 后端且无 KV 缓存的推理相比 Python 后端实现了 1.14 倍至 1.28 倍的加速,而使用 Pytorch AOTI 后端与 KV 缓存的推理在理想的全 GPU 缓存命中场景下实现了 2.20 倍至 2.38 倍的加速。

柱状图展示了在 Triton Inference Server 上服务多种尺寸 HSTU 模型时,使用 Python 后端、无缓存 AOTI 后端和有缓存 AOTI 后端时请求延迟的下降情况。图 2. 使用 recsys-examples Triton 后端配置的 HSTU 性能

语义ID-GRhttps://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#semantic_id-gr

基于语义ID的 GR 引入了一种与基于聊天的 LLM 推理截然不同的服务模式:长用户上下文、短自回归解码以及在受限项目token空间上的大集束宽度。在实际的语义ID工作负载中,一个请求可能包含数千个历史token,只解码2-3个语义IDtoken,并使用128或256等集束宽度以提高推荐多样性。

现有的 LLM 服务系统(如 vLLM、SGLang 和 TensorRT LLM)主要针对具有分页 KV 缓存、动态批处理和长解码的多用户聊天式服务进行了优化。它们是强大的通用框架,但自然无法暴露语义ID-GR所需的核心抽象:共享的请求级上下文 KV、每个集束的短解码 KV、集束路径跟踪、动态集束宽度和项目受限生成。

为解决此问题,recsys-examples 提供了一个面向 GR 优化的推理框架示例 (https://github.com/NVIDIA/recsys-examples/tree/main/examples/sid-gr-inference),适用于基于 Qwen 的语义ID模型。该框架分离 KV 缓存以隔离集束相关和无关组件:将长共享上下文放入 ContextKV,短解码历史放入 BeamKV,逻辑集束祖先关系放入 BeamPath。这避免了将每个集束视为一个独立的长序列。运行时还包括 GR 原生的连续批处理、直接的池视图 CUDA 图重放、项目受限的 topK 以及一个直接操作 GR KV 布局的专用 gr-decode-atten 后端。

在单块 NVIDIA H100 80GB GPU 上,使用 Qwen3-1.7B,上下文长度分别为 1,000 和 5,000 tokens,集束宽度为 256,输出 3 个token的情况下,GR 优化路径在测量的离线和在线基准测试中始终优于 SGLang 的集束搜索。

指标工作负载GR 结果基线提升
离线延迟ctx=1000, batch=4, beam=256, output=347.736 ms102.318 ms快 2.14 倍
离线延迟ctx=5000, batch=4, beam=256, output=3154.224 ms349.857 ms快 2.27 倍
离线延迟ctx=5000, batch=8, beam=256, output=3307.917 ms685.354 ms快 2.23 倍
在线服务吞吐量ctx=5000, concurrency=4, beam=256, output=3~19.7 req/s~10.7 req/s高 ~1.85 倍
在线中位数延迟ctx=5000, concurrency=4, beam=256, output=3~198 ms~370 ms降低 ~46%
离线正确性ctx=1000/5000, batch=1/2/4/8, beam=256Top1 精确度 1.000SGLang 比较TopK 重叠均值 0.960
表 2. recsys-examples 语义ID-GR 推理框架的服务结果

这使得语义ID-GR 服务成为 recsys-examples 中 HSTU 和嵌入组件的自然补充:HSTU 和 DynamicEmb 解决了生产规模的训练和嵌入密集型推理,而语义ID-GR 推理路径则针对具有长上下文、大集束搜索和严格推荐系统延迟要求的自回归语义ID生成。

nv-embedding-cachehttps://developer.nvidia.com/blog/how-generative-recommenders-are-redefining-recsys-at-scale/#nv-embedding-cache

nv-embedding-cache(NVE)是一个 SDK,用于加速推荐系统推理中的大规模嵌入表查找和操作。它提供模块化组件,包括优化的内核、软件缓存原语和 PyTorch 兼容绑定,从而能够低延迟地访问超出单个 GPU HBM 容量的嵌入表。

生产环境的推荐器嵌入表通常会超过单个 GPU 的 HBM 容量。

相似文章