optimizing-attention-flash

用途

该技能用于为 Transformer 模型启用 Flash Attention，以加速注意力计算并降低显存占用。资料称 Flash Attention 通过 IO 感知的分块（tiling）与重计算（recomputation），可为 Transformer 注意力带来 2-4 倍加速和 10-20 倍内存降低。

主要能力

使用 PyTorch 原生 torch.nn.functional.scaled_dot_product_attention（PyTorch 2.2+）自动启用 Flash Attention。
使用 flash-attn 库的 flash_attn_func，支持更多特性：多查询注意力（MQA）、滑动窗口注意力、H100 FP8（FlashAttention-3）。
提供性能基准测试与精度对比方法。
提供常见问题排查指引（导入错误、加速不明显、CUDA 错误、精度下降）。

使用方式

PyTorch 原生（最简单，PyTorch 2.2+）

import torch
import torch.nn.functional as F

q = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)  # [batch, heads, seq, dim]
k = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 8, 512, 64, device='cuda', dtype=torch.float16)

# Automatically uses Flash Attention if available
out = F.scaled_dot_product_attention(q, k, v)

强制使用 Flash Attention 后端：

with torch.backends.cuda.sdp_kernel(
    enable_flash=True,
    enable_math=False,
    enable_mem_efficient=False
):
    out = F.scaled_dot_product_attention(q, k, v)

flash-attn 库（更多特性）

安装：

pip install flash-attn --no-build-isolation

使用：

from flash_attn import flash_attn_func

# q, k, v: [batch, seqlen, nheads, headdim]
out = flash_attn_func(q, k, v, dropout_p=0.0, causal=True)

工作流 1：在现有 PyTorch 模型中启用

检查 PyTorch 版本（≥2.2）：python -c "import torch; print(torch.__version__)"，若低于 2.2 则 pip install --upgrade torch。
将标准注意力替换为 F.scaled_dot_product_attention。
用 profiling 验证加速效果。
测试精度与基线一致（float16 下最大差异应 <1e-3）。

工作流 2：使用 flash-attn 库的高级特性

安装：pip install flash-attn --no-build-isolation（NVIDIA GPU，CUDA 12.0+）。
修改注意力代码，注意输入格式为 [batch_size, seq_len, num_heads, head_dim]，必要时转置。
启用高级特性：多查询注意力（共享 K/V）、滑动窗口注意力（如 window_size=(256, 256)）。
基准测试性能。

工作流 3：H100 FP8 优化（FlashAttention-3）

确认 H100/H800 GPU：nvidia-smi --query-gpu=name --format=csv。
安装 flash-attn（H100 已包含 FP8 支持）。
将输入转换为 FP8：q.to(torch.float8_e4m3fn) 等。
运行 FP8 注意力，资料称约 1.2 PFLOPS，比 FP16 快 1.5-2 倍。

适用场景

训练序列长度 >512 token 的 Transformer。
长上下文（>2K token）推理。
GPU 显存受限（标准注意力 OOM）。
需要 2-4 倍加速且不损失精度。
使用 PyTorch 2.2+ 或可安装 flash-attn。

限制与注意事项

序列 <256 token 时标准注意力更合适（开销不划算）；<512 token 加速有限（10-20%），512-2K token 为 2-3 倍，>2K token 为 3-4 倍。
需要 GPU：NVIDIA Ampere+（A100、A10、A30）或 AMD MI200+；Turing（T4）支持；Volta（V100）不支持；不支持 CPU 推理。
CUDA 12.0+（最低 11.8），PyTorch 2.2+ 原生支持。
仅支持 float16 或 bfloat16，不支持 float32。
若需更多注意力变体（而非仅加速）可考虑 xFormers；CPU 推理可用内存高效注意力。
常见问题：ImportError: cannot import flash_attn 需加 --no-build-isolation；CUDA 错误需确认 GPU 计算能力 ≥(7, 5)；精度下降需检查 dtype。

硬件要求

GPU：NVIDIA Ampere+（A100、A10、A30）或 AMD MI200+。
显存：与标准注意力相同（Flash Attention 不增加内存）。
CUDA：12.0+（最低 11.8）。
PyTorch：2.2+ 原生支持。
不支持：V100（Volta）、CPU 推理。

参考资料

论文："FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness"（NeurIPS 2022）
论文："FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning"（ICLR 2024）
博客：https://tridao.me/blog/2024/flash3/
GitHub：https://github.com/Dao-AILab/flash-attention
PyTorch 文档：https://pytorch.org/docs/stable/generated/torch.nn.functional.scaled_dot_product_attention.html

技能介绍

用途

主要能力

使用方式

PyTorch 原生（最简单，PyTorch 2.2+）

flash-attn 库（更多特性）

工作流 1：在现有 PyTorch 模型中启用

工作流 2：使用 flash-attn 库的高级特性

工作流 3：H100 FP8 优化（FlashAttention-3）

适用场景

限制与注意事项

硬件要求

参考资料

相关技能