Flash Attention

# Cách dùng Flash Attention trong PyTorch
import torch
import torch.nn.functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

# PyTorch 2.0+ có built-in Flash Attention!
# Tự động chọn Flash Attention khi có thể.

q = torch.randn(1, 8, 4096, 64, device="cuda", dtype=torch.float16)  # (B, heads, N, d_k)
k = torch.randn(1, 8, 4096, 64, device="cuda", dtype=torch.float16)
v = torch.randn(1, 8, 4096, 64, device="cuda", dtype=torch.float16)

# Tự động dùng Flash Attention nếu available
output = F.scaled_dot_product_attention(q, k, v, is_causal=True)
# Memory: O(N) thay vì O(N²)
# Speed: 2–4× nhanh hơn naive implementation

# Ép dùng Flash backend (báo lỗi nếu không được) để debug
with sdpa_kernel(SDPBackend.FLASH_ATTENTION):
    output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# Hoặc dùng thư viện flash-attn trực tiếp (chi tiết hơn)
# pip install flash-attn --no-build-isolation
from flash_attn import flash_attn_func
output = flash_attn_func(q, k, v, causal=True, dropout_p=0.0, softmax_scale=1.0 / (64 ** 0.5))

# So sánh memory:
# N=4096:  Standard ~64MB,  Flash ~0.5MB (giảm 128×)
# N=128K:  Standard ~64GB,  Flash ~0.5MB (giảm 130,000×)

# Pseudocode Flash Attention forward pass
# Tham khảo: Dao et al. 2022, Algorithm 1
# N: seq length, d: head dim, Br/Bc: block row/col sizes

import math
import torch

def flash_attention_forward(Q, K, V, Br=64, Bc=64, causal=False):
    """
    Q, K, V: (N, d) tensors trên HBM
    Trả về O: (N, d) — tương đương softmax(QK^T / sqrt(d)) V
    Bộ nhớ thêm: O(N) (running stats), KHÔNG phải O(N^2)
    """
    N, d = Q.shape
    scale = 1.0 / math.sqrt(d)

    # Output + running stats, chỉ O(N) bộ nhớ
    O = torch.zeros_like(Q)             # (N, d)
    l = torch.zeros(N, device=Q.device) # running sum of exp
    m = torch.full((N,), -float("inf"), device=Q.device)  # running max

    Tr = (N + Br - 1) // Br  # số block của Q
    Tc = (N + Bc - 1) // Bc  # số block của K, V

    # OUTER loop trên Q blocks (FA2 đổi thứ tự: outer là Q)
    for i in range(Tr):
        qi_s, qi_e = i * Br, min((i + 1) * Br, N)
        Qi = Q[qi_s:qi_e]                        # (Br, d) — load SRAM
        Oi = torch.zeros(qi_e - qi_s, d, device=Q.device)
        li = torch.zeros(qi_e - qi_s, device=Q.device)
        mi = torch.full((qi_e - qi_s,), -float("inf"), device=Q.device)

        # INNER loop trên K,V blocks
        for j in range(Tc):
            if causal and j * Bc > qi_e - 1:
                break
            kj_s, kj_e = j * Bc, min((j + 1) * Bc, N)
            Kj = K[kj_s:kj_e]   # (Bc, d)
            Vj = V[kj_s:kj_e]   # (Bc, d)

            # Attention scores cho tile này — hoàn toàn trong SRAM
            Sij = Qi @ Kj.transpose(0, 1) * scale  # (Br, Bc)
            if causal:
                mask = torch.arange(kj_s, kj_e, device=Q.device)                      > torch.arange(qi_s, qi_e, device=Q.device)[:, None]
                Sij.masked_fill_(mask, -float("inf"))

            # ── Online softmax ───────────────────────────────
            mij_local = Sij.max(dim=-1).values          # (Br,)
            mi_new = torch.maximum(mi, mij_local)
            Pij = torch.exp(Sij - mi_new[:, None])      # (Br, Bc)
            lij_local = Pij.sum(dim=-1)                 # (Br,)
            alpha = torch.exp(mi - mi_new)              # rescale factor cho O cũ
            li = alpha * li + lij_local
            Oi = alpha[:, None] * Oi + Pij @ Vj
            mi = mi_new
            # ─────────────────────────────────────────────────

        # Ghi kết quả block Q này ra HBM
        O[qi_s:qi_e] = Oi / li[:, None]
        l[qi_s:qi_e] = li
        m[qi_s:qi_e] = mi

    return O, l, m  # l, m dùng lại ở backward