Enhancing Performance and Scalability of Large-Scale Recommendation Systems with Jagged Flash Attention

硬件加速器的整合显著提升了现代推荐系统的能力，使得探索先前被认为不切实际的复杂排名范例成为可能。然而，基于GPU的计算成本带来了巨大的挑战。在本文中，我们展示了一种以效率为驱动的方法，探索这些范例，超越了传统对本地PyTorch模块的依赖。我们解决了排名模型对分类特征的依赖所带来的具体挑战，这些特征在长度上有所不同，使GPU的利用变得复杂。我们引入了Jagged Feature Interaction Kernels，这是一种新颖的方法，旨在通过有效处理动态大小的张量，从长分类特征中提取细粒度的见解。我们进一步通过将Jagged张量与Flash Attention集成来增强注意机制的性能。我们的新型Jagged Flash Attention相比密集注意力实现了高达9倍的加速和22倍的内存减少。值得注意的是，它也优于密集Flash Attention，速度提高了3倍，内存效率提高了53%。在生产模型中，我们观察到10%的QPS改进和18%的内存节省，使我们能够以更长的特征和更复杂的架构扩展我们的推荐系统。

本论文旨在解决现代推荐系统中GPU计算成本高昂的问题，并提出一种高效的方法来处理分类特征的依赖性和长度变化。

论文提出了Jagged Feature Interaction Kernels的新方法来有效地处理动态大小的张量，并通过将Jagged张量与Flash Attention相结合来提高注意机制的性能。

论文的实验结果表明，Jagged Flash Attention相比传统的dense attention和dense flash attention可以实现更快的速度和更高的内存效率，同时也可以提高生产模型的性能。

最近的相关研究包括《Neural Collaborative Filtering with Embedding Regeneration》和《Deep Learning based Recommender System: A Survey and New Perspectives》。