Blog / AI / Deep Article
FlashAttention tăng tốc attention bằng cách giảm truy cập bộ nhớ GPU. Tìm hiểu tiling, online softmax, giới hạn thực tế và cách đo hiệu năng.
Bạn gửi một tài liệu dài cho chatbot và chờ câu trả lời. GPU rất mạnh, vậy vì sao máy vẫn chậm? Một nguyên nhân là mô hình phải chuyển quá nhiều dữ liệu qua bộ nhớ, dù bản thân các phép tính diễn ra nhanh. FlashAttention giải quyết đúng điểm nghẽn ấy: tổ chức lại phép attention để không phải ghi các ma trận trung gian khổng lồ vào bộ nhớ GPU, trong khi vẫn giữ attention đầy đủ. Vấn đề là không phải cứ tăng số phép tính mỗi giây thì câu trả lời sẽ đến nhanh hơn. Cần nhìn vào cả những lần đọc và ghi dữ liệu không cần thiết.
Bài viết phân tích cơ chế, bằng chứng, giới hạn và cách kiểm chứng hiệu năng thực tế; không coi mức tăng tốc trong phòng thí nghiệm là lời hứa cho mọi mô hình.
Trong Transformer, attention giúp mỗi vị trí tổng hợp thông tin từ các vị trí liên quan. Đầu vào được biến thành ba nhóm vector: query (Q), key (K) và value (V). Query so với các key để tạo điểm tương quan; softmax biến điểm thành trọng số; các trọng số được dùng để kết hợp value. Công thức rút gọn:
Attention(Q, K, V) = softmax(QKᵀ / √d) V
Với N token, ma trận QKᵀ có N × N phần tử cho mỗi đầu attention. Chuỗi dài gấp đôi tạo khoảng bốn lần số phần tử. Ví dụ, 8.192 token tạo ma trận 67.108.864 phần tử; lưu mỗi phần tử bằng hai byte cần khoảng 128 MiB cho một tensor, chưa tính ma trận xác suất, các lớp và đầu attention khác. Đây là phép tính minh họa, không phải mức VRAM của một mô hình cụ thể.
Quan trọng hơn, số phép tính và kích thước dữ liệu trung gian là hai vấn đề khác nhau. FlashAttention cải thiện mạnh vấn đề thứ hai, nhưng không xóa bỏ số cặp token phải xét trong attention đầy đủ.
GPU có nhiều tầng nhớ. HBM chứa lượng dữ liệu lớn, còn SRAM trên chip nhỏ hơn nhưng có thể truy cập nhanh hơn. Một cách triển khai attention thông thường tạo QKᵀ, ghi ma trận điểm ra HBM, đọc lại để tính softmax, ghi ma trận xác suất rồi đọc tiếp để nhân với V. Tính toán có thể nhanh, nhưng các lần truyền dữ liệu lớn trở thành nút thắt.
Thực tế là hai chương trình thực hiện cùng phép toán có thể chạy nhanh chậm rất khác nhau. Nghiên cứu của Tri Dao và cộng sự gọi cách thiết kế chú ý đến lưu lượng giữa các tầng nhớ là IO-aware. Nhóm tác giả không chỉ đếm phép nhân cộng mà còn xem xét những lần đọc/ghi HBM. Cách nhìn này giải thích vì sao một phương pháp attention xấp xỉ với ít phép toán hơn chưa chắc nhanh hơn khi đo thời gian thực: cấu trúc truy cập bộ nhớ, việc khởi chạy kernel và điều phối cũng có chi phí.
Ý tưởng cốt lõi là chia khối (tiling). Thay vì tạo toàn bộ ma trận N × N rồi đưa nó qua nhiều phép toán tách rời, FlashAttention lấy những khối nhỏ của Q, K và V có thể xử lý trong vùng nhớ nhanh trên chip. Với từng khối, nó tính điểm attention, cập nhật kết quả và chỉ lưu những thông tin cần thiết. Không cần ghi toàn bộ ma trận điểm và ma trận xác suất ra HBM.
Hãy hình dung bạn phải tổng hợp số liệu từ hàng triệu dòng. Cách thông thường là lập bảng phụ khổng lồ, lưu bảng, mở lại, chuẩn hóa từng hàng rồi mới cộng ra kết quả. Cách tổ chức theo khối là xử lý từng phần nhỏ và duy trì một bản tổng hợp có thể cập nhật. Hình dung này giúp hiểu dòng dữ liệu, nhưng không thay thế được toán học softmax: nếu chỉ cộng các phần lại một cách ngây thơ, kết quả sẽ sai vì mỗi phần có hệ số chuẩn hóa khác nhau.
Câu hỏi khó nằm ở đây: nếu không nhìn thấy cả hàng điểm cùng lúc, làm sao chuẩn hóa softmax đúng? FlashAttention cần một kỹ thuật thứ hai, quan trọng không kém tiling: cập nhật softmax dần mà vẫn giữ đúng phép chuẩn hóa toàn hàng.
Softmax biến một dãy điểm sᵢ thành trọng số bằng exp(sᵢ) chia cho tổng các giá trị mũ. Để ổn định số học, ta thường trừ điểm lớn nhất trước khi tính hàm mũ. Nhưng nếu chỉ có một khối điểm, làm sao biết cực đại của cả hàng? Đây là chỗ khó nhất khi chia khối.
FlashAttention giữ ba thông tin gọn: cực đại đang thấy, tổng chuẩn hóa đã điều chỉnh và đầu ra tích lũy. Khi khối mới đến, thuật toán tìm cực đại mới, điều chỉnh đóng góp cũ theo hệ số tương ứng rồi cộng phần mới. Nhờ vậy, kết quả tương ứng với softmax trên cả hàng mà không cần lưu cả ma trận điểm trong HBM.
“Exact” ở đây nghĩa là không chủ ý bỏ qua tương tác token như attention thưa hoặc xấp xỉ. Nó không có nghĩa mọi kết quả số thực đều giống nhau từng bit: thứ tự phép tính có thể gây sai khác làm tròn nhỏ.
Khi huấn luyện mô hình, giai đoạn lan truyền ngược cần các thông tin liên quan đến attention để tính gradient. Một cách làm là lưu nhiều tensor trung gian trong lượt đi tới, rồi đọc chúng trở lại ở lượt đi lui. Nhưng chính những tensor lớn ấy chiếm bộ nhớ và tạo thêm lưu lượng đọc/ghi.
FlashAttention chọn một đánh đổi khác: lưu đầu ra cùng các thống kê chuẩn hóa cần thiết, rồi tính lại một số khối attention trong lượt lan truyền ngược. Tính lại nghĩa là tăng một phần số phép toán. Điều nghe có vẻ nghịch lý là tổng thời gian vẫn có thể giảm nếu chi phí tính lại thấp hơn chi phí ghi rồi đọc ma trận lớn từ HBM. Đây là ví dụ rõ về việc tối ưu thời gian thực tế thay vì tối thiểu hóa mọi phép nhân cộng.
Đánh đổi này phụ thuộc phần cứng và khối lượng công việc. Nếu một tác vụ vốn đã bị giới hạn bởi năng lực tính toán, thêm phép tính không chắc có lợi. Còn với suy luận, vốn không có lượt lan truyền ngược, lợi ích chính đến từ đường tính attention ở lượt đi tới, không phải từ chiến lược tính lại gradient.
Một hiểu lầm phổ biến là FlashAttention biến attention đầy đủ thành thuật toán tuyến tính theo độ dài chuỗi. Không phải vậy. Với attention dày đặc, số cặp query–key cần xét vẫn tăng theo N². Cải thiện chính là không lưu ma trận trung gian N × N trong HBM và giảm nhiều lần truyền dữ liệu. Bộ nhớ phụ cho phép attention có thể tăng tuyến tính theo chiều dài chuỗi, nhưng không có nghĩa toàn bộ mô hình chỉ cần bộ nhớ tuyến tính trong mọi trường hợp.
Cũng cần phân biệt bản FlashAttention dày đặc với block-sparse FlashAttention: biến thể thưa bỏ qua một số khối theo cấu trúc và có đánh đổi riêng. Kết quả của biến thể xấp xỉ không thể gán cho thuật toán dày đặc. Nếu chatbot chậm vì tải trọng số, KV cache, mạng hoặc hàng đợi, thay kernel attention không tự giải quyết các vấn đề đó.
Bài báo năm 2022 báo cáo 15% cải thiện thời gian huấn luyện đầu-cuối trên BERT-large, 3 lần trên một thiết lập GPT-2 và 2,4 lần trên Long Range Arena. Những kết quả này thuộc các mô hình, GPU, độ dài chuỗi và đối chứng cụ thể. Không nên biến chúng thành lời hứa “mọi chatbot nhanh gấp ba”.
Mức tăng tốc riêng của phép attention còn có thể cao hơn số đo toàn hệ thống. Nhưng kernel chỉ là một phần của hệ thống. Nếu thời gian chủ yếu nằm ở tải mô hình, truy xuất dữ liệu hay sinh token nối tiếp, tăng tốc một kernel sẽ không làm toàn bộ ứng dụng nhanh tương ứng. Đây là trực giác của định luật Amdahl.
Một hệ thống có thể xử lý chuỗi dài hơn nhờ giảm bộ nhớ trung gian; điều đó không tự chứng minh mô hình hiểu tốt hơn, trích dẫn chính xác hơn hoặc ít bịa thông tin hơn.
Tài liệu Hugging Face Transformers mô tả các attention backend như flash_attention_2 và sdpa. Nhà phát triển có thể chọn backend phù hợp mà không nhất thiết viết kernel CUDA từ đầu. Nhưng tên backend không phải bảo đảm tốc độ: cần xem GPU, phiên bản thư viện, dtype, kiến trúc mô hình và attention mask có được hỗ trợ không.
Một chi tiết dễ bỏ sót: các backend dùng định dạng mask khác nhau. Tài liệu cảnh báo khi tự đăng ký hàm attention mà không đăng ký hàm tạo mask tương ứng, một số ràng buộc có thể bị bỏ qua. Đây là vấn đề tính đúng đắn, không chỉ hiệu năng.
Thực tế là hệ thống có thể chuyển sang backend dự phòng khi kernel mong muốn không phù hợp. Vì vậy, trước khi công bố benchmark, hãy kiểm tra triển khai thực tế đã được dùng.
Giả sử nhóm phát triển chatbot hỏi đáp tài liệu nhận thấy VRAM tăng mạnh khi prompt dài. Họ nghi attention là một nút thắt. Cách kiểm tra là giữ nguyên mô hình, prompt, độ dài đầu ra và số yêu cầu đồng thời, chỉ đổi backend nếu được hỗ trợ. Sau đó đo thời gian xử lý prompt, tốc độ sinh token, VRAM cực đại, tổng độ trễ và thông lượng khi nhiều người dùng cùng truy cập.
Câu hỏi quan trọng: cải thiện có còn khi chạy tải thật, hay chỉ xuất hiện trong phép đo một kernel? Cần ghi rõ GPU, dtype, phiên bản thư viện và backend thực sự được chọn.
Đây là kịch bản hướng dẫn đo lường, không phải benchmark AIWiki đã chạy. Không có số liệu giả định nào được trình bày như kết quả thực nghiệm.
FlashAttention cho thấy cách dữ liệu di chuyển có thể quan trọng không kém lượng tính toán. Chia khối, chuẩn hóa softmax theo dòng và tránh lưu ma trận trung gian giúp attention đầy đủ dùng bộ nhớ hiệu quả hơn; trong huấn luyện, tính lại một số giá trị còn có thể rẻ hơn lưu rồi đọc.
Bài học quan trọng: số phép tính attention dày đặc vẫn tăng theo bình phương độ dài chuỗi; tốc độ kernel không đồng nghĩa tốc độ toàn sản phẩm; và mức cải thiện phải được đo trên phần cứng, mô hình, dữ liệu thật. Hiểu ba điều này giúp đánh giá các tuyên bố hiệu năng mà không bị cuốn theo một con số benchmark.