A percepção-chave (Dao et al., 2022): a atenção padrão materializa a matriz completa N×N de atenção na HBM (high bandwidth memory) da GPU, o que é intensivo em memória (quadrático no comprimento da sequência) e lento (a largura de banda da HBM é o gargalo). O Flash Attention nunca materializa essa matriz. Em vez disso, calcula a atenção em blocos, carregando pequenos blocos de Q, K, V na SRAM rápida on-chip, computando resultados parciais e acumulando-os — uma técnica chamada "tiling" ou "fusão de kernels".
A atenção padrão armazena a matriz N×N de atenção, exigindo O(N²) de memória. Para um contexto de 128K com 128 cabeças de atenção, isso são centenas de gigabytes. O Flash Attention usa O(N) de memória ao computar o softmax incrementalmente e nunca armazenar a matriz completa. Isso é o que tornou janelas de contexto de 128K–1M viáveis no hardware existente. O FlashAttention-2 melhorou ainda mais o throughput ao paralelizar melhor entre os blocos de threads da GPU.
O Flash Attention exemplifica um princípio mais amplo: no hardware moderno, o gargalo frequentemente é a largura de banda de memória, não a computação. GPUs podem realizar trilhões de operações por segundo, mas só conseguem ler/escrever centenas de gigabytes por segundo da HBM. Algoritmos que minimizam o tráfego de memória (mesmo ao custo de computação extra) frequentemente vencem. Essa abordagem "ciente de IO" está influenciando como todo o campo pensa sobre design de algoritmos para cargas de trabalho de IA.