6.2 FlashAttention V2 chi tiết

Bài này phân tích sâu các cải tiến cốt lõi của FlashAttention V2 so với V1: đổi thứ tự vòng lặp ngoài/trong, phân công việc tốt hơn trong thread block, giảm phép không phải nhân ma trận, và đạt 50–73% FLOPS lý thuyết trên A100.

📑 Mục lục

• 1. Nhắc lại nút thắt hiệu năng của V1
• 2. Cải tiến 1: đổi thứ tự vòng lặp
• 3. Cải tiến 2: giảm phép không phải nhân ma trận
• 4. Cải tiến 3: phân công ở cấp Warp
• 5. Tối ưu Causal Mask
• 6. Thuật toán forward đầy đủ
• 7. Tối ưu backward
• 8. Phân tích hiệu năng và so sánh
• Tóm tắt
• Checklist tự kiểm
• Tài liệu tham khảo

---

1. Nhắc lại nút thắt hiệu năng của V1

1.1 Mức tận dụng thực tế của V1

FlashAttention V1 trên A100 chỉ đạt 25–40% FLOPS lý thuyết. Với một toán tử lẽ ra đã Compute-Bound sau khi giảm IO, mức tận dụng đó không lý tưởng.

Nút thắt đến từ ba phía (ảnh hưởng lớn → nhỏ):

1.2 Bối cảnh phần cứng

Thông số then chốt của A100:

• Đỉnh FP16 Tensor Core: 312 TFLOPS
• Băng thông HBM: 2 TB/s
• Dung lượng SRAM: 192 KB/SM (Shared Memory; tối đa cấu hình được 164 KB cho user)
• Số Warp: mỗi Thread Block có thể chứa nhiều Warp (thường 4–8)

Muốn tận dụng Tensor Core, pipeline tính toán không được bị cắt bởi phép non-GEMM.

---

2. Cải tiến 1: đổi thứ tự vòng lặp

2.1 Vấn đề cấu trúc vòng của V1

Ở V1 vòng ngoài duyệt khối K/V, vòng trong duyệt khối Q. Hai hệ quả:

1. Đọc/ghi Q và O lặp lại: mỗi khối K/V mới buộc mọi khối Q và O tương ứng load lại và ghi về HBM
2. Độ song song bị hạn: giữa các khối K/V vòng ngoài có phụ thuộc tuần tự qua running max/sum của online softmax, nên V1 chỉ song song theo batch × heads. Khi batch và heads nhỏ, chuỗi dài (ngữ cảnh dài), không tận dụng được song song theo chiều sequence, SM tận dụng kém

2.2 Cấu trúc vòng mới của V2

V2 đổi thành vòng ngoài duyệt khối Q, vòng trong duyệt khối K/V:

Outer loop: for i = 1 to T_r (over Q tiles) — mapped to different Thread Blocks
    Load Q_i from HBM into SRAM (once)
    Init O_i = 0, m_i = -inf, l_i = 0
    Inner loop: for j = 1 to T_c (over K/V tiles)
        Load K_j, V_j from HBM into SRAM
        Compute S_ij, update stats and O_i (all in SRAM)
    Write final O_i back to HBM (once)

2.3 Lợi ích

Đổi tưởng chừng đơn giản mang lại ba lợi ích:

Lợi ích 1: O chỉ ghi một lần cuối

Ở V1, mỗi khối K/V đều đọc rồi ghi O — tổng $2T_c$ lần truy cập HBM ($T_c$ đọc + $T_c$ ghi). Ở V2, O chỉ ghi một lần sau vòng trong.

Lợi ích 2: vòng ngoài song song tự nhiên

Các khối Q độc lập hoàn toàn (mỗi khối tính hàng output của mình), map thẳng lên Thread Block khác nhau để song song cấp GPU.

Lợi ích 3: Q chỉ load một lần

Mỗi Thread Block phụ trách một khối Q; Q nằm trong Shared Memory suốt vòng trong.

📌 Điểm then chốt: V1 để K/V ngoài vì logic tích lũy của online softmax (cập nhật chuẩn hóa O sau mỗi cột khối K/V). Cái giá là O phải đọc+ghi lặp. V2 nhận ra: load Q một lần tốt hơn, vì giảm đọc/ghi O mang lại nhiều hơn giảm đọc K/V.

---

3. Cải tiến 2: giảm phép không phải nhân ma trận

3.1 Vấn đề: nút thắt lệnh non-GEMM

Trong Attention tiled, ngoài hai nhân ma trận ($QK^\top$ và $PV$), còn nhiều phép từng phần tử:

• Trừ max: $S_{ij} - m_i$
• Exp: $\exp(\cdot)$
• Tổng hàng: $\text{rowsum}(\cdot)$
• Rescaling: nhân hệ số hiệu chỉnh

Các phép này không dùng được Tensor Core, chỉ CUDA Core, trở thành nút thắt trong pipeline tính toán.

3.2 Cách V2 tối ưu: trì hoãn Rescaling

V2 trì hoãn bước rescaling Softmax đến sau vòng trong, làm một lần:

Cách V1 (rescale mọi khối K/V):

$$O_i^{(j)} = \frac{l_i^{(j-1)} \cdot e^{m_i^{(j-1)} - m_i^{(j)}}}{l_i^{(j)}} \cdot O_i^{(j-1)} + \frac{e^{S_{ij} - m_i^{(j)}}}{l_i^{(j)}} \cdot V_j$$

Cách V2 (chỉ giữ tích lũy chưa chuẩn hóa):

$$\tilde{O}_i^{(j)} = e^{m_i^{(j-1)} - m_i^{(j)}} \cdot \tilde{O}_i^{(j-1)} + e^{S_{ij} - m_i^{(j)}} \cdot V_j$$

Chuẩn hóa một lần cuối:

$$O_i = \tilde{O}_i^{(T_c)} / l_i^{(T_c)}$$

3.3 Định lượng lợi ích

Tối ưu này bỏ ba phép từng phần tử mỗi vòng lặp trong:

• Bỏ chuẩn hóa O cũ: nhân $l^{(j-1)}$
• Chuẩn hóa lại O cũ: chia $l^{(j)}$
• Chuẩn hóa đóng góp mới: chia $l^{(j)}$

V2 gỡ hết 「2 chia + 1 nhân」 khỏi vòng trong, chỉ chia một lần cuối. Khi số vòng $T_c$ lớn, tiết kiệm khá rõ.

---

4. Cải tiến 3: phân công ở cấp Warp

4.1 Phương án V1: cắt theo chiều K/V

Ở V1, nhiều Warp trong Thread Block cắt cùng một khối $S_{ij}$ theo chiều K/V.

Giả sử 4 Warp, kích thước khối K/V $B_c = 128$:

• Warp 0 phụ trách cột K/V 0–31
• Warp 1 phụ trách cột 32–63
• Warp 2 phụ trách cột 64–95
• Warp 3 phụ trách cột 96–127

Mỗi Warp chỉ tính được một phần cột của $S_{ij}$, nên:

• Tính row-max / row-sum phải gộp kết quả cục bộ giữa các Warp
• Tính $P \cdot V$ thì nhiều Warp cùng cộng vào cùng một hàng O, phải sync lại

Kiểu phân 「split-K」 này gây sync và giao tiếp shared memory giữa Warp thường xuyên.

4.2 Phương án V2: cắt theo chiều Q

V2 chia tiếp khối Q cho các Warp khác nhau:

Giả sử 4 Warp, kích thước khối Q $B_r = 128$:

• Warp 0 phụ trách hàng Q 0–31
• Warp 1 phụ trách hàng 32–63
• Warp 2 phụ trách hàng 64–95
• Warp 3 phụ trách hàng 96–127

Mỗi Warp độc lập tính attention giữa các hàng Q của mình với toàn bộ khối K/V, không xen nhau.

4.3 So hai chiến lược

💡 Gợi ý: Nếu cắt theo K/V cho Warp khác nhau, mỗi Warp chỉ thấy một phần đầu vào Softmax, phải giao tiếp giữa Warp để gộp max và sum — đó là một lý do V1 chậm.

---

5. Tối ưu Causal Mask

5.1 Vai trò của Causal Mask

Trong mô hình tự hồi quy (như GPT), token thứ $i$ khi sinh chỉ được “nhìn” chính nó và các token trước — vị trí $i$ chỉ attend tới $j \leq i$. Đó là ràng buộc causality then chốt khi huấn luyện LM tự hồi quy: mô hình phải dự đoán token tiếp theo từ lịch sử, không được nhìn đáp án.

Trong hiện thực, dùng ma trận mask tam giác dưới, cộng vào điểm attention trước Softmax: $\tilde{S} = QK^\top / \sqrt{d} + M$. Vị trí bị che cộng $-\infty$, sau $\exp$ thành 0, không đóng góp vào Softmax.

Hình dung: ma trận attention $N \times N$ chỉ tam giác dưới (gồm đường chéo) có hiệu lực, tam giác trên toàn 0. Về lý thuyết khoảng nửa phép tính vô ích — đúng cấu trúc chiến lược block-skip của V2 khai thác.

5.2 Lãng phí của hiện thực ngây thơ

Nếu không xử lý đặc biệt, vị trí mask = 0 vẫn chạy đủ GEMM và Softmax, lãng phí khoảng nửa công tính.

5.3 Chiến lược bỏ khối cấp block của V2

V2 thêm tối ưu Causal Mask cấp block:

Với khối Q thứ $i$ (hàng $[iB_r, (i+1)B_r)$) và khối K thứ $j$ (cột $[jB_c, (j+1)B_c)$):

1. Hoàn toàn dưới mask ($(i+1)B_r - 1 \geq (j+1)B_c - 1$ và $iB_r \geq (j+1)B_c - 1$, tức min hàng Q ≥ max cột K): mọi vị trí trong khối thỏa row ≥ col, thấy hết, tính bình thường
2. Hoàn toàn trên mask ($(i+1)B_r - 1 < jB_c$, tức max hàng Q < min cột K): mọi vị trí thỏa row < col, không thấy gì, bỏ hẳn khối
3. Cắt đường chéo: trong khối vừa có vị trí thấy vừa có không thấy — áp mask từng phần tử

Below diagonal (all visible)    Diagonal block (partial mask)    Above diagonal (all skip)
┌─────────────┐      ┌─────────────┐      ┌─────────────┐
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ □ □ □ │      │ □ □ □ □ □ □ │
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ ■ □ □ │      │ □ □ □ □ □ □ │
│ ■ ■ ■ ■ ■ ■ │      │ ■ ■ ■ ■ ■ □ │      │ □ □ □ □ □ □ │
└─────────────┘      └─────────────┘      └─────────────┘

5.4 Hiệu quả tăng tốc

Với Causal Attention, bỏ khối của V2 giảm khoảng 50% tính vô ích. Kết hợp phân vòng ngoài theo Q, mỗi Thread Block chỉ cần duyệt K/V tới vị trí đường chéo rồi thoát vòng trong sớm.

---

6. Thuật toán forward đầy đủ

6.1 Pseudocode

def flash_attention_v2_forward(Q, K, V, B_r, B_c, causal=False):
    N, d = Q.shape
    O = zeros(N, d)
    L = zeros(N)  # logsumexp, for backward

    # Outer loop (parallelized across Thread Blocks)
    parallel for i in range(0, N, B_r):
        Q_i = Q[i:i+B_r]  # load into SRAM; stays for whole inner loop

        # Init accumulators (SRAM/registers)
        m_i = full(B_r, -inf)
        l_i = zeros(B_r)
        O_tilde = zeros(B_r, d)  # unnormalized accumulated output

        # Inner loop: walk K/V tiles
        kv_end = min(N, (i + B_r)) if causal else N
        for j in range(0, kv_end, B_c):
            K_j = K[j:j+B_c]  # load from HBM
            V_j = V[j:j+B_c]  # load from HBM

            # Local attention scores
            S_ij = Q_i @ K_j.T / sqrt(d)

            # Apply causal mask (only needed on diagonal-crossing tiles)
            # i, j are absolute start coords; when K max col (j+B_c-1) > Q min row (i), this tile crosses the diagonal
            if causal and i < j + B_c - 1:
                apply_causal_mask(S_ij, i, j, B_r, B_c)

            # Update statistics
            m_new = max(m_i, rowmax(S_ij))
            P_ij = exp(S_ij - m_new[:, None])
            l_new = exp(m_i - m_new) * l_i + rowsum(P_ij)

            # Update unnormalized output
            O_tilde = exp(m_i - m_new)[:, None] * O_tilde + P_ij @ V_j

            m_i = m_new
            l_i = l_new

        # Final normalize and write back
        O[i:i+B_r] = O_tilde / l_i[:, None]
        L[i:i+B_r] = m_i + log(l_i)  # save logsumexp for backward

    return O, L

6.2 Khác biệt then chốt so với V1

---

7. Tối ưu backward

7.1 Lưu logsumexp thay vì tách m và l

V2 trong forward lưu $L_i = m_i + \log(l_i)$ (logsumexp theo hàng) thay vì lưu riêng $m_i$ và $l_i$. Hai lợi ích:

• Giảm một nửa lưu thống kê
• Backward khôi phục trực tiếp $P_{ij}$: $P_{ij} = \exp(S_{ij} - L_i)$

7.2 Chiến lược song song của backward

Backward của V2 dùng vòng ngoài duyệt khối K/V (giống backward V1, ngược với forward V2 dùng Q ngoài), để cập nhật $dK$ và $dV$ không cần atomic:

• $dK_j$ và $dV_j$ do một Thread Block độc quyền tích lũy (Block đó duyệt mọi khối Q)
• $dQ_i$ phải cộng giữa các Thread Block (atomic add hoặc quét hai lượt)

7.3 Vector D trong backward

Gradient Softmax cần $dS_{ij} = P_{ij} \odot (dP_{ij} - D_i)$, trong đó

$$D_i = \text{rowsum}(dP_i \odot P_i) = \text{rowsum}(dO_i \odot O_i) \in \mathbb{R}^N$$

Phép đổi đẳng thức này đã có trong paper V1 (phụ lục B.3): vế trái phải reduce cả hàng dài $N$, $P_i$, $dP_i$ có thể không vừa SRAM; đổi sang vế phải chỉ cần tích vô hướng vector dài $d$, chi phí không phụ thuộc độ dài sequence.

Cải tiến của V2 nằm ở tầng hiện thực — đưa $D$ ra một preprocess kernel riêng, tính một lần cả vector $\mathbb{R}^N$, vòng chính backward đọc trực tiếp; còn V1 tính từng khối ngay trong vòng chính backward.

---

8. Phân tích hiệu năng và so sánh

8.1 Hiệu năng trên A100

V2 nhanh hơn V1 khoảng 1.6x.

8.2 So với các hiện thực khác

Trên A100-80GB, độ dài sequence 2048, head dim 128:

8.3 Lợi thế VRAM chuỗi dài (ước lượng tổng activation lúc train)

Bảng dưới ước lượng tổng VRAM activation liên quan attention khi train (gồm tensor trung gian và buffer gradient, ước theo cấu hình multi-head điển hình, để thấy khác biệt scale $N^2$ vs $N$), không phải chi phí lưu một ma trận đơn:

---

📝 Tóm tắt

FlashAttention V2 nâng hiệu suất trên A100 từ 25–40% lên 50–73% nhờ ba cải tiến cốt lõi + một tối ưu phụ:

1. Đổi thứ tự vòng: vòng ngoài duyệt khối Q (song song), vòng trong duyệt khối K/V — giảm đọc/ghi HBM của O và song song tự nhiên
2. Trì hoãn chuẩn hóa (giảm non-GEMM): vòng trong chỉ scale bằng exp, chuẩn hóa thống nhất cuối cùng
3. Tối ưu cấp Warp: mỗi Warp phụ trách hàng Q khác nhau — tránh sync và giao tiếp giữa Warp
4. Phụ: bỏ khối Causal: dùng cấu trúc tam giác của mask để bỏ khối vô ích — Causal Attention giảm thêm ~50% tính toán

---

🎯 Checklist tự kiểm

• Giải thích vì sao V2 đổi vòng ngoài từ K/V sang Q
• Nói rõ vì sao số lần truy cập HBM của O giảm từ $2T_c$ xuống 1
• Mô tả chiến lược phân Warp của V2 và ưu thế so với V1
• Chứng minh tính đúng của trì hoãn chuẩn hóa (tích lũy chưa chuẩn hóa rồi chia $l$ tương đương chuẩn hóa từng bước)
• Vẽ được ba trường hợp bỏ khối dưới Causal Mask
• Giải thích logsumexp $L = m + \log(l)$ đơn giản hóa backward thế nào
• So sánh khác biệt tận dụng TFLOPS của V1 và V2 trên A100

---

📚 Tài liệu tham khảo

• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• FlashAttention-1: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention hiện thực chính thức - GitHub
• Tri Dao - FlashAttention-2 Blog Post
• NVIDIA A100 Tensor Core GPU Architecture