Trong bài viết này
Từ được đánh dấu có giải thích. Di chuột hoặc chạm để xem, hay mở từ điển AI.
Một phép toán có thể nhanh hơn dù không giảm bậc số phép nhân. Nếu nó ghi ít dữ liệu trung gian xuống bộ nhớ lớn rồi đọc lại, GPU dành ít thời gian chờ dữ liệu hơn. FlashAttention đưa việc di chuyển dữ liệu vào trung tâm thiết kế attention.
Bạn nên hiểu công thức softmax(QK^T)V, ma trận và sự khác nhau giữa bộ nhớ trên chip với bộ nhớ thiết bị. Bài này bỏ qua nhiều chi tiết kernel để giải thích nguyên lý; các kích thước và điểm số là ví dụ giả định.
Ma trận trung gian là một khoản chi phí
Với chuỗi dài n, attention dày tạo n nhân n điểm trên mỗi đầu. Một cách triển khai trực tiếp tính ma trận điểm, lưu nó, tính softmax rồi nhân với V. Các tensor trung gian lớn phải đi qua bộ nhớ thiết bị, thường gọi là HBM trong bối cảnh GPU phù hợp.
Lấy n = 4096. Một ma trận điểm có 16.777.216 phần tử; nếu mỗi phần tử 2 byte thì riêng ma trận ấy chiếm 32 MiB cho một đầu, chưa tính batch và những tensor khác. Khi n gấp đôi, kích thước ma trận tăng bốn lần.
SRAM trên chip nhanh nhưng nhỏ hơn nhiều so với bộ nhớ thiết bị. Không thể chỉ đặt mọi ma trận vào SRAM. Thiết kế phải chia công việc để từng khối vừa với vùng nhớ nhanh và dùng lại dữ liệu trong lúc khối còn ở đó.
Tiling giữ công việc ở gần nơi tính
FlashAttention chia Q, K và V thành các khối, tính từng phần attention và cập nhật kết quả cần thiết. Mục tiêu là tránh vật chất hóa toàn bộ ma trận điểm và xác suất trong bộ nhớ lớn, đồng thời giảm đọc ghi giữa các tầng nhớ.
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness trình bày thuật toán attention chính xác theo hướng tối ưu IO. “Chính xác” ở đây phân biệt với việc xấp xỉ cơ chế attention; số học dấu phẩy động vẫn có thể tạo sai khác nhỏ do thứ tự tính.
Điều này không biến attention dày thành thuật toán tuyến tính theo chiều dài chuỗi. Các cặp query-key vẫn được xử lý theo thiết lập dày, nhưng dữ liệu trung gian được quản lý khác. Đừng dùng biểu đồ peak memory để suy ra trực tiếp bậc tính toán đã đổi.
Softmax qua nhiều khối cần thống kê chạy
Softmax ổn định thường trừ điểm lớn nhất trước khi lấy hàm mũ. Nhưng khi các điểm đến theo từng khối, điểm lớn nhất toàn hàng chưa biết ngay. Ta giữ cực đại m, tổng chuẩn hóa l và tổng có trọng số u để cập nhật dần.
Khi khối mới có cực đại m_b, đặt m_new = max(m, m_b). Phần tích lũy cũ được nhân exp(m − m_new); các điểm khối mới dùng exp(s − m_new). Cập nhật l bằng tổng hệ số và u bằng tổng hệ số nhân value. Kết quả cuối là u/l.
Với điểm giả định [0, ln2] và value [1, 3], softmax đầy đủ cho trọng số [1/3, 2/3], kết quả 7/3. Xử lý điểm đầu trước: m = 0, l = 1, u = 1. Khi gặp ln2, phần cũ co còn một nửa; l mới = 1,5 và u mới = 3,5. Tỉ số vẫn 7/3.
Một query được so với các key để tạo trọng số. Attention dùng những trọng số đó để tổng hợp các value thành biểu diễn mới của vị trí đang xét.
↔ Vuốt ngang để xem toàn bộ sơ đồ.
01 / Query
Vị trí đang xét có một vector truy vấn, gọi là query.
02 / Key
Query được so với các vector key ở những vị trí mà mô hình được phép tham chiếu.
03 / Trọng số
Điểm tương hợp được chuẩn hóa thành trọng số attention cho từng vị trí.
04 / Tổng hợp
Các vector value được cộng theo trọng số đó, tạo biểu diễn có thông tin từ ngữ cảnh.
Sơ đồ giản lược một đầu attention, bỏ qua kích thước vector và phép chiếu. Các đường nối không phải lời giải thích đầy đủ cho quyết định của mô hình.
Sơ đồ trên mô tả ý nghĩa kết hợp ngữ cảnh. Nó không biểu diễn tiling, SRAM hay thống kê online softmax; phép tính hai khối mới là phần dùng để kiểm tra cách giữ cùng kết quả.
Kernel đúng còn phải hợp với tải
Lợi ích thực phụ thuộc loại GPU, chiều head, độ dài chuỗi, dtype và chế độ tính. Kernel có thể có điều kiện hỗ trợ cụ thể; một đường chạy fallback không nhất thiết dùng thuật toán bạn đang nghĩ đến. Khi benchmark, cần xác nhận backend thực sự được chọn.
Tách thời gian khởi động, biên dịch và làm nóng khỏi các lượt đo ổn định, nhưng vẫn báo chúng nếu ứng dụng quan tâm thời gian khởi động. Đo cả forward lẫn backward nếu mục tiêu là huấn luyện. Một kết quả chỉ đo forward không chứng minh toàn bộ bước học tăng tốc cùng mức.
FlashAttention và KV cache xử lý những chi phí khác nhau. Cache tái sử dụng K, V của tiền tố khi decode; FlashAttention tối ưu thực hiện attention và dữ liệu trung gian. Có thể kết hợp trong một hệ thống, nhưng tên gọi không thay thế việc đo từng giai đoạn.
Bài tập kiểm tính đúng trước khi đo tốc độ
Tính ví dụ hai khối ở trên theo thứ tự ngược: xử lý ln2 trước rồi điểm không. Kết quả vẫn phải là 7/3 trong số học chính xác. Nếu khác, bạn có thể đã quên đổi thang phần tích lũy cũ.
Sau đó viết kế hoạch so đầu ra và gradient giữa hai triển khai với ngưỡng sai số phù hợp dtype, rồi mới đo thời gian và peak memory. Tiêu chí tự kiểm là báo rõ kernel, kích thước tensor và sai số, không chỉ một hệ số nhanh hơn. Đọc self-attention từng bước để kiểm lại phép toán nền.