"""
flash_lab.py —— 第 18 讲上半的共享模块：在线 softmax、FlashAttention 的前向与反向、"显存读写"记账、重计算
=====================================================================

【这个文件要回答的问题】
    标准注意力 O = softmax(QKᵀ/√d)·V 要把 N×N 的分数矩阵 S 和概率矩阵 P 写进显存（HBM）、再读出来。
    FlashAttention 把 Q、K、V 切成小块，每次只把几块搬进片上内存（SRAM），在片上算完、
    用"在线 softmax"一边算一边修正，最后只把 N×d 的输出写回去——N×N 的矩阵从来不落地。
    这个文件把两种算法都**真的实现一遍**（用 float64，可以和 PyTorch 逐位比较），
    并且在每一次"从显存读 / 往显存写"时记账，于是可以数出两者各访问了多少个数。

【本文件提供】
    1. 一维 softmax 的三种写法（逐个元素地读写，数访存次数）：
           Memory、softmax_naive、softmax_safe、softmax_online
    2. 分块 softmax 的"充分统计量"（m, l, o）与合并：block_state、merge_states、finalize
    3. 记账器 HBM（读 / 写了多少个元素、做了多少 FLOPs、片上同时放了多少个元素）
    4. 注意力：standard_attention（论文的算法 0）、standard_attention_backward（算法 3）、flash_attention_v1（算法 1：外层循环 K、V）、
           flash_attention_v2（FlashAttention-2 的循环顺序：外层循环 Q）、flash_attention_backward（带重计算的反向）、
           io_formula（三种做法的读写次数的闭式公式，和实测核对）、block_sizes_v1（论文的块大小）
    5. 重计算（激活检查点）：mlp_stack_forward、checkpointed_backward

【记号】和论文一致：N = 序列长度，d = 头维度，M = SRAM 能放的元素个数，B_r / B_c = Q / K 的块大小，
    T_r = N / B_r、T_c = N / B_c 是块数。单个头、形状 (N, d) 的二维张量（多头、多批只是外面再套一层循环）。

【边界】
    - 这里的 "HBM" 和 "SRAM" 是**记账的模型**，不是真的显存：CPU 上没有可编程的片上内存。
      我们只保证"算法访问了哪些数"和真实的 GPU 实现一致，耗时用前置讲 P3 的 roofline 模型估计。
    - 为了能和 PyTorch 逐位比较，全部用 float64；真实的 FlashAttention 用 FP16/BF16 输入、FP32 累加。
    - 要求 N 能被块大小整除（真实实现会处理最后一个不完整的块，这里为了代码清楚从简）。

【GenAI 联系】
    PyTorch 的 F.scaled_dot_product_attention、vLLM、Hugging Face 的 attn_implementation="flash_attention_2"
    背后都是这个算法；它让长上下文（第 19 讲）的训练在显存上成为可能。
"""

import math

import numpy as np
import torch


# =====================================================================
# 1. 一维 softmax 的三种写法：数一数访存次数
# =====================================================================

class Memory:
    """
    一段"慢速内存"：每读一个元素 loads 加 1，每写一个元素 stores 加 1。

    【参数】values：初始内容（列表或一维数组）；dtype：用哪种浮点数存（np.float32 / np.float16 ……）
    【用法】x = Memory([1.0, 2.0]); v = x.load(0); y.store(0, v)
    【为什么要它】Milakov & Gimelshein (2018) 比较 softmax 的几种写法时，数的正是"每个元素访问内存几次"：
        softmax 是访存受限的运算（P3 §5），访存次数基本决定了耗时。
    """

    def __init__(self, values, dtype=np.float32):
        self.data = np.array(values, dtype=dtype)
        self.loads = 0
        self.stores = 0

    def __len__(self):
        return len(self.data)

    def load(self, i):
        self.loads += 1
        return self.data[i]

    def store(self, i, v):
        self.stores += 1
        self.data[i] = v


def softmax_naive(x, y):
    """
    朴素 softmax（Milakov 论文的算法 1）：y_i = e^{x_i} / Σ_j e^{x_j}。
    【思路】第 1 遍读 x 求分母 d；第 2 遍再读 x、写 y。每个元素：读 2 次、写 1 次 = 3 次访存。
    【参数】x：输入 Memory；y：输出 Memory（长度相同）
    【边界】x 里有大数时 e^{x} 会溢出成 inf（float32 在 x > 88.7 时、float16 在 x > 11.1 时），
        结果变成 nan——所以没有哪个深度学习框架真的这样写。
    【Python 小知识】np.errstate(over="ignore", invalid="ignore") 让 numpy 在溢出时安静地给出 inf / nan，
        而不是发出警告（我们正是要展示这个溢出）。
    """
    one = y.data.dtype.type
    with np.errstate(over="ignore", invalid="ignore"):
        d = one(0)
        for j in range(len(x)):
            d = d + np.exp(x.load(j))
        for i in range(len(x)):
            y.store(i, np.exp(x.load(i)) / d)


def softmax_safe(x, y):
    """
    安全 softmax（算法 2）：先减去最大值 m = max_j x_j，y_i = e^{x_i − m} / Σ_j e^{x_j − m}。
    【为什么正确】分子分母同乘 e^{−m}，比值不变（softmax 对平移不变）；减去 m 之后最大的指数是 e^0 = 1，不会溢出。
    【代价】要 3 遍：求 m、求 d、写 y。每个元素：读 3 次、写 1 次 = 4 次访存（比朴素版多 1/3）。
    """
    one = y.data.dtype.type
    m = one(-np.inf)
    for j in range(len(x)):
        m = max(m, x.load(j))
    d = one(0)
    for j in range(len(x)):
        d = d + np.exp(x.load(j) - m)
    for i in range(len(x)):
        y.store(i, np.exp(x.load(i) - m) / d)


def softmax_online(x, y):
    """
    在线 softmax（算法 3，Milakov & Gimelshein 2018）：一遍同时求出最大值 m 和分母 d。

    【思路】读到第 j 个元素时，维护"到目前为止"的最大值 m_j 和分母 d_j = Σ_{k≤j} e^{x_k − m_j}：
            m_j = max(m_{j−1}, x_j)
            d_j = d_{j−1} · e^{m_{j−1} − m_j} + e^{x_j − m_j}
        第二行的第一项是"修正"：最大值变大了，之前加进去的每一项都要乘上 e^{m_{j−1} − m_j}（≤ 1）。
        最后 d_N 恰好等于安全 softmax 的分母。第 2 遍再读 x、写 y。
    【代价】2 遍；每个元素读 2 次、写 1 次 = 3 次访存，和朴素版一样少，又和安全版一样不溢出。
    【GenAI 联系】FlashAttention 正是把这个"边读边修正"推广到"一次读一整块"。
    """
    one = y.data.dtype.type
    m, d = one(-np.inf), one(0)
    for j in range(len(x)):
        xj = x.load(j)
        m_new = max(m, xj)
        d = d * np.exp(m - m_new) + np.exp(xj - m_new)
        m = m_new
    for i in range(len(x)):
        y.store(i, np.exp(x.load(i) - m) / d)


# =====================================================================
# 2. 分块 softmax 的充分统计量：(m, l, o)
# =====================================================================

def _exp_diff(a, b):
    """
    逐元素算 e^{a − b}，并约定 a = −∞ 时结果为 0（即使 b 也是 −∞）。
    【为什么】因果掩码下，某一行在某一块里可能一个能看的键都没有，最大值是 −∞；
        直接算 e^{−∞ − (−∞)} = e^{nan} = nan。约定成 0 的意思是"这一块对这一行没有贡献"。
    """
    out = torch.exp(a - torch.where(torch.isinf(b), torch.zeros_like(b), b))
    return torch.where(torch.isinf(a), torch.zeros_like(out), out)


def block_state(S, V):
    """
    一块分数 S（形状 (r, c)）和对应的值 V（形状 (c, d)）的"充分统计量"：
        m = 每行的最大值（r,）
        l = Σ_j e^{S_ij − m_i}（r,）              —— 这一块对分母的贡献（以 m 为基准）
        o = Σ_j e^{S_ij − m_i} V_j（r, d）         —— 这一块对"未归一化的输出"的贡献
    【边界】S 里可以有 −∞（被掩码的位置）；整行都是 −∞ 时 m = −∞、l = 0、o = 0。
    """
    m = S.max(dim=1).values
    P = _exp_diff(S, m.unsqueeze(1))
    return m, P.sum(dim=1), P @ V


def merge_states(a, b):
    """
    合并两块的统计量（softmax 不能直接分块求和，但 (m, l, o) 可以）：
        m = max(m_a, m_b)
        l = l_a e^{m_a − m} + l_b e^{m_b − m}
        o = o_a e^{m_a − m} + o_b e^{m_b − m}
    【性质】这个运算满足结合律和交换律（01 号脚本验证）：不管按什么顺序、怎样分组合并，结果都一样。
        这正是"分块"需要的性质——矩阵乘法里的加法天生满足它，softmax 加上 (m, l) 之后也满足了。
    """
    ma, la, oa = a
    mb, lb, ob = b
    m = torch.maximum(ma, mb)
    ea, eb = _exp_diff(ma, m), _exp_diff(mb, m)
    return m, la * ea + lb * eb, oa * ea.unsqueeze(1) + ob * eb.unsqueeze(1)


def finalize(state):
    """由统计量得到真正的输出 o / l，以及每行的 logsumexp L = m + log l（反向传播要用）。"""
    m, l, o = state
    return o / l.unsqueeze(1), m + torch.log(l)


# =====================================================================
# 3. 记账器
# =====================================================================

class HBM:
    """
    记录一个算法和"显存"之间的数据搬运。

    【属性】
        reads、writes : 从 HBM 读出 / 写入的元素个数
        flops         : 矩阵乘法的浮点运算次数（每次乘加算 2 次；softmax 里的 exp 等不计——它们相对很少）
        sram_peak     : 片上同时存放的元素个数的最大值（由算法自己在每个内层步骤报告）
    【用法】hbm.read(t) 返回 t 本身并记账；hbm.write(t) 同理；hbm.mm(a, b) 返回 a @ b 并记 FLOPs
    【边界】只是计数器，不会真的把数据搬到别处。
    """

    def __init__(self):
        self.reads = 0
        self.writes = 0
        self.flops = 0
        self.sram_peak = 0

    def read(self, t):
        self.reads += t.numel()
        return t

    def write(self, t):
        self.writes += t.numel()
        return t

    def mm(self, a, b):
        self.flops += 2 * a.shape[0] * a.shape[1] * b.shape[1]
        return a @ b

    def on_chip(self, *tensors):
        self.sram_peak = max(self.sram_peak, sum(t.numel() for t in tensors))

    @property
    def accesses(self):
        return self.reads + self.writes


def _causal_block(S, row0, col0):
    """给一块分数加因果掩码：全局行号 row0+i 只能看全局列号 ≤ row0+i 的键。"""
    r = torch.arange(row0, row0 + S.shape[0]).unsqueeze(1)
    c = torch.arange(col0, col0 + S.shape[1]).unsqueeze(0)
    return S.masked_fill(c > r, float("-inf"))


# =====================================================================
# 4. 注意力：标准做法、FlashAttention（两种循环顺序）、带重计算的反向
# =====================================================================

def standard_attention(Q, K, V, causal=False, hbm=None):
    """
    标准注意力（FlashAttention 论文的算法 0）：
        1. 从 HBM 读 Q、K，算 S = QKᵀ/√d，把 S 写回 HBM；
        2. 读 S，算 P = softmax(S)，把 P 写回 HBM；
        3. 读 P、V，算 O = PV，把 O 写回 HBM。
    【参数】Q、K、V：(N, d)；causal：是否加因果掩码；hbm：记账器（None 表示不记）
    【返回】(O, P)；P 是 N×N 的概率矩阵——反向传播要用它，所以标准实现会一直存着它。
    【访存】读 3Nd + 2N²、写 2N² + Nd，共 4N² + 4Nd 个元素（Θ(Nd + N²)，论文定理 2）。
    """
    hbm = hbm or HBM()
    N, d = Q.shape
    S = hbm.mm(hbm.read(Q), hbm.read(K).T) / math.sqrt(d)
    if causal:
        S = _causal_block(S, 0, 0)
    hbm.write(S)
    P = torch.softmax(hbm.read(S), dim=-1)
    hbm.write(P)
    O = hbm.mm(hbm.read(P), hbm.read(V))
    hbm.write(O)
    return O, P


def standard_attention_backward(Q, K, V, P, dO, hbm=None):
    """
    标准注意力的反向（论文的算法 3）：用前向存下的 N×N 矩阵 P。
        读 P、dO，写 dV = Pᵀ dO；读 dO、V，写 dP = dO Vᵀ（N×N）；读 P、dP，写 dS = P ∘ (dP − rowsum(P∘dP))（N×N）；
        读 dS、K，写 dQ = dS K/√d；读 dS、Q，写 dK = dSᵀ Q/√d。
    【返回】(dQ, dK, dV)
    【访存】读 5N² + 5Nd、写 2N² + 3Nd，共 7N² + 8Nd 个元素。
    """
    hbm = hbm or HBM()
    d = Q.shape[1]
    scale = 1.0 / math.sqrt(d)
    dV = hbm.write(hbm.mm(hbm.read(P).T, hbm.read(dO)))
    dP = hbm.write(hbm.mm(hbm.read(dO), hbm.read(V).T))
    Pr, dPr = hbm.read(P), hbm.read(dP)
    dS = hbm.write(Pr * (dPr - (Pr * dPr).sum(dim=1, keepdim=True)))
    dQ = hbm.write(hbm.mm(hbm.read(dS), hbm.read(K)) * scale)
    dK = hbm.write(hbm.mm(hbm.read(dS).T, hbm.read(Q)) * scale)
    return dQ, dK, dV


def block_sizes_v1(M, d):
    """论文算法 1 第 1 行的块大小：B_c = ⌈M / 4d⌉，B_r = min(⌈M / 4d⌉, d)。"""
    bc = math.ceil(M / (4 * d))
    return min(bc, d), bc


def flash_attention_v1(Q, K, V, M=None, block=None, causal=False, hbm=None):
    """
    FlashAttention（论文算法 1）：外层循环遍历 K、V 的块，内层循环遍历 Q 的块。

    【思路】
        O、l、m 放在 HBM 里，初值 0、0、−∞。对每一块 K_j、V_j（搬进 SRAM，只搬一次）：
            对每一块 Q_i：把 Q_i、O_i、l_i、m_i 搬进 SRAM，
                S_ij = Q_i K_jᵀ/√d；m̃ = 行最大值；P̃ = e^{S_ij − m̃}；l̃ = 行和
                m_new = max(m_i, m̃)；l_new = e^{m_i − m_new} l_i + e^{m̃ − m_new} l̃
                O_i ← (l_i e^{m_i − m_new} O_i + e^{m̃ − m_new} P̃ V_j) / l_new       （先撤销旧的归一化，再加新的一块）
                把 O_i、l_new、m_new 写回 HBM。
        S_ij、P̃ 只在片上存在，用完就丢——N×N 的矩阵从来不写进 HBM。
    【参数】
        M      : SRAM 能放的元素个数；给了 M 就按论文的公式取块大小
        block  : 也可以直接给 (B_r, B_c)
        causal : 因果掩码；完全在对角线上方的块直接跳过（不读、不算）
    【返回】(O, L)：O 是 (N, d) 的输出，L = m + log l 是每行的 logsumexp（反向要用）
    【访存】K、V 各读一次（2Nd）；每一块 K_j 都要把全部的 Q、O 读一遍、O 写一遍（加上 l、m）：
        T_c(3Nd + 4N) + 2Nd + 初始化 (Nd + 2N)。B_c = M/4d 时 T_c = 4Nd/M，主项约 12N²d²/M——论文定理 2 的 Θ(N²d²/M)。
    """
    hbm = hbm or HBM()
    N, d = Q.shape
    br, bc = block if block is not None else block_sizes_v1(M, d)
    if N % br or N % bc:
        raise ValueError(f"N={N} 必须能被块大小 B_r={br}、B_c={bc} 整除")
    scale = 1.0 / math.sqrt(d)
    O = hbm.write(torch.zeros(N, d, dtype=Q.dtype))
    l = hbm.write(torch.zeros(N, dtype=Q.dtype))
    m = hbm.write(torch.full((N,), float("-inf"), dtype=Q.dtype))
    for j0 in range(0, N, bc):
        Kj, Vj = hbm.read(K[j0:j0 + bc]), hbm.read(V[j0:j0 + bc])
        for i0 in range(0, N, br):
            if causal and j0 > i0 + br - 1:        # 这一块的所有键都在所有查询的"未来"：整块跳过
                continue
            Qi, Oi = hbm.read(Q[i0:i0 + br]), hbm.read(O[i0:i0 + br])
            li, mi = hbm.read(l[i0:i0 + br]), hbm.read(m[i0:i0 + br])
            S = hbm.mm(Qi, Kj.T) * scale
            if causal:
                S = _causal_block(S, i0, j0)
            mt, lt, PV = block_state_with_flops(S, Vj, hbm)
            m_new = torch.maximum(mi, mt)
            a, b = _exp_diff(mi, m_new), _exp_diff(mt, m_new)
            l_new = a * li + b * lt
            O[i0:i0 + br] = hbm.write((li * a).unsqueeze(1) * Oi / l_new.unsqueeze(1)
                                      + b.unsqueeze(1) * PV / l_new.unsqueeze(1))
            l[i0:i0 + br], m[i0:i0 + br] = hbm.write(l_new), hbm.write(m_new)
            hbm.on_chip(Kj, Vj, Qi, Oi, S, li, mi)
    return O, m + torch.log(l)


def block_state_with_flops(S, V, hbm):
    """和 block_state 相同，只是把 P̃V 的矩阵乘法记进 hbm.flops。"""
    m = S.max(dim=1).values
    P = _exp_diff(S, m.unsqueeze(1))
    return m, P.sum(dim=1), hbm.mm(P, V)


def flash_attention_v2(Q, K, V, block=64, causal=False, hbm=None):
    """
    FlashAttention-2（Dao, 2023）的循环顺序：外层循环 Q 的块，内层循环 K、V 的块。

    【和 v1 的区别】
        - Q_i 和它的输出 O_i、统计量 m_i、l_i 一直留在片上，内层循环结束才写回一次；
          代价是 K、V 要被每一块 Q 各读一遍。
        - 不在每一步都除以 l（只在最后除一次），少做一些非矩阵乘法的运算（GPU 上它们比矩阵乘法慢 16 倍左右）。
        - 每一块 Q 的计算互不依赖，可以分给不同的 SM 并行——长序列、小批量时更能占满 GPU。
    【参数】block：B_r = B_c 的块大小（整数）或 (B_r, B_c)
    【返回】(O, L)，同 v1
    【访存】Q 读一次、O 写一次、L 写一次，K、V 各读 T_r 次：Nd + 2Nd·T_r + Nd + N。
    """
    hbm = hbm or HBM()
    N, d = Q.shape
    br, bc = (block, block) if isinstance(block, int) else block
    if N % br or N % bc:
        raise ValueError(f"N={N} 必须能被块大小 {br}、{bc} 整除")
    scale = 1.0 / math.sqrt(d)
    O = torch.empty(N, d, dtype=Q.dtype)
    L = torch.empty(N, dtype=Q.dtype)
    for i0 in range(0, N, br):
        Qi = hbm.read(Q[i0:i0 + br])
        state = (torch.full((br,), float("-inf"), dtype=Q.dtype), torch.zeros(br, dtype=Q.dtype),
                 torch.zeros(br, d, dtype=Q.dtype))
        for j0 in range(0, N, bc):
            if causal and j0 > i0 + br - 1:
                break                                # 后面的块全在未来：内层循环提前结束
            Kj, Vj = hbm.read(K[j0:j0 + bc]), hbm.read(V[j0:j0 + bc])
            S = hbm.mm(Qi, Kj.T) * scale
            if causal:
                S = _causal_block(S, i0, j0)
            state = merge_states(state, block_state_with_flops(S, Vj, hbm))
            hbm.on_chip(Qi, Kj, Vj, S, state[2], state[0], state[1])
        Oi, Li = finalize(state)
        O[i0:i0 + br], L[i0:i0 + br] = hbm.write(Oi), hbm.write(Li)
    return O, L


def flash_attention_backward(Q, K, V, O, dO, L, block=64, causal=False, hbm=None):
    """
    FlashAttention 的反向（论文算法 4 / FlashAttention-2 的算法 2）：不存 P，用 L 把它**重新算出来**。

    【公式】记 P = softmax(S)，前向 O = PV。给定 dO = ∂ℓ/∂O：
        dV = Pᵀ dO
        dP = dO Vᵀ
        dS = P ∘ (dP − D)，其中 D_i = Σ_j P_ij dP_ij = Σ_k dO_ik O_ik（softmax 的雅可比；D 只要 O 和 dO 就能算，是 N 个数）
        dQ = dS K / √d，dK = dSᵀ Q / √d
    【重计算】P_ij 的任意一块都可以由 Q_i、K_j 和这一行的 L_i 现算：P_ij = e^{Q_i K_jᵀ/√d − L_i}——
        多做一次 QKᵀ 的矩阵乘法，换来不必存 N×N 的 P。
    【参数】Q、K、V、O、dO：(N, d)；L：前向返回的 logsumexp (N,)；block、causal 同 v2
    【返回】(dQ, dK, dV)
    【循环顺序】外层 K、V 的块（dK_j、dV_j 在片上累加，最后写一次），内层 Q 的块（dQ_i 每次读出、加上、写回）。
    """
    hbm = hbm or HBM()
    N, d = Q.shape
    br, bc = (block, block) if isinstance(block, int) else block
    scale = 1.0 / math.sqrt(d)
    D = hbm.write((hbm.read(dO) * hbm.read(O)).sum(dim=1))
    dQ = hbm.write(torch.zeros(N, d, dtype=Q.dtype))
    dK = torch.empty(N, d, dtype=Q.dtype)
    dV = torch.empty(N, d, dtype=Q.dtype)
    for j0 in range(0, N, bc):
        Kj, Vj = hbm.read(K[j0:j0 + bc]), hbm.read(V[j0:j0 + bc])
        dKj = torch.zeros(bc, d, dtype=Q.dtype)
        dVj = torch.zeros(bc, d, dtype=Q.dtype)
        for i0 in range(0, N, br):
            if causal and j0 > i0 + br - 1:
                continue
            Qi, dOi = hbm.read(Q[i0:i0 + br]), hbm.read(dO[i0:i0 + br])
            Li, Di, dQi = hbm.read(L[i0:i0 + br]), hbm.read(D[i0:i0 + br]), hbm.read(dQ[i0:i0 + br])
            S = hbm.mm(Qi, Kj.T) * scale                     # 重计算：再做一次 QKᵀ
            if causal:
                S = _causal_block(S, i0, j0)
            P = _exp_diff(S, Li.unsqueeze(1))
            dVj += hbm.mm(P.T, dOi)
            dP = hbm.mm(dOi, Vj.T)
            dS = P * (dP - Di.unsqueeze(1))
            dQ[i0:i0 + br] = hbm.write(dQi + hbm.mm(dS, Kj) * scale)
            dKj += hbm.mm(dS.T, Qi) * scale
            hbm.on_chip(Kj, Vj, dKj, dVj, Qi, dOi, dQi, S, Li, Di)
        dK[j0:j0 + bc], dV[j0:j0 + bc] = hbm.write(dKj), hbm.write(dVj)
    return dQ, dK, dV


def io_formula(kind, N, d, br=None, bc=None):
    """
    三种做法（无因果掩码）读写 HBM 的元素个数的闭式公式，和上面代码的实际计数逐项对应：
        standard : 4N² + 4Nd
        v1       : (Nd + 2N) + 2Nd + T_c·(3Nd + 4N)              T_c = N / B_c
        v2       : Nd + 2Nd·T_r + Nd + N                          T_r = N / B_r
    【返回】整数
    """
    if kind == "standard":
        return 4 * N * N + 4 * N * d
    if kind == "v1":
        return (N * d + 2 * N) + 2 * N * d + (N // bc) * (3 * N * d + 4 * N)
    if kind == "v2":
        return N * d + 2 * N * d * (N // br) + N * d + N
    raise ValueError(kind)


# =====================================================================
# 5. 重计算（激活检查点）：一个 L 层的 MLP
# =====================================================================

def mlp_stack_forward(x, weights, keep="all"):
    """
    L 层 MLP 的前向：h_{ℓ} = tanh(h_{ℓ−1} W_ℓ)，ℓ = 1..L。

    【参数】
        x       : (B, d)；weights：L 个 (d, d) 的权重（requires_grad=True）
        keep    : "all" —— 普通做法，autograd 自动保存每一层的输入和输出；
                  整数 k —— 激活检查点：只在每 k 层的边界保存 h（不建计算图），其余层的中间结果用完就丢
    【返回】(输出 h_L, 存下来的检查点列表)；keep="all" 时检查点列表为空（中间结果都在计算图里）
    【边界】keep 为整数时，输出不带计算图，必须用 checkpointed_backward 求梯度。
    """
    if keep == "all":
        h = x
        for W in weights:
            h = torch.tanh(h @ W)
        return h, []
    checkpoints = [x]
    h = x
    with torch.no_grad():                            # 不建计算图：中间结果算完即可释放
        for idx, W in enumerate(weights, start=1):
            h = torch.tanh(h @ W)
            if idx % keep == 0 and idx < len(weights):
                checkpoints.append(h)
    return h, checkpoints


def checkpointed_backward(weights, checkpoints, k, grad_out):
    """
    激活检查点的反向：从最后一段开始，每段从它的检查点**重新做一遍前向**（这次建计算图），
    再把上一段传回来的梯度往回传。

    【参数】weights：L 个权重；checkpoints：mlp_stack_forward(keep=k) 返回的检查点；k：每段的层数；
        grad_out：∂ℓ/∂h_L
    【返回】(每个权重的梯度列表, 反向时同时存在的激活值元素个数的最大值, 重算的层数)
    【直觉】存 L/k 个检查点 + 一段之内的 k 层中间结果；k ≈ √L 时显存 ∝ √L（Chen et al., 2016），
        代价是几乎每一层的前向都多做了一遍。
    """
    L = len(weights)
    grads = [None] * L
    g = grad_out
    peak, recomputed = 0, 0
    stored_ckpt = sum(c.numel() for c in checkpoints)
    seg_starts = list(range(0, L, k))
    for s in reversed(seg_starts):
        layer_ids = range(s, min(s + k, L))
        h0 = checkpoints[s // k].detach().requires_grad_(True)
        h, saved = h0, 0
        for i in layer_ids:                          # 重计算这一段（建计算图）
            h = torch.tanh(h @ weights[i])
            saved += h.numel()                       # 每层多存一个 tanh 的输出（它同时是下一层 matmul 的输入）
            recomputed += 1
        peak = max(peak, stored_ckpt + saved)
        seg_w = [weights[i] for i in layer_ids]
        outs = torch.autograd.grad(h, [h0] + seg_w, grad_outputs=g)
        g = outs[0]
        for i, gw in zip(layer_ids, outs[1:]):
            grads[i] = gw
    return grads, peak, recomputed