Flash Attention 优化

nanoGPT 中的注释

> "flash attention make GPU go brrrrr but support is only in PyTorch >= 2.0"

---

问题：标准注意力的瓶颈

标准注意力计算：


1. 计算 Q × K^T → 得到 (T×T) 的注意力矩阵      ← 写入 HBM（显存）
2. 对注意力矩阵做 Softmax                        ← 从 HBM 读出
3. 乘以 V                                        ← 从 HBM 读出

对于序列长度 T=1024：注意力矩阵大小 = 1024 × 1024 × 4 bytes ≈ 4MB

对于 T=8192（长文档）：注意力矩阵 = 256MB

瓶颈在于 IO：GPU 的计算速度很快，但内存带宽（HBM ↔ SRAM）是瓶颈。

---

Flash Attention 的解决方案

核心思路：分块（Tiling）计算，避免将完整注意力矩阵写入 HBM。


将 Q, K, V 分成小块
    ↓
每次只在 SRAM（片上缓存）中处理一小块
    ↓
在线计算 Softmax（不需要看完整矩阵）
    ↓
直接输出结果，不写中间矩阵到 HBM

---

性能提升

| 指标 | 标准注意力 | Flash Attention |

|------|-----------|----------------|

| HBM读写 | 存取大注意力中间矩阵 | 减少中间矩阵存取；复杂度还依赖头维度及SRAM容量 |

| 速度 | 基准 | 依硬件和工作负载测定，非固定倍数 |

| 显存占用 | O(T²) | O(T) |

| 支持更长序列 | 受限 | 可处理更长输入 |

---

nanoGPT 中的使用


# 检测SDPA接口是否存在，不是确认Flash内核已启用
self.flash = hasattr(torch.nn.functional, 'scaled_dot_product_attention')

if self.flash:
    # 调用SDPA，由后端选择可用实现
    y = torch.nn.functional.scaled_dot_product_attention(
        q, k, v,
        attn_mask=None,
        dropout_p=self.dropout if self.training else 0,
        is_causal=True  # 因果掩码由内核自动处理
    )
else:
    # 退回手动实现
    att = (q @ k.transpose(-2, -1)) * (1.0 / math.sqrt(k.size(-1)))
    att = att.masked_fill(self.bias[:,:,:T,:T] == 0, float('-inf'))
    att = F.softmax(att, dim=-1)
    att = self.attn_dropout(att)
    y = att @ v

---

与律师工作的关联

Flash Attention可降低注意力中间存储开销，但不能单独决定模型上下文长度或保证完整分析：

  • 是否容纳长合同取决于具体模型窗口、分词和输出预算。
  • 能放入上下文不等于已经完整读取、正确引用或分析。
  • 汉字与Token没有通用固定换算；必要时仍需分段并核对覆盖。

    • ---

    相关页面

  • 01-概念/注意力机制 — 注意力原理
  • 02-技术/nanoGPT代码解析 — 完整代码
  • 03-法律AI应用/AI合同审查 — 长文档处理应用

    • 来源

  • Dao et al., "FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness" (2022)
  • [nanoGPT model.py](https://github.com/karpathy/nanoGPT/blob/master/model.py)

    • 本轮核查来源与限制（2026-09-17）

  • https://docs.pytorch.org/docs/2.14/generated/torch.nn.functional.scaled_dot_product_attention.html
  • https://arxiv.org/abs/2205.14135

    • 只修本轮明确列出的陈述，不将旧页整篇升级为已验证。themis_only为本地保存和禁止外发的兼容标记，不代表本地模型推理。