# 因果 FlashAttention

## 实现目标

`flash_attention.py` 为形状 `(B, H, N, Dh)` 的 Q、K、V 实现因果注意力前向和
反向传播。实现按序列维度分块，使用在线 softmax 避免在 DRAM 中物化完整的
`N x N` 注意力矩阵。

## 核心数据流

1. `attention_forward()` 为每个 Q 块遍历 K/V 块，分开处理因果对角块与无需掩码
   的下三角块。
2. 在线 softmax 持续更新每行最大值、归一化分母和输出累加器；使用 `exp2` 与预缩放
   降低指数计算成本。
3. 前向保存每行的 log-sum-exp（LSE），供反向重建注意力概率。
4. `_attention_backward_preprocess()` 计算
   `delta = sum(output * grad_output)`。
5. `attention_backward()` 分两个阶段遍历分块：先计算 `grad_k`、`grad_v`，再计算
   `grad_q`，并对对角块应用因果掩码。
6. `_FlashAttention` 将这些内核封装为 `torch.autograd.Function`，向 PyTorch
   返回 Q、K、V 的梯度。

## 关键细节与限制

- Q、K、V 必须形状、设备、dtype 和 stride 一致；当前包装器要求 FP32。
- 头维 `Dh` 不得超过 128，以控制单个程序的 SRAM 使用量。
- 非块大小整数倍的序列长度通过 load/store 掩码处理。
- 实现只支持因果注意力，不接受额外的 attention mask、dropout 或偏置。
- 块大小和流水线配置由 autotune 选择，首次运行新形状时会产生调优开销。

## 测试与基准

```bash
python3 flash_attention.py
python3 flash_attention.py --benchmark
```

测试将前向输出及 `dQ`、`dK`、`dV` 与 PyTorch
`scaled_dot_product_attention()` 比较；基准测试分别测量前向和反向 TFLOPS。

## 配套视频

[bilibili](https://www.bilibili.com/video/BV1fMyWBgERM?spm_id_from=333.788.videopod.episodes&vd_source=e98b669ccbafff4b5aa59dd6303b722f&p=9)
