On Fine-Grained I/O Complexity of Attention Backward Passes

TL;DR

This paper analyzes the I/O complexity of attention mechanisms in large language models, deriving tight bounds for various cache scenarios and proposing optimal algorithms for small-cache environments.

Contribution

It provides a comprehensive theoretical framework for I/O complexity in attention, including new algorithms for small-cache settings and bounds for sparse attention.

Findings

FlashAttention is optimal in large-cache scenarios.

A new algorithm outperforms existing methods in small-cache environments.

Established lower bounds for sparse attention across cache sizes.

Abstract

Large Language Models (LLMs) exhibit exceptional proficiency in handling extensive context windows in natural language. Nevertheless, the quadratic scaling of attention computation relative to sequence length creates substantial efficiency bottlenecks, necessitating the development of I/O-optimized algorithms. In this work, we conduct a systematic examination of the I/O complexity inherent in attention mechanisms, with a specific emphasis on the backward pass under both small and large cache settings. By leveraging the red-blue pebble game framework, we derive tight bounds for I/O complexity across the full spectrum of cache sizes. We validate that FlashAttention, one of the current industry standards, achieves optimality in the large-cache scenario for both forward and backward passes. Conversely, for small-cache environments, we introduce a novel algorithm that outperforms contemporary methods and successfully attains theoretical tight bounds. Furthermore, we expand our investigation to include sparse attention by establishing granular lower bounds for both forward and backward passes across all cache configurations. Ultimately, our results solidify the theoretical framework regarding I/O complexity in attention mechanisms, providing critical guidance for the development of efficient LLM training and inference systems.