利用不规则闪存注意力提升大规模推荐系统的性能与可扩展性
机器学习
2024-09-25 v1 人工智能
信息检索
摘要
硬件加速器的集成显著推动了现代推荐系统的能力，使能够探索此前被视为不切实际的复杂排序范式。然而，基于GPU的计算成本仍然是面临的主要挑战。本文展示了我们开发一种以效率为导向的方法，超越传统依赖本地PyTorch模块的做法。我们针对排序模型对类别特征的依赖，这些特征长度变化，使GPU利用率受限，提出了具体挑战。我们引入了不规则特征交互核（Jagged Feature Interaction Kernels），一种新颖的方法，通过高效处理动态大小张量来从长类别特征中提取细粒度见解。我们进一步通过将不规则张量与Flash Attention集成来增强注意力机制的性能。我们新开发的Jagged Flash Attention相较于稠密注意力实现了最高达9倍的加速和22倍的内存降低。值得注意的是，它还优于稠密flash注意力，实现了最高3倍的加速和53%以上的内存效率提升。在实际生产模型中，我们观察到10%的QPS提升和18%的内存节省，使我们能够扩展推荐系统以支持更长的特征和更复杂的架构。
引用
@article{arxiv.2409.15373,
title = {Enhancing Performance and Scalability of Large-Scale Recommendation Systems with Jagged Flash Attention},
author = {Rengan Xu and Junjie Yang and Yifan Xu and Hong Li and Xing Liu and Devashish Shankar and Haoci Zhang and Meng Liu and Boyang Li and Yuxi Hu and Mingwei Tang and Zehua Zhang and Tunhou Zhang and Dai Li and Sijia Chen and Gian-Paolo Musumeci and Jiaqi Zhai and Bill Zhu and Hong Yan and Srihari Reddy},
journal= {arXiv preprint arXiv:2409.15373},
year = {2024}
}
备注
3 pages, 2 figures