Flash Attention 是一种快速且显存高效的注意力实现，通过 IO 感知的分块（tiling）与重计算（recomputation），为 Transformer 注意力带来 2-4 倍加速与 10-20 倍显存占用降低。该技能围绕如何在训练/推理中启用与优化它展开，适用于长序列 Transformer 场景。
【主要能力】 提供三条使用路径：一是 PyTorch 原生方式（PyTorch 2.2+），通过 F.scaled_dot_product_attention 自动使用 Flash Attention，并可用 torch.backends.cuda.sdp_kernel 强制指定后端；二是 flash-attn 库方式，支持多查询注意力（MQA）、滑动窗口注意力、自定义 softmax_scale 等更多特性；三是 H100 FP8 优化（FlashAttention-3），将输入转为 float8_e4m3fn 后调用 flash_attn_func，在 H100 上约 1.2 PFLOPS、比 FP16 快 1.5-2 倍。
【使用方式】 原生路径：先检查 PyTorch 版本（≥2.2，低于则 pip install --upgrade torch），把标准注意力（softmax(q @ k^T / sqrt(d_k)) @ v）替换为 F.scaled_dot_product_attention(q, k, v, attn_mask=mask)；可选用 torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=False, enable_mem_efficient=False) 固定后端。 flash-attn 路径：pip install flash-attn --no-build-isolation 安装（NVIDIA GPU，CUDA 12.0+），输入形状为 [batch, seq, heads, dim]（如原为 [batch, heads, seq, dim] 需 transpose(1,2)），调用 flash_attn_func(q, k, v, dropout_p, causal, window_size, softmax_scale)，支持 MQA（K/V 头数少于 Q 头数）与滑动窗口 window_size=(256,256)。 FP8 路径：确认 GPU 为 H100/H800（nvidia-smi 查询），安装支持 FP8 的 flash-attn，将 q/k/v 转为 torch.float8_e4m3fn 后调用 flash_attn_func。
【验证与工作流】 集成流程为：检查 PyTorch 版本 → 启用后端 → 用 torch.utils.benchmark 对比 Flash 与标准注意力的耗时（预期序列 >512 token 时 2-4 倍加速）→ 比对输出最大差值（float16 下应 <1e-3）。
【适用场景】 训练序列长度 >512 token 的 Transformer；长上下文（>2K token）推理；显存受限（标准注意力 OOM）；需要 2-4 倍加速且不损失精度；使用 PyTorch 2.2+ 或可安装 flash-attn。序列 <256 token 时建议用标准注意力；需要更多注意力变体可用 xFormers；CPU 推理用 memory-efficient attention（Flash Attention 需 GPU）。
【限制与硬件要求】 GPU：NVIDIA Ampere+（A100、A10、A30）或 AMD MI200+；CUDA 12.0+（最低 11.8）；PyTorch 2.2+ 原生支持。不支持 V100（Volta）与 CPU 推理。数据类型需为 float16 或 bfloat16，不支持 float32。加速效果随序列长度增加而提升：<512 token 仅 10-20%，512-2K 为 2-3 倍，>2K 为 3-4 倍。
【常见问题】 ImportError 时用 pip install flash-attn --no-build-isolation 或先 conda install cuda -c nvidia；CUDA 报错时用 torch.cuda.get_device_capability() 确认算力 ≥(7,5)；精度下降时检查 dtype 是否为 float16/bfloat16。
【参考资源】 论文「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 。进阶内容（HuggingFace Transformers 集成、基准对比、算法细节、旋转位置编码/ALiBi/paged KV cache/自定义掩码）见技能内 references 目录文档。