Ý Tưởng Tiling & Online Softmax Trong FlashAttention (Dao et al.)
Thay vì đợi toàn bộ hàng hoàn thành để tìm max và tổng exp, ta duy trì running maximum mi và running sum li:
Tiến độ của tôi ở bài này
Điểm được lưu vào tài khoản sau khi chấm bài.
Đang tải điểm của bạn…
Kiến thức và chủ đề
Kiến thức tiên quyết: multi-head-attention-from-scratch.
Nội dung đề bài
Mục tiêu kiến thức
- Trong Attention chuẩn, ma trận S = Q KT kích thước N × N được ghi vào bộ nhớ GPU (HBM) rồi đọc lại để tính Softmax. Với N = 32k, ma trận này tốn hàng chục GB VRAM.
- FlashAttention (Dao et al., 2022) chia Q, K, V thành các khối nhỏ vừa với bộ nhớ SRAM nhanh (Block Tiling) và dùng kỹ thuật Online Softmax:
Thay vì đợi toàn bộ hàng hoàn thành để tìm max và tổng exp, ta duy trì running maximum mi và running sum li: Khi duyệt qua khối mới có điểm cực đại mblock: mnew = max(mold, mblock) lnew = lold · emold - mnew + ∑ eSblock - mnew Onew = ( Oold · lold · emold - mnew + eSblock - mnew Vblock ) / lnew Thuật toán không bao giờ cần cấp phát ma trận N × N trong bộ nhớ!
Yêu cầu
Viết hàm:
def flash_attention_cpu_tiling(
q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
block_size: int = 4
) -> torch.Tensor:
passq, k, v: Tensor 3D shape(B, seq_len, d_head).block_size: Kích thước khối theo chiều sequence length.- Trả về tensor output cùng shape
(B, seq_len, d_head)khớp với attention chuẩn.
Input
- Hàm
flash_attention_cpu_tiling(q,k,v,block_size): Các tham số đầu vào chứa dữ liệu Tensor/mảng NumPy hoặc giá trị siêu tham số tương ứng.
Output
- Hàm
flash_attention_cpu_tiling: Trả về kết quả kiểutorch.Tensortheo đúng đặc tả kỹ thuật và kích thước quy định.
Ràng buộc
- Thời gian chạy tối đa: 3000ms.
- Giới hạn bộ nhớ: 512MB.
- Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.
Ví dụ 1
Input
torch.manual_seed(42)
B, N, D = (2, 12, 16)
q = torch.randn(B, N, D, dtype=torch.float64)
k = torch.randn(B, N, D, dtype=torch.float64)
v = torch.randn(B, N, D, dtype=torch.float64)
scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(D)
probs = F.softmax(scores, dim=-1)
flash_attention_cpu_tiling(q, k, v, block_size=4)Output
Tensor shape: (2, 12, 16), dtype=torch.float64Giải thích
Hàm/lớp được gọi với các tham số mẫu trên và trả về kết quả số học / kích thước tensor tương ứng theo đúng thiết kế.
Ví dụ 2
Input
B, N, D = (1, 7, 8)
q = torch.randn(B, N, D, dtype=torch.float64)
k = torch.randn(B, N, D, dtype=torch.float64)
v = torch.randn(B, N, D, dtype=torch.float64)
flash_attention_cpu_tiling(q, k, v, block_size=3)Output
Tensor shape: (1, 7, 8), dtype=torch.float64Giải thích
Hàm/lớp được gọi với các tham số mẫu trên và trả về kết quả số học / kích thước tensor tương ứng theo đúng thiết kế.
Gợi ý và lời giải chỉ mở sau khi bạn bấm Nộp bài. Giáo viên và quản trị viên mở được ngay.
Góp ý & báo lỗi bài tập
Đề bài chưa rõ, test có vấn đề hay bạn có ý tưởng giúp bài tốt hơn? Gửi cho đội ngũ AI Empire nhé — mỗi góp ý đều được đọc.