3.3 Hiểu sâu cơ chế Self-Attention
Self-Attention là trái tim của Transformer, cũng là module tập trung nhiều compute nhất và có nhiều kỹ thuật tối ưu nhất trong các LLM hiện đại
Self-Attention là trái tim của Transformer, cũng là module tập trung nhiều compute nhất và có nhiều kỹ thuật tối ưu nhất trong các LLM hiện đại. Dù bạn muốn hiểu ý tưởng tối ưu IO đằng sau FlashAttention, hay muốn nắm vì sao các biến thể Attention như GQA, MLA giảm được chi phí suy luận, bạn đều cần hiểu sâu Self-Attention. Bài này bắt đầu từ lịch sử Attention, lần lượt tháo từng bước toán của Scaled Dot-Product Attention, tự viết lại bằng PyTorch, phân tích bottleneck tính toán, rồi kéo sang FlashAttention và các biến thể Attention chính, để bạn dựng được chuỗi nhận thức từ trực giác đến công thức rồi đến hiện thực kỹ thuật.
📑 Mục lục
- 1. Lịch sử tiến hóa của Attention
- 2. Hiểu QKV qua góc nhìn truy xuất thông tin
- 3. Scaled Dot-Product Attention chi tiết
- 4. Tự viết bằng PyTorch
- 5. Theo dõi chiều chi tiết
- 6. Vì sao chia cho : suy luận ổn định phương sai
- 7. Ổn định số học của Softmax
- 8. Multi-Head Attention
- 9. MHA vs MQA vs GQA vs MLA
- 10. Vai trò và cách làm Causal Mask
- 11. Phân tích bottleneck tính toán của Self-Attention
- 12. Ý tưởng cốt lõi của FlashAttention
- Checklist tự kiểm
- Tài liệu tham khảo
1. Lịch sử tiến hóa của Attention
Trước khi Transformer xuất hiện năm 2017, cơ chế Attention đã được ươm trong mô hình hóa chuỗi vài năm. Theo dõi mạch tiến hóa đó giúp nắm Self-Attention thực sự giải quyết vấn đề gì, và vì sao thiết kế của nó thành hình như hôm nay.
1.1 Bahdanau Attention (2014): cột mốc mở đầu
Kiến trúc Encoder-Decoder truyền thống (ví dụ RNN cho dịch máy) có một lỗ hổng chết người: Encoder nén cả chuỗi đầu vào thành một vector độ dài cố định rồi đưa cho Decoder. Khi chuỗi dài, vector cố định đó không chứa hết thông tin, chất lượng dịch sụp nhanh.
Bahdanau và cộng sự năm 2014 đề xuất một cải tiến then chốt: đừng cố nén nữa. Khi Decoder sinh từng từ, hãy để nó nhìn lại mọi hidden state của Encoder và lấy đúng phần cần. Cụ thể, hidden state hiện tại của Decoder đóng vai “query”, được ghép với từng hidden state của Encoder (một mạng feed-forward nhỏ tính điểm khớp), rồi dùng điểm khớp đó để tổng có trọng số các hidden state Encoder thành “context vector”, làm thêm đầu vào cho bước giải mã hiện tại.
Đó là phôi thai của Attention: dạy mô hình “chú ý” các phần khác nhau của chuỗi đầu vào. Nhưng hàm khớp của Bahdanau Attention vẫn là một feed-forward phụ (additive attention), nên hiệu năng tính toán không cao.
1.2 Luong Attention (2015): đơn giản hóa tính toán
Năm 2015 Luong đơn giản hóa Bahdanau Attention theo hai hướng: hàm khớp rẻ hơn (tích vô hướng thay cho feed-forward), và thử nhiều cách căn chỉnh (global vs local). Trong đó, dot-product Attention (đo độ giống bằng tích trong của hai vector) sau này trở thành nền của Transformer.
Tích vô hướng rẻ hơn nhiều so với feed-forward: nó là nhân ma trận, vốn hợp GPU song song. Cải tiến trông đơn giản đó đã mở đường cho Self-Attention quy mô lớn.
1.3 Self-Attention (2017): Attention Is All You Need
Paper Transformer 2017 đưa ra quyết định táo bạo: bỏ hẳn RNN, chỉ dùng Attention để mô hình hóa chuỗi.
Attention trước đó là “chéo chuỗi”: Decoder nhìn Encoder. Self-Attention là “tự chú ý”: mỗi vị trí trong chuỗi nhìn mọi vị trí cùng chuỗi (kể cả chính nó). Hai lợi thế lớn:
- Song song hóa: RNN phải tính từng bước (bước phụ thuộc ). Trong Self-Attention mọi vị trí tính cùng lúc, hợp GPU.
- Phụ thuộc khoảng cách dài: Trong RNN, hai từ xa nhau phải truyền qua nhiều bước, tín hiệu suy dần. Trong Self-Attention hai vị trí bất kỳ tương tác trong một bước, về lý thuyết không có suy giảm theo khoảng cách.
Cái giá? Self-Attention có độ phức tạp : mỗi cặp vị trí đều cần điểm khớp, trong khi RNN là . Độ phức tạp bình phương đó trở thành điểm xuất phát của vô số công việc tối ưu về sau.
2. Hiểu QKV qua góc nhìn truy xuất thông tin
Ba khái niệm cốt lõi trong Attention là Query (Q), Key (K) và Value (V). Ba từ này mượn thẳng từ lĩnh vực truy xuất thông tin. Một phép so sánh đời thường giúp dựng trực giác.
2.1 So sánh với tìm sách trong thư viện
Bạn bước vào thư viện với một nhu cầu mơ hồ: “Tôi muốn tài liệu về tính toán song song.” Nhu cầu đó chính là Query.
Gáy mỗi cuốn sách dán nhãn: “Hệ điều hành”, “Mạng máy tính”, “Nhập môn tính toán song song”, “Cấu trúc dữ liệu”… Những nhãn đó là Key của từng cuốn.
Bạn đem Query đi so với từng Key. “Nhập môn tính toán song song” khớp cao; “Hệ điều hành” có liên quan vừa; “Cấu trúc dữ liệu” gần như không liên quan.
Sau khi có điểm khớp, bạn không lấy nhãn gáy (Key), mà lấy nội dung thật (Value) của từng cuốn, cộng có trọng số theo điểm. Nội dung parallel computing chiếm phần lớn, OS góp một ít, data structures gần như không, và bạn nhận được một bản tổng hợp lấy parallel computing làm trọng tâm, kèm chút kiến thức OS.
2.2 Định nghĩa hình thức của bộ ba
Quay lại Self-Attention. Cho chuỗi đầu vào hình dạng , ba phép biến đổi tuyến tính đưa vào các không gian khác nhau:
- : mỗi token tạo “vector truy vấn” của mình: “tôi cần thông tin gì”
- : mỗi token tạo “vector chỉ mục” của mình: “tôi cung cấp manh mối gì về thông tin mình có”
- : mỗi token tạo “vector nội dung” của mình: “thông tin tôi thực sự mang”
Q, K, V phải lấy từ ba phép chiếu tuyến tính khác nhau của cùng , chứ không dùng thẳng , vì cần tách vai trò. Những gì một token “cần” và những gì nó “cung cấp được” thường khác nhau. Trong một câu, động từ có thể cần nhìn chủ ngữ và tân ngữ (hướng Query), còn khi bị token khác chú ý, nó chủ yếu cung cấp ngữ nghĩa hành động (hướng Key/Value). Ba ma trận chiếu độc lập cho mô hình tự do học các ánh xạ khác nhau đó.
3. Scaled Dot-Product Attention chi tiết
Có Q, K, V rồi, quá trình Attention tách thành năm bước rõ ràng.
3.1 Bước 1: chiếu tuyến tính
Đầu vào có hình . Ba ma trận trọng số , , đều có hình (trường hợp một head):
Ba phép nhân ma trận này là ba GEMM (General Matrix Multiply) độc lập. Trong thực tế, để tăng mức dùng GPU, thường gộp , , thành một ma trận lớn hình , làm một GEMM rồi split. Cách đó tận dụng sức tính của GPU tốt hơn.
3.2 Bước 2: điểm Attention thô
Nhân với chuyển vị của để đo mức khớp mỗi cặp token:
Mỗi phần tử là tích vô hướng giữa query của token và key của token : mức chú ý thô của token tới token .
Bước này tạo ma trận . Đó chính là gốc của độ phức tạp bình phương của Self-Attention.
3.3 Bước 3: scale
Chia điểm thô cho ( là chiều của vector Key):
Vì sao cần scale? Nói ngắn: khi lớn, tích vô hướng phình to, gradient Softmax về sau tiến gần không. Suy luận toán chi tiết ở mục 6.
3.4 Bước 4: chuẩn hóa Softmax
Softmax theo hàng trên ma trận điểm đã scale, biến điểm thô thành phân phối xác suất:
Softmax làm hai việc: ánh xạ số thực bất kỳ vào để dùng làm trọng số, và bảo đảm tổng trọng số mỗi hàng bằng 1, thành phân phối hợp lệ.
3.5 Bước 5: tổng có trọng số
Dùng trọng số Attention để tổng có trọng số ma trận Value:
Mỗi vị trí token nhận một vector chiều, đã trộn thông tin của mọi token nó “chú ý”, với mức chú ý do trọng số trong quyết định.
3.6 Công thức đầy đủ
Gộp các bước trên, Self-Attention gói trong một dòng:
Cuối cùng còn qua một ma trận chiếu đầu ra hình để ánh xạ về không gian ẩn của mô hình. Phép chiếu này đặc biệt quan trọng trong Multi-Head Attention: nó trộn lại biểu diễn sau khi nối nhiều head.
4. Tự viết bằng PyTorch
Lý thuyết nhiều đến đâu cũng không bằng tự viết code một lần. Dưới đây lần lượt là Single-Head và Multi-Head Self-Attention.
4.1 Single-Head Self-Attention
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
class SingleHeadSelfAttention(nn.Module):
"""Tự viết single-head Self-Attention, không dùng MHA sẵn của PyTorch"""
def __init__(self, d_model: int):
super().__init__()
self.d_model = d_model
# Ba phép chiếu tuyến tính: vào d_model, ra d_model
self.W_Q = nn.Linear(d_model, d_model, bias=False)
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)
# Chiếu đầu ra
self.W_O = nn.Linear(d_model, d_model, bias=False)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
# x: (batch, seq_len, d_model)
# Bước 1: chiếu tuyến tính lấy Q, K, V
Q = self.W_Q(x) # (batch, seq_len, d_model)
K = self.W_K(x) # (batch, seq_len, d_model)
V = self.W_V(x) # (batch, seq_len, d_model)
# Bước 2: điểm Attention = Q @ K^T
scores = torch.matmul(Q, K.transpose(-2, -1)) # (batch, seq_len, seq_len)
# Bước 3: scale, tránh tích vô hướng quá lớn làm gradient softmax biến mất
scores = scores / math.sqrt(self.d_model)
# Bước 4: nếu có mask, đặt điểm các vị trí bị che thành -inf
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# Bước 5: Softmax để ra trọng số Attention
attn_weights = F.softmax(scores, dim=-1) # (batch, seq_len, seq_len)
# Bước 6: tổng có trọng số trên V
context = torch.matmul(attn_weights, V) # (batch, seq_len, d_model)
# Bước 7: chiếu đầu ra
output = self.W_O(context) # (batch, seq_len, d_model)
return output
4.2 Multi-Head Self-Attention
class MultiHeadSelfAttention(nn.Module):
"""Tự viết multi-head Self-Attention: cắt d_model thành num_heads không gian con"""
def __init__(self, d_model: int, num_heads: int):
super().__init__()
assert d_model % num_heads == 0, "d_model phải chia hết cho num_heads"
self.d_model = d_model
self.num_heads = num_heads
self.head_dim = d_model // num_heads # chiều mỗi head
# Chiếu QKV: vào d_model, ra d_model (gom chiếu của mọi head bên trong)
self.W_Q = nn.Linear(d_model, d_model, bias=False)
self.W_K = nn.Linear(d_model, d_model, bias=False)
self.W_V = nn.Linear(d_model, d_model, bias=False)
self.W_O = nn.Linear(d_model, d_model, bias=False)
def forward(self, x: torch.Tensor, mask: torch.Tensor = None):
batch_size, seq_len, _ = x.shape
# Bước 1: chiếu tuyến tính
Q = self.W_Q(x) # (batch, seq_len, d_model)
K = self.W_K(x)
V = self.W_V(x)
# Bước 2: reshape multi-head, transpose để head đứng trước seq_len
# (batch, seq_len, d_model) -> (batch, seq_len, num_heads, head_dim)
# -> (batch, num_heads, seq_len, head_dim)
Q = Q.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
K = K.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
V = V.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
# Bước 3: mỗi head tính điểm Attention riêng
# Q @ K^T: (batch, num_heads, seq_len, head_dim)
# @ (batch, num_heads, head_dim, seq_len)
# = (batch, num_heads, seq_len, seq_len)
scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(self.head_dim)
# Bước 4: causal mask tùy chọn
if mask is not None:
scores = scores.masked_fill(mask == 0, float('-inf'))
# Bước 5: Softmax + tổng có trọng số
attn_weights = F.softmax(scores, dim=-1)
# (batch, num_heads, seq_len, seq_len) @ (batch, num_heads, seq_len, head_dim)
# = (batch, num_heads, seq_len, head_dim)
context = torch.matmul(attn_weights, V)
# Bước 6: nối nhiều head: transpose về (batch, seq_len, num_heads, head_dim) rồi gộp hai chiều cuối
context = context.transpose(1, 2).contiguous().view(batch_size, seq_len, self.d_model)
# Bước 7: chiếu đầu ra, trộn biểu diễn đã nối của nhiều head
output = self.W_O(context) # (batch, seq_len, d_model)
return output
Đoạn code trên đi hết từ đầu vào đến đầu ra. Đáng chú ý: .transpose(1, 2) biến (batch, seq_len, num_heads, head_dim) thành (batch, num_heads, seq_len, head_dim), nhờ đó các phép nhân ma trận sau chạy độc lập trong từng head và tận dụng song song theo chiều batch để tính hết mọi head một lần.
5. Theo dõi chiều chi tiết
Ký hiệu chiều trừu tượng dễ làm rối. Dưới đây dùng một bộ số cụ thể để theo dõi hình dạng tensor ở từng bước trong Multi-Head Self-Attention.
5.1 Thiết lập tham số
seq_len = 6 # độ dài chuỗi (6 token)
d_model = 512 # chiều ẩn của mô hình
num_heads = 8 # số head Attention
head_dim = 512 / 8 = 64 # chiều mỗi head
batch = 1 # đơn giản hóa batch=1
5.2 Theo dõi từng bước
Đầu vào
X: (1, 6, 512) # 1 mẫu, 6 token, mỗi token 512 chiều
Chiếu tuyến tính
W_Q: (512, 512) # ma trận trọng số của nn.Linear
Q = X @ W_Q^T: (1, 6, 512) @ (512, 512) = (1, 6, 512)
K = X @ W_K^T: (1, 6, 512) @ (512, 512) = (1, 6, 512)
V = X @ W_V^T: (1, 6, 512) @ (512, 512) = (1, 6, 512)
Cắt nhiều head + transpose
Q.view(1, 6, 8, 64): (1, 6, 8, 64)
Q.transpose(1, 2): (1, 8, 6, 64) # 8 head, mỗi head thấy 6 token 64 chiều
K reshape tương tự: (1, 8, 6, 64)
V reshape tương tự: (1, 8, 6, 64)
Tính điểm Attention
K.transpose(-2, -1): (1, 8, 64, 6) # chuyển vị hai chiều cuối
Q @ K^T: (1, 8, 6, 64) @ (1, 8, 64, 6) = (1, 8, 6, 6)
# mỗi head tạo một ma trận điểm 6x6
Scale
scores / sqrt(64): (1, 8, 6, 6) # sqrt(64) = 8, mỗi điểm chia cho 8
Softmax
softmax(scores, dim=-1): (1, 8, 6, 6) # chuẩn hóa mỗi hàng (chiều cuối) thành phân phối
# mỗi hàng 6 chiều tổng bằng 1
Tổng có trọng số
attn_weights @ V: (1, 8, 6, 6) @ (1, 8, 6, 64) = (1, 8, 6, 64)
# mỗi head xuất 6 token 64 chiều
Nối nhiều head
transpose(1, 2): (1, 6, 8, 64) # đưa chiều head về lại
view(1, 6, 512): (1, 6, 512) # 8 * 64 = 512, nối lại thành d_model
Chiếu đầu ra
W_O: (512, 512)
context @ W_O^T: (1, 6, 512) @ (512, 512) = (1, 6, 512)
Đầu ra cuối
Output: (1, 6, 512) # trùng hình với đầu vào X
Từ đầu đến cuối, đầu vào và đầu ra đều , nên nhiều Transformer Block xếp chồng được như xếp Lego.
6. Vì sao chia cho : suy luận ổn định phương sai
Bước chia cho dễ bị coi là “mẹo chỉnh tham số”, nhưng đằng sau có lý do toán nghiêm ngặt.
6.1 Đặt vấn đề
Giả sử mỗi thành phần của và là biến ngẫu nhiên độc lập kỳ vọng 0, phương sai 1. Tính phương sai của tích vô hướng .
Đặt , , mỗi và i.i.d., , .
Định nghĩa tích vô hướng:
6.2 Quá trình suy luận
Kỳ vọng và phương sai của tích một thành phần:
Vì và độc lập:
Vậy:
Tích vô hướng là tổng biến ngẫu nhiên độc lập. Theo tính cộng của phương sai:
6.3 Kết luận
Phương sai của là . Nghĩa là khi , độ lệch chuẩn của tích vô hướng khoảng 8; khi , khoảng 11.3. Chiều càng lớn, giá trị tuyệt đối càng lớn, phân phối càng trải.
Hàm Softmax rất nhạy với thang đo đầu vào. Khi khoảng cách giữa các giá trị lớn (ví dụ 50 và -10), Softmax gần như one-hot, gradient gần không, tham số không cập nhật được, huấn luyện đứng.
Sau khi chia cho :
Phương sai kéo về 1. Dù thế nào, đầu vào Softmax dao động trong khoảng hợp lý, gradient còn khỏe. Đó là nguồn gốc chữ “Scaled” trong Scaled Dot-Product Attention.
7. Ổn định số học của Softmax
Softmax trông như công thức đơn giản, nhưng trên GPU thật nó giấu một bẫy dễ gây overflow.
7.1 Vấn đề overflow
Công thức Softmax chuẩn:
Vấn đề là tăng cực nhanh. Khi lớn (ví dụ ), vượt luôn float32 lẫn float16, thành inf. Dù tử và mẫu cùng tràn, inf / inf ra NaN, cả phép tính hỏng.
7.2 Mẹo trừ max
Cách xử lý dựa trên bất biến tịnh tiến của Softmax: trừ max khỏi mọi đầu vào, kết quả không đổi.
Chứng minh ngắn: nhân tử và mẫu với , tương đương chia cả hai cho . Sau khi trừ max, số mũ lớn nhất là 0, nên tối đa bằng 1, không overflow.
Nhưng điều đó tạo vấn đề hiệu năng: bản chuẩn cần ba lượt quét dữ liệu:
- Lượt 1: duyệt mọi phần tử, tìm max
- Lượt 2: duyệt mọi phần tử, tính và tổng của chúng
- Lượt 3: duyệt mọi phần tử, chia mỗi cho tổng
Mỗi lượt đọc từ HBM (bộ nhớ băng thông cao). Ba lượt là ba lần đọc HBM. Khi lớn (ví dụ độ dài chuỗi 128K), chi phí IO đó rất nặng.
7.3 Online Softmax
Milakov và Gimelshein năm 2018 đề xuất Online Softmax, gộp ba lượt thành một lượt. Ý chính là vừa duyệt vừa duy trì động max và mẫu số chuẩn hóa.
Khi đến phần tử thứ , giữ hai biến:
- : max của phần tử đầu
- : tổng mũ của phần tử đầu (lấy làm gốc)
Hệ thức đệ quy:
Mẹo quan trọng ở dòng hai: khi max đổi từ sang , tổng mũ cũ nhân hệ số sửa để “đổi gốc”. Một lượt vừa có max vừa có tổng mũ; thêm một lượt nữa để chuẩn hóa. Tổng hai lượt, ít hơn bản chuẩn một lượt.
Ý tưởng này then chốt với FlashAttention: FlashAttention cắt ma trận Attention thành tile và xử lý từng khối, mỗi khối cần cập nhật kết quả Softmax trung gian. Nếu không cập nhật Softmax online được, phải ghi cả ma trận ra HBM rồi Softmax toàn cục, mất hết ý nghĩa tính theo khối. Online Softmax khiến tính theo khối và Softmax chính xác sống chung được.
8. Multi-Head Attention
8.1 Vì sao cần nhiều head
Single-head Attention chỉ có một bộ chiếu QKV, nghĩa là mô hình chỉ học một “kiểu chú ý”. Nhưng quan hệ giữa token trong ngôn ngữ là đa chiều: cùng một từ có thể đồng thời mang quan hệ cú pháp (chủ-vị khớp), quan hệ ngữ nghĩa (đồng nghĩa), quan hệ vị trí (mẫu cục bộ của từ liền kề)…
Ví dụ: trong buổi review dự án, chỉ một reviewer thì chỉ nhìn được góc họ giỏi. Nếu có hội đồng (kiến trúc, chất lượng code, coverage test, tài liệu) chấm độc lập rồi gộp báo cáo, độ phủ rộng hơn nhiều.
Multi-Head Attention chính là cơ chế “hội đồng” đó. Mỗi head có , , riêng và bắt các kiểu quan hệ khác nhau trong các không gian con khác nhau.
8.2 Nguyên lý toán
Giả sử , , thì .
Công thức đầy đủ Multi-Head Attention:
Mỗi head có , , hình . Thực tế không giữ ma trận nhỏ: chiếu một lần bằng lớn hình , rồi reshape thành . Tương đương: ma trận lớn chính là 8 ma trận nhỏ ghép theo chiều ra.
8.3 Phân tích số tham số
Với :
| Ma trận tham số | Hình | Số tham số |
|---|---|---|
| 262,144 | ||
| 262,144 | ||
| 262,144 | ||
| 262,144 | ||
| Tổng |
Công thức chung: số tham số MHA là (không tính bias). Dù bao nhiêu head, tổng tham số không đổi: số head chỉ đổi cách cắt, không đổi tổng. Thêm một head thì chiều mỗi head giảm tương ứng, tích (tổng chiều chiếu) luôn bằng .
8.4 Ý nghĩa kỹ thuật của nhiều head
Cấu trúc nhiều head vốn hợp song song. Tám head hoàn toàn độc lập, có thể:
- Song song trong GPU: tận dụng chiều batch, xử lý mọi head trong một lần launch CUDA kernel
- Tensor Parallelism nhiều card: chia head khác nhau sang GPU khác nhau. Ví dụ 8 head trên 4 GPU thì mỗi card 2 head. Cuối cùng một AllReduce gom kết quả chiếu đầu ra
9. MHA vs MQA vs GQA vs MLA
Khi LLM bước vào thời đại ưu tiên hiệu năng suy luận, cấu trúc head Attention cũng tiến hóa. Động lực cốt lõi là một câu hỏi: KV Cache quá lớn.
9.1 MHA (Multi-Head Attention)
Trong MHA chuẩn, mỗi head Attention có Q, K, V riêng. Có head thì phải cache bộ K và bộ V.
- Q head, K head, V head
- Mỗi bộ Q/K/V độc lập, không chia sẻ
9.2 MQA (Multi-Query Attention)
Shazeer năm 2019 đề xuất: mọi Q head chia sẻ cùng một bộ K và V. Vẫn Q head, nhưng chỉ 1 K head và 1 V head.
- Q head, 1 K head, 1 V head
- Mọi Q head dùng chung một bản K và V
Lợi ích: KV Cache thu về , tốc độ suy luận tăng mạnh. Cái giá: năng lực biểu đạt giảm. Mọi head buộc lấy thông tin từ cùng một bộ KV, đa dạng bị hạn chế.
9.3 GQA (Grouped-Query Attention)
GQA là phương án trung gian giữa MHA và MQA (Ainslie et al., 2023). Chia Q head thành nhóm (mỗi nhóm Q head), mỗi nhóm chia sẻ một bộ K và V.
- Q head, K head, V head
- Mỗi Q head dùng chung một bộ K và V
- Khi thoái hóa về MHA; khi thoái hóa về MQA
GQA cân bằng chất lượng mô hình và hiệu năng suy luận tốt hơn. LLaMA-2-70B, Mistral-7B và nhiều mô hình主流 đều dùng GQA.
9.4 Bảng so sánh
Với , num_heads = 32, head_dim = 128, độ dài chuỗi , FP16:
| Chỉ số | MHA () | GQA () | MQA () |
|---|---|---|---|
| Số Q head | 32 | 32 | 32 |
| Số KV head | 32 | 8 | 1 |
| Tham số | |||
| Tham số | |||
| KV Cache / token | |||
| KV Cache / lớp @ 4096 token | |||
| Chất lượng mô hình | Tốt nhất | Gần MHA | Có giảm |
Ghi chú: Cache size = 2 × B × L × num_heads × head_dim × d_type_bytes (2 là K và V, B là batch_size, L là độ dài chuỗi đã sinh) Có thể thấy, từ MHA sang GQA(), KV Cache còn 1/4, tham số giảm, nhưng chất lượng gần như không đổi. Đó là lý do GQA trở thành lựa chọn mainstream.
9.5 MLA (Multi-Latent Attention)
DeepSeek-V2 đề xuất MLA (Multi-Latent Attention), một hướng nén KV Cache khác.
Ý tưởng cốt lõi: không cache trực tiếp vector K và V đầy đủ, mà nén chúng vào không gian tiềm ẩn (latent space) chiều thấp, chỉ cache biểu diễn thấp chiều đó. Lúc suy luận mới giải nén lại thành K và V đầy đủ.
- MHA chuẩn: (cache , chiều = num_kv_heads head_dim)
- MLA: (cache , chiều nhỏ hơn nhiều so với ), lúc suy luận giải nén khi cần
MLA dùng nén hạng thấp để thu nhỏ mạnh KV Cache (DeepSeek-V2 báo khoảng 93.3% tỷ lệ nén), đồng thời các ma trận chiếu lên được thiết kế kỹ để giữ chất lượng. Khác GQA (giảm số KV head), MLA giảm chiều mỗi biểu diễn. Mục tiêu giống nhau: làm KV Cache lúc suy luận càng nhỏ càng tốt.
10. Vai trò và cách làm Causal Mask
10.1 Vì sao cần che
Trong mô hình ngôn ngữ tự hồi quy (họ GPT, họ LLaMA), mục tiêu huấn luyện là “dựa vào phần trước để dự đoán token tiếp theo”. Khi tính Attention, mỗi vị trí chỉ được thấy chính nó và các token trước đó, không được “nhìn lén” tương lai. Nếu không, lúc train mô hình đã thấy đáp án, chẳng học được năng lực dự đoán có nghĩa.
10.2 Cách hiện thực
Causal Mask (che nhân quả) là một ma trận tam giác trên phủ lên ma trận điểm Attention, đặt điểm các vị trí tương lai thành âm vô cùng:
# Xây ma trận Causal Mask
def create_causal_mask(seq_len: int) -> torch.Tensor:
"""
Tạo ma trận tam giác dưới (gồm đường chéo); phần tam giác trên là False.
Vị trí True được giữ; vị trí False bị đặt thành -inf trong điểm Attention.
"""
mask = torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool))
return mask
# Ví dụ dùng: seq_len = 4
# mask:
# [[True, False, False, False],
# [True, True, False, False],
# [True, True, True, False],
# [True, True, True, True ]]
#
# Ma trận điểm Attention sau khi che (tam giác trên = -inf):
# [[s00, -inf, -inf, -inf],
# [s10, s11, -inf, -inf],
# [s20, s21, s22, -inf],
# [s30, s31, s32, s33 ]]
#
# Sau Softmax, vị trí -inf có trọng số 0, thông tin tương lai bị chặn hoàn toàn
Một ý nghĩa kỹ thuật khác của Causal Mask: ma trận Attention chỉ phần tam giác dưới có hiệu lực, về lý thuyết có thể bỏ tính phần tam giác trên. FlashAttention-2 tận dụng đúng điểm này: khi gặp tile vùng tam giác trên thì bỏ qua, cắt gần một nửa khối lượng tính toán.
11. Phân tích bottleneck tính toán của Self-Attention
Hiểu bottleneck hiệu năng của Self-Attention cần tách hai kịch bản khác nhau.
11.1 Phân tích khối lượng tính toán
Self-Attention có bốn phép nhân ma trận chính:
| Phép toán | Hình | FLOPs |
|---|---|---|
Tổng: .
Khi nhỏ (ví dụ ), hạng chiếm ưu thế: bottleneck nằm ở GEMM chiếu QKV. Khi lớn (ví dụ ), hạng chiếm ưu thế: bottleneck nằm ở tính ma trận Attention.
11.2 Compute Bound vs Memory Bound
Hiệu năng GPU bị giới hạn bởi hai chỉ số:
- Compute (sức tính): bao nhiêu phép số thực mỗi giây (FLOPS)
- Bandwidth (Memory Bandwidth): bao nhiêu dữ liệu đọc/ghi từ HBM mỗi giây (GB/s)
Một phép toán là Compute Bound hay Memory Bound phụ thuộc cường độ số học (Arithmetic Intensity) = FLOPs / Bytes.
- Nếu cường độ > tỉ số compute/bandwidth của GPU, phép toán Compute Bound: ALU no việc, phần lớn thời gian đang tính
- Nếu cường độ < tỉ số đó, phép toán Memory Bound: ALU không no, phần lớn thời gian chờ chuyển dữ liệu
Lấy A100: khoảng 312 TFLOPS FP16, khoảng 2 TB/s HBM, tỉ số khoảng 156 FLOP/Byte.
Giai đoạn Prefill (xử lý cả prompt):
Chiếu QKV là nhân ma trận lớn, chiều batch (seq_len) lớn, cường độ cao, thường Compute Bound. cũng là nhân ma trận lớn, cũng Compute Bound.
Giai đoạn Decode (sinh từng token):
Mỗi bước chỉ 1 token. Chiếu QKV thoái hóa thành nhân ma trận-vector (GEMV), cường độ cực thấp, Memory Bound. Attention trở thành một vector với cả KV Cache, cũng Memory Bound. Phần lớn thời gian dành để kéo trọng số mô hình và KV Cache từ HBM.
11.3 Bottleneck bộ nhớ
Bản Attention chuẩn phải tính và lưu tường minh cả ma trận Attention . Với seq_len = 128, num_heads = 32:
Rõ ràng vượt xa bộ nhớ mọi GPU đơn. Ngay cả khi độ dài chuỗi “chỉ” 8K, ma trận Attention cũng cần , đã là chi phí không thể bỏ qua.
Bottleneck bộ nhớ đó chính là bài toán cốt lõi FlashAttention muốn giải.
12. Ý tưởng cốt lõi của FlashAttention
FlashAttention (Dao et al., 2022) là một trong những công trình ảnh hưởng nhất vài năm gần đây trong tối ưu Attention. Nó không đổi kết quả tính (tính chính xác, không xấp xỉ), nhưng sắp xếp lại thứ tự tính toán để giảm mạnh lượng truy cập HBM.
12.1 Nhắc lại phân cấp bộ nhớ GPU
Để hiểu FlashAttention cần nắm phân cấp bộ nhớ GPU:
- HBM (High Bandwidth Memory): bộ nhớ chính của GPU, dung lượng lớn (ví dụ A100 80 GB) nhưng truy cập tương đối chậm (~2 TB/s)
- SRAM (cache trên chip): shared memory và thanh ghi trên mỗi SM (Streaming Multiprocessor), dung lượng rất nhỏ (A100 khoảng 192 KB shared memory mỗi SM, cả chip khoảng 20 MB) nhưng cực nhanh (~19 TB/s)
Luồng Attention chuẩn: hoàn thành trên HBM rồi ghi cả ma trận về HBM; đọc lại để Softmax rồi ghi về HBM; cuối cùng đọc kết quả Softmax và để nhân. Mỗi bước đều đọc/ghi ma trận lớn trên HBM. Đó là lãng phí băng thông khổng lồ.
12.2 Tiling: tính theo khối
Kỹ thuật then chốt đầu tiên của FlashAttention là Tiling (chia khối): không tính một lần cả ma trận Attention , mà cắt Q, K, V thành nhiều khối nhỏ (tile). Mỗi lần chỉ nạp một khối Q và một khối K, V vào SRAM, tính Attention của khối đó trong SRAM, rồi ghi kết quả về HBM.
Bản chuẩn:
Nạp hết Q, K vào HBM → tính đủ ma trận N x N → ghi HBM → đọc Softmax → ghi HBM → ...
FlashAttention:
Cắt Q thành T_r khối, K/V thành T_c khối
For mỗi khối Q_i:
For mỗi khối K_j, V_j:
Nạp Q_i, K_j, V_j từ HBM vào SRAM (mỗi khối nhỏ, vừa khít)
Trong SRAM: Q_i @ K_j^T → Softmax cục bộ → nhân V_j
Cộng dồn vào đầu ra tương ứng của khối Q_i
Ghi đầu ra cuối của khối Q_i về HBM
Mỗi lần chỉ xử lý một tile; phần nhỏ của ma trận Attention luôn nằm trong SRAM, không bao giờ cần ghi cả ma trận ra HBM.
12.3 Vai trò then chốt của Online Softmax
Tiling mang theo một bài toán khó: Softmax cần chuẩn hóa mọi phần tử của mỗi hàng (cần max toàn cục và tổng toàn cục), nhưng khi tính theo khối mỗi lần chỉ thấy một phần hàng. Làm sao tính Softmax toàn cục khi chỉ thấy dữ liệu cục bộ?
Đó là chỗ Online Softmax vào cuộc. Nhớ hệ thức đệ quy mục 7.3: khi khối dữ liệu mới đến, có thể cập nhật động max toàn cục và mẫu số chuẩn hóa, rồi sửa kết quả khối trước.
Cụ thể, sau Attention của với khối ta có đầu ra cục bộ và thống kê Softmax (max cục bộ , mẫu số cục bộ ). Tiếp tục với khối được và . Sửa bằng Online Softmax:
Cộng dồn liên tục đến khi hết mọi khối . Kết quả cuối trùng Softmax toàn cục của Attention chuẩn về mặt toán.
12.4 Vì sao giảm được truy cập HBM
Lượng truy cập HBM của Attention chuẩn:
- Ghi (): lần ghi
- Đọc để Softmax: lần đọc
- Ghi : lần ghi
- Đọc để : lần đọc
- Tổng truy cập HBM: (cộng thêm đọc của chính Q, K, V)
Lượng truy cập HBM của FlashAttention:
- Đọc Q, K, V:
- Ghi đầu ra cuối O:
- Ma trận Attention trung gian mỗi lần chỉ tính một tile trong SRAM rồi bỏ; ma trận đầy đủ không bao giờ ghi vào HBM
- Nhưng các khối Q, K, V, O cần đọc/ghi lặp: lấy V1 làm ví dụ, vòng ngoài duyệt khối K/V, mỗi lần quét hết Q và O ()
- Tổng truy cập HBM: ( là kích thước SRAM tính theo phần tử)
Vì nhỏ hơn nhiều so với (ví dụ thì , trong khi 192KB SRAM chứa khoảng 98304 phần tử FP16), lượng truy cập HBM nhỏ hơn vài lần so với của bản chuẩn. Nó vẫn tăng theo , chỉ là hằng số giảm mạnh; thứ thực sự xuống là bộ nhớ phụ thêm (không còn lưu ma trận ). Lưu ý, khối lượng tính (FLOPs) không đổi, vẫn ; chỉ đổi mô hình IO. Bản chất tăng tốc của FlashAttention đến từ giảm truy cập HBM chậm, để tính toán càng nhiều càng tốt trong SRAM nhanh.
12.5 FlashAttention-2 và FlashAttention-3
FlashAttention-2 (Dao, 2023) tối ưu thêm chiến lược song song và lịch warp trên nền FlashAttention:
- Đổi thứ tự vòng trong/ngoài (vòng ngoài duyệt khối Q, vòng trong duyệt khối K/V), giảm đọc/ghi shared memory
- Phân việc cấp warp tốt hơn, nâng occupancy GPU
- Dùng Causal Mask để bỏ tile toàn không, giảm thêm tính toán
FlashAttention-3 (Dao et al., 2024) tối ưu cho kiến trúc Hopper (H100), dùng TMA (Tensor Memory Accelerator) và wgmma (Warp Group MMA) để nâng mức dùng phần cứng thêm nữa.
Hiện FlashAttention đã là trang bị chuẩn của hầu hết framework suy luận và huấn luyện. PyTorch 2.0+ có sẵn torch.nn.functional.scaled_dot_product_attention, mặc định gọi FlashAttention kernel bên dưới.
🎯 Checklist tự kiểm
Sau khi học xong bài này, dùng các câu hỏi sau để kiểm độ sâu hiểu biết:
- Kể được mạch tiến hóa từ Bahdanau Attention đến Self-Attention, nói rõ mỗi bước giải quyết vấn đề gì
- Viết thuộc công thức đầy đủ Scaled Dot-Product Attention, giải thích vai trò từng phần của bộ ba Q, K, V
- Suy ra từ không vì sao chia cho : giả sử chuẩn tắc độc lập, suy ra phương sai tích vô hướng bằng
- Tự viết Single-Head và Multi-Head Self-Attention bằng PyTorch, nói rõ thay đổi chiều tensor từng bước
- Giải thích vì sao Softmax phải trừ max, và ý tưởng đệ quy cốt lõi của Online Softmax
- Cho số cụ thể seq_len, , num_heads, theo dõi trọn hình dạng tensor từng bước
- Vẽ được khác biệt cấu trúc MHA / MQA / GQA và ước lượng kích thước KV Cache tương ứng
- Nói rõ vai trò và cách làm Causal Mask (tam giác trên đặt -inf, Softmax về sau thành 0)
- Phân biệt Compute Bound và Memory Bound, nói rõ bottleneck Prefill và Decode thuộc loại nào
- Nói rõ ý tưởng cốt lõi FlashAttention: Tiling + Online Softmax làm thế nào để bộ nhớ phụ từ xuống , truy cập HBM từ xuống
📚 Tài liệu tham khảo
Paper
- Neural Machine Translation by Jointly Learning to Align and Translate (Bahdanau et al., 2014): https://arxiv.org/abs/1409.0473 — Bahdanau Attention, cột mốc mở đầu của Attention
- Effective Approaches to Attention-based Neural Machine Translation (Luong et al., 2015): https://arxiv.org/abs/1508.04025 — Luong Attention, đưa vào dot-product Attention
- Attention Is All You Need (Vaswani et al., 2017): https://arxiv.org/abs/1706.03762 — paper gốc Transformer
- Fast Transformer Decoding: One Write-Head is All You Need (Shazeer, 2019): https://arxiv.org/abs/1911.02150 — Multi-Query Attention
- GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints (Ainslie et al., 2023): https://arxiv.org/abs/2305.13245 — Grouped-Query Attention
- DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model (DeepSeek-AI, 2024): https://arxiv.org/abs/2405.04434 — Multi-Latent Attention
- FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (Dao et al., 2022): https://arxiv.org/abs/2205.14135 — FlashAttention v1
- FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning (Dao, 2023): https://arxiv.org/abs/2307.08691 — FlashAttention v2
- Online Normalizer Calculation for Softmax (Milakov & Gimelshein, 2018): https://arxiv.org/abs/1805.02867 — thuật toán Online Softmax
Tutorial và blog
- The Illustrated Transformer (Jay Alammar): https://jalammar.github.io/illustrated-transformer/ — nhập môn Transformer nhiều hình
- The Annotated Transformer (Harvard NLP): https://nlp.seas.harvard.edu/annotated-transformer/ — paper khớp từng dòng với PyTorch
- Andrej Karpathy: Let’s build GPT from scratch: https://www.youtube.com/watch?v=kCc8FmEb1nY — tự viết GPT từ không
- ELI5: FlashAttention (Aleksa Gordic): https://gordicaleksa.medium.com/eli5-flash-attention-5c44017022ad — giải thích FlashAttention dễ hiểu