"""
=============================================================================
CS349E · 第 11 讲共享模块：flash_lab.py —— FlashAttention、flash-decoding 与 megakernel 的 numpy 版本
=============================================================================

【这个模块要解决什么问题】
    第 11 讲有三件事，都是"同样的数学，换一种算法顺序，让数据少进出显存"：
      1. FlashAttention：注意力的 N×N 分数矩阵不写回显存——分块 + 在线 softmax；反向传播时重算而不是存。
      2. flash-decoding：decode 时只有 1 个 query，FlashAttention 沿 query 切不出足够的块；改沿 KV 的长度切，
         每块算出局部结果和 log-sum-exp，最后合并。
      3. megakernel：一次 decode 有上百个小内核，内核之间的空隙（启动、等待第一批数据、尾部）成了主要的时间。
    这个模块用 numpy **真的算**前两件事（和标准注意力比较到 1e-15），并数出每种算法读写显存的元素个数；
    第三件事和 decode 注意力的速度用显式写出假设的时间模型估计。

【主要内容】
    - 在线 softmax：softmax_3pass / softmax_online / merge_states（局部结果的合并，满足结合律）
    - 标准注意力 standard_attention（+ 显存读写计数）、standard_attention_bwd
    - FlashAttention：flash_fwd（order="fa1" 是论文算法 1 的循环顺序，"fa2" 是 FlashAttention-2 的顺序）、
      flash_bwd（重算 S、P 的反向）、block_sizes（论文的 Bc = ⌈M/4d⌉、Br = min(Bc, d)）
    - flash-decoding：split_kv_decode（分块 + LSE 合并）、DecodeGPU / decode_attention_time（时间模型）
    - megakernel：LLAMA_1B_OPS（Llama-3.2-1B 一次 decode 的各个算子和权重字节）、timeline（多内核 vs megakernel 的时间线模型）

【约定】注意力的输入是一个头：Q [N, d]、K [Nk, d]、V [Nk, dv]；缩放 1/√d。
    causal=True 时第 i 个 query（位置 Nk − N + i）只看位置 ≤ 它的 key（N = Nk 时就是普通的因果掩码；
    N < Nk 时相当于"前面有 Nk − N 个缓存的 token"）。
    显存读写按**元素个数**计（乘以每个元素的字节数就是字节）。

【运行】python3 flash_lab.py    # 自检：各种算法和标准注意力一致、梯度和有限差分一致、计数和公式一致
【依赖】numpy。
"""

import math
from dataclasses import dataclass, field

import numpy as np


# =============================================================================
# 1. 在线 softmax
# =============================================================================
def softmax_3pass(x):
    """
    "安全"的 softmax：先扫一遍求最大值 m，再扫一遍求 Σ exp(x − m)，最后扫一遍除。

    【返回】(y, 扫数据的遍数 = 3)
    【为什么减最大值】exp(x) 在 FP32 里 x > 88 就溢出成 inf；减掉最大值后每个指数 ≤ 0（前置讲 P3 §5）。
    """
    m = np.max(x, axis=-1, keepdims=True)
    e = np.exp(x - m)
    return e / np.sum(e, axis=-1, keepdims=True), 3


def softmax_online(x, block=128):
    """
    在线 softmax（Milakov & Gimelshein 2018）：一遍同时求出最大值 m 和 Σ exp(x − m)，第二遍再除。

    【思路】按块扫：手里维护 (m, l)，m = 目前见过的最大值，l = Σ exp(x_j − m)。来了一块新数据（块内最大值 m_b）：
        m_new = max(m, m_b)；l_new = l · exp(m − m_new) + Σ_块 exp(x_j − m_new)。
        旧的 l 是按旧的 m 算的，乘 exp(m − m_new) 就"换算"成按新的 m 算的——这个换算因子是整个 FlashAttention 的关键。
    【返回】(y, 遍数 = 2)
    【边界】block 可以不整除行长。
    """
    n = x.shape[-1]
    m = np.full(x.shape[:-1] + (1,), -np.inf)
    l = np.zeros(x.shape[:-1] + (1,))
    for s in range(0, n, block):
        xb = x[..., s:s + block]
        m_new = np.maximum(m, np.max(xb, axis=-1, keepdims=True))
        l = l * np.exp(m - m_new) + np.sum(np.exp(xb - m_new), axis=-1, keepdims=True)
        m = m_new
    return np.exp(x - m) / l, 2


def partial_state(x, v=None):
    """
    一段数据的"局部状态"：(m, l, o)。m = 最大值；l = Σ exp(x − m)；o = Σ exp(x − m)·v（v 为 None 时 o 为 None）。

    【用途】把一行切成任意几段，各算一个局部状态，再用 merge_states 合并——这就是 flash-decoding 和 split-K 的数学。
    """
    m = np.max(x, axis=-1)
    p = np.exp(x - m[..., None])
    l = np.sum(p, axis=-1)
    o = None if v is None else p @ v
    return m, l, o


def merge_states(a, b):
    """
    合并两个局部状态：结果和"把两段拼起来再算局部状态"完全一样。

    【公式】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 脚本演示 2）。
    【边界】一段全是 −∞（被掩码掉）时 m = −∞、l = 0：这里把 e^(−∞ − m) 当作 0 处理。
    """
    ma, la, oa = a
    mb, lb, ob = b
    m = np.maximum(ma, mb)
    with np.errstate(invalid="ignore"):
        fa = np.where(np.isneginf(ma), 0.0, np.exp(ma - m))
        fb = np.where(np.isneginf(mb), 0.0, np.exp(mb - m))
    l = la * fa + lb * fb
    o = None if oa is None else oa * fa[..., None] + ob * fb[..., None]
    return m, l, o


# =============================================================================
# 2. 标准注意力（+ 显存读写计数）
# =============================================================================
@dataclass
class IO:
    """
    显存读写计数（元素个数）和矩阵乘法的运算量。

    【字段】read、write：读 / 写的元素数；flops：矩阵乘法的 FLOPs（2 × 乘加，只算 QKᵀ、PV 及反向里对应的乘法）
    """
    read: int = 0
    write: int = 0
    flops: int = 0
    notes: dict = field(default_factory=dict)

    @property
    def total(self):
        return self.read + self.write


def _causal_mask(nq, nk, q0=0, k0=0, offset=0):
    """[nq, nk] 的布尔数组：True = 允许。第 q0+r 个 query 的位置是 offset + q0 + r，key 的位置是 k0 + c。"""
    qpos = offset + q0 + np.arange(nq)[:, None]
    kpos = k0 + np.arange(nk)[None, :]
    return kpos <= qpos


def standard_attention(Q, K, V, causal=False, scale=None):
    """
    标准注意力：S = QKᵀ·scale → 写回显存；P = softmax(S) → 写回显存；O = PV。

    【显存读写（FlashAttention 论文算法 0 的写法）】
        读 Q、K，写 S（N×Nk）；读 S，写 P；读 P、V，写 O。
        = 读 N·d + Nk·d + 2·N·Nk + Nk·dv，写 2·N·Nk + N·dv —— Θ(Nd + N²)。
    【返回】(O, P, io)：P 留给反向传播用（标准写法要存下 N×Nk 的 P）
    """
    N, d = Q.shape
    Nk, dv = V.shape
    scale = 1.0 / math.sqrt(d) if scale is None else scale
    S = (Q @ K.T) * scale
    if causal:
        S = np.where(_causal_mask(N, Nk, offset=Nk - N), S, -np.inf)
    P, _ = softmax_3pass(S)
    O = P @ V
    io = IO(read=N * d + Nk * d + N * Nk + N * Nk + Nk * dv, write=N * Nk + N * Nk + N * dv,
            flops=2 * N * Nk * d + 2 * N * Nk * dv)
    return O, P, io


def standard_attention_bwd(Q, K, V, P, dO, scale=None):
    """
    标准注意力的反向：用前向存下的 P。

    【公式】dV = Pᵀ·dO；dP = dO·Vᵀ；dS = P ∘ (dP − D)，D_i = Σ_j dP_ij P_ij = Σ dO_i ∘ O_i；
        dQ = dS·K·scale；dK = dSᵀ·Q·scale。（softmax 的雅可比矩阵乘向量，推导见讲义 §5.1）
    【显存读写】每个算子单独一个内核：读 P 两次（dV、dS）、写 dP 再读、写 dS 再读两次（dQ、dK）——
        N×Nk 的矩阵共读 5 次、写 2 次，外加 Q、K、V、O、dO、dQ、dK、dV 各一次。
    【返回】(dQ, dK, dV, io)
    """
    N, d = Q.shape
    Nk, dv = V.shape
    scale = 1.0 / math.sqrt(d) if scale is None else scale
    O = P @ V
    D = np.sum(dO * O, axis=1)
    dV = P.T @ dO
    dP = dO @ V.T
    dS = P * (dP - D[:, None])
    dQ = dS @ K * scale
    dK = dS.T @ Q * scale
    io = IO(read=5 * N * Nk + 2 * N * dv + N * dv + Nk * dv + Nk * d + N * d + N,
            write=2 * N * Nk + N * d + Nk * d + Nk * dv + N,
            flops=2 * N * Nk * (dv + dv + d + d))
    return dQ, dK, dV, io


# =============================================================================
# 3. FlashAttention
# =============================================================================
def block_sizes(M, d):
    """
    FlashAttention 论文算法 1 的块大小：Bc = ⌈M / (4d)⌉，Br = min(⌈M / (4d)⌉, d)。

    【参数】M：片上 SRAM 能放多少个**元素**；d：头维度
    【为什么是 4d】片上要同时放 Q_i（Br×d）、K_j、V_j（Bc×d）、O_i（Br×d）四块，每块不超过 M/4。
    """
    bc = math.ceil(M / (4 * d))
    return min(bc, d), bc


def flash_fwd(Q, K, V, Br, Bc, causal=False, scale=None, order="fa2"):
    """
    FlashAttention 前向：分块 + 在线 softmax，S、P 只在"片上"（这里是局部变量）存在。

    【两种循环顺序】
        order="fa1"（论文算法 1）：外层走 K、V 的块 j，内层走 Q 的块 i。O_i、l_i、m_i 每次都要从显存读进来、
            更新后写回去——K、V 只读一遍，但 Q、O 被读写 Tc = ⌈Nk/Bc⌉ 遍。
        order="fa2"（FlashAttention-2）：外层走 Q 的块 i，内层走 K、V 的块 j。O_i 的累加器一直留在片上，
            最后除以 l、写一次——Q、O 各读写一遍，K、V 被读 Tr = ⌈N/Br⌉ 遍。FA2 还把"除以 l"推迟到最后，
            省掉每一步的除法；不同的 Q 块之间没有依赖，可以分给不同的 SM 并行（FA1 的顺序要在 SM 之间同步 O_i）。
    【causal】整块都被掩码掉的 (i, j) 直接跳过（不读、不算）；对角线上的块在块内加掩码。
    【返回】(O, L, io)：L = m + log(l) 是每行的 log-sum-exp（反向传播只需要存它，N 个数，而不是 N×Nk 的 P）
    【边界】N、Nk 不必是块大小的倍数（最后一块小一些）。
    """
    N, d = Q.shape
    Nk, dv = V.shape
    scale = 1.0 / math.sqrt(d) if scale is None else scale
    off = Nk - N
    io = IO()
    O = np.zeros((N, dv))
    L = np.zeros(N)

    def block_scores(i0, br, j0, bc):
        S = (Q[i0:i0 + br] @ K[j0:j0 + bc].T) * scale
        io.flops += 2 * br * bc * d
        if causal:
            S = np.where(_causal_mask(br, bc, i0, j0, off), S, -np.inf)
        return S

    def skip(i0, br, j0):
        # 这一块里最靠后的 query 的位置 < 这一块最靠前的 key → 整块被掩码
        return causal and (off + i0 + br - 1) < j0

    if order == "fa2":
        for i0 in range(0, N, Br):
            br = min(Br, N - i0)
            io.read += br * d                                        # 读 Q_i（一次）
            m = np.full(br, -np.inf)
            l = np.zeros(br)
            acc = np.zeros((br, dv))
            for j0 in range(0, Nk, Bc):
                if skip(i0, br, j0):
                    break                                            # 后面的块更靠后，也都被掩码
                bc = min(Bc, Nk - j0)
                io.read += bc * d + bc * dv                          # 读 K_j、V_j
                S = block_scores(i0, br, j0, bc)
                m_new = np.maximum(m, S.max(axis=1))
                alpha = np.where(np.isneginf(m), 0.0, np.exp(m - np.where(np.isneginf(m_new), 0.0, m_new)))
                P = np.exp(S - np.where(np.isneginf(m_new), 0.0, m_new)[:, None])
                l = alpha * l + P.sum(axis=1)
                acc = alpha[:, None] * acc + P @ V[j0:j0 + bc]
                io.flops += 2 * br * bc * dv
                m = m_new
            O[i0:i0 + br] = acc / l[:, None]
            L[i0:i0 + br] = m + np.log(l)
            io.write += br * dv + br                                 # 写 O_i、L_i（一次）
    elif order == "fa1":
        m_all = np.full(N, -np.inf)
        l_all = np.zeros(N)
        io.write += N * dv + 2 * N                                   # 初始化 O、l、m（写进显存）
        for j0 in range(0, Nk, Bc):
            bc = min(Bc, Nk - j0)
            io.read += bc * d + bc * dv                              # 读 K_j、V_j（一次）
            for i0 in range(0, N, Br):
                br = min(Br, N - i0)
                if skip(i0, br, j0):
                    continue
                io.read += br * d + br * dv + 2 * br                 # 读 Q_i、O_i、l_i、m_i
                S = block_scores(i0, br, j0, bc)
                mt = S.max(axis=1)
                mt_safe = np.where(np.isneginf(mt), 0.0, mt)
                Pt = np.exp(S - mt_safe[:, None])
                lt = Pt.sum(axis=1)
                mi, li = m_all[i0:i0 + br], l_all[i0:i0 + br]
                m_new = np.maximum(mi, mt)
                m_new_safe = np.where(np.isneginf(m_new), 0.0, m_new)
                fa = np.where(np.isneginf(mi), 0.0, np.exp(mi - m_new_safe))
                fb = np.where(np.isneginf(mt), 0.0, np.exp(mt - m_new_safe))
                l_new = fa * li + fb * lt
                PV = Pt @ V[j0:j0 + bc]
                io.flops += 2 * br * bc * dv
                with np.errstate(invalid="ignore", divide="ignore"):
                    O[i0:i0 + br] = np.where(l_new[:, None] > 0,
                                             ((li * fa)[:, None] * O[i0:i0 + br] + fb[:, None] * PV) / l_new[:, None], 0.0)
                m_all[i0:i0 + br], l_all[i0:i0 + br] = m_new, l_new
                io.write += br * dv + 2 * br                         # 写回 O_i、l_i、m_i
        L = m_all + np.log(l_all)
    else:
        raise ValueError("order 只能是 'fa1' 或 'fa2'")
    return O, L, io


def flash_bwd(Q, K, V, O, dO, L, Br, Bc, causal=False, scale=None):
    """
    FlashAttention 反向（FlashAttention-2 的写法）：不读 P，而是用 Q、K 和前向存下的 L 在片上**重算** S 和 P。

    【步骤】先算 D_i = Σ dO_i ∘ O_i（每行一个数）。外层走 K、V 的块 j（dK_j、dV_j 的累加器留在片上），内层走 Q 的块 i：
        S_ij = Q_i K_jᵀ·scale（+ 掩码）；P_ij = exp(S_ij − L_i)        ← 重算，不需要 m、l 分开存
        dV_j += P_ijᵀ dO_i；dP_ij = dO_i V_jᵀ；dS_ij = P_ij ∘ (dP_ij − D_i)
        dQ_i += dS_ij K_j·scale（累加回显存）；dK_j += dS_ijᵀ Q_i·scale
    【运算量】每个 (i, j) 5 次矩阵乘法（前向 2 次），所以反向 ≈ 前向的 2.5 倍；比标准反向多一次（重算 S）。
    【返回】(dQ, dK, dV, io)
    """
    N, d = Q.shape
    Nk, dv = V.shape
    scale = 1.0 / math.sqrt(d) if scale is None else scale
    off = Nk - N
    io = IO()
    D = np.sum(dO * O, axis=1)
    io.read += 2 * N * dv
    io.write += N
    dQ = np.zeros_like(Q)
    dK = np.zeros_like(K)
    dV = np.zeros_like(V)
    io.write += N * d                                                # dQ 清零
    for j0 in range(0, Nk, Bc):
        bc = min(Bc, Nk - j0)
        io.read += bc * d + bc * dv
        dKj = np.zeros((bc, d))
        dVj = np.zeros((bc, dv))
        for i0 in range(0, N, Br):
            br = min(Br, N - i0)
            if causal and (off + i0 + br - 1) < j0:
                continue
            io.read += br * d + br * dv + 2 * br + br * d            # Q_i、dO_i、L_i、D_i、dQ_i
            S = (Q[i0:i0 + br] @ K[j0:j0 + bc].T) * scale
            if causal:
                S = np.where(_causal_mask(br, bc, i0, j0, off), S, -np.inf)
            P = np.exp(S - L[i0:i0 + br, None])
            dVj += P.T @ dO[i0:i0 + br]
            dP = dO[i0:i0 + br] @ V[j0:j0 + bc].T
            dS = P * (dP - D[i0:i0 + br, None])
            dQ[i0:i0 + br] += dS @ K[j0:j0 + bc] * scale
            dKj += dS.T @ Q[i0:i0 + br] * scale
            io.flops += 2 * br * bc * (d + dv + dv + d + d)
            io.write += br * d                                       # 写回 dQ_i
        dK[j0:j0 + bc] = dKj
        dV[j0:j0 + bc] = dVj
        io.write += bc * d + bc * dv
    return dQ, dK, dV, io


def fa_io_formula(N, d, M, order="fa2"):
    """
    FlashAttention 前向读写元素数的公式（不含 causal、N = Nk、dv = d），和 flash_fwd 的计数对照用。

    fa1：K、V 各读一次 + Q 读 Tc 次 + O 读写 Tc 次（+ 初始化写一次）+ l、m 的读写
    fa2：Q 读一次 + K、V 各读 Tr 次 + O 写一次 + L 写一次
    """
    Br, Bc = block_sizes(M, d)
    Tr, Tc = math.ceil(N / Br), math.ceil(N / Bc)
    if order == "fa1":
        return 2 * N * d + Tc * (N * d + N * d + 2 * N) + Tc * (N * d + 2 * N) + N * d + 2 * N
    return N * d + Tr * 2 * N * d + N * d + N


# =============================================================================
# 4. flash-decoding
# =============================================================================
def split_kv_decode(q, K, V, n_splits, scale=None):
    """
    flash-decoding 的数学：一个（或几个）query 对很长的 K、V 做注意力，把 K、V 切成 n_splits 段，
    每段独立算出 (o_s, lse_s)，最后合并。

    【参数】q：[R, d]（R 个 query 头共享这组 K、V——GQA 时把同一组的 query 头"打包"在一起）；K：[S, d]；V：[S, dv]
    【每段的输出】o_s = softmax(q K_sᵀ)·V_s（这一段内部归一化过的输出），lse_s = log Σ_段 exp(q k·scale)
    【合并】lse = log Σ_s exp(lse_s)；o = Σ_s exp(lse_s − lse)·o_s。和 merge_states 是同一件事，换成了 lse 的写法
        （每段只需要额外存 1 个数，而不是 m 和 l 两个）。
    【返回】(o, lse, 各段的 (o_s, lse_s) 列表)
    【边界】n_splits 可以不整除 S；n_splits > S 时多出来的段为空，跳过。
    """
    S, d = K.shape
    scale = 1.0 / math.sqrt(d) if scale is None else scale
    bounds = np.linspace(0, S, n_splits + 1).round().astype(int)
    parts = []
    for a, b in zip(bounds[:-1], bounds[1:]):
        if b <= a:
            continue
        s = (q @ K[a:b].T) * scale
        m = s.max(axis=1, keepdims=True)
        p = np.exp(s - m)
        l = p.sum(axis=1, keepdims=True)
        parts.append(((p @ V[a:b]) / l, (m + np.log(l))[:, 0]))
    lses = np.stack([p[1] for p in parts])                # [段数, R]
    lse = np.log(np.sum(np.exp(lses - lses.max(axis=0)), axis=0)) + lses.max(axis=0)
    o = sum(np.exp(ls - lse)[:, None] * os for os, ls in parts)
    return o, lse, parts


@dataclass(frozen=True)
class DecodeGPU:
    """
    decode 注意力时间模型用的 GPU 参数。

    【字段】
        name；n_sm：SM 个数；tc_flops：Tensor Core 的 FP16 稠密峰值；hbm：显存带宽（字节/秒）
        tc_eff：注意力内核在 Tensor Core 上能达到的峰值比例（**假设**：FlashAttention-2 论文报告前向在 A100 上
            达到理论峰值的 50%–73%，这里取 0.7）
        sm_bw_cap：一个 SM 自己最多能拉多少带宽（**假设**：按 Little 定律，在途约 64 KB、延迟约 0.8 µs → 约 80 GB/s）
        launch：每个内核的固定开销（**假设**：几微秒）
    """
    name: str
    n_sm: int
    tc_flops: float
    hbm: float
    tc_eff: float = 0.7
    sm_bw_cap: float = 80e9
    launch: float = 5e-6


A100_DEC = DecodeGPU("A100", 108, 312e12, 1.555e12)
H100_DEC = DecodeGPU("H100", 132, 989e12, 3.35e12)


def decode_attention_time(B, Hq, Hkv, S, d, method, gpu=A100_DEC, splits=1, block_m=None, elem=2):
    """
    decode 注意力（每个序列 1 个新 token）的时间模型：按"线程块（CTA）× 波次"估计。

    【两种方法】
        "fa2"：FlashAttention-2 的并行方式——每个 (批, query 头) 一个 CTA，沿 query 切块；只有 1 个 query 时，
            一个 CTA 的 Tensor Core 仍按 block_m（默认 128）行算——其余行是补齐的 0。每个 query 头各自读一遍
            它那组的 K、V（不同 CTA 之间不共享）。
        "fd"：flash-decoding——每个 (批, KV 头, 段) 一个 CTA；同一组的 R = Hq/Hkv 个 query 头打包成 R 行
            （补齐到 block_m，默认 64）；K、V 每组只读一遍；最后多一个合并的小内核。
    【每个 CTA 的时间】max(补齐后的运算量 / (每 SM 的 Tensor Core 峰值 × tc_eff)，
                         读的字节 / min(sm_bw_cap, 显存带宽 / 同时活跃的 SM 数))
    【总时间】⌈CTA 数 / SM 数⌉ 波 × 每个 CTA 的时间 + 内核固定开销（fd 两个内核）
    【返回】字典：t（秒）、ctas、waves、bound（"compute" / "memory"）、util（同时活跃的 SM 比例）
    【局限】不模拟 L2 的复用、负载不均、合并内核的读写（很小）。只比较趋势；参数是假设的（见 DecodeGPU）。
    """
    R = Hq // Hkv
    if method == "fa2":
        ctas = B * Hq
        rows = block_m or 128
        keys = S
        byts = 2 * S * d * elem
        n_kernels = 1
    elif method == "fd":
        ctas = B * Hkv * splits
        rows = max(R, block_m or 64)
        keys = math.ceil(S / splits)
        byts = 2 * keys * d * elem
        n_kernels = 2 if splits > 1 else 1
    else:
        raise ValueError("method 只能是 'fa2' 或 'fd'")
    active = min(ctas, gpu.n_sm)
    bw = min(gpu.sm_bw_cap, gpu.hbm / active)
    t_comp = 4 * rows * keys * d / (gpu.tc_flops / gpu.n_sm * gpu.tc_eff)
    t_mem = byts / bw
    waves = math.ceil(ctas / gpu.n_sm)
    t = waves * max(t_comp, t_mem) + n_kernels * gpu.launch
    return {"t": t, "ctas": ctas, "waves": waves, "bound": "compute" if t_comp > t_mem else "memory",
            "util": active / gpu.n_sm, "t_comp": t_comp, "t_mem": t_mem}


def best_splits(B, Hq, Hkv, S, d, gpu=A100_DEC, candidates=(1, 2, 4, 8, 16, 32, 64, 128, 256)):
    """
    在候选的段数里选模型时间最短的那个（flash-decoding 的实现用类似的启发式：段数让 CTA 大约填满所有 SM）。

    【规则】只考虑 CTA 数不超过 4 波的段数（段数 = 1 总是候选）；时间相同（差 < 1%）时取段数少的。
    【返回】(段数, decode_attention_time 的结果)
    """
    res = [(s, decode_attention_time(B, Hq, Hkv, S, d, "fd", gpu, splits=s)) for s in candidates
           if s <= S and (s == 1 or B * Hkv * s <= 4 * gpu.n_sm)]
    best_t = min(r[1]["t"] for r in res)
    return min((r for r in res if r[1]["t"] <= best_t * 1.01), key=lambda r: r[0])


# =============================================================================
# 5. megakernel：一次 decode 的时间线模型
# =============================================================================
@dataclass(frozen=True)
class Op:
    """
    一次 decode 里的一个算子（或融合后的一组算子）。

    【字段】name；bytes：要从显存读的字节（主要是权重；注意力是 KV cache）；ctas：多内核写法里这个内核的线程块数
        （0 表示"很小的内核"——只有 1 个块、几乎不读数据，如单独的 RMSNorm、RoPE）
    """
    name: str
    bytes: float
    ctas: int


def llama_1b_ops(context=160, rows_per_cta=16, kernels="7"):
    """
    Llama-3.2-1B（d = 2048、16 层、32 个 query 头 / 8 个 KV 头、头维度 64、FFN 8192、词表 128,256、
    词嵌入和 lm_head 共享）batch 1 的一次 decode：各算子读的字节（BF16）。

    【kernels】
        "7"：Hazy 的博客估计"每层 7 个内核"：RMSNorm、QKV（+RoPE）、注意力、O 投影（+残差）、RMSNorm、gate/up（+SiLU）、down（+残差）
        "12"：更朴素的写法，RoPE、写 KV、残差、SiLU 各自一个内核
        "fused"：Hazy 的 megakernel 里的指令粒度：RMSNorm + QKV + RoPE、注意力、O 投影 + 残差、
                 RMSNorm + gate/up + SiLU、down + 残差（每层 5 个）
    【返回】Op 的列表（最后是 final norm + lm_head）
    【边界】context = 注意力要读的 KV cache 长度（Hazy 的测试：32 个提示 token + 128 个生成 token）。
    """
    d, L, Hq, Hkv, dh, f, vocab, eb = 2048, 16, 32, 8, 64, 8192, 128256, 2
    qkv = d * (Hq + 2 * Hkv) * dh * eb
    o = Hq * dh * d * eb
    gu = 2 * d * f * eb
    dn = f * d * eb
    kv = 2 * context * Hkv * dh * eb
    ct = lambda byts, rows: max(1, rows // rows_per_cta)  # noqa: E731
    ops = []
    for _ in range(L):
        if kernels == "fused":
            ops += [Op("norm+qkv+rope", qkv, ct(qkv, (Hq + 2 * Hkv) * dh)), Op("attention", kv, Hq),
                    Op("o_proj+resid", o, ct(o, d)), Op("norm+gate_up+silu", gu, ct(gu, 2 * f)),
                    Op("down+resid", dn, ct(dn, d))]
        elif kernels == "7":
            ops += [Op("rmsnorm", 2 * d * eb, 0), Op("qkv+rope", qkv, ct(qkv, (Hq + 2 * Hkv) * dh)),
                    Op("attention", kv, Hq), Op("o_proj+resid", o, ct(o, d)), Op("rmsnorm", 2 * d * eb, 0),
                    Op("gate_up+silu", gu, ct(gu, 2 * f)), Op("down+resid", dn, ct(dn, d))]
        elif kernels == "12":
            ops += [Op("rmsnorm", 2 * d * eb, 0), Op("qkv", qkv, ct(qkv, (Hq + 2 * Hkv) * dh)), Op("rope", 0, 0),
                    Op("kv_write", 0, 0), Op("attention", kv, Hq), Op("o_proj", o, ct(o, d)), Op("residual", 0, 0),
                    Op("rmsnorm", 2 * d * eb, 0), Op("gate_up", gu, ct(gu, 2 * f)), Op("silu_mul", 0, 0),
                    Op("down", dn, ct(dn, d)), Op("residual", 0, 0)]
        else:
            raise ValueError("kernels 只能是 '7'、'12' 或 'fused'")
    lm = vocab * d * eb
    ops += [Op("final_norm", 2 * d * eb, 0), Op("lm_head", lm, ct(lm, vocab))] if kernels != "fused" else \
        [Op("norm+lm_head", lm, ct(lm, vocab))]
    return ops


@dataclass(frozen=True)
class TimelineParams:
    """
    时间线模型的参数（单位：秒、字节/秒）。

    【字段】
        n_sm、hbm：SM 个数、显存带宽
        eff：长时间连续读时能达到的带宽比例（**假设** 0.9）
        launch：相邻两个内核之间的启动间隙（Hazy 的博客：H100 上空内核在 stream 上约 2.1 µs，CUDA Graphs 约 1.3 µs）
        ramp：一个内核开始后、第一批权重到达之前的等待（**假设** 1 µs：显存延迟 + 地址计算等"序幕"）
        small：一个"很小的内核"（RMSNorm 这类）的执行时间（**假设** 2 µs：读几 KB 的激活、写回）
        sync：megakernel 里一个算子的输入依赖满足的代价——前一个算子写激活、计数器更新、这个算子看到计数器、读激活
             （Hazy：每次读或写激活约 500 ns，所以**假设** 1 µs）
        instr_overhead：megakernel 里每条指令的固定开销（取指令、设置、warp 间同步；**假设** 0.5 µs）
        smem_buf：megakernel 每个 SM 可以预取的权重缓冲（Hazy：H100 上 13 页 × 16 KiB = 208 KiB）
    """
    n_sm: int = 132
    hbm: float = 3.35e12
    eff: float = 0.9
    launch: float = 1.3e-6
    ramp: float = 1.0e-6
    small: float = 2.0e-6
    sync: float = 1.0e-6
    instr_overhead: float = 0.5e-6
    smem_buf: float = 13 * 16 * 1024


def timeline(ops, mode, p=TimelineParams()):
    """
    一次 decode 的时间线模型。

    【mode="kernels"（一个算子一个内核）】
        每个内核：等上一个内核的**所有**块结束 → 启动间隙 launch → 等 ramp（第一批数据到达；之后的波次和前一波重叠，只付一次）
        → 每个块以"每 SM 一份"的带宽读完它那份字节。
        块数多于 SM 数时分几波，最后一波没占满的 SM 空闲（尾部效应）。很小的内核按 small 计。
    【mode="megakernel"】
        一个常驻的内核：每个 SM 按顺序执行"指令"（每个算子在每个 SM 上一条）。权重**不依赖**上一个算子的结果，
        所以在缓冲区（smem_buf）允许的范围内可以提前读；只有激活的依赖要等：上一个算子在所有 SM 上结束 + sync。
        每条指令还有 instr_overhead 的固定开销。
    【返回】字典：t（秒）、bw_util（理想时间 / t）、parts（时间的分解：读权重、启动、序幕、尾部、小内核、同步等）
    【局限】带宽在 SM 之间平均分；不模拟 L2、预取的细节、注意力内核的并行度不足（注意力只读几百 KB，影响很小）。
    """
    bw_sm = p.hbm * p.eff / p.n_sm
    ideal = sum(o.bytes for o in ops) / p.hbm
    parts = {"读权重（按满带宽）": 0.0, "带宽效率损失": 0.0, "启动间隙": 0.0, "序幕（等第一批数据）": 0.0,
             "尾部（最后一波没占满）": 0.0, "小内核": 0.0, "同步与指令开销（没被预取藏住的）": 0.0}
    t = 0.0
    if mode == "kernels":
        for o in ops:
            t += p.launch
            parts["启动间隙"] += p.launch
            if o.ctas == 0:
                t += p.small
                parts["小内核"] += p.small
                continue
            waves = math.ceil(o.ctas / p.n_sm)
            per_cta = o.bytes / o.ctas
            dur = p.ramp + waves * per_cta / bw_sm          # 序幕只付一次：之后的波次和前一波重叠
            full = o.bytes / (p.n_sm * bw_sm)
            t += dur
            parts["序幕（等第一批数据）"] += p.ramp
            parts["读权重（按满带宽）"] += o.bytes / p.hbm
            parts["带宽效率损失"] += full - o.bytes / p.hbm
            parts["尾部（最后一波没占满）"] += waves * per_cta / bw_sm - full
    elif mode == "megakernel":
        weight_ready = 0.0                      # 每个 SM 的权重"读到哪儿了"（所有 SM 对称，只算一个）
        dep_done = 0.0                          # 上一个算子在所有 SM 上结束的时刻
        for o in ops:
            share = o.bytes / p.n_sm            # 这条指令在每个 SM 上要读的字节
            start_load = weight_ready
            dep_ready = dep_done + p.sync
            prefetched = min(share, p.smem_buf, max(0.0, dep_ready - start_load) * bw_sm)
            begin = max(dep_ready, start_load)
            load_end = max(start_load + share / bw_sm, begin + (share - prefetched) / bw_sm)
            finish = max(load_end, begin) + p.instr_overhead
            weight_ready = load_end
            dep_done = finish
        t = dep_done
        parts["读权重（按满带宽）"] = ideal
        parts["带宽效率损失"] = ideal / p.eff - ideal
        parts["同步与指令开销（没被预取藏住的）"] = t - ideal / p.eff
    else:
        raise ValueError("mode 只能是 'kernels' 或 'megakernel'")
    return {"t": t, "ideal": ideal, "bw_util": ideal / t, "parts": parts}


# =============================================================================
# 自检
# =============================================================================
def _selftest():
    rng = np.random.default_rng(0)
    # 在线 softmax 与合并
    x = rng.standard_normal((5, 300)) * 4
    y3, _ = softmax_3pass(x)
    yo, _ = softmax_online(x, block=37)
    assert np.abs(y3 - yo).max() < 1e-15
    a = partial_state(x[:, :100])
    b = partial_state(x[:, 100:])
    m, l, _ = merge_states(a, b)
    mf, lf, _ = partial_state(x)
    assert np.allclose(m, mf) and np.allclose(l, lf, rtol=1e-14)
    # 前向：两种顺序、causal、不整除
    for (N, Nk, d) in ((64, 64, 16), (50, 50, 8), (7, 33, 8)):
        Q, K, V = (rng.standard_normal((n, d)) for n in (N, Nk, Nk))
        for causal in (False, True):
            O_ref, _, _ = standard_attention(Q, K, V, causal)
            for order in ("fa1", "fa2"):
                O, L, io = flash_fwd(Q, K, V, 16, 12, causal, order=order)
                assert np.abs(O - O_ref).max() < 1e-12, (N, Nk, causal, order)
    # 不带 causal 时计数等于公式
    N, d, M = 256, 16, 2048
    Q, K, V = (rng.standard_normal((N, d)) for _ in range(3))
    Br, Bc = block_sizes(M, d)
    for order in ("fa1", "fa2"):
        _, _, io = flash_fwd(Q, K, V, Br, Bc, order=order)
        assert io.total == fa_io_formula(N, d, M, order), (order, io.total, fa_io_formula(N, d, M, order))
    # 反向：和标准反向一致，和有限差分一致
    N, d = 12, 4
    Q, K, V = (rng.standard_normal((N, d)) for _ in range(3))
    dO = rng.standard_normal((N, d))
    for causal in (False, True):
        O, P, _ = standard_attention(Q, K, V, causal)
        dQs, dKs, dVs, _ = standard_attention_bwd(Q, K, V, P, dO)
        O2, L, _ = flash_fwd(Q, K, V, 4, 5, causal)
        dQ, dK, dV, _ = flash_bwd(Q, K, V, O2, dO, L, 4, 5, causal)
        assert max(np.abs(dQ - dQs).max(), np.abs(dK - dKs).max(), np.abs(dV - dVs).max()) < 1e-12
        eps = 1e-6
        Qp = Q.copy()
        Qp[3, 1] += eps
        Qm = Q.copy()
        Qm[3, 1] -= eps
        num = (np.sum(standard_attention(Qp, K, V, causal)[0] * dO) - np.sum(standard_attention(Qm, K, V, causal)[0] * dO)) / (2 * eps)
        assert abs(num - dQ[3, 1]) < 1e-7
    # flash-decoding
    q = rng.standard_normal((4, 32))
    K, V = rng.standard_normal((1000, 32)), rng.standard_normal((1000, 32))
    o_ref, _, _ = standard_attention(q, K, V)
    for ns in (1, 3, 16, 999):
        o, _, _ = split_kv_decode(q, K, V, ns)
        assert np.abs(o - o_ref).max() < 1e-13
    # 时间线：理想 ≤ 任何写法；Llama-1B 的权重约 2.47 GB
    ops = llama_1b_ops(kernels="7")
    tot = sum(o.bytes for o in ops)
    assert abs(tot - 2.47e9) < 0.02e9, tot
    for mode in ("kernels", "megakernel"):
        assert timeline(ops, mode)["bw_util"] < 1.0
    print("flash_lab 自检通过。")


if __name__ == "__main__":
    print("== flash_lab 自检 ==")
    _selftest()