"""
attention_kernels.py —— 注意力的几种算法实现（PyTorch 张量版）与它们的读写量公式（第 5 讲共用）
======================================================================================

【这个文件要回答的问题】
    FlashAttention 不是一个"近似"，它和标准注意力在数学上**完全相同**，只是换了计算顺序：
        标准：先把整个 S = QKᵀ 算出来写进显存，再 softmax，再乘 V；
        Flash：把 K、V 切成块，一块一块地流过片上内存，用"在线 softmax"边走边更新答案。
    这个文件把几种顺序都写成**普通的 PyTorch 代码**（在 CPU 上跑，用来验证正确性、接进模型），
    第 5 讲的 03 号脚本再把同样的算法写成 gpusim 上的 CUDA 风格内核，数清楚读写了多少字节。

【内容】
    1. 在线 softmax 的"状态"与合并规则：SoftmaxState、merge_states
    2. 标准注意力 standard_attention（会返回 logsumexp，方便和 Flash 版对照）
    3. FlashAttention 前向 flash_forward（FA2 的循环顺序：外层 Q 块，内层 K/V 块）
    4. FlashAttention 反向 flash_backward（重计算 P，不存 N×N 矩阵）
    5. Flash-Decoding：flash_decode（沿 K/V 方向切成几段并行，最后合并）
    6. 读写量与 FLOPs 的闭式公式：io_standard、io_flash、attention_flops

【记号】（和 FlashAttention 论文一致）
    N  序列长度（query 数 = key 数，除非特别说明）     d  每个头的维度
    Br 一个 Q 块的行数                                   Bc 一个 K/V 块的行数
    S = QKᵀ·scale（N×N 分数），P = softmax(S)（按行），O = PV
    m  每行的最大值；ℓ（代码里写 l）每行 exp(S - m) 的和；L = m + log ℓ 是 logsumexp

【Python 小知识】
    - `x.amax(dim=-1, keepdim=True)`：沿最后一维取最大值，keepdim=True 保留这一维（长度变 1），方便广播。
    - `torch.where(cond, a, b)`：逐元素"cond 为真取 a，否则取 b"。
    - 张量切片 `k[:, j0:j1]` 不复制数据，只是"看"原张量的一部分（view）。
"""

import math
from dataclasses import dataclass
from typing import List, Optional, Tuple

import torch

NEG_INF = float("-inf")


# ---------------------------------------------------------------------------
# 1. 在线 softmax 的状态与合并
# ---------------------------------------------------------------------------
@dataclass
class SoftmaxState:
    """一段分数 s_1..s_k 的"摘要"：只要两个数就能代表它们对 softmax 分母的贡献。

        m = max_i s_i              （这一段的最大值）
        l = Σ_i exp(s_i - m)        （以 m 为基准的指数和）

    整段的 softmax 分母 Σ exp(s_i) = exp(m) · l。存 (m, l) 而不是直接存 Σ exp(s_i)，是为了不溢出。
    做注意力时再加一个 o = Σ_i exp(s_i - m) · v_i（未归一化的加权和），最后输出 o / l。

    【边界】空的一段：m = -inf, l = 0, o = 0。和任何状态合并都不改变对方（它是合并运算的"单位元"）。
    """
    m: torch.Tensor
    l: torch.Tensor
    o: Optional[torch.Tensor] = None


def merge_states(a: SoftmaxState, b: SoftmaxState) -> SoftmaxState:
    """把两段的摘要合并成一段：(m, l, o) = a ⊕ b。

    【思路】新的最大值 m = max(m_a, m_b)；两段原来的基准不同，各自乘一个"换基准"的系数：
                l = l_a · exp(m_a - m) + l_b · exp(m_b - m)
                o = o_a · exp(m_a - m) + o_b · exp(m_b - m)
           这个运算满足**结合律**和**交换律**——所以可以像第 4 讲的归约一样任意分组、并行计算，
           这就是 FlashAttention（按块顺序合并）和 Flash-Decoding（并行算各段再合并）的数学基础。
    【边界】m_a = m_b = -inf（两段都空）时 exp(-inf - (-inf)) = exp(nan)，要特判，结果仍是空。
    """
    m = torch.maximum(a.m, b.m)
    # 两段都是 -inf 时把基准设成 0，避免 -inf - (-inf) = nan；此时 l_a = l_b = 0，结果不受影响
    safe = torch.where(torch.isinf(m), torch.zeros_like(m), m)
    ca, cb = torch.exp(a.m - safe), torch.exp(b.m - safe)
    l = a.l * ca + b.l * cb
    o = None
    if a.o is not None and b.o is not None:
        o = a.o * ca[..., None] + b.o * cb[..., None]     # o 比 m、l 多最后一维 d，系数要补一维才能广播
    return SoftmaxState(m, l, o)


def state_of_scores(s: torch.Tensor, v: Optional[torch.Tensor] = None) -> SoftmaxState:
    """直接从一段分数（最后一维）算出它的摘要；给了 v 就同时算未归一化的输出 o = exp(s - m) @ v。"""
    m = s.amax(dim=-1)
    safe = torch.where(torch.isinf(m), torch.zeros_like(m), m)
    p = torch.exp(s - safe[..., None])
    return SoftmaxState(m, p.sum(-1), None if v is None else p @ v)


# ---------------------------------------------------------------------------
# 2. 标准注意力（参考实现）
# ---------------------------------------------------------------------------
def _scale(d: int, scale: Optional[float]) -> float:
    return 1.0 / math.sqrt(d) if scale is None else scale


def _mask(q_start: int, tq: int, tk: int, k_start: int = 0) -> torch.Tensor:
    """因果掩码 (tq, tk)：key 位置 > query 位置的地方为 True（屏蔽）。"""
    qpos = torch.arange(q_start, q_start + tq)[:, None]
    kpos = torch.arange(k_start, k_start + tk)[None, :]
    return kpos > qpos


def standard_attention(q, k, v, causal: bool = False, q_start: int = 0, scale: Optional[float] = None):
    """标准注意力：S → P → O，三个中间结果都是完整的矩阵。

    【参数】q: (..., Nq, d)；k, v: (..., Nk, d)；causal=True 时第 i 个 query（绝对位置 q_start + i）
            只看位置 ≤ 它的 key；scale 默认 1/√d。
    【返回】(O, L)：O 形状 (..., Nq, d)；L = logsumexp(S) 形状 (..., Nq)。
    """
    sc = _scale(q.shape[-1], scale)
    s = (q @ k.transpose(-1, -2)) * sc
    if causal:
        s = s.masked_fill(_mask(q_start, q.shape[-2], k.shape[-2]), NEG_INF)
    lse = torch.logsumexp(s, dim=-1)
    p = torch.exp(s - lse[..., None])
    return p @ v, lse


# ---------------------------------------------------------------------------
# 3. FlashAttention 前向（FA2 的循环顺序）
# ---------------------------------------------------------------------------
def flash_forward(q, k, v, Br: int = 64, Bc: int = 64, causal: bool = False, q_start: int = 0,
                  scale: Optional[float] = None, stats: Optional[dict] = None):
    """分块 + 在线 softmax 的注意力前向。数学上与 standard_attention 相同（只差浮点舍入）。

    【思路】（FlashAttention-2 论文的算法 1，去掉了 GPU 细节）
        for 每个 Q 块 Q_i（Br 行）：               ← GPU 上：每个线程块负责一个 Q 块，彼此独立、并行
            m = -inf, l = 0, acc = 0               ← 这三个量一直留在片上（寄存器）
            for 每个 K/V 块 K_j, V_j（Bc 行）：     ← GPU 上：从显存读进共享内存
                S_ij = Q_i K_jᵀ · scale            （Br × Bc，只在片上存在）
                m_new = max(m, rowmax(S_ij))
                P̃ = exp(S_ij - m_new)             （未归一化）
                l   = l · exp(m - m_new) + rowsum(P̃)
                acc = acc · exp(m - m_new) + P̃ V_j
                m = m_new
            O_i = acc / l；   L_i = m + log l     ← 只在最后除一次（FA2 相对 FA1 的改进之一）
    【参数】
        Br, Bc : 块大小；不要求整除 N（最后一块可以短一些）
        causal : 因果掩码。整块都在"未来"的 K/V 块直接跳过（FA2 的另一个省时技巧）
        q_start: 第一个 query 的绝对位置（解码时 query 在序列的末尾）
        stats  : 传入一个 dict 时，记录 "tiles"（真正计算的块数）、"skipped"（因果跳过的块数）、
                 "rescales"（m 变大、需要把旧的 acc 缩放的行次）
    【返回】(O, L)，与 standard_attention 相同。
    【边界】因果模式下要求每个 query 至少能看到一个 key（q_start + i ≥ 0 总成立，所以没问题）。
    """
    sc = _scale(q.shape[-1], scale)
    Nq, Nk = q.shape[-2], k.shape[-2]
    O = torch.empty_like(q)
    L = torch.empty(q.shape[:-1], dtype=q.dtype)
    for i0 in range(0, Nq, Br):
        i1 = min(i0 + Br, Nq)
        qi = q[..., i0:i1, :]
        m = torch.full(qi.shape[:-1], NEG_INF, dtype=q.dtype)
        l = torch.zeros(qi.shape[:-1], dtype=q.dtype)
        acc = torch.zeros_like(qi)
        for j0 in range(0, Nk, Bc):
            j1 = min(j0 + Bc, Nk)
            if causal and j0 > q_start + i1 - 1:        # 这一块的第一个 key 比块里最后一个 query 还靠后
                if stats is not None:
                    stats["skipped"] = stats.get("skipped", 0) + 1
                continue
            s = (qi @ k[..., j0:j1, :].transpose(-1, -2)) * sc
            if causal:
                s = s.masked_fill(_mask(q_start + i0, i1 - i0, j1 - j0, k_start=j0), NEG_INF)
            m_new = torch.maximum(m, s.amax(-1))
            alpha = torch.exp(m - m_new)                # 旧基准换到新基准的系数（第一块时 exp(-inf)=0）
            p = torch.exp(s - m_new[..., None])
            l = l * alpha + p.sum(-1)
            acc = acc * alpha[..., None] + p @ v[..., j0:j1, :]
            if stats is not None:
                stats["tiles"] = stats.get("tiles", 0) + 1
                stats["rescales"] = stats.get("rescales", 0) + int(((m_new > m) & torch.isfinite(m)).sum())
            m = m_new
        O[..., i0:i1, :] = acc / l[..., None]
        L[..., i0:i1] = m + torch.log(l)
    return O, L


# ---------------------------------------------------------------------------
# 4. FlashAttention 反向（重计算）
# ---------------------------------------------------------------------------
def flash_backward(q, k, v, O, L, dO, Br: int = 64, Bc: int = 64, causal: bool = False,
                   scale: Optional[float] = None, stats: Optional[dict] = None):
    """FlashAttention 反向：前向只留下 O 和 L（每行一个数），P 在反向里**重新算一遍**。

    【思路】（FlashAttention-2 论文的算法 2）
        D = rowsum(dO ∘ O)                         ← 每行一个数，先算好
        for 每个 K/V 块 j（外层；GPU 上每个线程块负责一个 K/V 块）：
            dK_j = dV_j = 0
            for 每个 Q 块 i（内层）：
                S = Q_i K_jᵀ · scale;  P = exp(S - L_i)        ← 重计算：用存下的 L 直接得到归一化的 P
                dV_j += Pᵀ dO_i
                dP = dO_i V_jᵀ
                dS = P ∘ (dP - D_i)                            ← softmax 的反向，D 让它不需要整行
                dQ_i += dS K_j · scale                         ← 多个线程块会同时加同一个 dQ_i → GPU 上用原子加
                dK_j += dSᵀ Q_i · scale
    【为什么 dS = P ∘ (dP - D)】softmax 的雅可比是 diag(p) - p pᵀ，所以 dS_ij = P_ij (dP_ij - Σ_k P_ik dP_ik)；
        而 Σ_k P_ik dP_ik = Σ_k P_ik (dO_i · V_k) = dO_i · (Σ_k P_ik V_k) = dO_i · O_i = D_i。
        于是"整行的和"可以用 dO 和 O 提前算好——这正是能分块的关键。
    【返回】(dQ, dK, dV)。stats 里记录 "dq_atomic_adds"（GPU 上对 dQ 的原子加元素数）。
    """
    sc = _scale(q.shape[-1], scale)
    Nq, Nk = q.shape[-2], k.shape[-2]
    D = (dO * O).sum(-1)
    dQ = torch.zeros_like(q)
    dK = torch.zeros_like(k)
    dV = torch.zeros_like(v)
    for j0 in range(0, Nk, Bc):
        j1 = min(j0 + Bc, Nk)
        kj, vj = k[..., j0:j1, :], v[..., j0:j1, :]
        dKj = torch.zeros_like(kj)
        dVj = torch.zeros_like(vj)
        for i0 in range(0, Nq, Br):
            i1 = min(i0 + Br, Nq)
            if causal and j0 > i1 - 1:
                continue
            s = (q[..., i0:i1, :] @ kj.transpose(-1, -2)) * sc
            if causal:
                s = s.masked_fill(_mask(i0, i1 - i0, j1 - j0, k_start=j0), NEG_INF)
            p = torch.exp(s - L[..., i0:i1, None])
            dVj += p.transpose(-1, -2) @ dO[..., i0:i1, :]
            dp = dO[..., i0:i1, :] @ vj.transpose(-1, -2)
            ds = p * (dp - D[..., i0:i1, None])
            dQ[..., i0:i1, :] += (ds @ kj) * sc
            dKj += (ds.transpose(-1, -2) @ q[..., i0:i1, :]) * sc
            if stats is not None:
                stats["dq_atomic_adds"] = stats.get("dq_atomic_adds", 0) + (i1 - i0) * q.shape[-1]
        dK[..., j0:j1, :] = dKj
        dV[..., j0:j1, :] = dVj
    return dQ, dK, dV


# ---------------------------------------------------------------------------
# 5. Flash-Decoding
# ---------------------------------------------------------------------------
def flash_decode(q, k, v, n_split: int, scale: Optional[float] = None) -> Tuple[torch.Tensor, List[SoftmaxState]]:
    """解码阶段的注意力：query 很少（通常每个序列 1 个），K/V 很长。

    【思路】把 K/V 沿序列方向切成 n_split 段，每段独立算出 (m, l, o) 摘要（GPU 上每段一个线程块，
           于是同一个头也能用上很多个 SM），最后用 merge_states 合并（一个很小的归约内核）。
           解码时 query 是序列里最新的位置，能看到全部 key，不需要因果掩码。
    【参数】q: (..., Nq, d)（Nq 通常是 1）；k, v: (..., Nk, d)；n_split ≥ 1。
    【返回】(O, 各段的摘要列表)。
    【边界】n_split 大于 Nk 时有的段是空的：摘要为 (m=-inf, l=0, o=0)，合并时自动被忽略。
    """
    if n_split < 1:
        raise ValueError("n_split 至少为 1")
    sc = _scale(q.shape[-1], scale)
    Nk = k.shape[-2]
    bounds = [round(Nk * s / n_split) for s in range(n_split + 1)]
    parts = []
    for s in range(n_split):
        j0, j1 = bounds[s], bounds[s + 1]
        if j1 == j0:
            shape = q.shape[:-1]
            parts.append(SoftmaxState(torch.full(shape, NEG_INF, dtype=q.dtype), torch.zeros(shape, dtype=q.dtype),
                                      torch.zeros_like(q)))
            continue
        scores = (q @ k[..., j0:j1, :].transpose(-1, -2)) * sc
        parts.append(state_of_scores(scores, v[..., j0:j1, :]))
    total = parts[0]
    for p in parts[1:]:
        total = merge_states(total, p)
    return total.o / total.l[..., None], parts


# ---------------------------------------------------------------------------
# 6. 闭式公式：FLOPs 与显存读写量（单位：元素个数）
# ---------------------------------------------------------------------------
def attention_flops(N: int, d: int, causal: bool = False, backward: bool = False, recompute: bool = True) -> float:
    """一个头的注意力里矩阵乘法的 FLOPs（乘加算 2 次；softmax 的 exp 等不计入）。

        前向：S = QKᵀ（2N²d）+ O = PV（2N²d）= 4N²d
        反向：dV = PᵀdO、dP = dO Vᵀ、dQ = dS K、dK = dSᵀ Q 各 2N²d = 8N²d；重计算 S 再加 2N²d
        因果：只算一半（按大 N 的近似）
    """
    f = 4.0 * N * N * d
    if backward:
        f = (10.0 if recompute else 8.0) * N * N * d
    return f / 2 if causal else f


def io_standard(N: int, d: int) -> float:
    """标准注意力前向（3 个内核）在"理想缓存"下的显存读写元素数：
        内核 1：读 Q、K（2Nd），写 S（N²）
        内核 2：读 S，写 P（2N²）
        内核 3：读 P、V（N² + Nd），写 O（Nd）
        合计 4Nd + 4N²（不论块多大，都要把 N×N 的 S、P 各写一次、读一次）
    """
    return 4.0 * N * d + 4.0 * N * N


def io_flash(N: int, d: int, Br: int, kv_cached: bool = False) -> float:
    """FlashAttention 前向（FA2 循环顺序）的显存读写元素数：
        读 Q 一次（Nd），写 O 一次（Nd），写 L（N）；
        每个 Q 块都要把整个 K、V 读一遍：2Nd × ⌈N/Br⌉ 次。
        kv_cached=True 表示 K、V 被 L2 缓存命中、实际只从显存读一次（真实 GPU 上常常接近这种情况）：
        这时只剩 4Nd + N。
    """
    kv_reads = 2.0 * N * d * (1 if kv_cached else math.ceil(N / Br))
    return 2.0 * N * d + kv_reads + N