Fine-grained Attention I/O Complexity: Comprehensive Analysis for Backward Passes
This work addresses efficiency problems for LLM training and inference by providing a theoretical foundation for I/O complexity, though it is incremental as it builds on existing frameworks like the red-blue pebble game.
This paper tackled the computational challenges of attention mechanisms in Large Language Models by analyzing I/O complexity for backward passes, establishing tight bounds across cache sizes and confirming FlashAttention's optimality for large caches while providing an improved algorithm for small caches.
Large Language Models (LLMs) have demonstrated remarkable capabilities in processing long-context information. However, the quadratic complexity of attention computation with respect to sequence length poses significant computational challenges, and I/O aware algorithms have been proposed. This paper presents a comprehensive analysis of the I/O complexity for attention mechanisms, focusing on backward passes by categorizing into small and large cache scenarios. Using the red-blue pebble game framework, we establish tight bounds on I/O complexity across all cache sizes. We confirm that the de facto standard I/O aware algorithm FlashAttention is optimal for both forward and backward passes for the large cache size scenario. For small cache sizes, we provide an algorithm that improves over existing methods and achieves the tight bounds. Additionally, we extend our analysis to sparse attention, a mainstream speeding-up approach, deriving fine-grained lower bounds for both forward and backward passes and both small and large caches. Our findings complete the theoretical foundation for I/O complexity in attention mechanisms, offering insights for designing efficient algorithms of LLM training and inference.