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

TL;DR

This paper introduces Jagged Flash Attention, a novel method that significantly improves the efficiency and scalability of large-scale recommendation systems by optimizing attention mechanisms for variable-length categorical features.

Contribution

We develop Jagged Feature Interaction Kernels and integrate them with Flash Attention to handle dynamic tensor sizes, achieving substantial speedups and memory savings in recommendation models.

Findings

Up to 9x speedup over dense attention

22x memory reduction compared to dense attention

10% QPS improvement in production models

Abstract

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.