Inference Cookbook

常见问题

目录

安装与环境

Q1: 安装时编译失败，报内存不足

问题：pip install flash-attn 时编译器因内存不足被 OOM Kill。

解决方案：

# 限制并行编译任务数
MAX_JOBS=2 pip install flash-attn --no-build-isolation

# 或者使用预编译 wheel（推荐）
pip install flash-attn --no-build-isolation

Flash Attention 的编译需要大量内存（每个编译进程约 8-16GB），推荐至少 64GB RAM。如果机器内存有限，设置 MAX_JOBS=1 或 MAX_JOBS=2。

Q2: 如何确认 Flash Attention 已正确安装？ 

import flash_attn
print(flash_attn.__version__)

# 验证 CUDA 内核
import torch
from flash_attn import flash_attn_func
q = torch.randn(1, 64, 4, 64, device='cuda', dtype=torch.float16)
k = torch.randn(1, 64, 4, 64, device='cuda', dtype=torch.float16)
v = torch.randn(1, 64, 4, 64, device='cuda', dtype=torch.float16)
out = flash_attn_func(q, k, v)
print(f"成功! 输出形状: {out.shape}")

Q3: Flash Attention 2 和 Flash Attention 3 有什么区别？

特性 FA2 (flash-attn) FA3 (hopper/)
架构支持 SM80+ (Ampere/Ada/Hopper) SM90 (Hopper only)
安装方式 pip install flash-attn cd hopper && python setup.py install
性能 基准 前向 1.5-2× 提升
FP8 不支持 支持
特有优化 - Warp Specialization, TMA, GMMA

FA3 是 FA2 在 Hopper 架构上的深度优化版本。非 Hopper GPU 只能使用 FA2。

Q4: 支持哪些 GPU？

  • 完全支持：NVIDIA Ampere (A100, A6000)、Ada (RTX 4090, L40S)、Hopper (H100, H200)
  • AMD ROCm：MI200x, MI250x, MI300x（通过 Triton 后端）
  • 不支持：Turing (T4)、Volta (V100) 及更早架构

API 使用

Q5: Flash Attention 的张量布局与 PyTorch 不同吗？

是的。Flash Attention 使用 (batch, seqlen, num_heads, head_dim) 布局，而 PyTorch 的 nn.MultiheadAttention 默认使用 (seqlen, batch, embed_dim)。

# PyTorch 风格 → Flash Attention 风格
# 假设 x: (seqlen, batch, embed_dim)
x = x.transpose(0, 1)  # → (batch, seqlen, embed_dim)
x = x.view(batch, seqlen, num_heads, head_dim)  # → Flash Attention 格式

Q6: 因果遮蔽在 seqlen_q != seqlen_k 时如何对齐？

Flash Attention 的因果遮蔽对齐到注意力矩阵的右下角，而非左上角。这意味着：

  • 当 seqlen_q < seqlen_k 时，前缀部分的 K 对所有 Q 可见（适用于 prefix-filling）
  • 当 seqlen_q > seqlen_k 时，前面的 Q 位置没有可注意的 K

这是有意的设计决策，支持 KV Cache 场景中 seqlen_q=1 而 seqlen_k 很大的情况。

Q7: 如何实现交叉注意力？ 

from flash_attn import flash_attn_func

# Q 来自解码器，K/V 来自编码器
q = decoder_output  # (batch, seqlen_q, num_heads, head_dim)
k = encoder_output  # (batch, seqlen_k, num_heads, head_dim)
v = encoder_output  # (batch, seqlen_k, num_heads, head_dim)

out = flash_attn_func(q, k, v, causal=False)  # 交叉注意力不用因果遮蔽

也可以使用 KV 打包格式以获得更快的反向传播：

from flash_attn import flash_attn_kvpacked_func

kv = torch.stack([k, v], dim=2)  # (batch, seqlen_k, 2, num_heads, head_dim)
out = flash_attn_kvpacked_func(q, kv)

Q8: 如何处理变长序列的 padding？ 

from flash_attn import flash_attn_varlen_func
from flash_attn.bert_padding import unpad_input, pad_input

# 方法一：手动构造 cu_seqlens
seqlens = [100, 200, 150]  # 各序列长度
cu_seqlens = torch.tensor([0] + list(torch.cumsum(torch.tensor(seqlens), 0).numpy()),
                          dtype=torch.int32, device='cuda')

# 方法二：使用 unpad_input 工具
x_unpad, indices, cu_seqlens, max_seqlen, _ = unpad_input(x_padded, attention_mask)
# ... 执行变长注意力 ...
out_padded = pad_input(out_unpad, indices, batch_size, max_seqlen_padded)

Q9: return_attn_probs=True 是否影响性能？

是的，会显著影响。此选项会分配一个完整的 (batch, heads, seqlen_q, seqlen_k) 注意力矩阵，占用 $O(N^2)$ 内存，完全抵消了 Flash Attention 的内存优势。仅用于调试和测试，生产环境中不应使用。

数值与正确性

Q10: Flash Attention 的结果与标准注意力完全一致吗？

不完全一致，但数值差异在浮点精度范围内。Flash Attention 是 exact attention 的近似实现（不是算法层面的近似，而是浮点运算顺序的差异）。

参考容差：

数据类型 绝对误差 (atol) 相对误差 (rtol)
FP16 1e-3 1e-3
BF16 5e-3 5e-3
FP8 1e-2 1e-2

Q11: 反向传播的结果是否确定性的？

默认不是。Flash Attention 使用原子操作（atomicAdd）累加 dQ 梯度，不同运行的结果可能有微小差异。

启用确定性模式：

out = flash_attn_func(q, k, v, deterministic=True)

确定性模式通过避免原子操作来保证结果一致，但性能会下降约 10-20%。

Q12: 训练中出现 NaN/Inf 怎么办？

排查步骤：

  1. 检查输入是否包含 NaN：torch.isnan(q).any()
  2. 检查 softmax_scale 是否过大（默认 $1/\sqrt{d}$）
  3. 启用 residual_in_fp32=True 保持残差路径精度
  4. 使用 BF16 替代 FP16（更大的动态范围）
  5. 降低学习率

性能问题

Q13: Flash Attention 比标准实现慢？

可能的原因：

  1. 序列太短（< 128）：内核启动开销占比大
  2. Head dimension 不是 8 的倍数：需要 padding
  3. 输入不连续：需要先 .contiguous()
  4. GPU 利用率不足：batch × heads 太小
  5. 首次调用开销：CUDA 内核编译/加载，后续调用会快很多

验证 Flash Attention 确实被使用：

import os
os.environ['FLASH_ATTN_DEBUG'] = '1'  # 启用调试输出

Q14: 如何选择最优的 num_splits？

推荐使用 num_splits=0（自动决策）。自动策略考虑了 SM 数量、L2 Cache 大小和序列长度。手动调优：

  • batch × heads >= SM_count：num_splits=1
  • batch × heads < SM_count / 2：增加 splits 直到接近 SM_count
  • 超长序列（KV > L2）：增加 splits 使每个 split 的 KV 能放入 L2

Q15: GQA 比 MHA 慢？

当 seqlen_q 足够长且 num_heads_q 已经接近 kBlockM 的倍数时，PackGQA 可能比非打包模式略慢（额外的 divmod 和 __shfl_sync 开销）。此时可以手动禁用：

out = flash_attn_func(q, k, v, pack_gqa=False)

Q16: 如何测量 Flash Attention 的实际吞吐量？ 

from triton.testing import do_bench

# 正确的测量方式（包含 warmup）
time_ms = do_bench(lambda: flash_attn_func(q, k, v, causal=True),
                   warmup=5, rep=30)

# 计算 TFLOPS
total_flops = batch * nheads * 2 * seqlen_q * (seqlen_k / 2) * (headdim * 2)  # causal
tflops = total_flops / (time_ms * 1e-3 * 1e12)

详见 性能测试

推理与部署

Q17: KV Cache 需要手动管理 cache_seqlens 吗？

是的。flash_attn_with_kvcache 不会自动更新 cache_seqlens。每次追加新 token 后需要手动递增：

out = flash_attn_with_kvcache(q, k_cache, v_cache, k=k_new, v=v_new,
                              cache_seqlens=cache_seqlens, causal=True)
cache_seqlens += k_new.shape[1]  # 手动更新

Q18: Paged KV Cache 的 page_size 如何选择？

Flash Attention 内部的 page_block_size 为 256 token。用户侧的 page_size 应该是 kBlockN 的倍数以获得最佳性能（通常 64 或 128 的倍数）。与 vLLM 集成时需要匹配 vLLM 的 block_size。

Q19: 可以使用 CUDA Graph 加速推理吗？

可以。Flash Attention 兼容 CUDA Graph：

# 预热
for _ in range(3):
    out = flash_attn_with_kvcache(q, k_cache, v_cache, ...)
torch.cuda.synchronize()

# 捕获 Graph
g = torch.cuda.CUDAGraph()
with torch.cuda.graph(g):
    out = flash_attn_with_kvcache(q, k_cache, v_cache, ...)

# 重放（极低延迟）
g.replay()

注意：使用 CUDA Graph 时，输入张量的地址不能改变（只能改变内容）。

Q20: FP8 推理的输出为什么是 BF16？

这是设计选择。FP8 的精度不足以作为最终输出，因此 Flash Attention 内部使用 FP8 计算但以 BF16 输出。这在 softmax 归一化和 V 的 descale 应用后已经需要更高精度的中间结果。

高级功能

Q21: 如何同时使用因果遮蔽和滑动窗口？ 

# 因果 + 向左 256 token 的滑动窗口
out = flash_attn_func(q, k, v, window_size=(256, 0))

# 等价于 causal=True 并限制左侧窗口
# 位置 i 只能看到 [max(0, i-256), i] 范围的 K

window_size=(left, right) 中 right=0 即因果，left 控制滑动窗口大小。

Q22: Softcap 和因果遮蔽可以同时使用吗？

可以。Softcap 在 masking 之前应用：

out = flash_attn_func(q, k, v, causal=True, softcap=50.0)
# 计算顺序: QK^T → softcap(tanh) → causal mask → softmax → PV

Q23: 如何在不支持 Flash Attention 的环境中回退？ 

try:
    from flash_attn import flash_attn_func
    USE_FLASH = True
except ImportError:
    USE_FLASH = False

def attention(q, k, v, causal=False):
    if USE_FLASH:
        return flash_attn_func(q, k, v, causal=causal)
    else:
        # PyTorch 原生 SDPA
        q = q.transpose(1, 2)  # (b, h, s, d)
        k = k.transpose(1, 2)
        v = v.transpose(1, 2)
        out = torch.nn.functional.scaled_dot_product_attention(
            q, k, v, is_causal=causal
        )
        return out.transpose(1, 2)  # 回到 (b, s, h, d)

Q24: torch.compile 与 Flash Attention 兼容吗？

兼容。Flash Attention 使用 torch.library.custom_op 注册，支持 torch.compile 的 tracing。

model = MyModel()  # 内含 flash_attn_func 调用
compiled = torch.compile(model, mode='reduce-overhead')
out = compiled(x)  # 正常工作

限制：return_attn_probs=True 不兼容 compile。

导航

2026年2月24日
GitHub
论文阅读指南