Chuyển đến nội dung chính
Kiến thức nền AI Infra

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

Transformer Self-Attention Multi-Head Attention FlashAttention GQA

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

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 tt phụ thuộc t−1t-1). 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 O(N2)O(N^2): mỗi cặp vị trí đều cần điểm khớp, trong khi RNN là O(N)O(N). Độ 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 XX hình dạng (N,dmodel)(N, d_{model}), ba phép biến đổi tuyến tính đưa XX vào các không gian khác nhau:

  • Q=X⋅WQQ = X \cdot W_Q: mỗi token tạo “vector truy vấn” của mình: “tôi cần thông tin gì”
  • K=X⋅WKK = X \cdot W_K: 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ó”
  • V=X⋅WVV = X \cdot W_V: 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 XX, chứ không dùng thẳng XX, 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 XX có hình (N,dmodel)(N, d_{model}). Ba ma trận trọng số WQW_Q, WKW_K, WVW_V đều có hình (dmodel,dmodel)(d_{model}, d_{model}) (trường hợp một head):

Q=X⋅WQ(N,dmodel)×(dmodel,dmodel)=(N,dmodel)\begin{aligned} Q &= X \cdot W_Q \quad & (N, d_{model}) \times (d_{model}, d_{model}) = (N, d_{model}) \\ \end{aligned} K=X⋅WK(N,dmodel)×(dmodel,dmodel)=(N,dmodel)\begin{aligned} K &= X \cdot W_K \quad & (N, d_{model}) \times (d_{model}, d_{model}) = (N, d_{model}) \\ \end{aligned} V=X⋅WV(N,dmodel)×(dmodel,dmodel)=(N,dmodel)\begin{aligned} V &= X \cdot W_V \quad & (N, d_{model}) \times (d_{model}, d_{model}) = (N, d_{model}) \end{aligned}

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 WQW_Q, WKW_K, WVW_V thành một ma trận lớn WQKVW_{QKV} hình (dmodel,3×dmodel)(d_{model}, 3 \times d_{model}), 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 QQ với chuyển vị của KK để đo mức khớp mỗi cặp token:

S=QK⊤(N,dmodel)×(dmodel,N)=(N,N)S = Q K^\top \quad (N, d_{model}) \times (d_{model}, N) = (N, N)

Mỗi phần tử S[i][j]S[i][j] là tích vô hướng giữa query của token ii và key của token jj: mức chú ý thô của token ii tới token jj.

Bước này tạo ma trận N×NN \times 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 dk\sqrt{d_k} (dkd_k là chiều của vector Key):

Sscaled=SdkS_{\text{scaled}} = \frac{S}{\sqrt{d_k}}

Vì sao cần scale? Nói ngắn: khi dkd_k 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:

A=softmax(Sscaled)∈RN×N, moˆ˜i haˋng tổng ba˘ˋng 1A = \text{softmax}(S_{\text{scaled}}) \quad \in \mathbb{R}^{N \times N}, \text{ mỗi hàng tổng bằng 1}

Softmax làm hai việc: ánh xạ số thực bất kỳ vào (0,1)(0, 1) để 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:

Output=A⋅V(N,N)×(N,dmodel)=(N,dmodel)\text{Output} = A \cdot V \quad (N, N) \times (N, d_{model}) = (N, d_{model})

Mỗi vị trí token nhận một vector dmodeld_{model} chiều, đã trộn thông tin của mọi token nó “chú ý”, với mức chú ý do trọng số trong AA 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:

Attention(Q,K,V)=softmax(QKTdk)V\text{Attention}(Q, K, V) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}}\right) V

Cuối cùng còn qua một ma trận chiếu đầu ra WOW_O hình (dmodel,dmodel)(d_{model}, d_{model}) để á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 (batch,seq_len,dmodel)(batch, seq\_len, d_{model}), nên nhiều Transformer Block xếp chồng được như xếp Lego.


6. Vì sao chia cho dk\sqrt{d_k}: suy luận ổn định phương sai

Bước chia cho dk\sqrt{d_k} 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 QQ và KK 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 q⋅kq \cdot k.

Đặt q=(q1,q2,…,qdk)q = (q_1, q_2, \ldots, q_{d_k}), k=(k1,k2,…,kdk)k = (k_1, k_2, \ldots, k_{d_k}), mỗi qiq_i và kjk_j i.i.d., E[qi]=0E[q_i] = 0, Var(qi)=1\text{Var}(q_i) = 1.

Định nghĩa tích vô hướng:

q⋅k=∑i=1dkqi⋅kiq \cdot k = \sum_{i=1}^{d_k} q_i \cdot k_i

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:

E[qi⋅ki]=E[qi]⋅E[ki]=0×0=0E[q_i \cdot k_i] = E[q_i] \cdot E[k_i] = 0 \times 0 = 0 Var(qi⋅ki)=E[qi2⋅ki2]−(E[qi⋅ki])2\text{Var}(q_i \cdot k_i) = E[q_i^2 \cdot k_i^2] - (E[q_i \cdot k_i])^2

Vì qiq_i và kik_i độc lập:

E[qi2⋅ki2]=E[qi2]⋅E[ki2]=Var(qi)⋅Var(ki)=1×1=1E[q_i^2 \cdot k_i^2] = E[q_i^2] \cdot E[k_i^2] = \text{Var}(q_i) \cdot \text{Var}(k_i) = 1 \times 1 = 1

Vậy:

Var(qi⋅ki)=1−0=1\text{Var}(q_i \cdot k_i) = 1 - 0 = 1

Tích vô hướng là tổng dkd_k biến ngẫu nhiên độc lập. Theo tính cộng của phương sai:

Var(q⋅k)=∑i=1dkVar(qi⋅ki)=dk\text{Var}(q \cdot k) = \sum_{i=1}^{d_k} \text{Var}(q_i \cdot k_i) = d_k

6.3 Kết luận

Phương sai của q⋅kq \cdot k là dkd_k. Nghĩa là khi dk=64d_k = 64, độ lệch chuẩn của tích vô hướng khoảng 8; khi dk=128d_k = 128, 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 softmax(zi)=exp⁡(zi)/∑exp⁡(zj)\text{softmax}(z_i) = \exp(z_i) / \sum \exp(z_j) 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 dk\sqrt{d_k}:

Var(q⋅kdk)=Var(q⋅k)dk=dkdk=1\text{Var}\left(\frac{q \cdot k}{\sqrt{d_k}}\right) = \frac{\text{Var}(q \cdot k)}{d_k} = \frac{d_k}{d_k} = 1

Phương sai kéo về 1. Dù dkd_k 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:

softmax(zi)=ezi∑j=1Nezj\text{softmax}(z_i) = \frac{e^{z_i}}{\sum_{j=1}^{N} e^{z_j}}

Vấn đề là exp⁡\exp tăng cực nhanh. Khi ziz_i lớn (ví dụ zi=1000z_i = 1000), exp⁡(1000)\exp(1000) 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.

softmax(zi)=ezi−max⁡(z)∑j=1Nezj−max⁡(z)\text{softmax}(z_i) = \frac{e^{z_i - \max(z)}}{\sum_{j=1}^{N} e^{z_j - \max(z)}}

Chứng minh ngắn: nhân tử và mẫu với exp⁡(−max⁡(z))\exp(-\max(z)), tương đương chia cả hai cho exp⁡(max⁡(z))\exp(\max(z)). Sau khi trừ max, số mũ lớn nhất là 0, nên exp⁡\exp 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:

  1. Lượt 1: duyệt mọi phần tử, tìm max mm
  2. Lượt 2: duyệt mọi phần tử, tính exp⁡(zi−m)\exp(z_i - m) và tổng của chúng
  3. Lượt 3: duyệt mọi phần tử, chia mỗi exp⁡(zi−m)\exp(z_i - m) 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 NN 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ứ ii, giữ hai biến:

  • mim_i: max của ii phần tử đầu
  • did_i: tổng mũ của ii phần tử đầu (lấy mim_i làm gốc)

Hệ thức đệ quy:

mi=max⁡(mi−1,zi)\begin{aligned} m_i &= \max(m_{i-1}, z_i) \end{aligned} di=di−1⋅emi−1−mi+ezi−mi\begin{aligned} d_i &= d_{i-1} \cdot e^{m_{i-1} - m_i} + e^{z_i - m_i} \end{aligned}

Mẹo quan trọng ở dòng hai: khi max đổi từ mi−1m_{i-1} sang mim_i, tổng mũ cũ di−1d_{i-1} nhân hệ số sửa exp⁡(mi−1−mi)\exp(m_{i-1} - m_i) để “đổ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 N×NN \times 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ó WQW_Q, WKW_K, WVW_V 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ử dmodel=512d_{model} = 512, num_heads=8num\_heads = 8, thì head_dim=64head\_dim = 64.

Công thức đầy đủ Multi-Head Attention:

MultiHead(Q,K,V)=Concat(head1,…,headh)⋅WO\text{MultiHead}(Q, K, V) = \text{Concat}(\text{head}_1, \ldots, \text{head}_h) \cdot W_O headi=Attention(Q⋅WQi,K⋅WKi,V⋅WVi)\text{head}_i = \text{Attention}(Q \cdot W_Q^i, K \cdot W_K^i, V \cdot W_V^i)

Mỗi head có WQiW_Q^i, WKiW_K^i, WViW_V^i hình (dmodel,head_dim)=(512,64)(d_{model}, head\_dim) = (512, 64). Thực tế không giữ hh ma trận nhỏ: chiếu một lần bằng WQW_Q lớn hình (512,512)(512, 512), rồi reshape thành (seq_len,8,64)(seq\_len, 8, 64). 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 dmodel=512d_{model} = 512:

Ma trận tham sốHìnhSố tham số
WQW_Q(512,512)(512, 512)262,144
WKW_K(512,512)(512, 512)262,144
WVW_V(512,512)(512, 512)262,144
WOW_O(512,512)(512, 512)262,144
Tổng4×5122=1,048,5764 \times 512^2 = 1{,}048{,}576

Công thức chung: số tham số MHA là 4×dmodel24 \times d_{model}^2 (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 dmodeld_{model}.

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ó hh head thì phải cache hh bộ K và hh bộ V.

  • hh Q head, hh K head, hh 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 hh Q head, nhưng chỉ 1 K head và 1 V head.

  • hh 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ề 1/h1/h, 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 hh Q head thành gg nhóm (mỗi nhóm h/gh/g Q head), mỗi nhóm chia sẻ một bộ K và V.

  • hh Q head, gg K head, gg V head
  • Mỗi h/gh/g Q head dùng chung một bộ K và V
  • Khi g=hg = h thoái hóa về MHA; khi g=1g = 1 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 dmodel=4096d_{model} = 4096, num_heads = 32, head_dim = 128, độ dài chuỗi N=4096N = 4096, FP16:

Chỉ sốMHA (g=32g=32)GQA (g=8g=8)MQA (g=1g=1)
Số Q head323232
Số KV head3281
Tham số WKW_K32×128×4096=16M32 \times 128 \times 4096 = 16\text{M}8×128×4096=4M8 \times 128 \times 4096 = 4\text{M}1×128×4096=0.5M1 \times 128 \times 4096 = 0.5\text{M}
Tham số WVW_V16M16\text{M}4M4\text{M}0.5M0.5\text{M}
KV Cache / token2×32×128×2B=16KB2 \times 32 \times 128 \times 2\text{B} = 16\text{KB}2×8×128×2B=4KB2 \times 8 \times 128 \times 2\text{B} = 4\text{KB}2×1×128×2B=0.5KB2 \times 1 \times 128 \times 2\text{B} = 0.5\text{KB}
KV Cache / lớp @ 4096 token64MB64\text{MB}16MB16\text{MB}2MB2\text{MB}
Chất lượng mô hìnhTốt nhấtGần MHACó 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(g=8g=8), 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: X→WK→KX \to W_K \to K (cache KK, chiều = num_kv_heads ×\times head_dim)
  • MLA: X→WDKV→cX \to W_{DKV} \to c (cache cc, chiều nhỏ hơn nhiều so với KK), lúc suy luận c→WUK→Kc \to W_{UK} \to K 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ánHìnhFLOPs
Q=X⋅WQQ = X \cdot W_Q(N,d)×(d,d)(N, d) \times (d, d)2Nd22Nd^2
K=X⋅WKK = X \cdot W_K(N,d)×(d,d)(N, d) \times (d, d)2Nd22Nd^2
V=X⋅WVV = X \cdot W_V(N,d)×(d,d)(N, d) \times (d, d)2Nd22Nd^2
S=Q⋅KTS = Q \cdot K^T(N,d)×(d,N)(N, d) \times (d, N)2N2d2N^2d
O=S⋅VO = S \cdot V(N,N)×(N,d)(N, N) \times (N, d)2N2d2N^2d
Final=O⋅WOFinal = O \cdot W_O(N,d)×(d,d)(N, d) \times (d, d)2Nd22Nd^2

Tổng: 8Nd2+4N2d8Nd^2 + 4N^2d.

Khi NN nhỏ (ví dụ N<dN < d), hạng 8Nd28Nd^2 chiếm ưu thế: bottleneck nằm ở GEMM chiếu QKV. Khi NN lớn (ví dụ N>dN > d), hạng 4N2d4N^2d 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. Q⋅KTQ \cdot K^T 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ớ O(N2)O(N^2)

Bản Attention chuẩn phải tính và lưu tường minh cả ma trận Attention N×NN \times N. Với seq_len = 128K\text{K}, num_heads = 32:

32×128K×128K×2B=32×16,384M×2=1,048,576 MB=1 TB32 \times 128\text{K} \times 128\text{K} \times 2\text{B} = 32 \times 16{,}384\text{M} \times 2 = 1{,}048{,}576\text{ MB} = 1\text{ TB}

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 32×8K×8K×2B=4 GB32 \times 8\text{K} \times 8\text{K} \times 2\text{B} = 4\text{ GB}, đã là chi phí không thể bỏ qua.

Bottleneck bộ nhớ O(N2)O(N^2) đó 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 Q⋅KTQ \cdot K^T trên HBM rồi ghi cả ma trận N×NN \times N về HBM; đọc lại để Softmax rồi ghi về HBM; cuối cùng đọc kết quả Softmax và VV để nhân. Mỗi bước đều đọc/ghi ma trận N×NN \times 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 N×NN \times N, 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 N×NN \times 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 QiQ_i với khối K1K_1 ta có đầu ra cục bộ O1O_1 và thống kê Softmax (max cục bộ m1m_1, mẫu số cục bộ l1l_1). Tiếp tục với khối K2K_2 được m2m_2 và l2l_2. Sửa bằng Online Softmax:

mnew=max⁡(m1,m2)\begin{aligned} m_{\text{new}} &= \max(m_1, m_2) \end{aligned} lnew=l1⋅em1−mnew+l2⋅em2−mnew\begin{aligned} l_{\text{new}} &= l_1 \cdot e^{m_1 - m_{\text{new}}} + l_2 \cdot e^{m_2 - m_{\text{new}}} \end{aligned} Onew=O1⋅l1⋅em1−mnew+O2local⋅l2⋅em2−mnewlnew\begin{aligned} O_{\text{new}} &= \frac{O_1 \cdot l_1 \cdot e^{m_1 - m_{\text{new}}} + O_2^{\text{local}} \cdot l_2 \cdot e^{m_2 - m_{\text{new}}}}{l_{\text{new}}} \end{aligned}

Cộng dồn liên tục đến khi hết mọi khối KK. 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 SS (N×NN \times N): O(N2)O(N^2) lần ghi
  • Đọc SS để Softmax: O(N2)O(N^2) lần đọc
  • Ghi P=Softmax(S)P = \text{Softmax}(S): O(N2)O(N^2) lần ghi
  • Đọc PP để P⋅VP \cdot V: O(N2)O(N^2) lần đọc
  • Tổng truy cập HBM: O(N2)O(N^2) (cộng thêm O(Nd)O(Nd) đọc của chính Q, K, V)

Lượng truy cập HBM của FlashAttention:

  • Đọc Q, K, V: O(Nd)O(Nd)
  • Ghi đầu ra cuối O: O(Nd)O(Nd)
  • Ma trận Attention trung gian mỗi lần chỉ tính một tile trong SRAM rồi bỏ; ma trận N×NN \times 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 Tc=O(Nd/M)T_c = O(Nd/M) khối K/V, mỗi lần quét hết Q và O (O(Nd)O(Nd))
  • Tổng truy cập HBM: Θ(N2d2/M)\Theta(N^2 d^2 / M) (MM là kích thước SRAM tính theo phần tử)

Vì d2d^2 nhỏ hơn nhiều so với MM (ví dụ d=128d = 128 thì d2=16384d^2 = 16384, 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 Θ(N2)\Theta(N^2) của bản chuẩn. Nó vẫn tăng theo N2N^2, chỉ là hằng số giảm mạnh; thứ thực sự xuống O(N)O(N) là bộ nhớ phụ thêm (không còn lưu ma trận N×NN \times N). Lưu ý, khối lượng tính (FLOPs) không đổi, vẫn O(N2d)O(N^2d); 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 dk\sqrt{d_k}: giả sử qi,kiq_i, k_i chuẩn tắc độc lập, suy ra phương sai tích vô hướng bằng dkd_k
  • 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, dmodeld_{model}, 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ừ O(N2)O(N^2) xuống O(N)O(N), truy cập HBM từ Θ(N2)\Theta(N^2) xuống Θ(N2d2/M)\Theta(N^2 d^2 / M)

📚 Tài liệu tham khảo

Paper

Tutorial và blog