Flash Attention Là Gì?
Flash Attention là thuật toán tối ưu hóa attention computation trong Transformer, giúp giảm đáng kể việc sử dụng bộ nhớ GPU và tăng tốc độ huấn luyện. Được giới thiệu bởi Tri Dao năm 2022, Flash Attention nhanh hơn attention thông thường 2-4 lần và tiết kiệm bộ nhớ tới 10-20 lần cho các chuỗi dài.
Mình nghĩ đây là một trong những cải tiến hạ tầng AI quan trọng nhất gần đây. Không phô trương như các model mới, nhưng gần như mọi mô hình lớn hiện nay đều dùng nó, từ GPT-4 đến Llama 3.
Vấn Đề Của Attention Chuẩn
Trong Attention Mechanism chuẩn, mỗi token cần so sánh với mọi token khác. Nếu chuỗi đầu vào có N token, ma trận attention có kích thước N x N. Với N = 8192 (khá phổ biến hiện nay), ma trận đó chiếm hàng trăm MB bộ nhớ GPU.
Vấn đề lớn hơn là cách GPU hoạt động. GPU có hai loại bộ nhớ: HBM (bộ nhớ chính, dung lượng lớn nhưng chậm) và SRAM (bộ nhớ trên chip, rất nhanh nhưng dung lượng nhỏ). Attention chuẩn liên tục đọc ghi giữa HBM và SRAM, và đó chính là bottleneck, không phải bản thân phép tính.
Flash Attention Giải Quyết Thế Nào?
Flash Attention thay đổi cách tính attention để tối thiểu hóa việc truy cập HBM. Ý tưởng cốt lõi là tiling chia ma trận attention thành các block nhỏ vừa đủ để nằm trong SRAM. Thay vì tính toàn bộ ma trận rồi ghi lại HBM, Flash Attention tính từng block, xử lý xong rồi mới ghi kết quả cuối cùng.
Điều này giống như nấu ăn. Thay vì lấy từng nguyên liệu từ tủ lạnh (HBM) ra bếp (SRAM) rồi cất lại, lấy nguyên liệu khác ra, bạn lấy tất cả nguyên liệu cho một món ra một lần, nấu xong rồi mới lấy món tiếp theo. Ít chuyến đi lại hơn, nấu nhanh hơn.
Ba Kỹ Thuật Chính Của Flash Attention
1. Tiling (Chia Khối)
Thay vì tải toàn bộ ma trận Q, K, V vào bộ nhớ, Flash Attention chia chúng thành các block nhỏ. Mỗi block được xử lý độc lập trong SRAM. Kết quả từng block được gộp lại bằng một kỹ thuật gọi là online softmax, cho phép tính softmax từng phần mà không cần toàn bộ ma trận.
2. Recomputation (Tính Lại Trong Backprop)
Thông thường, huấn luyện mô hình cần lưu ma trận attention trung gian để dùng trong backward pass. Ma trận này rất lớn. Flash Attention không lưu nó. Thay vào đó, trong backward pass, nó tính lại ma trận attention từ Q, K, V. Nghe có vẻ phí, nhưng việc tính lại trong SRAM nhanh hơn nhiều so với đọc ma trận khổng lồ từ HBM.
3. Kernel Fusion (Gộp Phép Tính)
Attention chuẩn thực hiện nhiều phép tính tuần tự: nhân ma trận, scale, mask, softmax, rồi nhân ma trận nữa. Mỗi bước đều đọc ghi HBM. Flash Attention gộp tất cả thành một kernel duy nhất, chạy hoàn toàn trong SRAM, loại bỏ hoàn toàn việc đọc ghi HBM giữa các bước.
Flash Attention 2 và 3
Tri Dao đã liên tục cải tiến thuật toán. Flash Attention 2 (2023) tối ưu thêm việc phân bổ thread trên GPU, giảm register pressure và cải thiện parallelism. Nó nhanh hơn phiên bản đầu tiên khoảng 2 lần, đặc biệt hiệu quả với chuỗi dài.
Flash Attention 3 (2024) khai thác kiến trúc Hopper GPU của NVIDIA, sử dụng asynchronous data transfer và FP8 support. Nó đạt tốc độ 1.5-2 PFLOP/s trên H100 GPU, gần như tận dụng tối đa phần cứng. Đây là lý do chính khiến các mô hình AI năm 2025-2026 huấn luyện nhanh hơn nhiều so với trước.
Tại Sao Flash Attention Quan Trọng?
Không phải nói quá khi Flash Attention thay đổi cách ngành AI vận hành. Trước Flash Attention, huấn luyện mô hình với context length 32K token cực kỳ tốn kém, chỉ các phòng lab lớn mới làm được. Sau Flash Attention, context 128K, thậm chí 1 triệu token trở nên khả thi.
Mỗi lần tăng context length, mô hình hiểu được tài liệu dài hơn, duyệt codebase lớn hơn, và trả lời chính xác hơn. Không có Flash Attention, Context Window 1 triệu token của Gemini hay Claude sẽ tốn GPU gấp 10 lần và chậm gấp 5 lần.
Flash Attention Vs các kỹ Thuật Khác
Nhiều người nhầm Flash Attention với các phương pháp approximated attention như Linear Attention hay Sparse Attention. Khác biệt quan trọng: Flash Attention tính chính xác cùng kết quả với attention chuẩn, không xấp xỉ. Nó tối ưu cách tính, không thay đổi kết quả.
Đây là điểm cộng lớn. Bạn không cần lo về độ chính xác. Kết quả bằng hệt attention thường, chỉ là nhanh hơn và ít RAM hơn. Các phương pháp approximated attention thì đánh đổi accuracy lấy tốc độ, và thực tế cho thấy chúng thường kém hơn.
Ứng Dụng Thực Tế
Gần như mọi mô hình AI lớn hiện nay đều dùng Flash Attention. OpenAI, Anthropic, Google, Meta, tất cả đều tích hợp nó vào pipeline huấn luyện. PyTorch 2.0+ hỗ trợ Flash Attention mặc định qua scaled_dot_product_attention. Bạn chỉ cần gọi hàm, nó tự chọn backend tối ưu.
Nếu bạn tự huấn luyện mô hình, mình khuyên luôn bật Flash Attention từ đầu. Không có lý do gì để tắt nó cả. Nó compatible với mọi kiến trúc Transformer, không cần thay đổi code, và lợi ích quá rõ ràng.
Kết Luận
Flash Attention là một trong những tối ưu hạ tầng AI có tác động lớn nhất thập kỷ này. Nó chứng minh rằng đôi khi cải thiện cách tính toán quan trọng không kém gì thiết kế mô hình mới. Nhờ nó, mô hình AI xử lý được chuỗi dài hơn, huấn luyện nhanh hơn, và tốn ít GPU hơn.
Khi ngành AI tiến về phía các mô hình có context cực dài và multimodal, Flash Attention sẽ còn quan trọng hơn nữa. Hiểu được nó là hiểu được một phần lý do tại sao AI năm 2026 tiến bộ nhanh đến vậy.
