FlashAttention 4: Cơ chế Attention Nhanh và Tiết Kiệm Bộ Nhớ cho LLMs

Giới thiệu

Cơ chế chú ý nhân chấm được tỉ lệ hóa (Scaled-dot-product attention - SDPA) chiếm ưu thế trong thời gian suy luận và ngân sách năng lượng của các mô hình ngôn ngữ lớn (LLM): gần như tất cả các thực thi toán tử xảy ra trong một phép nguyên thuỷ duy nhất. Chú ý lấy các truy vấn (Q) và nhân chúng với các khóa (K). Sau đó, nó chuẩn hóa kết quả thông qua phép softmax trước khi nhân với các giá trị (V) để tạo ra đầu ra. Nó bị giới hạn bởi bộ nhớ: đọc Q/K/V nhiều lần từ bộ nhớ băng thông cao, sau đó liên tục ghi các ô trung gian.

Các dòng FlashAttention giảm tắc nghẽn này bằng cách giữ nhiều dữ liệu hơn trên chip trong thời gian dài hơn nhờ tận dụng các bộ xử lý đa luồng GPU (SMs). Lần đầu được công khai mã nguồn vào năm 2022, FlashAttention đã thiết kế lại quy trình attention để tính toán softmax trực tiếp. Chỉ một vài hàng của Q × K được lưu trên chip tại một thời điểm nhất định.

Bài viết này nhằm mục đích là một nghiên cứu chuyên sâu hướng tới người thực hành về FlashAttention 4 (FA4). Chúng tôi sẽ tóm tắt những gì đã thay đổi so với các phiên bản trước và cung cấp một số hướng dẫn về việc áp dụng/đánh giá hiệu năng. Mục tiêu của chúng tôi là cung cấp cho các kỹ sư hạ tầng LLM, kỹ sư kernel và các nhóm nền tảng ML đủ thông tin để biết khi nào – và liệu – FA4 nên được sử dụng trong hệ thống của họ.

Điểm chính

• A4 được thiết kế như một nhân chú ý Blackwell đầu tiên, nhằm triển khai trên GPU SM100. Nó có một quy trình 5 giai đoạn chuyên biệt theo warp, cho phép tăng khả năng chồng lấp và tái sử dụng trên chip so với các thế hệ FlashAttention trước.
• Hiệu suất được cải thiện trong softmax thông qua hai kỹ thuật chính: thực thi phần mềm exp2() ngay trên các lõi CUDA để giảm sự cạnh tranh trên SFU, và điều chỉnh lại tỉ lệ trực tuyến thích ứng để tránh các phép tính tỉ lệ không cần thiết trong khi vẫn duy trì độ ổn định.
• Vẫn còn một số tính năng thiếu: FA4 ưu tiên tiến, hiện chưa có các triển khai cho bước lùi, varlen và GQA/MQA, điều này giới hạn việc sử dụng nó cho huấn luyện và một số kiến trúc mô hình.
• Áp dụng FA4 một cách thực dụng: sử dụng trước cho suy luận Blackwell, đặt nó sau các bước kiểm tra đánh giá + kiểm tra độ chính xác + các phương án dự phòng, và giữ FA3 (Hopper)/FA2 (Ampere/Ada) làm mặc định ổn định cho các nhu cầu sản xuất rộng rãi.

Tại sao các lõi attention vẫn chiếm ưu thế trong chi phí của LLM

Mặc dù GPU cung cấp FLOPS cao, các mô hình transformer thường bị giới hạn băng thông bộ nhớ do cơ chế attention phải tải Q, K và V từ bộ nhớ toàn cục nhiều lần. Để tính toán đầu ra của attention cho mỗi token truy vấn, mỗi token phải nhân vô hướng với hàng nghìn token khóa, chuẩn hóa và điều chỉnh kết quả này, và cuối cùng truy xuất các giá trị. Các triển khai đơn giản tính toàn bộ ma trận Q×K và lưu trữ nó trong bộ nhớ. Điều này nhanh chóng vượt quá kích thước bộ nhớ trên chip, dẫn đến hiệu năng bị ảnh hưởng bởi các lần đọc ngoài chip.

FlashAttention giảm thiểu vấn đề này bằng cách chia ma trận attention thành các ô. Thay vì tạo ra toàn bộ Q×K, nó truyền các ô nhỏ của Q và K vào bộ nhớ chia sẻ, tính softmax trên ô đó, và chỉ giữ giá trị tối đa và tổng chạy cho mỗi hàng truy vấn. Điều này làm giảm đáng kể lưu lượng bộ nhớ, khiến kernel trở nên bị giới hạn bởi tính toán.

Tiến hóa FlashAttention tổng quan

Việc hiểu cách FlashAttention phát triển theo thời gian là rất hữu ích. Mỗi phiên bản đều tận dụng các tính năng phần cứng mới để cải thiện tốc độ hoặc hiệu quả. FA4 được tối ưu hóa đặc biệt cho GPU thế hệ Blackwell. Dưới đây là bảng các phiên bản FlashAttention cùng với những cải tiến chính mà chúng mang lại và tác động về hiệu suất:

Những gì mới trong FlashAttention 4

FlashAttention 4 tiếp tục mô hình tính toán theo ô mà các phiên bản FlashAttention trước đã thiết lập: các tensor Q/K/V được xử lý từng ô một (từng khối một) để tối đa hóa tính cục bộ dữ liệu trên chip và giảm thiểu lưu lượng tới bộ nhớ toàn cục. Các instance kernel riêng lẻ (gọi là mảng luồng hợp tác, hoặc CTA) làm việc cùng nhau để tính toán một hoặc nhiều ô đầu ra bằng cách truyền các ô tương ứng của Q, K và V từ bộ nhớ toàn cục, sau đó chiếu chúng lên các đầu ra attention cuối cùng.

Các thành phần bên trong của FA4 bao gồm một lịch trình được pipeline hóa cao, cho phép chồng lấp nhiều giai đoạn tính toán attention nhằm cải thiện độ trễ và tối đa hóa việc sử dụng tài nguyên GPU. Trái ngược với FlashAttention 3, vốn chỉ có sự chồng lấp giữa hai giai đoạn tải/tính toán, FA4 phân chia tính toán thành nhiều tập hợp các giai đoạn đồng thời, với các phần khác nhau của pipeline được xử lý bởi các nhóm warp riêng biệt. Mỗi nhóm warp này chuyên môn hóa vào một nhiệm vụ duy nhất trong suốt “vòng đời” của một ô trong hàm attention.

Mô hình chuyên môn hóa warp này phân công các warp 32 luồng vào các nhiệm vụ cụ thể, chẳng hạn như:

• Di chuyển dữ liệu (tải): lấy các ô dữ liệu của truy vấn, khóa và giá trị từ bộ nhớ toàn cục vào bộ nhớ trong nhanh hơn trên chip.
• Các giai đoạn tính toán: Thực hiện điểm chú ý từng phần bằng cách sử dụng phần cứng nhân ma trận-tích lũy với các lõi tensor.
• Logic Softmax và chuẩn hóa: Các phép toán mũ + chuẩn hóa cho trọng số softmax trong cơ chế attention.
• Logic tái tỉ lệ và giảm: áp dụng các điều chỉnh số học hoặc kết hợp các kết quả từng phần trước khi tích lũy.
• Lời kết/hoạt động lưu trữ: ghi các ô kết thúc đã hoàn tất của tensor đầu ra trở lại bộ nhớ toàn cục.

Trực giác đằng sau mô hình thực thi cơ bản là duy trì các đường ống không đồng bộ cho mỗi nhóm warp, sao cho bộ lập lịch warp của GPU có thể nhanh chóng chuyển đổi giữa các nhóm warp khi các toán hạng nguồn trở nên khả dụng.

Số mũ phần mềm qua nhân CUDA

Một mẹo mới đặc biệt gây chú ý trong FA4 là cách nó tính toán hàm mũ được sử dụng trong softmax. Thông thường, các kernel sẽ tận dụng các SFU chuyên dụng trên GPU để thực hiện các phép toán tốn kém như exp(). Tuy nhiên, vì số lượng SFU trên mỗi SM khá ít, chúng có thể dễ dàng trở thành nút thắt trong hàng đợi.

FA4 vượt qua giới hạn này bằng cách sử dụng các lõi CUDA thông thường để thực hiện các hàm số mũ gần đúng trong phần mềm. Cụ thể, FA4 chứa một triển khai tùy chỉnh của xấp xỉ đa thức exp2(x) với độ chính xác phần cứng. FA4 mô phỏng việc tính toán số mũ bằng đa thức bậc ba, sau đó có thể được thực thi song song trên nhiều lõi CUDA thay vì tập trung vào một vài SFU.

Kết luận: không còn phải chờ SFU để thực hiện softmax. Bằng cách xử lý exp() trên các nhân đa năng, FA4 cho phép warp tính toán softmax theo kịp với dòng xử lý.

Điều Chỉnh Lại Softmax Trực Tuyến Thông Minh Hơn (Mới trong FA4)

Một cải tiến khác trong FA4 liên quan đến cách mà tích lũy softmax duy trì sự ổn định về số học. Giống như tất cả các phiên bản FlashAttention đều thực hiện softmax "trực tuyến" – tích lũy các giá trị max() và tổng từng phần khi chúng đi qua chuỗi (để tránh tràn số). Các phiên bản trước liên tục cập nhật hệ số tỉ lệ đang chạy (thường là logit lớn nhất cho tới thời điểm đó). Chúng sẽ liên tục thay đổi tỷ lệ – bất cứ khi nào phát hiện giá trị lớn nhất mới, chúng sẽ điều chỉnh tỷ lệ để duy trì sự ổn định của softmax. FA4 giới thiệu một chính sách điều chỉnh tỷ lệ thích ứng. Nó chỉ thay đổi tỷ lệ nếu giá trị max mới mà bạn thấy lớn hơn đáng kể so với giá trị cũ. Theo báo cáo, điều này giảm tần suất các thao tác điều chỉnh tỷ lệ đi khoảng 10 lần. Hãy xem xét sơ đồ sau:

Sơ đồ so sánh hai phiên bản của việc tái tỉ lệ được sử dụng trong FlashAttention, cho phép tính toán attention ổn định và tiết kiệm bộ nhớ. Chiến lược gốc tái tỉ lệ đầu ra của mỗi khối nếu giá trị lớn nhất xuất hiện, đảm bảo tổng các giá trị nhỏ vẫn chính xác bằng cách chuẩn hóa các đóng góp từng bước. Với tái tỉ lệ thích ứng, chúng ta theo dõi các giá trị cực đại cục bộ và toàn cục, chỉ tái tỉ lệ khi có sự nhảy giá trị hoặc kích thước khối tăng lên. Điều này giảm tải trong khi vẫn ngăn chặn hiện tượng tràn số và cho phép tổng hợp các giá trị nhỏ một cách ổn định về mặt số học.

Khả năng tương thích và tình trạng hiện tại

Bảng này cung cấp tổng quan về khả năng tương thích dự kiến và "những gì hoạt động hiện nay" của FlashAttention 4 (FA4) trên toàn bộ hệ thống:

Phép truyền tiến (forward pass) của FA4 hiện đã được công khai rộng rãi (đã cam kết) trong cây mã nguồn mở. Có những bài đăng mở trong các dự án như kho FlashAttention của Lab, nơi người dùng yêu cầu hỗ trợ đầy đủ cho phép truyền ngược (backward pass)/varlen/group/GQA. Điều này ngụ ý rằng đây vẫn là một vấn đề công khai đang mở.

Hỗ trợ chuỗi có độ dài biến đổi (xử lý hiệu quả các chuỗi có độ dài khác nhau mà không cần điền đầy) và GQA (Grouped-Query Attention) hiện vẫn chưa được hỗ trợ trong quá trình lan truyền ngược. GQA là một thủ thuật attention, nơi nhiều đầu attention sử dụng cùng một phép chiếu key/value (đôi khi được sử dụng trong LLM để giảm sử dụng bộ nhớ). Các phiên bản trước của FlashAttention đã hỗ trợ các cấu hình kiểu MQA/GQA. Tuy nhiên, kernel chính của FA4 trong tiến trình lan truyền có thể sẽ giả định rằng số lượng đầu dùng cho mỗi truy vấn là giống nhau. Vì vậy, có lẽ nó chưa hỗ trợ đúng các đầu theo nhóm. Các mô hình với grouped-query attention (như một số biến thể của Llama2) sẽ không được FA4 hỗ trợ cho đến khi tính năng này được triển khai.

Tích hợp hệ sinh thái: FA4 vẫn chưa được nhiều framework/lớp cao hơn áp dụng vì nó còn mới. Một ví dụ là động cơ phục vụ vLLM (tối ưu hóa cao cho suy luận LLM), hiện tại không hỗ trợ FA4 – chỉ FlashAttention 2 và 3 được hỗ trợ. Chọn FA4 trong vLLM sẽ dẫn đến lỗi “phiên bản không được hỗ trợ” vì bộ chọn phiên bản của họ hiện tại chỉ cho phép bạn chọn giữa phiên bản 2 hoặc 3.

Cách Áp Dụng FlashAttention – Hướng Dẫn Quyết Định

Sử dụng bảng này để xác định phiên bản FlashAttention nào là tối ưu cho nhu cầu của bạn.

Câu hỏi thường gặp

FlashAttention 4 (FA4) đang cố gắng giải quyết vấn đề gì?

Nó giảm lưu lượng bộ nhớ HBM nặng của cơ chế chú ý bằng cách tính toán chú ý từng ô trên chip, tránh việc tạo ra toàn bộ ma trận QK⊤.

Thay đổi kiến trúc nổi bật nhất trong FA4 so với FA3 là gì?

FA4 sử dụng một đường ống ~5 giai đoạn chuyên dụng cho warp (tải, tính toán, softmax/chuẩn hóa, điều chỉnh/tổng hợp, lưu trữ) để tăng khả năng chồng lấp và tái sử dụng trên chip trên Blackwell (SM10.x).

Tại sao FA4 tính toán hàm mũ trên các lõi CUDA thay vì trên SFU?

Vì các SFU rất hiếm và có thể dễ dàng trở thành nút thắt cổ chai khi tính softmax, FA4 thay vào đó chuyển sang tính toán xấp xỉ bằng phần mềm của exp2() trên các lõi CUDA.

“Tái chuẩn trực tuyến thích ứng” thay đổi điều gì trong softmax?

Thay vì thay đổi tỉ lệ mỗi khi xuất hiện giá trị lớn nhất mới, FA4 chỉ thay đổi tỉ lệ khi giá trị lớn nhất đó tăng lên một lượng đáng kể. Điều này giảm bớt khối lượng công việc khi thay đổi tỉ lệ trong khi vẫn duy trì độ ổn định về mặt số học.

Khi nào một kỹ sư nên áp dụng FA4 trong sản xuất hiện nay?

Bắt đầu với Blackwell chỉ suy luận, độ dài cố định, attention chuẩn phía sau một cờ tính năng; giữ cuDNN/SDPA (và FA3/FA2 trên GPU cũ hơn) làm phương án dự phòng cho đến khi hỗ trợ backward/varlen/GQA trở nên ổn định.

Kết luận

FlashAttention 4 xây dựng trên mục tiêu của FlashAttention là giảm yêu cầu băng thông bộ nhớ của cơ chế attention bằng cách thực hiện nhiều công việc Q/K/V trực tiếp trên chip. Tuy nhiên, không giống như các phiên bản trước, FA4 thực hiện điều này với một pipeline 5 giai đoạn chuyên dụng theo warp dành riêng cho Blackwell và hai cải tiến chính về softmax (sử dụng hàm mũ phần mềm trên các nhân CUDA để giảm nghẽn SFU; và sử dụng điều chỉnh tỷ lệ trực tuyến thích ứng để loại bỏ các phép tính ổn định dư thừa). FA4 đã cho thấy sự cải thiện đáng kể về thông lượng khi thực hiện forward-pass so với cuDNN ở các chuỗi dài, nhưng vẫn giữ nguyên hướng forward-first và chỉ dành cho Blackwell. Hơn nữa, FA4 hiện nay còn thiếu nhiều tính năng quan trọng (thiếu backward, varlen, GQA, v.v.), và hỗ trợ trong hệ sinh thái chưa thực sự hoàn chỉnh. Điều rút ra hiện tại là nên xem xét các lộ trình FA4 cho suy luận Blackwell như một hướng có tiềm năng cao, nhưng cần áp dụng sau khi thực hiện các đánh giá nghiêm ngặt về hiệu năng, độ chính xác và cơ chế dự phòng. Đồng thời, các khối lượng công việc huấn luyện và phần cứng rộng hơn nên tiếp tục sử dụng các lộ trình mã FA3/FA2 đã trưởng thành cho môi trường sản xuất.

Tài liệu tham khảo

• FlashAttention-3: Chú ý nhanh và chính xác với bất đồng bộ và độ chính xác thấp
• FlashAttention 4: Phá vỡ rào cản Petaflop trong các nhân chú ý GPU
• Làm thế nào để áp dụng FA4 trên B200?
• Hỗ trợ cho quá trình truyền ngược FA4 trên Blackwell

Cảm ơn bạn đã học cùng cộng đồng Linux-Console.net.