"""
=============================================================================
CS349E · 第 11 讲 · 02_flash_forward.py —— FlashAttention 前向：分块、两种循环顺序与显存读写
=============================================================================

【这个脚本要回答的问题】
    FlashAttention 和标准注意力算的是同一个数（精确注意力，不是近似）、运算量也差不多，为什么快好几倍？
    因为它读写显存少得多。这个脚本真的把分块算法跑一遍，并数出每种算法读写了多少个元素。
      演示 1：正确性：两种循环顺序（论文算法 1 / FlashAttention-2）× causal × 不整除的长度，都等于标准注意力。
      演示 2：显存读写随序列长度 N 的变化：标准 Θ(Nd + N²) vs FlashAttention Θ(N²d²/M)。
      演示 3：显存读写随片上 SRAM 大小 M 的变化；论文命题 3 的下界。
      演示 4：论文图 2 的设置（GPT-2 medium：N = 1024、d = 64、16 个头、批 64，A100）：我们的计数 vs 论文的测量。

【运行】python3 02_flash_forward.py      （约 3 秒）
【依赖】numpy；本目录的 flash_lab；../lecture01 的 console。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from flash_lab import IO, block_sizes, fa_io_formula, flash_bwd, flash_fwd, standard_attention, standard_attention_bwd  # noqa: E402
from console import fmt_bytes, fmt_time, print_header, print_observations, print_table  # noqa: E402

M_100KB = 100 * 1024 // 2          # 论文说片上 SRAM 约 100 KB；FP16 每个元素 2 字节 → 约 51,200 个元素


def std_io(N, d):
    """标准注意力前向的读写元素数（flash_lab.standard_attention 的公式，不实际计算 N×N 的矩阵）。"""
    return IO(read=2 * N * d + 2 * N * N + N * d, write=2 * N * N + N * d, flops=4 * N * N * d)


def demo1_correctness():
    print_header("演示 1：正确性——和标准注意力（FP64）比较")
    rng = np.random.default_rng(0)
    rows = []
    for (N, Nk, d, Br, Bc) in ((256, 256, 64, 64, 64), (300, 300, 64, 64, 100), (77, 513, 32, 16, 50), (1, 1000, 64, 64, 128)):
        Q, K, V = rng.standard_normal((N, d)), rng.standard_normal((Nk, d)), rng.standard_normal((Nk, d))
        for causal in (False, True):
            O_ref, _, _ = standard_attention(Q, K, V, causal)
            errs = []
            for order in ("fa1", "fa2"):
                O, _, _ = flash_fwd(Q, K, V, Br, Bc, causal, order=order)
                errs.append(np.abs(O - O_ref).max())
            rows.append([f"N={N}, Nk={Nk}, d={d}", f"Br={Br}, Bc={Bc}", "是" if causal else "否",
                         f"{errs[0]:.1e}", f"{errs[1]:.1e}"])
    print_table(["形状", "块", "causal", "算法 1 顺序的最大误差", "FlashAttention-2 顺序"], rows)
    print_observations([
        "FlashAttention 是**精确**注意力：两种循环顺序、带不带因果掩码、长度是不是块的倍数，结果都和标准写法相同（1e-15 量级的舍入差）。",
        "N = 1、Nk = 1000 就是 decode 的情形（1 个新 token 看 1000 个缓存的 key）：算法照样对，但只有一个 query 块——"
        "在 GPU 上只能用一个 SM，这是 04 脚本 flash-decoding 要解决的问题。",
        "Nk ≠ N 且 causal 时，第 i 个 query 的位置是 Nk − N + i（前面有 Nk − N 个缓存的 token），掩码按位置算。",
    ])


def demo2_vs_N():
    print_header(f"演示 2：显存读写随序列长度变化（一个头，d = 64，片上 SRAM M = 100 KB = {M_100KB} 个 FP16 元素）")
    d = 64
    Br, Bc = block_sizes(M_100KB, d)
    rng = np.random.default_rng(1)
    rows = []
    for N in (512, 1024, 2048, 4096, 8192):
        Q, K, V = (rng.standard_normal((N, d)) for _ in range(3))
        s = std_io(N, d)
        _, _, io1 = flash_fwd(Q, K, V, Br, Bc, order="fa1")
        _, _, io2 = flash_fwd(Q, K, V, Br, Bc, order="fa2")
        _, _, io3 = flash_fwd(Q, K, V, 128, 64, order="fa2")
        assert io1.total == fa_io_formula(N, d, M_100KB, "fa1") and io2.total == fa_io_formula(N, d, M_100KB, "fa2")
        unit = N * N * d * d / M_100KB
        rows.append([N, fmt_bytes(s.total * 2), fmt_bytes(io1.total * 2), fmt_bytes(io2.total * 2), fmt_bytes(io3.total * 2),
                     f"{s.total / io1.total:.1f}", f"{io1.total / unit:.1f}", f"{io2.total / unit:.1f}",
                     f"{io2.flops / s.flops:.3f}"])
    print(f"论文算法 1 的块大小：Bc = ⌈M/4d⌉ = {Bc}，Br = min(Bc, d) = {Br}；最后一种是 FlashAttention-2 实际常用的 Br = 128、Bc = 64")
    print_table(["N", "标准（最少的读写，FP16）", f"FA 算法 1 顺序（{Br}×{Bc}）", f"FA-2 顺序（{Br}×{Bc}）", "FA-2 顺序（128×64）",
                 "标准 / 算法 1 顺序", "算法 1 ÷ (N²d²/M)", "FA-2 ÷ (N²d²/M)", "FA 的 FLOPs / 标准"], rows)
    print_observations([
        "标准注意力的读写由 N×N 的分数矩阵主导（写 S、读 S、写 P、读 P，4N² 个元素）；FlashAttention 的读写也随 N² 增长，"
        "系数是 d²/M 乘一个常数（\"÷ (N²d²/M)\"两列约 12–25，几乎不随 N 变——和论文定理 2 的 Θ(N²d²M⁻¹) 一致）。",
        "d = 64、M = 100 KB 时，FlashAttention 前向的读写只比\"最少读写的标准注意力\"少 3–4 倍（算法 1 顺序）——不是几十倍。"
        "省得多少取决于 M/d²：SRAM 越大、头维度越小越省。真实的标准实现还有单独的掩码、dropout、softmax 内核（演示 4），差距更大。",
        "两种循环顺序被反复读写的东西不同：算法 1 顺序反复读写 Q、O（次数 = N/Bc），FA-2 顺序反复读 K、V（次数 = N/Br）。"
        "论文的公式让 Br = min(Bc, d) = 64，对 FA-2 顺序不利（读写是算法 1 顺序的 2 倍）；FlashAttention-2 实际用 Br = 128，读写和算法 1 顺序相当，"
        "同时换来了不同 query 块之间没有依赖、可以分给不同 SM 的并行性（讲义 §4.3）。",
        "FLOPs 完全一样：分块不增加运算（causal 时还能跳过整块被掩码的部分，见 03 脚本）。快是因为搬得少。",
    ])


def demo3_vs_M():
    print_header("演示 3：显存读写随片上 SRAM 大小 M 变化（一个头，N = 4096，d = 64）")
    N, d = 4096, 64
    rows = []
    s = std_io(N, d)
    for kb in (16, 32, 64, 100, 192, 256):
        M = kb * 1024 // 2
        Br, Bc = block_sizes(M, d)
        t1, t2 = fa_io_formula(N, d, M, "fa1"), fa_io_formula(N, d, M, "fa2")
        lower = N * N * d * d / M
        rows.append([f"{kb} KB", f"{Br}×{Bc}", fmt_bytes(t1 * 2), fmt_bytes(t2 * 2), fmt_bytes(lower * 2),
                     f"{s.total / t2:.1f}"])
    print_table(["SRAM", "块 Br×Bc", "FA 算法 1 顺序", "FA-2 顺序", "N²d²/M（量级）", "标准 / FA-2"], rows)
    print(f"标准注意力：{fmt_bytes(s.total * 2)}（和 M 无关）")
    print_observations([
        "算法 1 顺序的读写约 ∝ 1/M：M 越大，Bc 越大，Q、O 被反复读写的次数（N/Bc）越少。A100 每个 SM 有 192 KB 的 SRAM（共享内存 + L1），"
        "H100 是 256 KB——这一项硬件参数直接决定 FlashAttention 能省多少。",
        "FA-2 顺序在 M ≥ 32 KB 之后不再下降：论文公式里 Br = min(Bc, d) 封顶在 d = 64，K、V 被读 N/64 次不变。"
        "真实的实现不完全按论文的公式选块，而是按寄存器、共享内存和 Tensor Core 的形状选（FlashAttention-2 在 A100 上常用 128×64 或 128×128）。",
        "论文的命题 3：对所有 M ∈ [d, Nd]，任何精确注意力算法都做不到 o(N²d²/M) 次显存访问——FlashAttention 在这个意义下是最优的（只差常数）。",
    ])


def demo4_gpt2():
    print_header("演示 4：论文图 2 的设置——GPT-2 medium（N = 1024，d = 64，16 个头，批 64），前向 + 反向，A100")
    N, d, heads = 1024, 64, 16 * 64
    rng = np.random.default_rng(2)
    Q, K, V = (rng.standard_normal((N, d)) for _ in range(3))
    dO = rng.standard_normal((N, d))
    _, P, sf = standard_attention(Q, K, V, causal=False)
    _, _, _, sb = standard_attention_bwd(Q, K, V, P, dO)
    rows = []
    res = {}
    for name, (Br, Bc) in (("FlashAttention（论文公式，M = 100 KB）", block_sizes(M_100KB, d)), ("FlashAttention（128×128 的块）", (128, 128))):
        O2, L, ff = flash_fwd(Q, K, V, Br, Bc, causal=False)
        _, _, _, fb = flash_bwd(Q, K, V, O2, dO, L, Br, Bc, causal=False)
        res[name] = (ff, fb)
    peak, bw = 312e12, 1.555e12
    entries = [("标准注意力（我们的计数）", sf, sb)] + [(k, v[0], v[1]) for k, v in res.items()]
    for name, f, b in entries:
        flops = (f.flops + b.flops) * heads
        byts = (f.total + b.total) * heads * 2
        t = max(flops / peak, byts / bw)
        rows.append([name, f"{flops / 1e9:.1f}", fmt_bytes(byts), fmt_time(flops / peak), fmt_time(byts / bw), fmt_time(t)])
    rows.append(["标准注意力（论文测量）", "66.6", "40.3 GB", "", "", "41.7 ms（实测）"])
    rows.append(["FlashAttention（论文测量）", "75.2", "4.4 GB", "", "", "7.3 ms（实测）"])
    print_table(["", "GFLOPs（矩阵乘法）", "显存读写（FP16）", "按 312 TFLOPS", "按 1.555 TB/s", "roofline 下界"], rows)
    print_observations([
        "我们的计数是\"每个算法最少要读写的元素\"：标准注意力 25 GB、FlashAttention 7.7–8.8 GB，约 3 倍。论文测到的是 40.3 GB vs 4.4 GB（约 9 倍）："
        "PyTorch 的标准实现还有单独的掩码、dropout、softmax 内核，每个都把 N×N 的矩阵再读写一遍；论文的 FlashAttention 实现的分块也和我们的不同。",
        "论文的 GFLOPs 数（66.6 / 75.2）比我们按\"批 64 × 16 个头\"算出的小一个数量级，论文没有说明它的 FLOP 计法和统计范围，"
        "我们没能对上绝对值；两者的**比值**接近：FlashAttention 因为反向重算 S，多做约 13%（论文）–17%（我们，只算矩阵乘法）的运算。",
        "按 roofline：标准注意力受限于带宽（读写的时间是计算的 6 倍），FlashAttention 仍受限于带宽（读写约是计算的 1.6–1.8 倍），但近了很多。运算多了一点，时间却少了几倍——"
        "我们的下界 16.3 → 4.9–5.6 ms，论文实测 41.7 → 7.3 ms（5.7 倍）。",
        "块从论文公式的 64×200 换成 128×128，读写还能少约 12%（反向传播里 Q、dO、dQ 被反复读写的次数随 Bc 变大而减少）。",
    ])


if __name__ == "__main__":
    demo1_correctness()
    demo2_vs_N()
    demo3_vs_M()
    demo4_gpt2()