Flash Attention

Flash Attention (FlashAttention) là một thuật toán và họ kernel GPU tính attention chính xác cho các model Transformer mà không bao giờ ghi toàn bộ ma trận attention vào bộ nhớ chính của GPU. Self-attention tiêu chuẩn trên N token tạo ma trận điểm N × N, ghi ma trận đó vào bộ nhớ băng thông cao (HBM), rồi đọc lại để tính softmax và tích với các giá trị; vì vậy, phần lớn thời gian chạy của attention dành cho việc di chuyển dữ liệu. Tri Dao và cộng sự tại Stanford giới thiệu phương pháp này trong bài báo năm 2022 FlashAttention: Attention chính xác nhanh và tiết kiệm bộ nhớ với nhận biết IO. Phương pháp chia công việc thành các tile vừa với SRAM trên chip tốc độ cao, nhờ đó mức sử dụng bộ nhớ tăng tuyến tính theo độ dài chuỗi và đầu ra khớp với attention tiêu chuẩn, ngoại trừ sai số làm tròn dấu phẩy động. Các kernel của phương pháp chạy trong PyTorch và trong các engine suy luận mô hình ngôn ngữ lớn (LLM) như vLLM.

Cách hoạt động của Flash Attention#

GPU có một lượng nhỏ bộ nhớ rất nhanh nằm cạnh các đơn vị tính toán và một lượng lớn bộ nhớ ngoài chip chậm hơn. Bài báo FlashAttention mô tả NVIDIA A100 có 40-80 GB HBM với băng thông 1.5-2.0 TB/s và 192 KB SRAM trong mỗi trong số 108 streaming multiprocessor, với băng thông ước tính 19 TB/s. Attention thực hiện ít phép toán trên mỗi byte được đọc, vì vậy tốc độ phụ thuộc vào tần suất dữ liệu di chuyển giữa hai cấp bộ nhớ; phương pháp này đếm số lần truyền đó, và đó là ý nghĩa của thuật ngữ "nhận biết IO".

FlashAttention nạp một tile query rồi lần lượt truyền các tile key và value qua SRAM. Với mỗi tile, phương pháp tính điểm, cập nhật giá trị cực đại đang chạy và tổng đang chạy cho từng hàng query, điều chỉnh lại đầu ra từng phần mỗi khi giá trị cực đại tăng, rồi cộng phần đóng góp của tile mới. Softmax trực tuyến này cho kết quả softmax chính xác mà không cần xử lý toàn bộ một hàng cùng lúc; chỉ đầu ra cuối cùng và hai con số trên mỗi hàng được ghi trở lại HBM. Lượt lan truyền ngược tính lại các tile điểm từ Q, K và V thay vì lưu ma trận N × N; cách này tăng số phép toán nhưng vẫn chạy nhanh hơn vì đọc ít bộ nhớ hơn nhiều. Vòng lặp tương tự trong PyTorch, xử lý tất cả query với từng tile key-value một:

import torch

torch.manual_seed(0)
n, d, block = 4096, 64, 256
q, k, v = torch.randn(3, n, d).unbind(0)

# Standard attention builds the full n x n score matrix in memory
reference = torch.softmax(q @ k.T / d**0.5, dim=-1) @ v

# FlashAttention-style: stream key/value tiles, keeping a running max and sum per query (online softmax)
row_max, row_sum, out = torch.full((n, 1), -torch.inf), torch.zeros(n, 1), torch.zeros(n, d)
for start in range(0, n, block):
    scores = q @ k[start : start + block].T / d**0.5  # only an n x block tile exists at once
    new_max = torch.maximum(row_max, scores.max(dim=-1, keepdim=True).values)
    rescale, probs = torch.exp(row_max - new_max), torch.exp(scores - new_max)
    row_sum = row_sum * rescale + probs.sum(dim=-1, keepdim=True)
    out = out * rescale + probs @ v[start : start + block]
    row_max = new_max
out = out / row_sum

print(f"Largest score tile: {tuple(scores.shape)} vs full matrix: ({n}, {n})")
print(f"Max difference from standard attention: {(out - reference).abs().max().item():.1e}")

Tile điểm lớn nhất có kích thước 4,096 × 256, bằng một phần mười sáu ma trận 4,096 × 4,096; kết quả khác attention tiêu chuẩn tối đa 4.3e-07, tương đương độ lớn của sai số làm tròn FP32. Kernel thực tế cũng chia query thành các tile, mỗi tile thuộc một khối luồng GPU, rồi hợp nhất toàn bộ vòng lặp vào một kernel duy nhất để các tile không rời khỏi SRAM.

FlashAttention-2, -3 và -4#

Mỗi phiên bản hướng đến thế hệ GPU mới hơn. FlashAttention-2 (2023) tổ chức lại cách phân chia công việc giữa các luồng và chạy nhanh gần gấp đôi phiên bản gốc, đạt 50-73% thông lượng cực đại của A100. FlashAttention-3 (2024) sử dụng Tensor Core bất đồng bộ của H100 và hỗ trợ FP8, đạt 740 TFLOPs/s ở FP16 và gần 1.2 PFLOPs/s ở FP8. FlashAttention-4 (2026) hướng đến GPU Blackwell và báo cáo hiệu suất tối đa 1,613 TFLOPs/s ở BF16 trên B200, bằng 71% hiệu suất cực đại. Các kernel được phân phối trong thư viện mã nguồn mở flash-attention, hỗ trợ đầu vào FP16 và BF16; hàm scaled_dot_product_attention của PyTorch tự động chọn kernel FlashAttention-2 khi GPU và đầu vào hỗ trợ, nếu không sẽ chuyển sang các cách triển khai tiết kiệm bộ nhớ hoặc thông thường. SAM 2 của Meta, chạy trong package Ultralytics, tính attention trong bộ mã hóa ảnh Hiera và attention bộ nhớ thông qua hàm đó; do vậy, PyTorch định tuyến các lớp này đến FlashAttention bất cứ khi nào phần cứng và độ chính xác đầu vào cho phép.

Flash Attention và các thuật ngữ liên quan#

• Attention tiêu chuẩn: Cùng phép toán toán học đó, được tính bằng cách ghi toàn bộ ma trận điểm và xác suất vào HBM.
• Attention thưa: Chỉ tính điểm cho các cặp query-key được chọn. Bài báo FlashAttention cũng giới thiệu biến thể thưa theo khối, bỏ qua các tile rỗng.
• Attention tuyến tính: Thay softmax bằng một feature map để chi phí tăng tuyến tính theo độ dài, làm thay đổi đầu ra của model.
• PagedAttention: Bộ quản lý bộ nhớ lưu KV cache thành các page kích thước cố định trong quá trình phục vụ model. Các kernel attention, bao gồm FlashAttention, đọc key và value từ những page đó.
• Ring attention: Chia một chuỗi rất dài giữa nhiều thiết bị; mỗi thiết bị tính attention theo khối trên phần dữ liệu của mình trong khi các khối key-value di chuyển theo vòng.

Tác động đến huấn luyện và ngữ cảnh dài#

Vì bộ nhớ tăng tuyến tính theo độ dài chuỗi, FlashAttention giúp việc huấn luyện với cửa sổ ngữ cảnh dài trở nên khả thi. Bài báo gốc báo cáo tốc độ toàn quy trình tăng 15% cho BERT-large với độ dài chuỗi 512 so với kỷ lục huấn luyện MLPerf 1.1, tăng tốc 3× cho GPT-2 với độ dài chuỗi 1,024 và 2.4× trên benchmark Long Range Arena. Mức tiết kiệm tương tự cũng áp dụng cho thị giác, nơi Vision Transformer (ViT) hoặc Transformer khuếch tán chuyển ảnh và video độ phân giải cao thành các chuỗi token dài. Chuỗi dài hơn cũng cải thiện độ chính xác: bài báo ghi nhận những Transformer đầu tiên đạt kết quả cao hơn ngẫu nhiên trên Path-X, một tác vụ yêu cầu xác định xem hai điểm trong ảnh 128 × 128, được làm phẳng thành chuỗi 16,384 pixel, có được nối với nhau bằng một đường đi hay không.