Flash Attention 3 原理：GPU 内存层次的最优利用

# 标准注意力（显存不友好）
def standard_attention(Q, K, V):
    # Q, K, V: (batch, heads, seq, dim)
    # 需要 O(N^2) 显存访问来生成 attention matrix
    
    S = Q @ K.T  # (N, N) 矩阵，需要全部存储在 HBM
    P = softmax(S)  # 需要读取整个 S 矩阵
    O = P @ V  # 需要读取 V 和 P
    
    # 总 HBM 访问：O(N^2 * d) 只读，O(N^2) 写
    return O

标准：  HBM[S(N,N)] ← 全部 attention 矩阵
Flash1: L2[S(M, d)] ← 只有 tile 大小的矩阵

H100 之前：
  Warp 0: [MMA] → [等待] → [MMA] → [等待]
  Warp 1: [等待] → [MMA] → [等待] → [MMA]
  → 串行执行，利用率低

H100 Hopper:
  Warpgroup 0 (4 warps): [MMA] ───────── [MMA]
  Warpgroup 1 (4 warps):      [MMA] ───────── [MMA]
  → 异步执行，流水线

# 伪代码：异步 MMA + Softmax Pipeline
def async_pipeline(Q_tile, K_tile, V_tile):
    """
    异步执行 pipeline：
    Warpgroup 0: 计算 GEMM (Q @ K^T)
    Warpgroup 1: 同时计算 Softmax (上一 tile 的结果)
    """
    # 假设有 2 个 warpgroup
    num_warpgroups = 2
    
    # 状态
    gemm_results = []  # Warpgroup 0 的 GEMM 结果
    softmax_results = []  # Warpgroup 1 的 Softmax 结果
    
    for i in range(num_tiles):
        # Warpgroup 0: 异步执行 GEMM
        # 使用 H100 的 TMA (Tensor Memory Accelerator)
        if i < len(K_tiles):
            Q_tile = load_tile_async(Q_tiles[i], target_mem='shared')
            K_tile = load_tile_async(K_tiles[i], target_mem='shared')
            
            # 异步 GEMM
            S_ij = warpgroup_mma_async(Q_tile, K_tile)
            gemm_results.append(S_ij)
        
        # Warpgroup 1: 同时处理上一个 GEMM 结果的 Softmax
        if i > 0:
            S_prev = gemm_results[i-1]
            P_prev = online_softmax(S_prev)
            O_ij = P_prev @ V_tiles[i-1]
            softmax_results.append(O_ij)
    
    # 等待所有异步操作完成
    synchronize_warpgroups()
    
    return combine_results(softmax_results)

def block_quantize_fp8(tensor, block_size=128):
    """
    块量化：每个块独立量化，保持局部精度
    """
    B, H, N, D = tensor.shape
    
    # 按块计算缩放因子
    num_blocks = (D + block_size - 1) // block_size
    scales = torch.zeros(B, H, N, num_blocks, device=tensor.device)
    
    quantized = torch.zeros_like(tensor, dtype=torch.int8)
    
    for block_idx in range(num_blocks):
        start = block_idx * block_size
        end = min(start + block_size, D)
        
        block = tensor[:, :, :, start:end]
        scale = block.abs().max(dim=-1, keepdim=True).values / 127
        scale = scale.clamp(min=1e-8)
        
        quantized[:, :, :, start:end] = (block / scale).clamp(-127, 127)
        scales[:, :, :, block_idx] = scale.squeeze(-1)
    
    return quantized, scales


def fp8_attention(Q_fp8, K_fp8, V_fp8, Q_scale, K_scale, V_scale):
    """
    FP8 注意力计算
    
    注意：QK^T 用 FP8 计算，但 Softmax 和输出投影用 FP16
    """
    # 反量化到 FP16 进行 Softmax
    Q = Q_fp8.to(torch.float16) * Q_scale.unsqueeze(-1)
    K = K_fp8.to(torch.float16) * K_scale.unsqueeze(-1)
    V = V_fp8.to(torch.float16) * V_scale.unsqueeze(-1)
    
    # 注意力计算（混合精度）
    S = torch.matmul(Q, K.transpose(-2, -1)) / (Q.shape[-1] ** 0.5)
    P = torch.softmax(S, dim=-1)
    O = torch.matmul(P, V)
    
    return O

Warp 分组：
  Warpgroup 0 (4 warps): GEMM QK^T
  Warpgroup 1 (4 warps): Softmax + 在线归约
  Warpgroup 2 (4 warps): GEMM PV
  
  总共 12 个 warps（H100 每个 SM 支持 4 个 warpgroup）

# H100 上 Flash Attention 3 的 warp 分配
WARPGROUP_SIZE = 4  # 每个 warpgroup 有 4 个 warp
NUM_WARPGROUPS = 3  # 总共 3 个 warpgroup

def fa3_warp_specialized(Q, K, V, block_size=128):
    """
    Warp 专用化的 Flash Attention 3
    """
    # 初始化共享内存
    shared_mem = allocate_shared_memory(block_size * D * 3)  # Q, K, V tiles
    
    # Warpgroup 0: 加载 Q, K tile 并计算 S = QK^T
    if warpgroup_id() == 0:
        for i in range(num_Q_tiles):
            Q_tile = load_Q_tile(i, shared_mem)
            for j in range(num_K_tiles):
                K_tile = load_K_tile(j, shared_mem)
                S_ij = mma(Q_tile, K_tile)  # Tensor Core MMA
                store_S_tile(S_ij, shared_mem)
    
    # Warpgroup 1: 读取 S_ij，计算在线 Softmax
    if warpgroup_id() == 1:
        for i in range(num_Q_tiles):
            for j in range(num_K_tiles):
                S_ij = load_S_tile(i, j, shared_mem)
                P_ij, m_new, l_new = online_softmax_update(
                    S_ij, m_old[i], l_old[i]
                )
                store_P_tile(P_ij, shared_mem)
                update_stats(i, m_new, l_new)
    
    # Warpgroup 2: 读取 P_ij 和 V_j，计算 O_i += P_ij * V_j
    if warpgroup_id() == 2:
        for i in range(num_Q_tiles):
            O_i = load_O_tile(i, shared_mem)
            for j in range(num_V_tiles):
                P_ij = load_P_tile(i, j, shared_mem)
                V_j = load_V_tile(j, shared_mem)
                O_i += mma(P_ij, V_j)
            store_O_tile(i, O_i, shared_mem)
    
    synchronize()  # 等待所有 warpgroup 完成
    return gather_output(shared_mem)

# 使用 TMA 异步加载（CUDA 伪代码）
__global__
void tma_load_kernel(
    const half* __restrict__ Q_global,  # HBM
    half* __restrict__ Q_shared,         # Shared Memory
    int num_tiles
) {
    // TMA 描述符
    CUtensorMap tma_desc;
    cuTensorMapEncodeTiled(
        &tma_desc,
        CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_FLOAT16,
        2,  // 维度
        Q_global,
        tensor_extent,
        tensor_stride,
        tile_size,
        CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
        CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE,
        CUtensorMapL2Promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
        CUtensorMapOOBfill::CU_TENSOR_MAP_OOB_FILL_NONE
    );
    
    // 启动 TMA 传输（异步）
    for (int i = 0; i < num_tiles; i++) {
        // TMA 在后台传输，warp 可以做其他事
        __cuda_tma_im2col(&tma_desc, Q_shared + i * TILE_SIZE,
                          0, 0, 0, 0);
    }
}

def online_softmax_fa3(S_block, m_prev, l_prev, safe_exp=True):
    """
    Flash Attention 3 的在线 Softmax
    
    改进点：
    1. 在 block 内使用 max subtraction 避免 exp 溢出
    2. 使用 Kahan 求和提高精度
    """
    # Block 内的最大值
    m_local = S_block.max(dim=-1, keepdim=True).values
    
    # 全局最大值
    m_new = torch.maximum(m_prev, m_local)
    
    # 修正之前的累积值（Kahan 求和）
    if safe_exp:
        # 使用 log-sum-exp 技巧避免数值溢出
        correction = torch.exp(m_prev - m_new)
        l_new = l_prev * correction + torch.sum(
            torch.exp(S_block - m_new), dim=-1, keepdim=True
        )
    else:
        correction = torch.exp(m_prev - m_new)
        l_new = l_prev * correction + torch.sum(
            torch.exp(S_block - m_new), dim=-1, keepdim=True
        )
    
    # P_block = exp(S_block - m_new) / l_new
    P_block = torch.exp(S_block - m_new) / (l_new + 1e-8)
    
    return P_block, m_new, l_new

# Flash Attention 3 需要 CUDA 12 + H100
pip install flash-attn --no-build-isolation

# 检查是否支持
python -c "import flash_attn; print(flash_attn.__version__)"
# 应该 >= 3.0

# 验证 FP8 支持
python -c "
import torch
print('FP8 supported:', torch.cuda.is_available() and 
      torch.cuda.get_device_capability()[0] >= 9)
"

import torch
import flash_attn_interface as fa3

# 标准使用
Q = torch.randn(2, 32, 4096, 64, dtype=torch.float16, device='cuda')
K = torch.randn(2, 32, 4096, 64, dtype=torch.float16, device='cuda')
V = torch.randn(2, 32, 4096, 64, dtype=torch.float16, device='cuda')

# FP16 模式
output = fa3.flash_attn_func(
    Q, K, V,
    causal=True,      # 因果注意力
    softmax_scale=1.0 / (64 ** 0.5)
)

# FP8 模式（需要 H100）
Q_fp8 = Q.to(torch.float8_e4m3fn)
K_fp8 = K.to(torch.float8_e4m3fn)
V_fp8 = V.to(torch.float8_e4m3fn)

output_fp8 = fa3.flash_attn_func(
    Q_fp8, K_fp8, V_fp8,
    causal=True,
    fp8=True,  # 启用 FP8
)

# 与 HuggingFace 集成
from transformers import LlamaForCausalLM
model = LlamaForCausalLM.from_pretrained(
    "meta-llama/Llama-3-8B",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_3"  # 需要最新 transformers
)

# 调优 block size（根据序列长度）
def auto_block_size(seq_len):
    if seq_len <= 2048:
        return 128
    elif seq_len <= 16384:
        return 64
    else:
        return 32  # 更小的 block 减少共享内存占用

# 混合精度策略
MIXED_PRECISION_CONFIG = {
    'prefill': 'fp16',   # 预填充阶段用 fp16 保证精度
    'decode': 'fp8',     # 解码阶段用 fp8 加速
}