5.2 Hiện thực và tối ưu CUDA Online Softmax

Safe Softmax truyền thống cần ba lần quét dữ liệu đầu vào. Online Softmax dùng một công thức đệ quy khéo léo để vừa tính max vừa tính sum trong một lần quét, giảm số lần đọc bộ nhớ từ 3 xuống 2. Bài này xuất phát từ suy diễn thuật toán, rồi hiện thực và tối ưu dần CUDA Kernel của Online Softmax — cũng là kiến thức nền cốt lõi để hiểu FlashAttention.

📑 Mục lục

• 1. Vì sao cần Online Softmax
• 2. Suy diễn thuật toán Online Softmax
• 3. Phiên bản V0: Online Softmax đơn luồng
• 4. Phiên bản V1: song song cấp Block + Warp Shuffle merge reduce
• 5. Phiên bản V2: tải vector hóa float4
• 6. Phiên bản V3: cache thanh ghi để bỏ lần đọc thứ hai
• 7. Phiên bản V4: song song nhiều hàng + Grid Stride
• 8. So sánh hiệu năng và liên hệ với FlashAttention
• Tóm tắt
• Checklist tự kiểm
• Tài liệu tham khảo

---

1. Vì sao cần Online Softmax

1.1 Cái giá của ba lần quét

Trong bài trước, Safe Softmax dùng ba lần quét:

$$\underbrace{m = \max_j(x_j)}_{\text{第1遍}} \quad \rightarrow \quad \underbrace{d = \sum_j e^{x_j - m}}_{\text{第2遍}} \quad \rightarrow \quad \underbrace{y_i = \frac{e^{x_i - m}}{d}}_{\text{第3遍}}$$

Mỗi lần quét đọc input một lần, tổng lượng đọc bộ nhớ là $3N$. Với thao tác Memory-Bound, bớt một lần quét gần như tăng hiệu năng 33%. Vấn đề: lần 2 (tính sum) phải biết kết quả lần 1 (max), nên trông như không gộp được.

Ví dụ: bạn muốn chuẩn hóa điểm cả lớp. Cách cũ: lật sổ điểm một lần tìm điểm cao nhất, lật lần nữa tính tổng đã điều chỉnh, rồi lật lần ba viết điểm đã chuẩn hóa — ba lần lật sổ. Có cách chỉ lật hai lần, thậm chí một lần?

1.2 Insight cốt lõi

Insight then chốt của Online Softmax: max và sum có thể duy trì trong một lần quét, với cái giá là “sửa” sum đã tích lũy mỗi khi xuất hiện max mới.

Hình dung bạn vừa lật sổ vừa cộng dồn “điểm đã điều chỉnh”. Đến người thứ 50 thấy điểm cao hơn max trước — mọi điểm đã điều chỉnh trước đó đều lệch! Nhưng lệch bao nhiêu thì biết chắc: mỗi $e^{x_i - m_{\text{old}}}$ trước đây phải thành $e^{x_i - m_{\text{new}}}$, tổng chỉ cần nhân hệ số sửa $e^{m_{\text{old}} - m_{\text{new}}}$.

Đó là tinh túy của Online Softmax — dùng phép nhân để sửa thay vì tính lại.

---

2. Suy diễn thuật toán Online Softmax

2.1 Suy diễn quan hệ đệ quy

Giả sử đã xử lý $k$ phần tử đầu, trạng thái đang giữ là:

• $m_k = \max(x_0, x_1, \ldots, x_{k-1})$ (max của $k$ phần tử đầu)
• $d_k = \sum_{j=0}^{k-1} e^{x_j - m_k}$ (tổng mũ theo max hiện tại)

Khi xử lý phần tử thứ $k+1$ là $x_k$:

Bước 1: cập nhật max

$$m_{k+1} = \max(m_k, x_k)$$

Bước 2: sửa sum đã có và thêm phần tử mới

$$d_{k+1} = d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}}$$

Suy diễn:

$$\begin{aligned} d_{k+1} &= \sum_{j=0}^{k} e^{x_j - m_{k+1}} \end{aligned}$$ $$\begin{aligned} &= \sum_{j=0}^{k-1} e^{x_j - m_{k+1}} + e^{x_k - m_{k+1}} \end{aligned}$$ $$\begin{aligned} &= \sum_{j=0}^{k-1} e^{x_j - m_k} \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}} \end{aligned}$$ $$\begin{aligned} &= d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}} \end{aligned}$$

💡 Gợi ý: then chốt là hệ số sửa $e^{m_k - m_{k+1}}$. Nếu phần tử mới không phải max mới ($m_{k+1} = m_k$), thì $e^{m_k - m_{k+1}} = e^0 = 1$, bước sửa rút gọn thành cộng đơn giản — không tốn thêm gì.

2.2 Pseudocode thuật toán đầy đủ

// Pass 1: Online compute max and sum together
m = -∞
d = 0
for j = 0 to N-1:
    m_new = max(m, x[j])
    d = d * exp(m - m_new) + exp(x[j] - m_new)
    m = m_new

// Pass 2: normalize output
for j = 0 to N-1:
    y[j] = exp(x[j] - m) / d

Ba lần xuống còn hai lần: lần 1 vừa có $m$ vừa có $d$, lần 2 chuẩn hóa. Số lần đọc bộ nhớ từ $3N$ xuống $2N$, lý thuyết tăng 33%.

2.3 Đảm bảo ổn định số

Ổn định số của Online Softmax giống hệt Safe Softmax:

• $m$ luôn là max các phần tử đã thấy, mọi đối số mũ $x_j - m \leq 0$, không tràn trên
• Hệ số sửa $e^{m_k - m_{k+1}} \leq 1$ (vì $m_{k+1} \geq m_k$), cũng không tràn trên
• $d$ cuối cùng tương đương toán học với kết quả ba lần của Safe Softmax

2.4 Môi trường kiểm thử

Giống bài trước: A100 80GB SXM4, ma trận kiểm thử $(M, N) = (4096, 4096)$.

---

3. Phiên bản V0: Online Softmax đơn luồng

3.1 Ý tưởng thuật toán

Hiện thực thẳng nhất: mỗi thread xử lý một hàng, tuần tự duy trì $(m, d)$, rồi làm một lần chuẩn hóa.

3.2 Hiện thực Kernel

// V0: single-thread Online Softmax (one thread per row)
__global__ void online_softmax_v0(float* input, float* output, int M, int N) {
    int row = blockIdx.x * blockDim.x + threadIdx.x;
    if (row >= M) return;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: Online compute max and sum together
    float m = -INFINITY;
    float d = 0.0f;

    for (int i = 0; i < N; i++) {
        float xi = x[i];
        float m_new = fmaxf(m, xi);
        d = d * expf(m - m_new) + expf(xi - m_new);
        m = m_new;
    }

    // Pass 2: normalize
    float inv_d = 1.0f / d;
    for (int i = 0; i < N; i++) {
        y[i] = expf(x[i] - m) * inv_d;
    }
}

3.3 Phân tích hiệu năng

Giống Safe Softmax V0, một thread xử lý cả hàng bị thiếu song song nghiêm trọng và truy cập không coalesced. Nhưng so với V0 ba lần quét, ở đây chỉ đọc input hai lần.

⚠️ Lưu ý: mỗi bước đệ quy Online cần gọi expf (cho hệ số sửa), tăng gánh tính toán. Nhưng Softmax vốn Memory-Bound, nên ở phiên bản hiệu năng cao phần tính thêm này bị latency bộ nhớ che hết.

V0 băng thông đo được: khoảng 7% (~143 GB/s) — chỉ dùng làm baseline kiểm đúng.

---

4. Phiên bản V1: song song cấp Block + Warp Shuffle merge reduce

4.1 Thách thức song song hóa

Đệ quy Online Softmax là tuần tự — $(m, d)$ mỗi bước phụ thuộc bước trước. Song song hóa thế nào?

Quan sát then chốt: cho mỗi thread xử lý một phần phần tử, tự duy trì $(m_{\text{local}}, d_{\text{local}})$, rồi gộp các kết quả cục bộ.

Công thức gộp hai kết quả cục bộ $(m_1, d_1)$ và $(m_2, d_2)$:

$$\begin{aligned} m_{\text{merged}} &= \max(m_1, m_2) \end{aligned}$$ $$\begin{aligned} d_{\text{merged}} &= d_1 \cdot e^{m_1 - m_{\text{merged}}} + d_2 \cdot e^{m_2 - m_{\text{merged}}} \end{aligned}$$

Phép gộp này thỏa kết hợp tính (associative — dùng được tree reduce), nên song song hóa hiệu quả.

4.2 Merge reduce bằng Warp Shuffle

Theo kinh nghiệm tối ưu Safe Softmax ở bài trước, reduce tốt nhất làm ở tầng thanh ghi: Warp Shuffle latency thấp, không Bank Conflict, Warp tự đồng bộ, hiệu quả hơn Shared Memory reduce.

Khác Reduce thường: reduce của Online Softmax không phải max hay sum đơn giản mà gộp đồng thời hai giá trị $(m, d)$. May là __shfl_down_sync truyền riêng từng giá trị, mỗi vòng dùng công thức gộp ở trên. Cấu trúc tổng thể vẫn là reduce hai cấp “Warp Shuffle trong Warp + Shared Memory giữa các Warp”.

4.3 Hiện thực Kernel

// Online merge reduce within a Warp
__device__ void warpReduceOnline(float& m, float& d, const int Warp_Threads = 32) {
    for (int offset = Warp_Threads / 2; offset > 0; offset >>= 1) {
        float m2 = __shfl_down_sync(0xffffffff, m, offset);
        float d2 = __shfl_down_sync(0xffffffff, d, offset);

        float m_new = fmaxf(m, m2);
        d = d * expf(m - m_new) + d2 * expf(m2 - m_new);
        m = m_new;
    }
}

// V1: Block-level parallel + Warp Shuffle two-level merge reduce
__global__ void online_softmax_v1(float* input, float* output, int M, int N, const int Warp_Size) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % Warp_Size;
    int wid  = tid / Warp_Size;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Pass 1: each thread Online-processes a segment
    float local_m = -INFINITY;
    float local_d = 0.0f;

    for (int i = tid; i < N; i += blockDim.x) {
        float xi = x[i];
        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // Level 1: merge within Warp
    warpReduceOnline(local_m, local_d, Warp_Size);

    // Exchange across Warps via Shared Memory
    __shared__ float warp_m[32];
    __shared__ float warp_d[32];

    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    // Level 2: Warp 0 does the final merge
    int num_warps = blockDim.x / Warp_Size;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d, num_warps);
    }

    // Broadcast final result
    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // Pass 2: normalize output
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - row_max) * inv_sum;
    }
}

Trong hiện thực, cần chú ý biên offset ở cấp Warp thứ hai: khi lane > num_warps thì local_m = -INFINITY. Giả sử num_warps = 8, khi offset = 16, thread 9 lấy được m2 = -INFINITY từ thread 25, trong khi thread 9 cũng có m = -INFINITY, nên m_new = -INFINITY, m - m_new = -INFINITY - (-INFINITY) = NaN, expf(NaN) = NaN, cuối cùng sum thành Nan, kết quả tính ra là NaN.

4.4 So với Safe Softmax ba lần

⚠️ Lưu ý: thứ tự gộp ảnh hưởng độ chính xác số thực (cộng không kết hợp chặt). Softmax không nhạy với sai số nhỏ, nên độ chính xác tree reduce là đủ.

💡 Gợi ý: reduce Online nặng hơn sum-reduce thường (mỗi bước 2 lần expf), nhưng băng thông bộ nhớ tiết kiệm được lớn hơn nhiều so với tính thêm. Với thao tác Memory-Bound, bỏ một lần đọc global là quyết định.

V1 băng thông đo được: khoảng 58% (~1183 GB/s), khoảng 8.3× so với V0 đơn luồng, khoảng 1.4× so với phiên bản tương ứng của Safe Softmax ba lần.

---

5. Phiên bản V2: tải vector hóa float4

5.1 Ý tưởng cải tiến

Cùng đường tối ưu với Safe Softmax: dùng float4 tải vector hóa để giảm số lệnh. Ở giai đoạn đệ quy Online, mỗi lần tải 4 phần tử rồi cập nhật $(m, d)$ lần lượt.

5.2 Hiện thực Kernel

// V2: float4 vectorized load + Online Softmax
__global__ void online_softmax_v2(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    float* x = input  + row * N;
    float* y = output + row * N;

    float4* x4 = reinterpret_cast<float4*>(x);
    float4* y4 = reinterpret_cast<float4*>(y);
    int N4 = N / 4;

    // Pass 1: vectorized Online recurrence
    float local_m = -INFINITY;
    float local_d = 0.0f;

    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];

        // Update (m, d) for 4 elements in order
        float vals[4] = {data.x, data.y, data.z, data.w};
        for (int k = 0; k < 4; k++) {
            float m_new = fmaxf(local_m, vals[k]);
            local_d = local_d * expf(local_m - m_new) + expf(vals[k] - m_new);
            local_m = m_new;
        }
    }
    // Tail handling
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        float xi = x[i];
        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // Two-level Warp Shuffle merge (same as V1)
    warpReduceOnline(local_m, local_d);

    __shared__ float warp_m[32];
    __shared__ float warp_d[32];
    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d);
    }

    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // Pass 2: vectorized normalize
    for (int i = tid; i < N4; i += blockDim.x) {
        float4 data = x4[i];
        float4 result;
        result.x = expf(data.x - row_max) * inv_sum;
        result.y = expf(data.y - row_max) * inv_sum;
        result.z = expf(data.z - row_max) * inv_sum;
        result.w = expf(data.w - row_max) * inv_sum;
        y4[i] = result;
    }
    for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
        y[i] = expf(x[i] - row_max) * inv_sum;
    }
}

5.3 float4 kết hợp đệ quy Online thế nào

float4 tải một lần 4 phần tử, nhưng đệ quy Online vẫn cập nhật $(m, d)$ từng phần tử. Trông như không lợi, thực tế:

1. Bộ nhớ: float4 chuyển cùng lượng dữ liệu với ít lệnh hơn, giảm nghẽn phát lệnh
2. Tính toán: bốn lần expf có thể pipeline; compiler sắp xếp lệnh tốt hơn
3. Tổng thể: số vòng lặp giảm 4 lần, overhead vòng (nhánh, tính địa chỉ) giảm mạnh

V2 băng thông đo được: khoảng 70% (~1427 GB/s), khoảng 1.21× so với V1.

---

6. Phiên bản V3: cache thanh ghi để bỏ lần đọc thứ hai

6.1 Ý tưởng tối ưu tận cùng

Hai lần quét nghĩa là input vẫn bị đọc hai lần. Có đọc một lần thôi không?

Có, nếu mỗi hàng không quá dài ($N / \text{blockDim} \leq$ vài chục): số phần tử mỗi thread xử lý có hạn, dữ liệu lần 1 có thể giữ tạm trong thanh ghi, lần 2 lấy từ thanh ghi để chuẩn hóa — không đọc global memory lần hai.

Đó mới là Softmax một lần (One-Pass) đúng nghĩa — global memory chỉ đọc một lần, ghi một lần.

6.2 Hiện thực Kernel

// V3: register cache for One-Pass Softmax
// For small N / blockDim.x (elements per thread <= MAX_ELEMS)
#define MAX_ELEMS_PER_THREAD 32

__global__ void online_softmax_v3(float* input, float* output, int M, int N) {
    int row = blockIdx.x;
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    float* x = input  + row * N;
    float* y = output + row * N;

    // Register array caches input data
    float reg_cache[MAX_ELEMS_PER_THREAD];

    // Pass 1 (only global read): load + Online recurrence
    float local_m = -INFINITY;
    float local_d = 0.0f;
    int count = 0;

    for (int i = tid; i < N; i += blockDim.x) {
        float xi = x[i];
        reg_cache[count] = xi;  // cache into registers
        count++;

        float m_new = fmaxf(local_m, xi);
        local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
        local_m = m_new;
    }

    // Two-level Warp Shuffle merge
    warpReduceOnline(local_m, local_d);

    __shared__ float warp_m[32];
    __shared__ float warp_d[32];
    if (lane == 0) {
        warp_m[wid] = local_m;
        warp_d[wid] = local_d;
    }
    __syncthreads();

    int num_warps = blockDim.x / 32;
    if (wid == 0) {
        local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
        local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
        warpReduceOnline(local_m, local_d);
    }

    __shared__ float final_m, final_d;
    if (tid == 0) {
        final_m = local_m;
        final_d = local_d;
    }
    __syncthreads();

    float row_max = final_m;
    float inv_sum = 1.0f / final_d;

    // Read from register cache and write out (no second global read)
    int idx = 0;
    for (int i = tid; i < N; i += blockDim.x) {
        y[i] = expf(reg_cache[idx] - row_max) * inv_sum;
        idx++;
    }
}

6.3 Phân tích áp lực thanh ghi

⚠️ Lưu ý: hạn chế then chốt là dung lượng thanh ghi. Mỗi thread dùng MAX_ELEMS_PER_THREAD thanh ghi float để cache:

A100 mỗi SM có 65536 thanh ghi 32-bit, mỗi thread tối đa 255. MAX_ELEMS_PER_THREAD quá lớn thì thanh ghi spill sang Local Memory (thực chất là global memory), càng chậm hơn.

📌 Điểm then chốt: V3 phù hợp $N \leq 8192$ với blockDim=256. Với $N$ lớn hơn, quay lại V2 (hai lần).

V3 băng thông đo được: khoảng 82% (~1672 GB/s) ($N=4096$, blockDim=256), khoảng 1.17× so với V2.

---

7. Phiên bản V4: song song nhiều hàng + Grid Stride

7.1 Trường hợp áp dụng

Các phiên bản trước đều “một Block một hàng”, “khởi động $M$ Block”. Khi $M$ rất lớn (vài chục nghìn hàng), có hai vấn đề: overhead lập lịch Grid tăng tuyến tính; mỗi Block chạy xong một hàng rồi thoát, Shared Memory và trạng thái thanh ghi không tái dùng qua hàng.

Giải pháp là Grid Stride: số Block cố định (thường gấp vài lần số SM) vòng qua mọi hàng, vẫn chồng các tối ưu trước (tải vector hóa + Warp Shuffle merge reduce). Về thuật toán V4 giống hệt V2 (hai lần quét); khác chỉ ở ánh xạ Block–hàng — V2 là 1:1, V4 là 1:nhiều.

⚠️ Lưu ý: Grid Stride không giải được trường hợp ”$M$ rất nhỏ (ví dụ batch=1), $N$ rất lớn” khiến thiếu Block — khi đó cần nhiều Block cùng xử lý một hàng (reduce phân đoạn trong hàng + Kernel hai pha hoặc atomic merge).

7.2 Hiện thực Kernel

// V4: Grid Stride multi-row + all optimizations
__global__ void online_softmax_v4(float* input, float* output, int M, int N) {
    int tid = threadIdx.x;
    int lane = tid % 32;
    int wid  = tid / 32;

    // Grid Stride over all rows
    for (int row = blockIdx.x; row < M; row += gridDim.x) {
        float* x = input  + row * N;
        float* y = output + row * N;

        float4* x4 = reinterpret_cast<float4*>(x);
        float4* y4 = reinterpret_cast<float4*>(y);
        int N4 = N / 4;

        // Pass 1: vectorized Online recurrence
        float local_m = -INFINITY;
        float local_d = 0.0f;

        for (int i = tid; i < N4; i += blockDim.x) {
            float4 data = x4[i];
            float vals[4] = {data.x, data.y, data.z, data.w};
            for (int k = 0; k < 4; k++) {
                float m_new = fmaxf(local_m, vals[k]);
                local_d = local_d * expf(local_m - m_new) + expf(vals[k] - m_new);
                local_m = m_new;
            }
        }
        for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
            float xi = x[i];
            float m_new = fmaxf(local_m, xi);
            local_d = local_d * expf(local_m - m_new) + expf(xi - m_new);
            local_m = m_new;
        }

        // Two-level Warp Shuffle merge
        warpReduceOnline(local_m, local_d);

        __shared__ float warp_m[32];
        __shared__ float warp_d[32];
        if (lane == 0) {
            warp_m[wid] = local_m;
            warp_d[wid] = local_d;
        }
        __syncthreads();

        int num_warps = blockDim.x / 32;
        if (wid == 0) {
            local_m = (lane < num_warps) ? warp_m[lane] : -INFINITY;
            local_d = (lane < num_warps) ? warp_d[lane] : 0.0f;
            warpReduceOnline(local_m, local_d);
        }

        __shared__ float final_m, final_d;
        if (tid == 0) {
            final_m = local_m;
            final_d = local_d;
        }
        __syncthreads();

        float row_max = final_m;
        float inv_sum = 1.0f / final_d;

        // Pass 2: vectorized normalize
        for (int i = tid; i < N4; i += blockDim.x) {
            float4 data = x4[i];
            float4 result;
            result.x = expf(data.x - row_max) * inv_sum;
            result.y = expf(data.y - row_max) * inv_sum;
            result.z = expf(data.z - row_max) * inv_sum;
            result.w = expf(data.w - row_max) * inv_sum;
            y4[i] = result;
        }
        for (int i = N4 * 4 + tid; i < N; i += blockDim.x) {
            y[i] = expf(x[i] - row_max) * inv_sum;
        }

        __syncthreads();  // Ensure Shared Memory is ready before the next row
    }
}

Cách gọi:

int num_sms;
cudaDeviceGetAttribute(&num_sms, cudaDevAttrMultiProcessorCount, 0);
int grid_size = min(M, num_sms * 4);  // maximize GPU utilization
int block_size = 256;
online_softmax_v4<<<grid_size, block_size>>>(d_input, d_output, M, N);

7.3 Lợi ích của Grid Stride

Khi $M$ lớn, Grid Stride để Block vòng qua nhiều hàng:

• Grid cố định tránh overhead lập lịch Grid quá lớn
• Shared Memory và trạng thái thanh ghi trong Block tái dùng qua hàng
• GPU luôn đầy tải

Khi $M$ nhỏ, grid_size = M, hành vi rút về V2 — mỗi Block một hàng.

V4 băng thông đo được: khoảng 72% (~1468 GB/s) ($M=4096, N=4096$), gần ngang V2 (trần lý thuyết của lược đồ hai lần với $N$ lớn). Lợi ích Grid Stride chủ yếu khi $M$ cực lớn nhờ tiết kiệm overhead lập lịch; với quy mô kiểm thử bài này ($M=4096$) cải thiện có hạn.

---

8. So sánh hiệu năng và liên hệ với FlashAttention

8.1 Tổng hợp hiệu năng các phiên bản

📌 Điểm then chốt: One-Pass của V3 tốt nhất khi $N \leq 8192$, nhưng bị giới hạn dung lượng thanh ghi. Two-Pass của V2/V4 phổ quát hơn, với nhiều kích thước $N$ vẫn đạt 70%+ băng thông.

8.2 Liên hệ với FlashAttention

Ý tưởng cốt lõi của FlashAttention chính là tổng quát hóa Online Softmax. Trong Self-Attention, tính toán đầy đủ là:

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

Hiện thực truyền thống phải tính hết ma trận $QK^T$ ($N \times N$), Softmax, rồi nhân $V$. Điều đó đòi hỏi giữ cả ma trận Attention $N \times N$ trong HBM — khi $N=128K$ đó là 64GB.

FlashAttention dùng ý đệ quy Online Softmax: tải $K$ và $V$ theo khối vào SRAM (Shared Memory), mỗi lần chỉ tính một phần $QK^T$ và cập nhật dần trạng thái Softmax $(m, d)$, rồi sửa output. Cả quá trình không cần lưu ma trận Attention đầy đủ.

graph LR
    A["Online Softmax<br>công thức đệ quy"] --> B["Tính QK^T theo khối<br>cập nhật dần (m, d)"]
    B --> C["Sửa output O<br>O = O * correction + new_block"]
    C --> D["FlashAttention<br>IO-aware Attention"]

Hệ số sửa $e^{m_{\text{old}} - m_{\text{new}}}$ của Online Softmax trong FlashAttention dùng để sửa đóng góp output của các khối trước:

$$O_{\text{new}} = O_{\text{old}} \cdot \frac{d_{\text{old}}}{d_{\text{new}}} \cdot e^{m_{\text{old}} - m_{\text{new}}} + \frac{e^{m_{\text{block}} - m_{\text{new}}}}{d_{\text{new}}} \cdot P_{\text{block}} \cdot V_{\text{block}}$$

💡 Gợi ý: hiểu đệ quy và công thức gộp Online Softmax trong bài này thì chiến lược chia khối của FlashAttention chỉ là thêm cập nhật dần phép nhân ma trận — toán học cốt lõi hoàn toàn giống.

8.3 Gợi ý chọn phương án kỹ thuật

---

📝 Tóm tắt

Đóng góp cốt lõi của Online Softmax là một công thức đệ quy gọn, gộp ba lần quét của Safe Softmax thành hai:

1. Nền toán: dùng $d_{k+1} = d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}}$ để duy trì max và sum trong một lần
2. Then chốt song song: gộp $(m, d)$ thỏa kết hợp tính, song song hiệu quả bằng Warp Shuffle hai cấp (V1)
3. Tiến gần trần băng thông: tải vector hóa (V2) giảm số lệnh; cache thanh ghi (V3) đi tới One-Pass; Grid Stride (V4) khớp mọi batch
4. Đường tới FlashAttention: ý đệ quy Online Softmax là nền toán của tính theo khối trong FlashAttention

Nhìn kỹ thuật, Online Softmax đưa số lần truy cập bộ nhớ của thao tác Memory-Bound xuống mức tối thiểu lý thuyết; kết hợp Warp Shuffle và tải vector hóa, trên A100 đạt 70–82% băng thông.

---

🎯 Checklist tự kiểm

• Giải thích Online Softmax tiết kiệm lần quét nào so với Safe Softmax ba lần, và vì sao
• Tự suy diễn $d_{k+1} = d_k \cdot e^{m_k - m_{k+1}} + e^{x_k - m_{k+1}}$
• Viết được code gộp hai $(m_1, d_1)$ và $(m_2, d_2)$ cục bộ thành $(m, d)$ toàn cục
• Giải thích vì sao phép gộp $(m, d)$ thỏa kết hợp tính nên song song hóa được
• Hiện thực warpReduceOnline dùng __shfl_down_sync
• Nêu điều kiện áp dụng và giới hạn áp lực thanh ghi của V3
• Giải thích liên hệ toán giữa Online Softmax và chiến lược chia khối FlashAttention
• Chọn phiên bản Softmax phù hợp theo kích thước $M$ và $N$
• Dùng Nsight Compute so sánh throughput bộ nhớ của Two-Pass và Three-Pass
• Mở rộng ý đệ quy $(m, d)$ của Online Softmax sang các tình huống tính tăng dần khác

---

📚 Tài liệu tham khảo

• Online normalizer calculation for softmax - Milakov & Gimelshein, 2018
• FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness
• FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning
• From Online Softmax to FlashAttention - UW CSE 599M
• NVIDIA CUDA C++ Programming Guide - Warp Shuffle Functions
• Chuỗi tối ưu GPU dễ hiểu: tối ưu Softmax
• Giải thích nguyên lý cốt lõi FlashAttention
• NVIDIA CUB Library