第 3 章 进阶原理：性能模型与数值稳定性

第 3 章 进阶原理：性能模型与数值稳定性

本章导读：第二章我们看到了两个模块"怎么做"。这一章回答"为什么这样更快、会不会算错、怎么量化收益"。不少工程师把 FlashAttention 当成黑盒直接 enable，结果遇到精度对不齐、长序列收益不及预期时完全无从排查。本章会给出一套可落地的性能测试框架，并解释 FlashAttention 的数值等价性——它和朴素实现在数学上是否一致、在什么情况下会有微小差异。读完你能独立设计一次"加速是否真的有效"的对照实验。

3.1 性能模型：为什么长序列收益更大

GPU 上一次注意力内核的时间，可以粗略分解为：

T ≈ 计算时间（正比于 FLOPS）+ 显存搬运时间（正比于 HBM 字节数）
​

关键事实：A100 的算力约 312 TFLOPS（fp16），但 HBM 带宽只有约 2 TB/s。当算子的"算术强度"（每搬一字节能做的运算量）不足时，内核是**带宽受限（memory-bound）**的——它在等数据，不在算。

朴素注意力是典型带宽受限：它要把 O(n²) 的分数矩阵写出再读回做 softmax。FlashAttention-2 把数据留在 SRAM，把 HBM 访问从 O(n²) 降到 O(n)，于是搬运时间大幅下降。

为什么长序列更赚：n 越大，朴素实现的 O(n²) 显存压力和 HBM 往返越夸张，FA-2 节约的相对比例越高。短序列时，两者都可能在别处受限（如 kernel launch、KV 读取本身），收益有限。

可操作的测试清单：

1. 固定模型、固定 batch，扫不同的 max_seq_len（如 512 / 1024 / 2048 / 4096 / 8192）。
2. 每个长度下分别跑"朴素注意力"与"FlashAttention-2"，记录 tokens/s 与峰值显存。
3. 用 nvidia-smi dmon 或 PyTorch Profiler 看 SM 利用率与 HBM 带宽占用。
4. 结论不应只写"快了 X 倍"，要写"在 >2048 长度后收益显著，4096 时提速约 Y%"。

3.2 数值稳定性与精度对齐

FlashAttention-2 是否改变结果？答案是：数学上等价，实现上数值等价性由 online softmax 保证。因为 softmax 是"平移不变"的（分子分母同减最大值不影响结果），在线增量更新只是换了个顺序累计 m 和 l。

但你仍可能观察到微小差异，来源通常是：

• 浮点累加顺序不同（并行 reduce 的合并次序）。
• 是否使用不同精度的累加器（FA-2 内部常用 fp32 累加以提升稳定性）。
• 头部维度 d 不是 2 的幂或触发了不同的内核路径。

排查建议：

• 先确认两者在短序列、小 batch 下输出是否一致到 1e-3 量级。
• 若差异过大，优先怀疑自己是否混用了不同 attention 实现（例如某层没开 FA）。
• 涉及量化（INT8/FP8）时，精度差异主要来自量化本身，而非 FA。

PagedAttention 共享场景的写时复制：当多个序列共享前缀块时，读操作共用物理块；一旦某个序列要写入新 KV（生成新 token），才复制该块，避免互相污染。这正是操作系统中 copy-on-write 的翻版，也是 beam search 能省显存的原因。

本章给你的是"方法论武器"：下一章我们用真实框架（vLLM、HuggingFace Transformers）把这些原理落成可直接抄的作业。