📚 课程目录第4章使用FlashAttention官方库

4. 使用FlashAttention官方库

任务目标

本案例演示如何通过官方 flash-attn 库调用 FlashAttention-2 内核，替换标准 PyTorch 的 scaled_dot_product_attention，并验证输出正确性。你将得到一个可直接运行的脚本，对比两种注意力实现在相同输入下的数值误差。

环境依赖

Python 3.8+，PyTorch ≥ 2.0（需支持 torch.backends.cuda.flash_sdp_enabled）
GPU（Ampere 架构及以上，如 A100/A30/RTX 3090）
安装官方库：pip install flash-attn --no-build-isolation（源码编译需 nvcc 与匹配的 CUDA 工具链）

关键完整代码

以下脚本实现：构造随机 Q/K/V，分别用 PyTorch 原生 SDPA 与 FlashAttention 官方库计算，比较输出误差并打印耗时。

import torch
import torch.nn.functional as F
from flash_attn import flash_attn_func
import time

# 配置：batch=2, seqlen=1024, heads=8, head_dim=64
batch, seqlen, nheads, head_dim = 2, 1024, 8, 64
dtype = torch.float16
device = "cuda"

torch.manual_seed(42)
q = torch.randn(batch, seqlen, nheads, head_dim, device=device, dtype=dtype)
k = torch.randn(batch, seqlen, nheads, head_dim, device=device, dtype=dtype)
v = torch.randn(batch, seqlen, nheads, head_dim, device=device, dtype=dtype)

# 参考实现：PyTorch 原生 SDPA（内部自动选择 flash 内核）
# 需要将 (b, s, h, d) 转为 (b, h, s, d) 格式
q_sdpa = q.transpose(1, 2)
k_sdpa = k.transpose(1, 2)
v_sdpa = v.transpose(1, 2)

torch.cuda.synchronize()
start = time.time()
out_sdpa = F.scaled_dot_product_attention(q_sdpa, k_sdpa, v_sdpa)
torch.cuda.synchronize()
time_sdpa = time.time() - start

# FlashAttention 官方库：直接接受 (b, s, h, d) 布局
torch.cuda.synchronize()
start = time.time()
out_fa = flash_attn_func(q, k, v, dropout_p=0.0, causal=False)
torch.cuda.synchronize()
time_fa = time.time() - start

# 验证：比较输出形状与数值误差
assert out_fa.shape == (batch, seqlen, nheads, head_dim), "输出形状不匹配"
out_sdpa_perm = out_sdpa.transpose(1, 2)  # 转回 (b, s, h, d) 以便比较

# 计算最大误差（fp16 下允许较小误差）
diff = (out_sdpa_perm.float() - out_fa.float()).abs().max().item()
print(f"PyTorch SDPA 耗时: {time_sdpa*1000:.2f} ms")
print(f"FlashAttention 耗时: {time_fa*1000:.2f} ms")
print(f"最大误差: {diff:.6f}")

# 判定标准：误差小于 1e-2 即认为数值一致（fp16 精度范围）
if diff < 1e-2:
    print("✓ 验证通过：FlashAttention 输出与 PyTorch SDPA 一致")
else:
    print("✗ 验证失败：误差超出预期范围")

运行方式

# 保存为 flash_attn_demo.py 后执行
python flash_attn_demo.py

验证方法与判定标准

正确性验证：脚本自动比较两种实现的输出，最大误差 < 1e-2 时打印“验证通过”。由于 fp16 精度限制，该阈值已考虑舍入误差。
性能观察：打印的耗时仅作参考，实际加速比取决于 GPU 型号与序列长度。若你的 GPU 不支持 FlashAttention-2（如 V100），脚本会报错，此时需更换硬件或使用 flash_attn_varlen_func 的兼容版本。
预期输出示例（非实测值，仅示意格式）：最大误差: 0.000412，且验证通过。

常见问题

安装失败：flash-attn 需要与 PyTorch 匹配的 CUDA 版本。建议使用 Docker 镜像 pytorch/pytorch:2.1.0-cuda12.1-cudnn8-devel 后执行 pip install flash-attn。
布局差异：官方库使用 (batch, seqlen, nheads, head_dim) 布局，而 PyTorch SDPA 使用 (batch, nheads, seqlen, head_dim)，代码中已通过 transpose 处理。
因果掩码：本案例使用非因果注意力（causal=False）。若需因果掩码，将 flash_attn_func 的 causal 参数设为 True 即可，无需额外构造掩码矩阵。