用新方法处理变长类别特征，让推荐系统更快更省内存。
Enhancing Performance and Scalability of Large-Scale Recommendation Systems with Jagged Flash Attention
- 设计可处理动态张量的锯齿型特征交互核，高效提取长类别特征信息。
- 相比稠密注意力，实现最高9倍加速和22倍内存降低。
- 适合大规模推荐系统研发者，尤其关注性能与可扩展性的场景。
硬件加速器的集成显著提升了现代推荐系统的性能，使其能够探索以往被认为不切实际的复杂排序范式。然而，基于GPU的计算成本仍构成重大挑战。本文提出一种面向效率的方法，突破传统PyTorch原生模块的局限，针对排序模型依赖长度可变的类别特征导致的GPU利用率下降问题，提出锯齿型特征交互核（Jagged Feature Interaction Kernels），通过高效处理动态大小张量，从长类别特征中提取细粒度信息。进一步将锯齿张量与Flash Attention结合，提出锯齿型Flash注意力（Jagged Flash Attention），在性能上实现最高9倍加速和22倍内存减少；相较于稠密Flash注意力，更达到最高3倍加速和53%内存效率提升。在生产模型中，观察到10% QPS提升和18%内存节省，使我们得以支持更长特征和更复杂的架构，实现系统规模化。
原文摘要 · Abstract (English)
The integration of hardware accelerators has significantly advanced the capabilities of modern recommendation systems, enabling the exploration of complex ranking paradigms previously deemed impractical. However, the GPU-based computational costs present substantial challenges. In this paper, we demonstrate our development of an efficiency-driven approach to explore these paradigms, moving beyond traditional reliance on native PyTorch modules. We address the specific challenges posed by ranking models' dependence on categorical features, which vary in length and complicate GPU utilization. We introduce Jagged Feature Interaction Kernels, a novel method designed to extract fine-grained insights from long categorical features through efficient handling of dynamically sized tensors. We further enhance the performance of attention mechanisms by integrating Jagged tensors with Flash Attention. Our novel Jagged Flash Attention achieves up to 9x speedup and 22x memory reduction compared to dense attention. Notably, it also outperforms dense flash attention, with up to 3x speedup and 53% more memory efficiency. In production models, we observe 10% QPS improvement and 18% memory savings, enabling us to scale our recommendation systems with longer features and more complex architectures.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。