FlashAttention (TL)

FlashAttention — ito ay isang rebolusyonaryong algorithm para sa pagkalkula ng mekanismo ng attention, na binuo para sa makabuluhang pagpapabilis ng pagsasanay at inferens ng malalaking language model (LLM) habang pinapanatili ang buong katumpakan ng mga kalkulasyon. Ang algorithm ay unang ipinakita noong 2022 ng isang pangkat ng mga mananaliksik mula sa Stanford University sa ilalim ng pamumuno ni Tri Dao.

Ang pangunahing ideya ng FlashAttention ay ang muling pagsasaayos ng mga kalkulasyon na isinasaalang-alang ang hierarchy ng memorya ng GPU, na nagbibigay-daan sa pag-minimize ng bilang ng mga access sa mabagal na memorya at pag-aalis ng pangunahing bottleneck ng karaniwang mekanismo ng attention.

Problemang naroroon sa karaniwang attention

Ang karaniwang mekanismo ng self-attention sa mga transformer ay kinakalkula ayon sa formula: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ kung saan ang Q, K, V — ay ang mga matrix ng mga query, key, at value.

Ang pangunahing problema ng pamamaraang ito ay ang quadratic na kumplikasyon sa oras at memorya (O(N²)) kaugnay ng haba ng sequence na N. Sa isang simpleng implementasyon, kinakailangan na kalkulahin at itago sa memorya ng GPU ang buong matrix ng attention na S na may sukat na N×N, na humahantong sa dalawang kritikal na problema:

1. Malaking paggamit ng memorya: Ang pag-iimbak ng matrix na N×N ay nagiging imposible kapag nagtatrabaho sa mahahabang konteksto.
2. Mga operasyon ng input-output (IO): Ang pangunahing bottleneck ay hindi ang bilang ng mga aritmetikong operasyon, kundi ang patuloy na pag-access sa mabagal na memorya ng GPU.

Hierarchy ng memorya ng GPU

Para maunawaan ang problema, mahalagang makilala ang dalawang uri ng memorya sa GPU (gamit ang NVIDIA A100 bilang halimbawa):

• SRAM (static na memorya): Mabilis na on-chip na memorya na may maliit na kapasidad (~20 MB) na may napakalaking bandwidth (hanggang 19 TB/s).
• HBM (high-bandwidth memory): Mabagal na memorya na may malaking kapasidad (40–80 GB) na may mas mababang bandwidth (humigit-kumulang 1.5 TB/s).

Ang asymmetry na ito ay nagpapaging limitado sa bandwidth ng memorya (memory-bound) ang karaniwang algorithm ng attention, dahil patuloy itong nagbabasa at nagsusulat ng malalaking matrix mula sa mabagal na HBM, na siyang pangunahing pinagmumulan ng mga pagkaantala.

Mga pangunahing inobasyon ng FlashAttention

Ang FlashAttention ay isang IO-aware na algorithm na naglulutas ng problema sa pamamagitan ng pag-minimize ng mga access sa HBM. Ito ay nakamit sa pamamagitan ng tatlong pangunahing teknik.

Tiling at block na pagproseso

Sa halip na iproseso ang buong matrix nang sabay-sabay, hinahati ng FlashAttention ang mga input na matrix na Q, K, V sa maliliit na bloke (tiles), na kasyang-kasya sa mabilis na SRAM. Sunud-sunod na ino-load ng algorithm ang mga blokeng ito, isinasagawa ang lahat ng kalkulasyon ng attention para sa mga ito at ina-update ang panghuling resulta, nang hindi iniimbak ang buong matrix ng attention sa mabagal na HBM.

Online na pagkalkula ng Softmax

Ang pangunahing teknikal na tagumpay ay ang "online" na pagkalkula ng Softmax. Ang karaniwang Softmax ay nangangailangan ng kaalaman sa lahat ng elemento ng input vector para sa normalisasyon. Gumagamit ang FlashAttention ng binagong algorithm na nagbibigay-daan sa pagkalkula ng Softmax sa mga bahagi. Nagpapanatili ito ng dalawang intermediate na halaga (kasalukuyang maximum at kabuuan ng mga eksponente), na ina-update habang pinoproseso ang mga bagong bloke, na nagpapahintulot sa pagkuha ng tumpak na resulta nang hindi kino-access ang buong matrix nang sabay-sabay.

Pagsasama ng mga operasyon sa isang CUDA kernel

Lahat ng operasyon ng attention (matrix multiplication na QKᵀ, masking, Softmax, multiplication sa V) ay pinagsama sa isang pinagsamang CUDA kernel (fused kernel). Ito ay radikal na nagbabawas ng bilang ng mga operasyon ng pagbabasa/pagsulat sa HBM: sa halip na maraming beses na dumaan sa buong matrix, isang beses na ino-load ng algorithm ang bloke sa SRAM, isinasagawa ang lahat ng kalkulasyon at isinusulat lamang ang panghuling resulta.

Teoretikal at praktikal na kahusayan

Kumplikasyon at optimalidad

Binabawasan ng FlashAttention ang paggamit ng memorya mula O(N²) hanggang O(N), na nagbibigay ng linear na scalability. Napatunayan na ang IO-complexity ng algorithm ay teoretikal na optimal para sa pagkalkula ng attention sa dalawang antas na hierarchy ng memorya, ibig sabihin, imposibleng magsagawa ng tumpak na attention nang mas mabilis nang hindi binabago ang hardware.

Mga empirikal na resulta

Ipinakita ng unang bersyon ng FlashAttention ang mga makabuluhang pagpapabuti:

• Pagpapabilis:
  • BERT-large (haba na 512): 15% na pagpapabilis ng pagsasanay.
  • GPT-2 (haba na 1K): 3-beses na pagpapabilis.
  • Mga gawain sa Long-Range Arena (1K-4K): 2.4-beses na pagpapabilis.
• Pagtitipid ng memorya: Hanggang 20-beses na pagtitipid ng memorya kumpara sa mga tumpak na base na implementasyon.
• Pagpapabuti ng kalidad ng mga modelo: Salamat sa kakayahang magtrabaho sa mas mahahabang konteksto, ang FlashAttention ay hindi lamang hindi nawawalan, kundi nagpapabuti rin ng kalidad ng mga modelo. Halimbawa, ang perplexity ng GPT-2 ay bumuti ng 0.7 puntos, at ang katumpakan sa mga gawaing pag-uuri ng mahahabang dokumento ay tumaas ng 6.4 puntos.

Ebolusyon at karagdagang pag-unlad

Ang tagumpay ng FlashAttention ay nagbukas ng isang serye ng mga hardware-oriented na algorithm.

FlashAttention-2 (2023)

Ang ikalawang bersyon ay naglayong mas ganap na magamit ang mga resources ng GPU. Sa orihinal na FlashAttention, ang kahusayan sa NVIDIA A100 ay 25–40% lamang ng maximum. Nagpakilala ang FlashAttention-2 ng mga pagpapabuti sa parallelization ng mga kalkulasyon, na nagpahintulot:

• Na makamit ang dalawang beses na pagpapabilis kumpara sa unang bersyon.
• Na mapataas ang paggamit ng GPU hanggang 50–73% ng teoretikal na maximum.
• Na palawakin ang suporta para sa mga attention head na may sukat na 256, pati na rin para sa mga arkitektura ng Multi-Query Attention (MQA).

FlashAttention-3 (2024)

Ang ikatlong bersyon ay na-optimize nang espesyal para sa arkitektura ng GPU na NVIDIA Hopper (H100). Gumagamit ito ng mga bagong kakayahan ng hardware, tulad ng asynchrony ng Tensor Cores at suporta para sa FP8, na nagpahintulot:

• Na makamit ang karagdagang 1.5–2-beses na pagpapabilis kumpara sa FlashAttention-2.
• Na makamit ang performance na hanggang 740 TFLOPS sa FP16 at malapit sa 1.2 PFLOPS sa FP8.

Mga espesyal na solusyon

Ang mga ideya ng FlashAttention ay napalago sa iba pang mga proyekto:

• FlashInfer (2025): Isang nako-customize na attention engine na na-optimize nang espesyal para sa mga gawaing inferens ng LLM. Nakatuon ito sa mahusay na pagtatrabaho sa KV-cache sa streaming generation mode.
• FlashMLA (2024): Isang implementasyon ng attention na may compression ng context cache (latent attention), na nagbibigay-daan sa pagtitipid ng memorya sa napakahahabang sequence na may minimal na pagkawala ng impormasyon.

Impluwensya sa industriya at ecosystem

Naging pundamental na breakthrough ang FlashAttention at mabilis na naging pamantayan ng industriya para sa mahusay na pagsasanay at inferens ng LLM. Ito ay naisama sa mga pangunahing library tulad ng PyTorch at Hugging Face, at ginagamit sa karamihan ng malalaking language model (LLaMA, MPT, Falcon, Claude, atbp.).

Ang FlashAttention at ang mga kasunod na bersyon nito ay gumanap ng mapagpasyang papel sa pagpapalaki ng context window ng mga language model: mula 2–4 libong token (GPT-3) hanggang 128 libong token (GPT-4) at maging hanggang milyun-milyong token sa mga eksperimental na modelo. Tinanggal ng algorithm ang isa sa mga pangunahing hadlang sa scalability ng mga transformer, na nagbubukas ng mga bagong posibilidad para sa mga aplikasyon ng AI, mula sa pagsusuri ng mahahabang dokumento hanggang sa multimodal na pag-unawa.

Mga Sanggunian

• Opisyal na repository ng FlashAttention sa GitHub

Mga Babasahin

• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness. arXiv:2205.14135.
• Dao, T. (2023). FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning. arXiv:2307.08691.
• Shah, J. et al. (2024). FlashAttention-3: Fast and Accurate Attention with Asynchrony and Low-Precision. arXiv:2407.08608.
• Kwon, W. et al. (2023). Efficient Memory Management for Large Language Model Serving with PagedAttention. arXiv:2309.06180.
• Hong, K. et al. (2023). FlashDecoding++: Faster Large Language Model Inference on GPUs. arXiv:2311.01282.
• Ye, Z. et al. (2025). FlashInfer: Efficient and Customizable Attention Engine for LLM Inference Serving. arXiv:2501.01005.
• Dege, P. et al. (2025). FlashMLA-ETAP: Efficient Transpose Attention Pipeline for Accelerating MLA Inference on NVIDIA H20 GPUs. arXiv:2506.01969.
• Wang, G. et al. (2025). FlashMask: Efficient and Rich Mask Extension of FlashAttention. OpenReview wUtXB43Chi.
• Dao, T. et al. (2022). FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (OpenReview version). OpenReview H4DqfPSibmx.
• Gholami, A. et al. (2024). FlashAttention on a Napkin: A Diagrammatic Approach to Deep Learning IO-Awareness. OpenReview pF2ukh7HxA.