optimizing-attention-flash

Flash Attention - 快速且节省内存的注意力机制

快速开始

Flash Attention 通过感知 IO 的分块和重新计算，为 Transformer 注意力机制带来 2-4 倍的加速和 10-20 倍的内存占用降低。

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-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 模型中启用

复制此清单：

Flash Attention Integration:
- [ ] Step 1: Check PyTorch version (≥2.2)
- [ ] Step 2: Enable Flash Attention backend
- [ ] Step 3: Verify speedup with profiling
- [ ] Step 4: Test accuracy matches baseline

步骤 1：检查 PyTorch 版本

python -c "import torch; print(torch.__version__)"
# Should be ≥2.2.0

如果版本低于 2.2，请升级：

pip install --upgrade torch

步骤 2：启用 Flash Attention 后端

替换标准注意力机制：

# Before (standard attention)
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / math.sqrt(d_k), dim=-1)
out = attn_weights @ v

# After (Flash Attention)
import torch.nn.functional as F
out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask)

强制使用 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)

步骤 3：通过性能分析验证加速效果

import torch.utils.benchmark as benchmark

def test_attention(use_flash):
    q, k, v = [torch.randn(2, 8, 2048, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

    if use_flash:
        with torch.backends.cuda.sdp_kernel(enable_flash=True):
            return F.scaled_dot_product_attention(q, k, v)
    else:
        attn = (q @ k.transpose(-2, -1) / 8.0).softmax(dim=-1)
        return attn @ v

# Benchmark
t_flash = benchmark.Timer(stmt='test_attention(True)', globals=globals())
t_standard = benchmark.Timer(stmt='test_attention(False)', globals=globals())

print(f"Flash: {t_flash.timeit(100).mean:.3f}s")
print(f"Standard: {t_standard.timeit(100).mean:.3f}s")

预期效果：序列长度 >512 个 token 时可加速 2-4 倍。

步骤 4：测试准确率是否与基准一致

# Compare outputs
q, k, v = [torch.randn(1, 8, 512, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

# Flash Attention
out_flash = F.scaled_dot_product_attention(q, k, v)

# Standard attention
attn_weights = torch.softmax(q @ k.transpose(-2, -1) / 8.0, dim=-1)
out_standard = attn_weights @ v

# Check difference
diff = (out_flash - out_standard).abs().max()
print(f"Max difference: {diff:.6f}")
# Should be <1e-3 for float16

工作流 2：使用 flash-attn 库实现高级功能

适用于多查询注意力、滑动窗口注意力或 H100 FP8。

复制此清单：

flash-attn Library Setup:
- [ ] Step 1: Install flash-attn library
- [ ] Step 2: Modify attention code
- [ ] Step 3: Enable advanced features
- [ ] Step 4: Benchmark performance

步骤 1：安装 flash-attn 库

# NVIDIA GPUs (CUDA 12.0+)
pip install flash-attn --no-build-isolation

# Verify installation
python -c "from flash_attn import flash_attn_func; print('Success')"

步骤 2：修改注意力代码

from flash_attn import flash_attn_func

# Input: [batch_size, seq_len, num_heads, head_dim]
# Transpose from [batch, heads, seq, dim] if needed
q = q.transpose(1, 2)  # [batch, seq, heads, dim]
k = k.transpose(1, 2)
v = v.transpose(1, 2)

out = flash_attn_func(
    q, k, v,
    dropout_p=0.1,
    causal=True,  # For autoregressive models
    window_size=(-1, -1),  # No sliding window
    softmax_scale=None  # Auto-scale
)

out = out.transpose(1, 2)  # Back to [batch, heads, seq, dim]

步骤 3：启用高级功能

多查询注意力（各个头之间共享 K/V）：

from flash_attn import flash_attn_func

# q: [batch, seq, num_q_heads, dim]
# k, v: [batch, seq, num_kv_heads, dim]  # Fewer KV heads
out = flash_attn_func(q, k, v)  # Automatically handles MQA

滑动窗口注意力（局部注意力）：

# Only attend to window of 256 tokens before/after
out = flash_attn_func(
    q, k, v,
    window_size=(256, 256),  # (left, right) window
    causal=True
)

步骤 4：对性能进行基准测试

import torch
from flash_attn import flash_attn_func
import time

q, k, v = [torch.randn(4, 4096, 32, 64, device='cuda', dtype=torch.float16) for _ in range(3)]

# Warmup
for _ in range(10):
    _ = flash_attn_func(q, k, v)

# Benchmark
torch.cuda.synchronize()
start = time.time()
for _ in range(100):
    out = flash_attn_func(q, k, v)
    torch.cuda.synchronize()
end = time.time()

print(f"Time per iteration: {(end-start)/100*1000:.2f}ms")
print(f"Memory allocated: {torch.cuda.max_memory_allocated()/1e9:.2f}GB")

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

在 H100 GPU 上实现最高性能。

FP8 Setup:
- [ ] Step 1: Verify H100 GPU available
- [ ] Step 2: Install flash-attn with FP8 support
- [ ] Step 3: Convert inputs to FP8
- [ ] Step 4: Run with FP8 attention

步骤 1：验证 H100 GPU

nvidia-smi --query-gpu=name --format=csv
# Should show "H100" or "H800"

步骤 2：安装支持 FP8 的 flash-attn

pip install flash-attn --no-build-isolation
# FP8 support included for H100

步骤 3：将输入转换为 FP8

import torch

q = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)
k = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)
v = torch.randn(2, 4096, 32, 64, device='cuda', dtype=torch.float16)

# Convert to float8_e4m3 (FP8)
q_fp8 = q.to(torch.float8_e4m3fn)
k_fp8 = k.to(torch.float8_e4m3fn)
v_fp8 = v.to(torch.float8_e4m3fn)

步骤 4：使用 FP8 注意力运行

from flash_attn import flash_attn_func

# FlashAttention-3 automatically uses FP8 kernels on H100
out = flash_attn_func(q_fp8, k_fp8, v_fp8)
# Result: ~1.2 PFLOPS, 1.5-2x faster than FP16

适用场景与替代方案

以下情况适合使用 Flash Attention：

• 训练序列长度 >512 个 token 的 Transformer
• 使用长上下文（>2K 个 token）进行推理
• GPU 内存受限（标准注意力会导致 OOM）
• 需要提升 2-4 倍速度且不损失准确率
• 使用 PyTorch 2.2+ 或能够安装 flash-attn

以下情况适合使用替代方案：

• 标准注意力：序列长度 <256 个 token（开销不值得）
• xFormers：需要更多注意力变体（不只是提升速度）
• 节省内存的注意力机制：CPU 推理（Flash Attention 需要 GPU）

常见问题

问题：ImportError: cannot import flash_attn

使用 no-build-isolation 标志安装：

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

或者先安装 CUDA 工具包：

conda install cuda -c nvidia
pip install flash-attn --no-build-isolation

问题：速度比预期慢（没有加速）

序列越长，Flash Attention 的优势越明显：

• <512 个 token：加速效果很小（10-20%）
• 512-2K 个 token：加速 2-3 倍

• 2K 个 token：加速 3-4 倍

请检查序列长度是否足够。

问题：RuntimeError: CUDA error

验证 GPU 是否支持 Flash Attention：

import torch
print(torch.cuda.get_device_capability())
# Should be ≥(7, 5) for Turing+

Flash Attention 的要求：

• Ampere（A100、A10）：✅ 完全支持
• Turing（T4）：✅ 支持
• Volta（V100）：❌ 不支持

问题：准确率下降

请检查 dtype 是否为 float16 或 bfloat16（而非 float32）：

q = q.to(torch.float16)  # Or torch.bfloat16

Flash Attention 使用 float16/bfloat16 提升速度，不支持 float32。

高级主题

与 HuggingFace Transformers 集成：请参阅 references/transformers-integration.md，了解如何在 BERT、GPT、Llama 模型中启用 Flash Attention。

性能基准测试：请参阅 references/benchmarks.md，了解不同 GPU 和序列长度下的详细速度与内存对比。

算法详情：请参阅 references/algorithm.md，了解分块策略、重新计算和 IO 复杂度分析。

高级功能：请参阅 references/advanced-features.md，了解旋转位置嵌入、ALiBi、分页 KV 缓存和自定义注意力掩码。

硬件要求

• GPU：NVIDIA Ampere+（A100、A10、A30）或 AMD MI200+
• VRAM：与标准注意力相同（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

Agent 操作流程

1. 检查环境。 检查 GPU 类型、显存和驱动/CUDA 版本（nvidia-smi）、框架版本，以及数据集的位置和大小。
2. 确定输入。 根据数据或用户确认格式、标识符和参数。任何会改变结果的值都应询问，而不是猜测。
3. 先运行小规模版本。 进行冒烟测试：使用小型模型或子集，运行少量步骤，并确认损失下降且检查点能够保存。
4. 使用上述说明和参考资料执行完整任务。
5. 验证结果。 在留出数据上跟踪指标，与基线进行比较，并记录随机种子、配置和硬件。
6. 报告。 说明运行了什么（版本、命令、参数）、检查了什么，以及还有哪些不确定之处。

如果发生这种情况	请这样做
CUDA 显存不足	减小批次大小，启用梯度累积/检查点或混合精度，或者对模型进行分片。
损失为 NaN 或发散	降低学习率，检查数据中是否有无效值，并启用梯度裁剪。
已安装版本中缺少这些说明里的函数、标志或端点	查看已安装版本自身的文档（help()、--help、官方文档），调整做法，并告知用户。切勿编造 API。
必需的输入、标识符或参数含义不明确	询问用户，或在运行前明确说明所做的假设。

诚信准则

• 切勿捏造结果、参数、标识符、引用或统计数据。如果某项内容无法运行或验证，请如实说明。
• 没有日志时，切勿声称取得了训练结果；启动大型任务前应估算计算成本。
• 此处与版本相关的细节可能已过时：请根据已安装版本的官方文档进行确认。
• 涉及费用、占用共享 GPU 或云配额、触及个人或患者数据，或无法撤销的操作前，必须先征求同意。

相关技能

• rwkv-architecture：介绍 RWKV（Receptance Weighted Key Value）架构。这是一种 RNN/Transformer 混合架构，推理复杂度为 O(n)，且无需 KV 缓存，包括 RWKV……
• llama-cpp：在 CPU、Apple Silicon 和消费级 GPU 上运行 LLM 推理，无需 NVIDIA 硬件。
• tensorrt-llm：使用 NVIDIA TensorRT 优化 LLM 推理，实现最高吞吐量和最低延迟。