第12章 FlashAttention原理与实现¶

def flash_attention_forward(Q, K, V, B_r, B_c):
    """
    Q, K, V: (N, d) 存储在 HBM 中
    B_r: Q 的分块大小（行方向）
    B_c: K/V 的分块大小（列方向）
    """
    N, d = Q.shape
    O = zeros(N, d)          # 输出，存在 HBM
    l = zeros(N)             # softmax 分母（每行一个标量）
    m = full(N, -inf)        # 当前行最大值（每行一个标量）

    T_r = ceil(N / B_r)      # Q 的块数
    T_c = ceil(N / B_c)      # K/V 的块数

    # --- 外层循环: 遍历 K/V 的块（加载到 SRAM）---
    for j in range(T_c):
        # 从 HBM 加载 Kⱼ, Vⱼ 到 SRAM
        K_j = K[j*B_c : (j+1)*B_c, :]   # (B_c, d)
        V_j = V[j*B_c : (j+1)*B_c, :]   # (B_c, d)

        # --- 内层循环: 遍历 Q 的块 ---
        for i in range(T_r):
            # 从 HBM 加载 Qᵢ, Oᵢ, lᵢ, mᵢ 到 SRAM
            Q_i = Q[i*B_r : (i+1)*B_r, :]
            O_i = O[i*B_r : (i+1)*B_r, :]
            l_i = l[i*B_r : (i+1)*B_r]
            m_i = m[i*B_r : (i+1)*B_r]

            # Step 1: 计算局部注意力分数（在 SRAM 中）
            S_ij = Q_i @ K_j.T / sqrt(d)    # (B_r, B_c) — 小矩阵!

            # Step 2: Online Softmax 更新
            m_ij = rowmax(S_ij)              # (B_r,)
            m_new = max(m_i, m_ij)           # 新的全局最大值

            P_ij = exp(S_ij - m_new[:, None])  # (B_r, B_c)

            l_new = l_i * exp(m_i - m_new) + rowsum(P_ij)

            # Step 3: 更新输出（带修正因子）
            O_i = O_i * (l_i * exp(m_i - m_new))[:, None] / l_new[:, None] \
                + P_ij @ V_j / l_new[:, None]

            # Step 4: 写回 HBM
            m_i = m_new
            l_i = l_new
            O[i*B_r : (i+1)*B_r, :] = O_i
            l[i*B_r : (i+1)*B_r] = l_i
            m[i*B_r : (i+1)*B_r] = m_i

    return O

# FlashAttention-1: 每处理一个 K/V 块都要 rescale
O_i = O_i * (l_old / l_new) + P_ij @ V_j / l_new    # 多次除法

# FlashAttention-2: 延迟归一化
O_i = diag(exp(m_old - m_new)) @ O_i + P_ij @ V_j    # 只维护未归一化的累加
# ... 最后一步:
O_i = diag(1/l_final) @ O_i                           # 一次归一化

"""
使用 flash-attn 库的 API
安装: pip install flash-attn --no-build-isolation
要求: NVIDIA GPU (Ampere/Hopper), CUDA 11.6+, PyTorch 1.12+

> **编译环境要求**：flash-attn 需要从源码编译，需要：
> - CUDA 编译器 (nvcc) 在 PATH 中
> - 与 PyTorch 匹配的 CUDA 版本
> - C++ 编译器 (gcc/g++ 或 MSVC)
> - 充足的编译内存（建议 16GB+ RAM）
>
> 如遇编译问题，可尝试使用预编译 wheel：`pip install flash-attn --no-build-isolation` 或从 GitHub Releases 下载对应版本。
"""
import torch
from flash_attn import flash_attn_func

# 创建输入 (batch, seqlen, nheads, headdim) — 注意维度顺序!
batch_size, seq_len, n_heads, head_dim = 2, 4096, 32, 128
dtype = torch.float16
device = "cuda"

q = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=dtype, device=device)
k = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=dtype, device=device)
v = torch.randn(batch_size, seq_len, n_heads, head_dim, dtype=dtype, device=device)

# 前向传播
output = flash_attn_func(
    q, k, v,
    dropout_p=0.0,           # 推理时设为 0
    softmax_scale=None,      # 默认 1/sqrt(head_dim)
    causal=True,             # 因果 mask（decoder 模型使用）
)
# output shape: (batch, seqlen, nheads, headdim)
print(f"Output shape: {output.shape}")  # (2, 4096, 32, 128)

# 支持变长序列（packed/varlen）
from flash_attn import flash_attn_varlen_func
# 适用于 batch 内不同样本长度不同的情况，避免 padding 浪费

"""
PyTorch 2.0+ 内置 FlashAttention 后端
通过 torch.nn.functional.scaled_dot_product_attention 自动调用
"""
import torch
import torch.nn.functional as F

batch_size, n_heads, seq_len, head_dim = 2, 32, 4096, 128
dtype = torch.float16
device = "cuda"

# 注意: PyTorch API 的维度顺序是 (batch, nheads, seqlen, headdim)
q = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device=device)
k = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device=device)
v = torch.randn(batch_size, n_heads, seq_len, head_dim, dtype=dtype, device=device)

# 方式1: 自动选择最优后端（推荐）
with torch.backends.cuda.sdp_kernel(
    enable_flash=True,       # 启用 FlashAttention 后端
    enable_math=False,       # 禁用数学后端（标准实现）
    enable_mem_efficient=False  # 禁用 memory-efficient 后端
):
    output = F.scaled_dot_product_attention(  # F.xxx PyTorch函数式API
        q, k, v,
        attn_mask=None,
        dropout_p=0.0,
        is_causal=True,      # 因果 mask
        scale=None,          # 默认 1/sqrt(head_dim)
    )

print(f"Output shape: {output.shape}")  # (2, 32, 4096, 128)

# 方式2: 让 PyTorch 自动选择后端（最简单）
output = F.scaled_dot_product_attention(q, k, v, is_causal=True)

# 查看当前使用了哪个后端
# 可通过 torch.backends.cuda.flash_sdp_enabled() 检查
print(f"FlashSDP enabled: {torch.backends.cuda.flash_sdp_enabled()}")

"""
简化版 FlashAttention Forward — Triton 实现
核心展示 Tiling + Online Softmax 的逻辑
参考: Triton 官方教程 + FlashAttention 论文算法1

注意: 这是教学用简化版本，省略了 causal mask、dropout 等特性
"""
import torch
import triton
import triton.language as tl

@triton.jit
def _flash_attn_fwd_kernel(
    Q_ptr, K_ptr, V_ptr, O_ptr,
    stride_qb, stride_qh, stride_qm, stride_qd,  # Q 的 strides
    stride_kb, stride_kh, stride_kn, stride_kd,  # K 的 strides
    stride_vb, stride_vh, stride_vn, stride_vd,  # V 的 strides
    stride_ob, stride_oh, stride_om, stride_od,  # O 的 strides
    N: tl.constexpr,           # 序列长度
    D: tl.constexpr,           # 头维度
    BLOCK_M: tl.constexpr,     # Q 的分块大小
    BLOCK_N: tl.constexpr,     # K/V 的分块大小
    sm_scale,                  # softmax 缩放因子 = 1/sqrt(D)
):
    # ---- 确定当前 block 负责的区域 ----
    block_m_idx = tl.program_id(0)   # Q 的块索引
    batch_head_idx = tl.program_id(1)  # batch * num_heads
    # 从 batch_head_idx 解析 batch 和 head
    # （此处简化，假设 batch_head_idx 直接用于偏移）

    # ---- 计算基址偏移 ----
    q_offset = batch_head_idx * stride_qh
    k_offset = batch_head_idx * stride_kh
    v_offset = batch_head_idx * stride_vh
    o_offset = batch_head_idx * stride_oh

    # ---- 加载 Q 块到 SRAM ----
    offs_m = block_m_idx * BLOCK_M + tl.arange(0, BLOCK_M)   # 行偏移
    offs_d = tl.arange(0, D)                                   # 列偏移

    # Q_i: (BLOCK_M, D)
    q_ptrs = Q_ptr + q_offset + offs_m[:, None] * stride_qm + offs_d[None, :] * stride_qd
    q_mask = offs_m[:, None] < N
    q = tl.load(q_ptrs, mask=q_mask, other=0.0)

    # ---- 初始化 Online Softmax 状态 ----
    m_i = tl.full([BLOCK_M], value=float("-inf"), dtype=tl.float32)   # 行最大值
    l_i = tl.full([BLOCK_M], value=0.0, dtype=tl.float32)            # 行指数和
    o_i = tl.zeros([BLOCK_M, D], dtype=tl.float32)                   # 累加输出

    # ---- 内层循环: 遍历 K/V 的所有块 ----
    for block_n_start in range(0, N, BLOCK_N):
        offs_n = block_n_start + tl.arange(0, BLOCK_N)

        # 加载 K_j: (BLOCK_N, D)
        k_ptrs = K_ptr + k_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd
        k_mask = offs_n[:, None] < N
        k = tl.load(k_ptrs, mask=k_mask, other=0.0)

        # 加载 V_j: (BLOCK_N, D)
        v_ptrs = V_ptr + v_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd
        v = tl.load(v_ptrs, mask=k_mask, other=0.0)

        # Step 1: 计算 S_ij = Q_i @ K_j^T  — (BLOCK_M, BLOCK_N)
        s_ij = tl.dot(q, tl.trans(k)) * sm_scale

        # 对超出序列长度的位置 mask
        s_ij = tl.where(offs_n[None, :] < N, s_ij, float("-inf"))

        # Step 2: Online Softmax 更新
        m_ij = tl.max(s_ij, axis=1)              # 当前块的行最大值 (BLOCK_M,)
        m_new = tl.maximum(m_i, m_ij)             # 全局行最大值更新

        # 修正因子
        alpha = tl.exp(m_i - m_new)               # 旧最大值的修正
        p_ij = tl.exp(s_ij - m_new[:, None])      # 当前块的 exp 值 (BLOCK_M, BLOCK_N)

        l_new = l_i * alpha + tl.sum(p_ij, axis=1)

        # Step 3: 更新输出累加器
        # 先对旧的 O 做修正（旧 max 变化带来的 rescale）
        o_i = o_i * alpha[:, None]
        # 加上当前块的贡献
        o_i += tl.dot(p_ij.to(v.dtype), v)

        # Step 4: 更新统计量
        m_i = m_new
        l_i = l_new

    # ---- 最终归一化 ----
    o_i = o_i / l_i[:, None]

    # ---- 写回 HBM ----
    o_ptrs = O_ptr + o_offset + offs_m[:, None] * stride_om + offs_d[None, :] * stride_od
    o_mask = offs_m[:, None] < N
    tl.store(o_ptrs, o_i.to(O_ptr.dtype.element_ty), mask=o_mask)

def flash_attention_triton(q, k, v):
    """
    简化版 FlashAttention 包装函数
    q, k, v: (batch, nheads, seqlen, headdim) — float16/bfloat16
    """
    B, H, N, D = q.shape
    assert D in {16, 32, 64, 128}, "headdim 必须是 16/32/64/128"  # assert断言

    o = torch.empty_like(q)
    sm_scale = 1.0 / (D ** 0.5)

    BLOCK_M = 128
    BLOCK_N = 64

    grid = (triton.cdiv(N, BLOCK_M), B * H)

    _flash_attn_fwd_kernel[grid](
        q, k, v, o,
        q.stride(0), q.stride(1), q.stride(2), q.stride(3),
        k.stride(0), k.stride(1), k.stride(2), k.stride(3),
        v.stride(0), v.stride(1), v.stride(2), v.stride(3),
        o.stride(0), o.stride(1), o.stride(2), o.stride(3),
        N=N, D=D,
        BLOCK_M=BLOCK_M, BLOCK_N=BLOCK_N,
        sm_scale=sm_scale,
    )
    return o

# ---- 测试 ----
if __name__ == "__main__":
    torch.manual_seed(42)
    B, H, N, D = 2, 4, 1024, 64
    q = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    k = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    v = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)

    # FlashAttention (Triton 实现)
    o_flash = flash_attention_triton(q, k, v)

    # 标准 Attention (参考实现)
    scale = 1.0 / (D ** 0.5)
    attn = torch.softmax(q @ k.transpose(-2, -1) * scale, dim=-1)
    o_ref = attn @ v

    # 验证正确性
    max_diff = (o_flash.float() - o_ref.float()).abs().max().item()  # 将单元素张量转为Python数值
    print(f"Max absolute difference: {max_diff:.6f}")
    assert max_diff < 1e-2, f"差异过大: {max_diff}"
    print("✓ Triton FlashAttention 结果正确!")

"""
标准 Attention vs FlashAttention 性能对比
测量不同序列长度下的速度和显存使用
"""
import torch
import torch.nn.functional as F
import time

def benchmark_attention(func, q, k, v, name, warmup=10, repeat=50):
    """通用 benchmark 函数"""
    # Warmup
    for _ in range(warmup):
        _ = func(q, k, v)
    torch.cuda.synchronize()

    # 记录显存
    torch.cuda.reset_peak_memory_stats()
    mem_before = torch.cuda.memory_allocated()

    # 计时
    start = time.perf_counter()
    for _ in range(repeat):
        _ = func(q, k, v)
    torch.cuda.synchronize()
    elapsed = (time.perf_counter() - start) / repeat * 1000  # ms

    mem_peak = torch.cuda.max_memory_allocated()
    mem_used = (mem_peak - mem_before) / 1024**2  # MB

    print(f"  {name:30s} | Time: {elapsed:8.2f} ms | Peak Mem: {mem_used:8.1f} MB")
    return elapsed, mem_used

def standard_attention(q, k, v):
    """标准 Attention（显式计算 N×N 矩阵）"""
    scale = 1.0 / (q.shape[-1] ** 0.5)  # [-1]负索引取最后元素
    attn = torch.softmax(q @ k.transpose(-2, -1) * scale, dim=-1)
    return attn @ v

def flash_attention_pytorch(q, k, v):
    """PyTorch SDPA（自动选择 FlashAttention 后端）"""
    return F.scaled_dot_product_attention(q, k, v, is_causal=False)

print("=" * 80)
print("Attention Performance Benchmark")
print("=" * 80)

B, H, D = 2, 32, 128

for N in [512, 1024, 2048, 4096, 8192]:
    q = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    k = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)
    v = torch.randn(B, H, N, D, device="cuda", dtype=torch.float16)

    print(f"\n--- Seq Length = {N} ---")
    try:  # try/except捕获异常
        t1, m1 = benchmark_attention(standard_attention, q, k, v, "Standard Attention")
    except torch.cuda.OutOfMemoryError:
        print(f"  {'Standard Attention':30s} | OOM!")
        t1, m1 = float("inf"), float("inf")

    t2, m2 = benchmark_attention(flash_attention_pytorch, q, k, v, "FlashAttention (PyTorch SDPA)")

    if t1 != float("inf"):
        print(f"  Speedup: {t1/t2:.2f}x | Memory Saving: {m1/m2:.2f}x")

print("\n" + "=" * 80)