FlashAttention: Cơ chế Attention chính xác, nhanh và tối ưu bộ nhớ với nhận thức IO (IO-Awareness)

Hướng dẫn đọc

Vấn đề và động lực

Transformer có độ phức tạp thời gian và bộ nhớ bậc hai theo độ dài chuỗi $N$. Mặc dù các phương pháp attention xấp xỉ giảm được số lượng FLOP lý thuyết, chúng thường thất bại trong việc tăng tốc thời gian thực tế (wall-clock time) vì các phép tính trên GPU hiện đại chủ yếu bị nghẽn do quá trình truyền dữ liệu bộ nhớ (IO) chứ không phải các phép toán số học.

Trong attention tiêu chuẩn, việc tính $\mathbf{S}=\mathbf{Q}{\mathbf{K}}^{\top }$, $\mathbf{P}=\mathrm{softmax}(\mathbf{S})$, và $\mathbf{O}=\mathbf{PV}$ phải ghi các ma trận trung gian kích thước $N\times N$ vào High Bandwidth Memory (HBM) chậm và đọc lại chúng nhiều lần. Điều này làm phát sinh $\Theta (Nd+N^2)$ lượt truy cập HBM và dung lượng bộ nhớ $O(N^2)$, hạn chế kích thước cửa sổ ngữ cảnh trong thực tế.

Các ý tưởng cốt lõi

• IO-Awareness: Tính toán rõ ràng cấu trúc phân cấp bộ nhớ GPU: bộ nhớ SRAM trên chip nhỏ nhưng tốc độ cao (ví dụ: 192 KB/SM với băng thông ~19 TB/s trên A100) so với HBM lớn nhưng chậm (40–80 GB với băng thông 1.5–2.0 TB/s).
• Tiling (Tính toán Softmax theo khối): Chia nhỏ đầu vào $\mathbf{Q},\mathbf{K},\mathbf{V}$ thành các khối vừa vặn hoàn toàn trong SRAM. Bằng cách theo dõi các thống kê động về giá trị cực đại của hàng (row-max) và hệ số chuẩn hóa của softmax, phép giản lược softmax được tính toán tăng dần qua các khối mà không cần hiện thực hóa toàn bộ ma trận $N\times N$ trong HBM.
• Tính toán lại trong lượt lan truyền ngược (Recomputation in Backward Pass): Tránh lưu ma trận attention trung gian $N\times N$ cho quá trình lan truyền ngược. Thay vào đó, chỉ lưu đầu ra $\mathbf{O}$ và các giá trị thống kê softmax $(\ell ,m)$, rồi tính toán lại $\mathbf{S}$ và $\mathbf{P}$ ngay trên chip theo từng khối trong SRAM ở lượt truyền ngược.
• Kernel Fusion: Hợp nhất các phép nhân ma trận, tạo mặt nạ (masking), softmax, và dropout vào một CUDA kernel duy nhất để loại bỏ các lượt truy xuất vòng lặp dư thừa đến HBM.

Phương pháp (kèm các phương trình chính)

Cho $\mathbf{Q},\mathbf{K},\mathbf{V}\in {\mathbb{R}}^{N\times d}$, FlashAttention thiết lập kích thước khối $B_c=\lceil M/(4d)\rceil$ và $B_r=\min (\lceil M/(4d)\rceil ,d)$, trong đó $M$ là dung lượng SRAM.

Theo dõi Softmax tăng dần

Đối với một vector được phân rã thành các vector con $x^{(1)}$ và $x^{(2)}$, các thống kê động cập nhật thông qua: $$m(x)=\max (m(x^{(1)}),m(x^{(2)})),$$ $$\ell (x)=e^{m(x^{(1)})-m(x)}\ell (x^{(1)})+e^{m(x^{(2)})-m(x)}\ell (x^{(2)}).$$

Đối với khối $j$ của keys/values và khối $i$ của queries:

1. Tính điểm số chưa chuẩn hóa trên SRAM: ${\mathbf{S}}_{ij}={\mathbf{Q}}_i{\mathbf{K}}_j^{\top }\in {\mathbb{R}}^{B_r\times B_c}$.
2. Tính các giá trị thống kê cục bộ của khối: ${\tilde{m}}_{ij}=\mathrm{rowmax}({\mathbf{S}}_{ij})$, ${\tilde{\mathbf{P}}}_{ij}=\exp ({\mathbf{S}}_{ij}-{\tilde{m}}_{ij})$, và ${\tilde{\ell }}_{ij}=\mathrm{rowsum}({\tilde{\mathbf{P}}}_{ij})$.
3. Cập nhật các hệ số chuẩn hóa toàn cục động: $$m_i^{\mathrm{new}}=\max (m_i,{\tilde{m}}_{ij}),{\ell }_i^{\mathrm{new}}=e^{m_i-m_i^{\mathrm{new}}}{\ell }_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\ell }}_{ij}.$$
4. Cập nhật khối đầu ra động ${\mathbf{O}}_i\in {\mathbb{R}}^{B_r\times d}$: $${\mathbf{O}}_i\leftarrow \mathrm{diag}({\ell }_i^{\mathrm{new}}{)}^{-1}\left(\mathrm{diag}({\ell }_i)e^{m_i-m_i^{\mathrm{new}}}{\mathbf{O}}_i+e^{{\tilde{m}}_{ij}-m_i^{\mathrm{new}}}{\tilde{\mathbf{P}}}_{ij}{\mathbf{V}}_j\right).$$

Phân tích độ phức tạp

• Độ phức tạp IO: Attention tiêu chuẩn yêu cầu $\Theta (Nd+N^2)$ lượt truy cập HBM. FlashAttention chỉ yêu cầu $\Theta (N^2{d}^2{M}^{-1})$ lượt truy cập HBM. Các tác giả chứng minh rằng không có thuật toán exact attention nào có thể đạt mức tiệm cận $o(N^2{d}^2{M}^{-1})$ lượt truy cập trên mọi $M\in [d,Nd]$.
• FLOPs và Bộ nhớ: FlashAttention thực hiện $O(N^{2}d)$ FLOPs (với số FLOP nhiều hơn một chút trong lượt truyền ngược do phải tính toán lại) và chỉ yêu cầu $O(N)$ bộ nhớ phụ ngoài đầu vào và đầu ra.

Mở rộng: Block-Sparse FlashAttention

Bằng cách giới hạn việc tính toán trong các khối khác không theo một mặt nạ thưa định trước $\mathbf{M}\in \{0,1{\}}^{N/B_r\times N/B_c}$, thuật toán đạt độ phức tạp IO là $\Theta (Nd+N^2{d}^2{M}^{-1}s)$, trong đó $s$ là tỷ lệ các khối khác không.

Kết quả (kèm các số liệu quan trọng nhất)

• Tăng tốc huấn luyện:
  • BERT-large (độ dài chuỗi 512): Vượt qua kỷ lục tốc độ huấn luyện MLPerf 1.1 tới 15% (17.4 phút so với 20.0 phút trên cấu hình $8\times$ A100).
  • GPT-2 (độ dài chuỗi 1K): Đạt tốc độ nhanh hơn tới 3.5 lần so với HuggingFace và tới 1.8 lần so với Megatron-LM.
  • Long-Range Arena (LRA): Đạt tốc độ nhanh hơn 2.4 lần so với attention tiêu chuẩn; Block-sparse FlashAttention đạt mức tăng tốc 2.8 lần.
• Khả năng xử lý ngữ cảnh dài hơn:
  • GPT-2 với ngữ cảnh 4K: Huấn luyện nhanh hơn Megatron-LM ở ngữ cảnh 1K trong khi cải thiện perplexity tập validation thêm 0.7 điểm (17.5 so với 18.2).
  • Phân loại văn bản dài: Mở rộng ngữ cảnh từ 512 lên 16K mang lại mức tăng 4.3 điểm micro-$F_1$ trên MIMIC-III (52.8 lên 57.1) và tăng tới 8.5 điểm trên ECtHR (72.2 lên 80.7 ở mức 8K).
  • Path-X và Path-256: FlashAttention là Transformer đầu tiên giải quyết được Path-X (độ dài chuỗi 16K) với độ chính xác 61.4%. Block-sparse FlashAttention đạt độ chính xác 63.1% trên Path-256 (độ dài chuỗi 64K).

Hạn chế và các câu hỏi mở

• Độ phức tạp kỹ thuật CUDA: Việc triển khai các kernel có nhận thức IO đòi hỏi kỹ thuật lập trình CUDA cấp thấp tùy chỉnh thay vì dùng PyTorch chuẩn hay các lớp trừu tượng của trình biên dịch, khiến việc chuyển đổi giữa các kiến trúc phần cứng gặp nhiều thách thức.
• Multi-GPU Attention: Lịch trình IO tối ưu trên một node đơn lẻ không tự động mở rộng sang attention phân tán trên nhiều GPU qua các kết nối liên mạng (NVLink, InfiniBand).
• Cận dưới của Exact Attention: Dù đã tối ưu trên các giá trị $M$ nói chung, việc chứng minh cận dưới độ phức tạp tham số hóa chi tiết theo $M$ vẫn là một bài toán mở.