Flash Attention 3 指南：GPU IO 优化的终极武器

SRAM 容量: ~164 KB (A100)
Block 大小: B_r × B_c (需要满足 B_r × d + B_c × d + B_r × B_c ≤ SRAM)
典型值: B_r = 128, B_c = 128, d = 128 → 128×128 + 128×128 + 128×128 = 48KB ✓

def flash_attention_v1(Q, K, V, block_size=128):
    """
    Q: (n, d), K: (n, d), V: (n, d)
    """
    n, d = Q.shape
    O = torch.zeros_like(Q)
    l = torch.zeros(n)
    m = torch.full((n,), -float('inf'))

    for i in range(0, n, block_size):
        Qi = Q[i:i+block_size]  # 加载到 SRAM
        Oi = torch.zeros_like(Qi)
        li = torch.zeros(block_size)
        mi = torch.full((block_size,), -float('inf'))

        for j in range(0, n, block_size):
            Kj = K[j:j+block_size]  # 加载到 SRAM
            Vj = V[j:j+block_size]

            # 在 SRAM 中计算
            Sij = Qi @ Kj.T / math.sqrt(d)  # (B_r, B_c)
            
            # 在线 softmax 更新
            mij_new = torch.maximum(mi, Sij.max(dim=-1).values)
            Pij = torch.exp(Sij - mij_new.unsqueeze(-1))
            
            # 重缩放
            alpha = torch.exp(mi - mij_new)
            Oi = Oi * alpha.unsqueeze(-1)
            li = li * alpha
            
            Oi += Pij @ Vj
            li += Pij.sum(dim=-1)
            mi = mij_new

        O[i:i+block_size] = Oi / li.unsqueeze(-1)
    
    return O

def flash_attention_v2(Q, K, V, block_size=128):
    n, d = Q.shape
    O = torch.zeros_like(Q)
    
    for j in range(0, n, block_size):
        Kj = K[j:j+block_size]
        Vj = V[j:j+block_size]
        
        for i in range(0, n, block_size):
            Qi = Q[i:i+block_size]
            Sij = Qi @ Kj.T / math.sqrt(d)
            
            # 只需要当前块的 max
            mij = Sij.max(dim=-1).values
            Pij = torch.exp(Sij - mij.unsqueeze(-1))
            
            # 累积更新 (需要保存 mi, li 用于全局归一化)
            if j == 0:
                O[i:i+block_size] = Pij @ Vj
                m_global[i] = mij
                l_global[i] = Pij.sum(dim=-1)
            else:
                m_new = torch.maximum(m_global[i], mij)
                alpha = torch.exp(m_global[i] - m_new)
                O[i:i+block_size] = O[i:i+block_size] * alpha.unsqueeze(-1) \
                    + Pij @ Vj * torch.exp(mij - m_new).unsqueeze(-1)
                # ... 更新 m, l
    
    return O / l_global.unsqueeze(-1)

// FA v3 伪代码: 异步 GEMM + softmax 重叠
__global__ void flash_attn_v3_kernel(/* ... */) {
    // 1. TMA 异步加载 K, V 块到共享内存
    cp_async_bulk(shared_k, global_k);
    cp_async_bulk(shared_v, global_v);
    
    // 2. WGMMA 异步计算 Q @ K^T
    wgmma_async(C_shared, Q_shared, K_shared);
    
    // 3. 等待 GEMM 完成
    wgmma_wait_group(0);
    
    // 4. 在 GEMM 结果到达时，同时做 softmax
    softmax_in_place(C_shared, m, l);
    
    // 5. 异步启动 P @ V 的 GEMM
    wgmma_async(O_shared, P_shared, V_shared);
    
    // 6. 同时 TMA 加载下一块 K, V
    cp_async_bulk(shared_k_next, global_k_next);
}

# FA v3 FP8 路径
def flash_attn_v3_fp8(Q, K, V):
    # Q, K, V: FP16 输入
    Q_fp8 = Q.to(torch.float8_e4m3fn)
    K_fp8 = K.to(torch.float8_e4m3fn)
    V_fp8 = V.to(torch.float8_e4m3fn)
    
    # FP8 GEMM (2x 吞吐)
    S = fp8_gemm(Q_fp8, K_fp8.T)  # FP32 累积
    S = S / math.sqrt(d)
    P = softmax(S)
    P_fp8 = P.to(torch.float8_e4m3fn)
    O = fp8_gemm(P_fp8, V_fp8)    # FP32 累积
    
    return O.to(torch.float16)

def compute_block_size(sram_bytes, d, dtype_bytes=2):
    """计算最优 block size"""
    # SRAM 需要容纳: Q_block + K_block + V_block + P_block + O_block
    # B_r * d + B_c * d + B_c * d + B_r * B_c + B_r * d <= SRAM
    # 典型 SRAM: 228 KB (H100)
    
    available = sram_bytes - 3 * d * dtype_bytes  # 留余量
    # 最大化 B_r * B_c 同时满足约束
    # 启发式: B_r = B_c = sqrt(available / (2*d + 1))
    block = int(math.sqrt(available / (2 * d * dtype_bytes + 1)))
    # 对齐到 64 (GEMM 对齐要求)
    block = (block // 64) * 64
    return min(block, 256)  # 上限 256