FlashAttention (VI)

FlashAttention — là một thuật toán mang tính cách mạng để tính toán cơ chế attention, được phát triển nhằm tăng tốc đáng kể quá trình huấn luyện và inference của các mô hình ngôn ngữ lớn (LLM) trong khi vẫn đảm bảo độ chính xác tính toán đầy đủ. Thuật toán được giới thiệu lần đầu vào năm 2022 bởi nhóm các nhà nghiên cứu từ Đại học Stanford dưới sự dẫn dắt của Tri Dao.

Ý tưởng cốt lõi của FlashAttention là tái tổ chức các phép tính dựa trên hệ thống phân cấp bộ nhớ GPU, giúp giảm thiểu số lần truy cập vào bộ nhớ chậm và loại bỏ điểm nghẽn chính của cơ chế attention tiêu chuẩn.

Vấn đề của cơ chế attention tiêu chuẩn

Cơ chế self-attention tiêu chuẩn trong các transformer được tính theo công thức: $\text{Attention}(Q,K,V)=\text{softmax}\left(\frac{Q{K}^T}{\sqrt{d_k}}\right)V$ trong đó Q, K, V lần lượt là các ma trận truy vấn, khóa và giá trị.

Vấn đề chính của cách tiếp cận này là độ phức tạp bậc hai về thời gian và bộ nhớ (O(N²)) so với độ dài chuỗi N. Với cách triển khai đơn giản, cần phải tính toán và lưu trữ trong bộ nhớ GPU toàn bộ ma trận attention S có kích thước N×N, dẫn đến hai vấn đề nghiêm trọng:

1. Tiêu thụ bộ nhớ lớn: Việc lưu trữ ma trận N×N trở nên bất khả thi khi làm việc với các ngữ cảnh dài.
2. Các thao tác vào/ra (IO): Điểm nghẽn chính không phải là số lượng phép tính số học mà là các lần truy cập liên tục vào bộ nhớ chậm của GPU.

Hệ thống phân cấp bộ nhớ GPU

Để hiểu vấn đề, cần phân biệt hai loại bộ nhớ trong GPU (lấy ví dụ NVIDIA A100):

• SRAM (bộ nhớ tĩnh): Bộ nhớ nội trú nhanh, dung lượng nhỏ (~20 MB) với băng thông cực lớn (lên đến 19 TB/s).
• HBM (bộ nhớ băng thông cao): Bộ nhớ chậm, dung lượng lớn (40–80 GB) với băng thông thấp hơn nhiều (khoảng 1,5 TB/s).

Sự bất đối xứng này khiến thuật toán attention tiêu chuẩn bị giới hạn bởi băng thông bộ nhớ (memory-bound), vì nó liên tục đọc và ghi các ma trận lớn từ HBM chậm, đây chính là nguồn gốc chủ yếu gây ra độ trễ.

Những đổi mới chính của FlashAttention

FlashAttention là một thuật toán nhận thức IO (IO-aware), giải quyết vấn đề bằng cách giảm thiểu số lần truy cập HBM. Điều này đạt được nhờ ba kỹ thuật chính.

Tiling và xử lý theo khối

Thay vì xử lý toàn bộ ma trận một lúc, FlashAttention chia các ma trận đầu vào Q, K, V thành các khối nhỏ (tile) vừa khớp với SRAM nhanh. Thuật toán tuần tự nạp các khối này, thực hiện toàn bộ các phép tính attention cho chúng và cập nhật kết quả cuối cùng mà không lưu toàn bộ ma trận attention vào HBM chậm.

Tính toán Softmax trực tuyến

Đột phá kỹ thuật quan trọng là tính toán Softmax "trực tuyến". Softmax tiêu chuẩn yêu cầu biết tất cả các phần tử của vector đầu vào để chuẩn hóa. FlashAttention sử dụng một thuật toán sửa đổi cho phép tính Softmax theo từng phần. Nó duy trì hai giá trị trung gian (giá trị cực đại hiện tại và tổng các số mũ), được cập nhật khi các khối mới được xử lý, cho phép thu được kết quả chính xác mà không cần truy cập toàn bộ ma trận cùng một lúc.

Hợp nhất các thao tác thành một kernel CUDA duy nhất

Tất cả các thao tác attention (nhân ma trận QKᵀ, masking, Softmax, nhân với V) được tổng hợp thành một fused kernel CUDA duy nhất. Điều này giảm đáng kể số lần đọc/ghi vào HBM: thay vì nhiều lần duyệt qua toàn bộ ma trận, thuật toán nạp một khối vào SRAM một lần, thực hiện tất cả các phép tính và chỉ ghi kết quả cuối cùng.

Hiệu quả lý thuyết và thực tiễn

Độ phức tạp và tính tối ưu

FlashAttention giảm mức tiêu thụ bộ nhớ từ O(N²) xuống O(N), đảm bảo khả năng mở rộng tuyến tính. Đã được chứng minh rằng độ phức tạp IO của thuật toán là tối ưu về mặt lý thuyết cho việc tính toán attention trong hệ thống phân cấp bộ nhớ hai tầng, tức là không thể thực hiện attention chính xác nhanh hơn mà không thay đổi phần cứng.

Kết quả thực nghiệm

Phiên bản đầu tiên của FlashAttention đã cho thấy những cải tiến đáng kể:

• Tăng tốc:
  • BERT-large (độ dài 512): Tăng tốc huấn luyện 15%.
  • GPT-2 (độ dài 1K): Tăng tốc gấp 3 lần.
  • Các tác vụ Long-Range Arena (1K-4K): Tăng tốc gấp 2,4 lần.
• Tiết kiệm bộ nhớ: Tiết kiệm bộ nhớ lên đến 20 lần so với các triển khai cơ sở chính xác.
• Cải thiện chất lượng mô hình: Nhờ khả năng làm việc với các ngữ cảnh dài hơn, FlashAttention không chỉ không làm giảm mà còn cải thiện chất lượng mô hình. Ví dụ, perplexity của GPT-2 cải thiện 0,7 điểm, và độ chính xác trong các tác vụ phân loại tài liệu dài tăng 6,4 điểm.

Sự tiến hóa và các phát triển tiếp theo

Thành công của FlashAttention đã mở đầu cho một loạt thuật toán định hướng phần cứng.

FlashAttention-2 (2023)

Phiên bản thứ hai nhắm đến việc khai thác đầy đủ hơn tài nguyên GPU. Trong FlashAttention gốc, hiệu suất trên NVIDIA A100 chỉ đạt 25–40% mức tối đa. FlashAttention-2 đã đưa ra các cải tiến về song song hóa tính toán, cho phép:

• Đạt tốc độ tăng gấp đôi so với phiên bản đầu tiên.
• Tăng mức sử dụng GPU lên 50–73% mức tối đa lý thuyết.
• Mở rộng hỗ trợ lên đến các đầu attention kích thước 256, cũng như cho kiến trúc Multi-Query Attention (MQA).

FlashAttention-3 (2024)

Phiên bản thứ ba được tối ưu hóa đặc biệt cho kiến trúc GPU NVIDIA Hopper (H100). Nó tận dụng các khả năng phần cứng mới như tính bất đồng bộ của Tensor Cores và hỗ trợ FP8, cho phép:

• Đạt thêm tốc độ tăng 1,5–2 lần so với FlashAttention-2.
• Đạt hiệu suất lên đến 740 TFLOPS trên FP16 và gần 1,2 PFLOPS trên FP8.

Các giải pháp chuyên biệt

Các ý tưởng của FlashAttention đã được phát triển trong các dự án khác:

• FlashInfer (2025): Một engine attention có thể tùy chỉnh, được tối ưu hóa đặc biệt cho các tác vụ inference của LLM. Nó tập trung vào việc làm việc hiệu quả với KV-cache trong chế độ tạo sinh luồng.
• FlashMLA (2024): Triển khai attention với nén bộ nhớ cache ngữ cảnh (latent attention), cho phép tiết kiệm bộ nhớ trên các chuỗi rất dài với mức mất thông tin tối thiểu.

Tác động đến ngành công nghiệp và hệ sinh thái

FlashAttention đã trở thành một đột phá nền tảng và nhanh chóng trở thành tiêu chuẩn ngành cho việc huấn luyện và inference LLM hiệu quả. Nó đã được tích hợp vào các thư viện quan trọng như PyTorch và Hugging Face, và được sử dụng trong hầu hết các mô hình ngôn ngữ lớn (LLaMA, MPT, Falcon, Claude và các mô hình khác).

Chính FlashAttention và các phiên bản tiếp theo của nó đã đóng vai trò quyết định trong việc mở rộng cửa sổ ngữ cảnh của các mô hình ngôn ngữ: từ 2–4 nghìn token (GPT-3) lên đến 128 nghìn token (GPT-4) và thậm chí lên đến hàng triệu token trong các mô hình thực nghiệm. Thuật toán đã loại bỏ một trong những rào cản chính trên con đường mở rộng quy mô transformer, mở ra những khả năng mới cho các ứng dụng AI, từ phân tích tài liệu dài đến hiểu biết đa phương thức.

Liên kết

• Kho lưu trữ chính thức của FlashAttention trên GitHub

Tài liệu tham khảo

• 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.