FlashAttention 与高效注意力内核：IO 感知与分块计算

标准注意力三步：
Q, K, V ∈ R^(N×d)
S = Q @ K^T           # 形状 N×N
P = softmax(S)        # 还是 N×N
O = P @ V             # 形状 N×d

内存特征：
□ S 和 P 都是 N×N 的中间矩阵 → 必须写回 HBM
□ 大矩阵乘法要切块循环 → 中间结果反复写读 HBM
□ N 越大，S/P 越大，HBM 带宽成为真正的瓶颈

为什么「算力够用但跑不快」：
□ A100 HBM 带宽 ~2 TB/s，算力 312 TFLOPS
□ 注意力是「带宽受限」算子：IO 决定墙钟时间
□ 即使算力空闲，数据搬不动 → 长序列卡在 IO

IO 量估算：
S 写回 HBM + 读回：2 × N² × 4 bytes（FP32）
P 写回 + 读回：2 × N² × 4 bytes
N=4096：中间矩阵 IO ≈ 256 MB（一次注意力）
N=32768：≈ 16 GB → 光搬数据就要 8 秒（A100）
→ 这就是「长上下文跑不动」的直接原因

三个关键手段：
□ 分块（Tiling）：
  - 把 Q/K/V 切成小块，一次只算一块注意力
  - 小块计算结果留在 SRAM（片上），不写回 HBM
□ 重算（Recompute）：
  - 反向传播不保存 S、P（省显存）
  - 反向时用保存的 Q/K/V 重新算一遍
□ IO 感知（IO-aware）：
  - 目标从「减少 FLOPs」变成「减少 HBM 访问」
  - 块大小按 SRAM 容量决定，最大化片上复用

HBM 访问对比：
□ 标准：O(N²) 次 HBM 访问
□ FlashAttention：O(N²×d²/SRAM) 次 → 长序列下近似 O(N)

分块示意图：
Q 切 [Q₁, Q₂, ...]，K/V 也切块
对每个 Q 块：
  加载 Qᵢ、遍历所有 K/V 块
  在 SRAM 里算局部注意力、累积 O 与 softmax 统计量
  块间结果在片上合并，直到最后才写一次 O

标准 softmax：
□ 需要完整一行算 max，再算 exp，再归一化
□ 要求整行 S 都在手边 → 与分块矛盾

Online Softmax 技巧：
□ 每个块维护「局部 max（m）」与「局部 sum（l）」
□ 新块到来时更新：
  - 新 max = max(旧 max, 新块 max)
  - 用新 max 修正旧块的累积值（rescaling）
□ 保证数值稳定（每步减 max）同时不落盘

修正公式（直觉）：
旧块统计量按「新 max」重新缩放
l_new = l_old × exp(m_old - m_new) + l_block
m_new = max(m_old, m_block)

流式过程：
块1：max=2.0, sum=3.1
块2：max=3.5 → 修正块1：sum_old × exp(2.0-3.5)
      总 max=3.5，总 sum = 修正后 + 块2
→ 最终 softmax 分母 = 总 sum，分子逐块用全局 max 归一

前向：
□ 只保存输出 O 与每块的 softmax 统计量（m、l）
□ 不保存 S、P → 峰值显存从 O(N²) 降到 O(N)
□ 块内计算在 SRAM，最后写 O

反向：
□ 标准实现需要 S、P 算梯度 → 全存显存
□ FlashAttention 反向「重算」：
  - 用保存的 Q/K/V + 统计量，重新算一遍注意力
  - 重算的 FLOPs 多一点点，但省下 O(N²) 显存
□ 重算在训练/长上下文微调中同样关键

显存对比：
□ 标准：注意力中间矩阵 O(N²)
□ FlashAttention：O(N)（只存 Q/K/V/O + 统计量）

反向重算流程：
加载 Q/K/V 块 → 重新算 S、P（不落盘，在 SRAM）
→ 结合上游梯度 dO 算 dQ/dK/dV
→ 每个块算完即弃，峰值显存只与「单块」有关

FlashAttention 变体族谱：
□ FlashAttention-2：进一步优化并行策略与分块
  - 在 H100/A100 上达到接近峰值算力利用率
□ FlashDecoding：解码阶段的变体
  - 针对「一个 query 对一个长 KV」的 decode 特征
  - 把 KV 切块并行算部分和，再 reduce
□ PagedAttention（vLLM）：
  - FlashAttention 的显存管理 + 页式 KV 存储
  - 注意力计算支持「非连续块」的 KV 读取
□ 分页 FlashAttention（Paged + Flash 融合）：
  - 兼顾页式显存复用与 IO 优化
□ MLA（Multi-head Latent Attention）：
  - DeepSeek 的低秩 KV 压缩
  - 把 KV 压缩成潜在向量 → KV Cache 大幅缩小
  - 推理时再解压（需要内核级支持）

变体选择场景：
短上下文高吞吐      → 标准 FlashAttention-2
超长上下文          → FlashAttention + 页式/稀疏
解码带宽受限        → FlashDecoding（KV 并行分块）
显存极度紧张        → MLA 低秩压缩

融合的价值：
□ 每少一次 HBM 写读，就少一段 IO 时间
□ 注意力的上下游算子可以「粘」进同一个 kernel

常见融合组合：
□ Attention + 残差 + LayerNorm：
  - pre-norm：Q/K/V 投影 → 注意力 → 残差 → LayerNorm
  - 合成一个 kernel，中间结果不落盘
□ QKV 投影融合：把三个线性层合成一个矩阵乘
□ 输出投影 + 残差 + MLP 的连续融合
□ RoPE 位置编码融合进 Q/K 投影

与 MoE / MoA 的配合：
□ 注意力部分用 FlashAttention（IO 优化）
□ 专家部分单独内核（序列专家 GEMM）
→ 各算各的最优，避免强行融合反而退化

融合前后对比：
未融合：Q投影→写HBM → 注意力→写HBM → 残差→写HBM
融合后：Q投影+注意力+残差+LN 一个 kernel 完成
→ HBM 写读次数从 4 次降到 1 次

衡量指标：
□ HBM 访问量（越低越好）
□ 算力利用率（越接近峰值越好）
□ 端到端注意力耗时

FlashAttention 实测特征（社区公开数据口径）：
□ A100 上 FlashAttention 相比标准实现提速 2-4x
□ 长序列收益更大：序列越长，IO 优势越明显
□ FlashAttention-2 在 A100/H100 达到 ~70%-90% 算力利用率
□ FlashDecoding 在 decode 阶段比标准快数倍

瓶颈回归点：
□ 短序列（N<256）：算力与 IO 差距小，收益收窄
□ 极长序列：从「IO 受限」彻底变成「算力受限」

典型量级（相对值，供直觉参考）：
N=512    ：Flash ≈ 1.5-2x
N=2048   ：Flash ≈ 2-3x
N=8192   ：Flash ≈ 3-4x
N=32768  ：Flash 收益进一步拉大（IO 差 O(N²) vs O(N)）

手写内核的核心决策：
□ 分块大小：由 SRAM 容量决定（如 128×128）
  - 太大：SRAM 溢出 → 被迫落盘
  - 太小：块间开销高、复用差
□ 数据布局：Q/K/V 按块连续加载，支持张量核心
□ online softmax 统计量：每个线程块维护 m/l
□ 反向重算：不保存 S/P，用统计量重建

CUDA 层面：
□ 用共享内存（shared memory）模拟 SRAM 手动分块
□ 寄存器级切分：让每个线程处理一小块，最大化复用
□ 用 Tensor Core / cuBLAS 风格的矩阵乘原语
□ 注意 bank conflict：共享内存布局要对齐

伪代码骨架（前向，一块）：
load Q_block, K_block, V_block → shared memory
S_block = Q_block @ K_block^T      # 在片上
m_new = max(m_old, rowmax(S_block))
P_block = exp(S_block - m_new)
l_new = l_old * exp(m_old - m_new) + rowsum(P_block)
O_block = rescale(O_old) + P_block @ V_block
m_old, l_old, O_old = m_new, l_new, O_block

选型建议：
□ 生产推理引擎：直接用 vLLM/SGLang/TensorRT-LLM 内置内核
□ 需要特殊变体：FlashDecoding、Paged、MLA 看引擎支持
□ 自研内核：先跑通数学，再优化，再对比标准实现
□ 训练侧：PyTorch 的 flash_attn 库或 xformers 内存高效注意力

高频踩坑：
□ 直接手写而不对比数值 → 精度偏差（online softmax 细节错）
□ 分块大小拍脑袋 → 块太大落盘、太小复用差
□ 忽视 bank conflict → 共享内存命中率暴跌
□ 短序列强上 Flash → 收益小、调度开销反而拖慢
□ 把 FLOPs 当性能指标 → 真正瓶颈是 HBM 访问量
□ 反向重算的 FLOPs 上升没算进成本 → 训练侧意外变慢

数值验证：
□ 前向：输出与 float64 参考一致（相对误差 < 1e-5 量级）
□ 反向：梯度与标准实现一致（gradcheck）

验证清单：
□ max 值稳定性（数值溢出边界测试）
□ 掩码/因果注意力支持
□ 不同 dtype（FP16/BF16/FP8）的精度
□ 变长序列（分页/打包）支持