FlashAttention & Subquadratic Sparse Attention

FlashAttention & Subquadratic Sparse Attention

What it is

Standard attention mechanisms suffer from quadratic complexity, meaning if you double the input length, the computational cost quadruples. FlashAttention is an IO-aware algorithm that optimizes how data moves between GPU memory levels (SRAM and HBM) to speed up computation without changing the mathematical output. Subquadratic Sparse Attention (SSA) goes further by using routing mechanisms to ignore irrelevant tokens, processing only a subset of key-value pairs to keep costs linear or near-linear as context grows.

Why it matters

In production AI systems, memory bandwidth is the primary bottleneck, not raw compute power. Without these optimizations, processing long documents or massive codebases would require prohibitively expensive hardware and result in slow inference speeds. These techniques allow models to handle context windows of up to 12 million tokens, enabling the processing of entire libraries or hour-long videos in a single pass.

Core concepts to master

• Tiling: Breaking the large attention matrix into smaller blocks that fit into the GPU's fast SRAM to reduce slow reads/writes to HBM.
• Recomputation: Avoiding the storage of the massive attention matrix for the backward pass by recalculating it on the fly, trading cheap compute for expensive memory.
• Sparsity Routing: The process of selectively identifying which tokens are most relevant to the current query, allowing the model to skip the majority of the KV cache.
• IO-Awareness: Designing the algorithm based on the physical hardware constraints of the GPU rather than just the mathematical operations.

Common mistakes

• Confusing Approximation with Optimization: Thinking FlashAttention changes the result. It is a "lossless" optimization; SSA, however, is an approximation that may trade some accuracy for speed.
• Ignoring Memory Bandwidth: Focusing only on TFLOPS (compute speed) while ignoring the fact that the GPU is often idling while waiting for data to move from memory.

Track connection

This concept bridges the gap between high-level model architecture and low-level hardware acceleration. It connects directly to GPU kernel development and KV cache management, providing the foundation for scaling inference throughput in the AI Infrastructure track.