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

arXiv: https://arxiv.org/abs/2409.15373 · HTML: https://mlsanity.com/paper/2409.15373/

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 6 of its 13 references that are in the map.
- First posted on arXiv 19 Sep 2024.

## Builds on (most cited first, 6 of 6)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Deep Learning Recommendation Model for Personalization and Recommendation Systems](https://mlsanity.com/paper/1906.00091/) (2019)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Wukong: Towards a Scaling Law for Large-Scale Recommendation](https://mlsanity.com/paper/2403.02545/) (2024)
- [Async Learned User Embeddings for Ads Delivery Optimization](https://mlsanity.com/paper/2406.05898/) (2024)

## Built on by (most cited first, 1 of 1)

- [TransAct V2: Lifelong User Action Sequence Modeling on Pinterest Recommendation](https://mlsanity.com/paper/2506.02267/) (2025)

## Abstract (arXiv)

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.