"""
=============================================================================
CS349E · 第 2 讲共享模块：transformer_lab.py —— 一个能真正运行、还会"记账"的 numpy Transformer
=============================================================================

【这个模块要解决什么问题】
    第 1 讲的 llm_math.py 用**公式**回答了"一个 token 要做多少运算、读多少字节、KV cache 多大"。
    公式容易写错，也容易忘记它省略了什么。本模块写一个**真的能跑**的 Llama 风格 decoder-only
    Transformer（numpy 实现），并在每一次矩阵乘法、每一次读 KV cache 时**数一数**：
        - 做了多少次浮点运算（FLOPs）；
        - 读了多少字节的权重、多少字节的 KV cache、多少字节的激活。
    然后拿数出来的结果和公式逐项比较（03 脚本）。第 3 讲（量化）和第 4 讲（张量并行）也会导入这个模块。

【模型结构（和 Llama 3 一样，只是小得多）】
    tokens ─► 词嵌入（查表）─► [ 第 1 层 … 第 L 层 ] ─► RMSNorm ─► lm_head ─► logits（每个词一个分数）
    每一层（pre-norm，残差连接）：
        h = RMSNorm(x);  q, k, v = h·Wq, h·Wk, h·Wv;  q, k 加 RoPE 位置编码
        a = 因果注意力(q, k, v)（GQA：H 个 query 头共享 G 组 K/V）
        x = x + a·Wo
        h = RMSNorm(x);  x = x + ( SiLU(h·W_gate) ⊙ (h·W_up) ) · W_down          ← SwiGLU MLP
    权重矩阵统一写成 W[输入维, 输出维]，计算 y = x·W（和第 1 讲的 W[D, F] 记号一致）。

【"记账"的约定（Counter）】
    - 矩阵乘法 x[M, K]·W[K, N]：FLOPs = 2·M·K·N（每个乘加算 2 次）；权重字节 = K·N·每个数的字节；
      激活字节 = (M·K + M·N)·每个数的字节（读输入、写输出）。
    - 注意力核心（QKᵀ、softmax、PV）按"融合内核"记账（像 FlashAttention 那样，分数矩阵不写回显存，第 11 讲）：
      FLOPs = 4·B·T·S·H·dh（T 个 query × S 个 key，两次矩阵乘法）；KV 字节 = 读一遍 K 和 V。
      另外记一个 useful（"有用的"FLOPs）：因果掩码挡掉的那一半不算（第 i 个 query 只需要看前 i 个 key）。
      numpy 的朴素实现会把整个 T×S 都算出来再掩掉，所以 flops ≥ useful。
    - 逐元素运算（RMSNorm、RoPE、SiLU、残差加）的 FLOPs 用固定的"每个元素几次运算"粗略估计，
      它们在总量里很小，llm_math 的公式也不计它们。
    - 字节按 bytes_per_elem 计（默认 2，即 BF16），**和 numpy 实际用的 float64 无关**：
      我们用 float64 是为了让"带 KV cache"和"不带 KV cache"的结果能精确比较。

【运行】python3 transformer_lab.py   # 自检：参数量 = 公式；带 / 不带 KV cache 生成结果相同；记账 = 解析公式
【依赖】numpy；从 ../lecture01 导入 llm_math（公式版的账本）。
"""

import math
import os
import sys
from dataclasses import dataclass

import numpy as np

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

from llm_math import ModelConfig  # noqa: E402

# 逐元素运算的"每个元素几次运算"（粗略估计；只影响 other 类 FLOPs，不影响矩阵乘法的账）
NORM_FLOPS = 4      # RMSNorm：平方、累加、乘 rsqrt、乘权重
ROPE_FLOPS = 3      # RoPE：每对数 4 次乘 + 2 次加 → 每个数 3 次
SOFTMAX_FLOPS = 5   # softmax：减最大值、exp、累加、除（exp 按 1 次算，实际更贵，见第 1 讲 §5.5）
SWIGLU_FLOPS = 5    # SiLU(g)·u：sigmoid 约 3 次 + 乘 g + 乘 u
RESID_FLOPS = 1     # 残差加法


# ---------------------------------------------------------------------------
# 模型配置
# ---------------------------------------------------------------------------
@dataclass(frozen=True)
class Config:
    """
    一个 Llama 风格 Transformer 的结构参数（字段含义和第 1 讲 llm_math.ModelConfig 相同，多了几个运行时需要的）。

    【字段】
        vocab, d_model, n_layers, n_heads, n_kv_heads, d_head, d_ff：同 llm_math
        max_seq   ：KV cache 预分配的最大长度
        rope_theta：RoPE 的基数（Llama 3 用 500000，原始 RoFormer / Llama 2 用 10000）
        tied      ：lm_head 是否和词嵌入共享同一个矩阵
        norm_eps  ：RMSNorm 里防止除以 0 的小常数
        name      ：名字（打印用）
    【边界】n_heads 必须能被 n_kv_heads 整除（每组 K/V 头服务 n_heads / n_kv_heads 个 query 头）；
        d_head 必须是偶数（RoPE 把相邻两个数当成一个二维向量旋转）。不满足时在创建时就报错。
    【Python 小知识】dataclass 的 __post_init__ 在自动生成的 __init__ 之后运行，适合做参数检查。
    """
    vocab: int = 512
    d_model: int = 128
    n_layers: int = 4
    n_heads: int = 8
    n_kv_heads: int = 2
    d_head: int = 16
    d_ff: int = 344
    max_seq: int = 256
    rope_theta: float = 10000.0
    tied: bool = False
    norm_eps: float = 1e-5
    name: str = "tiny"

    def __post_init__(self):
        if self.n_heads % self.n_kv_heads != 0:
            raise ValueError(f"n_heads={self.n_heads} 不能被 n_kv_heads={self.n_kv_heads} 整除")
        if self.d_head % 2 != 0:
            raise ValueError(f"d_head={self.d_head} 必须是偶数（RoPE 两两旋转）")

    @property
    def group(self):
        """每组 K/V 头被几个 query 头共享（MHA = 1，MQA = n_heads）。"""
        return self.n_heads // self.n_kv_heads

    def to_model_config(self):
        """转成第 1 讲的 ModelConfig，好直接调用它的公式（参数量、KV/token、FLOPs/token）。"""
        return ModelConfig(self.name, n_layers=self.n_layers, d_model=self.d_model, n_heads=self.n_heads,
                           n_kv_heads=self.n_kv_heads, d_head=self.d_head, d_ff=self.d_ff, vocab=self.vocab,
                           tied=self.tied)


TINY = Config()   # 默认的小模型：4 层、d=128、8 个 query 头、2 组 KV 头（GQA，组大小 4）


def config_like(cfg, **changes):
    """
    复制一个配置并改几个字段，例如 config_like(TINY, n_kv_heads=8) 得到 MHA 版本。

    【Python 小知识】frozen dataclass 不能直接改字段；dataclasses.replace 生成一个改了字段的新对象。
    """
    from dataclasses import replace
    return replace(cfg, **changes)


# ---------------------------------------------------------------------------
# 参数初始化
# ---------------------------------------------------------------------------
def init_params(cfg, seed=0, dtype=np.float64, logit_std=2.0):
    """
    随机初始化一组参数（没有训练过——本讲关心的是"计算本身"，不是模型说得对不对）。

    【思路】
        - 词嵌入 ~ N(0, 1)：RMSNorm 之后每个数的均方根约为 1。
        - 每个线性层 W[in, out] ~ N(0, 1/in)：输出和输入的尺度差不多（"Xavier/He 式"缩放）。
        - 残差分支的输出矩阵（Wo、W_down）再乘 1/√(2L)：L 层、每层两个分支叠加之后，残差流不会越来越大。
        - lm_head ~ N(0, logit_std²/d)：让 logits 的标准差约为 logit_std，输出分布不至于是均匀分布
          （第 3 讲要比较"量化前后输出分布差多少"，均匀分布就看不出差别了）。
        - RMSNorm 的缩放系数初始化为 1。
    【参数】cfg：Config；seed：随机种子；dtype：numpy 数据类型（默认 float64，便于精确比较）；logit_std：见上
    【返回】字典：{"embed", "layers": [每层一个字典], "final_norm", "lm_head"}；tied 时 lm_head 为 None
    【边界】同样的 (cfg, seed) 永远得到同样的参数。
    【Python 小知识】np.random.default_rng(seed) 是 numpy 推荐的随机数生成器；standard_normal(shape) 生成标准正态。
    """
    rng = np.random.default_rng(seed)
    d, H, G, dh, f = cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.d_head, cfg.d_ff
    resid_scale = 1.0 / math.sqrt(2 * cfg.n_layers)

    def mat(n_in, n_out, scale=1.0):
        return (rng.standard_normal((n_in, n_out)) * (scale / math.sqrt(n_in))).astype(dtype)

    params = {"embed": rng.standard_normal((cfg.vocab, d)).astype(dtype), "layers": []}
    for _ in range(cfg.n_layers):
        params["layers"].append({
            "attn_norm": np.ones(d, dtype=dtype),
            "wq": mat(d, H * dh),
            "wk": mat(d, G * dh),
            "wv": mat(d, G * dh),
            "wo": mat(H * dh, d, resid_scale),
            "mlp_norm": np.ones(d, dtype=dtype),
            "w_gate": mat(d, f),
            "w_up": mat(d, f),
            "w_down": mat(f, d, resid_scale),
        })
    params["final_norm"] = np.ones(d, dtype=dtype)
    params["lm_head"] = None if cfg.tied else mat(d, cfg.vocab, logit_std)
    return params


def lm_head_matrix(params):
    """lm_head 用的矩阵 [d, vocab]：没有共享时是自己的矩阵，共享时是词嵌入表的转置。"""
    return params["lm_head"] if params["lm_head"] is not None else params["embed"].T


def count_params(params):
    """
    数一数参数的个数（直接把每个数组的元素个数加起来）。

    【边界】tied 时 lm_head 是 None，不重复计数——和 llm_math.ModelConfig.n_params 的约定一致。
    【Python 小知识】arr.size 是数组的元素总数。
    """
    n = params["embed"].size + params["final_norm"].size
    if params["lm_head"] is not None:
        n += params["lm_head"].size
    for lp in params["layers"]:
        n += sum(a.size for a in lp.values())
    return n


# ---------------------------------------------------------------------------
# 记账器
# ---------------------------------------------------------------------------
class Counter:
    """
    按"运算名"累计 FLOPs 和字节数。

    【字段（每个运算名一行）】
        flops  ：实际执行的浮点运算次数
        useful ：其中"有用的"运算（只有注意力核心会比 flops 少：因果掩码挡掉的部分不算）
        w      ：读权重的字节
        kv     ：读 / 写 KV cache 的字节
        act    ：读写激活的字节（假设每个运算单独一个内核、输入输出都过一遍显存）
        calls  ：调用了几次
    【用法】
        c = Counter(bytes_per_elem=2)
        forward(params, cfg, tokens, counter=c)
        c.total("flops")                       # 全部运算
        c.total("w", ops=MATMUL_OPS)           # 只看矩阵乘法读的权重
    【边界】bytes_per_elem 只影响字节，不影响 FLOPs。
    """

    FIELDS = ("flops", "useful", "w", "kv", "act", "calls")

    def __init__(self, bytes_per_elem=2.0):
        self.bpe = float(bytes_per_elem)
        self.rows = {}

    def add(self, op, flops=0, useful=None, w_elems=0, kv_elems=0, act_elems=0):
        """
        记一笔账。

        【参数】op：运算名；flops：运算次数；useful：有用的运算次数（None 表示和 flops 相同）；
            w_elems / kv_elems / act_elems：读写了多少个**元素**（乘以 bytes_per_elem 变成字节）
        【Python 小知识】dict.setdefault(key, 默认值) 在 key 不存在时先放进默认值，再返回它。
        """
        r = self.rows.setdefault(op, dict.fromkeys(self.FIELDS, 0))
        r["flops"] += flops
        r["useful"] += flops if useful is None else useful
        r["w"] += w_elems * self.bpe
        r["kv"] += kv_elems * self.bpe
        r["act"] += act_elems * self.bpe
        r["calls"] += 1

    def total(self, key, ops=None):
        """某一列的总和；ops 给出时只加这些运算名。"""
        return sum(r[key] for name, r in self.rows.items() if ops is None or name in ops)

    def merged(self, groups):
        """
        按分组合并：groups = {"注意力投影": ["q_proj", ...], ...}。返回 {组名: {字段: 值}}。

        【边界】不在任何组里的运算被忽略。
        """
        out = {}
        for g, names in groups.items():
            out[g] = {k: sum(self.rows[n][k] for n in names if n in self.rows) for k in self.FIELDS}
        return out


# 运算名的分组（打印表格和比较公式时用）
ATTN_PROJ_OPS = ("q_proj", "k_proj", "v_proj", "o_proj")
MLP_OPS = ("mlp_gate", "mlp_up", "mlp_down")
MATMUL_OPS = ATTN_PROJ_OPS + MLP_OPS + ("lm_head",)          # 和 llm_math.matmul_params 对应的那些矩阵
ATTN_CORE_OPS = ("attn_core",)
ELEMENTWISE_OPS = ("norm", "rope", "softmax", "swiglu", "residual")
OP_GROUPS = {
    "词嵌入（查表）": ["embed"],
    "注意力投影 Wq/Wk/Wv/Wo": list(ATTN_PROJ_OPS),
    "注意力核心 QKᵀ·PV": list(ATTN_CORE_OPS),
    "写 KV cache": ["kv_write"],
    "MLP（gate/up/down）": list(MLP_OPS),
    "lm_head": ["lm_head"],
    "逐元素（norm/rope/softmax/…）": list(ELEMENTWISE_OPS),
}


# ---------------------------------------------------------------------------
# 基本运算
# ---------------------------------------------------------------------------
def linear(x, W, op, counter=None):
    """
    y = x·W，并记账。

    【参数】x：[..., K]（前面任意多维，会被当成 M 行）；W：[K, N]；op：运算名；counter：Counter 或 None
    【返回】[..., N]
    【Python 小知识】x @ W 是矩阵乘法；x 是多维数组时，@ 作用在最后两维（前面的维度当作"批"）。
        x.size // x.shape[-1] 就是"前面所有维度的乘积" = 行数 M。
    """
    y = x @ W
    if counter is not None:
        M = x.size // x.shape[-1]
        K, N = W.shape
        counter.add(op, flops=2 * M * K * N, w_elems=K * N, act_elems=M * K + M * N)
    return y


def rms_norm(x, w, eps, counter=None):
    """
    RMSNorm：x / sqrt(mean(x²) + eps) · w（沿最后一维）。

    【为什么不用 LayerNorm】LayerNorm 还要减均值；RMSNorm 只做缩放，少一次规约（reduction），
        效果差不多，Llama 系列都用它。
    【边界】eps 防止全零向量除以 0。
    """
    if counter is not None:
        counter.add("norm", flops=NORM_FLOPS * x.size, act_elems=2 * x.size)
    return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + eps) * w


def rope_angles(positions, d_head, theta=10000.0):
    """
    RoPE 的旋转角：位置 p、第 i 对（i = 0 … d_head/2 − 1）的角度 = p · θ^(−2i/d_head)。

    【参数】positions：一维整数数组（每个 token 的绝对位置）；d_head：每头维度；theta：基数
    【返回】(cos, sin)，形状都是 [len(positions), d_head/2]
    【直觉】每一对数是一个"钟表指针"，不同的对转速不同（频率按几何级数递减）：前几对转得快，
        能分辨相邻位置；后几对转得慢，能分辨远距离。和原论文正弦位置编码的频率完全一样。
    """
    inv_freq = theta ** (-np.arange(0, d_head, 2) / d_head)
    ang = np.asarray(positions, dtype=np.float64)[:, None] * inv_freq[None, :]
    return np.cos(ang), np.sin(ang)


def apply_rope(x, cos, sin, counter=None):
    """
    把 RoPE 作用在 x 上：x 的形状 [B, T, 头数, d_head]，第 2i、2i+1 两个数当作二维向量旋转对应角度。

    【思路】(a, b) 旋转 φ：(a·cosφ − b·sinφ, a·sinφ + b·cosφ)。query 在位置 m 转 mφ，key 在位置 n 转 nφ，
        点积只和 (m − n)φ 有关——注意力分数只依赖**相对位置**（02 脚本演示 5 验证）。
    【边界】cos/sin 的第一维必须等于 T（每个 token 一行角度）。这里用"相邻两个数一对"（RoFormer 论文的写法）；
        Hugging Face 的 Llama 实现用"前一半和后一半配对"，数学上等价（只是维度的排列不同）。
    【Python 小知识】x[..., 0::2] 取最后一维的偶数位置，x[..., 1::2] 取奇数位置；
        cos[None, :, None, :] 加两个长度为 1 的维度，好和 [B, T, 头数, d/2] 广播（broadcast）。
    """
    if counter is not None:
        counter.add("rope", flops=ROPE_FLOPS * x.size, act_elems=2 * x.size)
    c = cos[None, :, None, :]
    s = sin[None, :, None, :]
    a, b = x[..., 0::2], x[..., 1::2]
    out = np.empty_like(x)
    out[..., 0::2] = a * c - b * s
    out[..., 1::2] = a * s + b * c
    return out


def softmax(x, axis=-1):
    """
    数值稳定的 softmax：先减去最大值再 exp（前置讲 P3 §5.1）。

    【边界】整行都是 −∞（全部被掩掉）时会得到 nan；因果掩码保证每个 query 至少能看到自己，不会出现。
    """
    m = np.max(x, axis=axis, keepdims=True)
    e = np.exp(x - m)
    return e / np.sum(e, axis=axis, keepdims=True)


def silu(x):
    """SiLU(x) = x · sigmoid(x)（也叫 swish）。"""
    return x / (1.0 + np.exp(-x))


def attention_core(q, k, v, q_pos0, counter=None, causal=True):
    """
    注意力的核心：softmax(Q·Kᵀ/√dh + 掩码)·V，支持 GQA。

    【参数】
        q：[B, T, H, dh]（T 个新 token 的 query，绝对位置是 q_pos0, q_pos0+1, …）
        k, v：[B, S, G, dh]（S 个 key/value：缓存里的旧 token + 这次的新 token，位置 0 … S−1）
        q_pos0：第一个 query 的绝对位置
        causal：True 时第 i 个 query 只能看位置 ≤ q_pos0 + i 的 key
    【返回】[B, T, H·dh]
    【思路（GQA）】把 H 个 query 头分成 G 组、每组 R = H/G 个：q 变形成 [B, T, G, R, dh]，
        每组的 R 个头和同一组 K/V 做注意力。G = H 时就是普通多头注意力（MHA），G = 1 时是多查询注意力（MQA）。
    【记账】见模块开头：flops = 4·B·T·S·H·dh（整个矩阵都算了），useful 只算掩码允许的部分；
        kv = 读一遍 K 和 V；act = 读 q、写输出。分数矩阵 [B, H, T, S] 不计入字节（假设融合内核）。
    【Python 小知识】np.einsum("btgrd,bsgd->bgrts", q, k) 用字母描述"哪些维度相乘、哪些维度求和"：
        出现在输入但不在输出里的字母（这里是 d）被求和。这就是 How to Scale Your Model 用的 einsum 记号（§5）。
    """
    B, T, H, dh = q.shape
    S, G = k.shape[1], k.shape[2]
    R = H // G
    qg = q.reshape(B, T, G, R, dh)
    scores = np.einsum("btgrd,bsgd->bgrts", qg, k) / math.sqrt(dh)
    if causal:
        qpos = q_pos0 + np.arange(T)
        allowed = np.arange(S)[None, :] <= qpos[:, None]          # [T, S]
        scores = np.where(allowed[None, None, None], scores, -np.inf)
    p = softmax(scores, axis=-1)
    out = np.einsum("bgrts,bsgd->btgrd", p, v).reshape(B, T, H * dh)
    if counter is not None:
        flops = 4 * B * T * S * H * dh
        if causal:
            n_allowed = int(np.sum(np.minimum(q_pos0 + np.arange(T) + 1, S)))
        else:
            n_allowed = T * S
        counter.add("attn_core", flops=flops, useful=4 * B * n_allowed * H * dh,
                    kv_elems=2 * B * S * G * dh, act_elems=2 * B * T * H * dh)
        counter.add("softmax", flops=SOFTMAX_FLOPS * B * H * T * S, useful=SOFTMAX_FLOPS * B * H * n_allowed)
    return out


# ---------------------------------------------------------------------------
# KV cache
# ---------------------------------------------------------------------------
class KVCache:
    """
    预分配的 KV cache：每层一块 [B, max_len, G, dh] 的 K 和一块同样大的 V。

    【为什么能缓存】因果掩码保证：位置 j 的 K、V 只依赖 token 0 … j，后来的 token 不会改变它们。
        所以生成第 t 个 token 时，前 t−1 个 token 的 K、V 可以直接复用，不必重算（04 脚本验证）。
    【字段】k, v：[L, B, max_len, G, dh]；length：已经写入了多少个位置（批里所有序列一样长）
    【边界】写入超过 max_len 时抛出 ValueError（真实引擎会换出、重算或拒绝请求，第 7 讲）。
    """

    def __init__(self, cfg, batch, max_len=None, dtype=np.float64):
        max_len = max_len or cfg.max_seq
        shape = (cfg.n_layers, batch, max_len, cfg.n_kv_heads, cfg.d_head)
        self.k = np.zeros(shape, dtype=dtype)
        self.v = np.zeros(shape, dtype=dtype)
        self.length = 0
        self.max_len = max_len
        self.cfg = cfg

    def write(self, layer, k_new, v_new):
        """把这一层新 token 的 K、V 写到位置 length … length+T−1（length 等整次前向结束再推进）。"""
        T = k_new.shape[1]
        s = self.length
        if s + T > self.max_len:
            raise ValueError(f"KV cache 满了：已有 {s}，再写 {T} 个超过 {self.max_len}")
        self.k[layer, :, s:s + T] = k_new
        self.v[layer, :, s:s + T] = v_new

    def read(self, layer, upto):
        """读这一层位置 0 … upto−1 的 K、V（numpy 切片是"视图"，不复制数据）。"""
        return self.k[layer, :, :upto], self.v[layer, :, :upto]

    def nbytes(self, bytes_per_elem=2.0):
        """已经写入的部分按 bytes_per_elem 算占多少字节 = 2 · L · B · length · G · dh · 字节数。"""
        L, B, _, G, dh = self.k.shape
        return 2 * L * B * self.length * G * dh * bytes_per_elem


# ---------------------------------------------------------------------------
# 前向计算
# ---------------------------------------------------------------------------
def forward(params, cfg, tokens, cache=None, counter=None, last_only=False, linear_fn=None, use_rope=True,
            causal=True, return_hidden=False, rope_pos0=None):
    """
    对一批 token 做一次前向计算。没有 cache 时就是"整段算"（prefill / 训练时的前向）；
    有 cache 时，新 token 接在 cache 里已有的 token 后面（decode 时 T = 1）。

    【参数】
        params, cfg：参数和配置
        tokens   ：整数数组 [B, T]
        cache    ：KVCache 或 None
        counter  ：Counter 或 None
        last_only：True 时只对最后一个位置算 lm_head（生成时只需要最后一个 token 的 logits，省掉 (T−1)/T 的 lm_head）
        linear_fn：可选，替换"x·W"的函数 linear_fn(x, W, op, layer) → y（第 3 讲用它插入量化的矩阵乘法）；
                   记账仍按普通矩阵乘法记
        use_rope ：False 时不加位置编码（02 脚本演示"没有位置信息会怎样"）
        causal   ：False 时用双向注意力（同上）
        return_hidden：True 时额外返回最后一层的隐藏状态（final norm 之前）
        rope_pos0：RoPE 用的起始位置；None 时等于 cache 里已有的长度（正确做法）。真实引擎会显式传入每个 token 的
                   位置（position ids）；04 脚本用它演示"位置编号写错"这个常见 bug。
    【返回】logits：[B, T, vocab]（last_only 时 [B, 1, vocab]）；return_hidden 时返回 (logits, hidden)
    【边界】有 cache 时 tokens 的 B 必须等于 cache 的 B；位置从 cache.length 开始编号。
    """
    B, T = tokens.shape
    pos0 = cache.length if cache is not None else 0
    d, H, G, dh = cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.d_head

    def lin(x, W, op, layer):
        if linear_fn is None:
            return linear(x, W, op, counter)
        y = linear_fn(x, W, op, layer)
        if counter is not None:
            M = x.size // x.shape[-1]
            counter.add(op, flops=2 * M * W.shape[0] * W.shape[1], w_elems=W.size, act_elems=x.size + y.size)
        return y

    x = params["embed"][tokens]                              # 查表：[B, T, d]
    if counter is not None:
        counter.add("embed", w_elems=B * T * d, act_elems=B * T * d)
    rp0 = pos0 if rope_pos0 is None else rope_pos0
    cos, sin = rope_angles(np.arange(rp0, rp0 + T), dh, cfg.rope_theta)

    for li, lp in enumerate(params["layers"]):
        h = rms_norm(x, lp["attn_norm"], cfg.norm_eps, counter)
        q = lin(h, lp["wq"], "q_proj", li).reshape(B, T, H, dh)
        k = lin(h, lp["wk"], "k_proj", li).reshape(B, T, G, dh)
        v = lin(h, lp["wv"], "v_proj", li).reshape(B, T, G, dh)
        if use_rope:
            q = apply_rope(q, cos, sin, counter)
            k = apply_rope(k, cos, sin, counter)
        if cache is not None:
            cache.write(li, k, v)
            if counter is not None:
                counter.add("kv_write", kv_elems=k.size + v.size)
            K, V = cache.read(li, pos0 + T)
        else:
            K, V = k, v
        a = attention_core(q, K, V, pos0, counter, causal=causal)
        x = x + lin(a, lp["wo"], "o_proj", li)
        if counter is not None:
            counter.add("residual", flops=RESID_FLOPS * x.size, act_elems=3 * x.size)
        h = rms_norm(x, lp["mlp_norm"], cfg.norm_eps, counter)
        g = lin(h, lp["w_gate"], "mlp_gate", li)
        u = lin(h, lp["w_up"], "mlp_up", li)
        m = silu(g) * u
        if counter is not None:
            counter.add("swiglu", flops=SWIGLU_FLOPS * m.size, act_elems=3 * m.size)
        x = x + lin(m, lp["w_down"], "mlp_down", li)
        if counter is not None:
            counter.add("residual", flops=RESID_FLOPS * x.size, act_elems=3 * x.size)

    if cache is not None:
        cache.length += T
    hidden = x
    if last_only:
        x = x[:, -1:, :]
    x = rms_norm(x, params["final_norm"], cfg.norm_eps, counter)
    logits = lin(x, lm_head_matrix(params), "lm_head", None)
    return (logits, hidden) if return_hidden else logits


def generate(params, cfg, prompt, n_new, use_cache=True, counter_prefill=None, counter_decode=None):
    """
    贪心生成：每一步取 logits 最大的 token。

    【参数】prompt：[B, T0] 整数数组；n_new：生成几个 token；use_cache：是否用 KV cache；
        counter_prefill / counter_decode：分别记 prefill 和 decode 阶段的账（可为 None）
    【返回】(tokens [B, T0 + n_new], step_logits：列表，每一步最后一个位置的 logits [B, vocab])
    【思路】
        有 cache：prefill 一次算完整段提示（并写入 cache），之后每步只算 1 个新 token。
        无 cache：每一步都把"到目前为止的整个序列"重新算一遍，只取最后一个位置的 logits。
        两者的结果应该完全相同（浮点误差以内），04 脚本验证；区别只在运算量。
    【边界】T0 + n_new 不能超过 cfg.max_seq（cache 的容量）。
    """
    B, T0 = prompt.shape
    if T0 + n_new > cfg.max_seq:
        raise ValueError(f"提示 {T0} + 生成 {n_new} 超过 max_seq={cfg.max_seq}")
    toks = prompt.copy()
    step_logits = []
    if use_cache:
        cache = KVCache(cfg, B, dtype=params["embed"].dtype)
        logits = forward(params, cfg, toks, cache=cache, counter=counter_prefill, last_only=True)[:, -1]
        for i in range(n_new):
            step_logits.append(logits)
            nxt = np.argmax(logits, axis=-1)[:, None]
            toks = np.concatenate([toks, nxt], axis=1)
            if i == n_new - 1:
                break
            logits = forward(params, cfg, nxt, cache=cache, counter=counter_decode, last_only=True)[:, -1]
    else:
        for i in range(n_new):
            c = counter_prefill if i == 0 else counter_decode
            logits = forward(params, cfg, toks, counter=c, last_only=True)[:, -1]
            step_logits.append(logits)
            nxt = np.argmax(logits, axis=-1)[:, None]
            toks = np.concatenate([toks, nxt], axis=1)
    return toks, step_logits


# ---------------------------------------------------------------------------
# 解析公式：不运行模型，直接算出每类运算的 FLOPs 和字节（和 Counter 的约定完全一致）
# ---------------------------------------------------------------------------
def analytic_ops(cfg, batch, n_new, pos0=0, bytes_per_elem=2.0, last_only=False, with_cache=True):
    """
    一次前向（B 个序列、每个序列 n_new 个新 token、已有 pos0 个缓存 token）的逐类账本，**纯公式**。

    【用途】03 脚本先在小模型上验证"公式 = Counter 实际数出来的"，再把公式用到 Llama-3.1-8B / 70B 这种
        跑不动的大模型上。
    【参数】cfg：Config 或 llm_math.ModelConfig（只用到结构字段）；batch = B；n_new = T；pos0：已有的上下文长度；
        bytes_per_elem：每个数几个字节；last_only：lm_head 是否只算最后一个位置；with_cache：是否写 KV cache
    【返回】{运算名: {"flops", "useful", "w", "kv", "act"}}，运算名和 Counter 相同（不含 calls）
    【边界】假设 causal=True、use_rope=True（和 forward 的默认值一致）。
    """
    L, d, H, G, dh, f, V = (cfg.n_layers, cfg.d_model, cfg.n_heads, cfg.n_kv_heads, cfg.d_head, cfg.d_ff,
                            cfg.vocab)
    B, T, S = batch, n_new, pos0 + n_new
    M = B * T
    bpe = bytes_per_elem
    rows = {}

    def put(op, flops=0, useful=None, w=0, kv=0, act=0):
        r = rows.setdefault(op, {"flops": 0, "useful": 0, "w": 0.0, "kv": 0.0, "act": 0.0})
        r["flops"] += flops
        r["useful"] += flops if useful is None else useful
        r["w"] += w * bpe
        r["kv"] += kv * bpe
        r["act"] += act * bpe

    def mm(op, m, k, n, times=1):
        for _ in range(times):
            put(op, flops=2 * m * k * n, w=k * n, act=m * k + m * n)

    put("embed", w=M * d, act=M * d)
    n_allowed = sum(min(pos0 + i + 1, S) for i in range(T))
    for _ in range(L):
        put("norm", flops=NORM_FLOPS * M * d, act=2 * M * d)
        mm("q_proj", M, d, H * dh)
        mm("k_proj", M, d, G * dh)
        mm("v_proj", M, d, G * dh)
        put("rope", flops=ROPE_FLOPS * M * H * dh, act=2 * M * H * dh)
        put("rope", flops=ROPE_FLOPS * M * G * dh, act=2 * M * G * dh)
        if with_cache:
            put("kv_write", kv=2 * M * G * dh)
        put("attn_core", flops=4 * B * T * S * H * dh, useful=4 * B * n_allowed * H * dh,
            kv=2 * B * S * G * dh, act=2 * B * T * H * dh)
        put("softmax", flops=SOFTMAX_FLOPS * B * H * T * S, useful=SOFTMAX_FLOPS * B * H * n_allowed)
        mm("o_proj", M, H * dh, d)
        put("residual", flops=RESID_FLOPS * M * d, act=3 * M * d)
        put("norm", flops=NORM_FLOPS * M * d, act=2 * M * d)
        mm("mlp_gate", M, d, f)
        mm("mlp_up", M, d, f)
        put("swiglu", flops=SWIGLU_FLOPS * M * f, act=3 * M * f)
        mm("mlp_down", M, f, d)
        put("residual", flops=RESID_FLOPS * M * d, act=3 * M * d)
    m_head = B if last_only else M
    put("norm", flops=NORM_FLOPS * m_head * d, act=2 * m_head * d)
    mm("lm_head", m_head, d, V)
    return rows


def sum_rows(rows, key, ops=None):
    """analytic_ops 结果的某一列求和（ops 给出时只加这些运算名）。"""
    return sum(r[key] for name, r in rows.items() if ops is None or name in ops)


# ---------------------------------------------------------------------------
# einsum 记号与 FLOPs 规则
# ---------------------------------------------------------------------------
def einsum_flops(spec, *shapes):
    """
    How to Scale Your Model 第 4 章的规则：一次缩并（contraction）的 FLOPs = 2 × 所有出现过的维度之积
    （同一个字母只算一次）；如果没有被求和的维度（纯逐元素乘法），就不乘 2。

    【参数】spec：einsum 字符串，如 "btd,df->btf"；shapes：每个输入的形状（元组）
    【返回】整数 FLOPs
    【边界】同一个字母在不同输入里的长度必须相同，否则抛出 ValueError。
    【例子】einsum_flops("btd,df->btf", (2, 8, 64), (64, 256)) = 2·2·8·64·256
    """
    ins, out = spec.replace(" ", "").split("->")
    ins = ins.split(",")
    sizes = {}
    for letters, shape in zip(ins, shapes):
        if len(letters) != len(shape):
            raise ValueError(f"{letters} 和形状 {shape} 的维数不同")
        for ch, n in zip(letters, shape):
            if sizes.setdefault(ch, n) != n:
                raise ValueError(f"维度 {ch} 的长度不一致：{sizes[ch]} vs {n}")
    total = 1
    for n in sizes.values():
        total *= n
    contracted = set("".join(ins)) - set(out)
    return 2 * total if contracted else total


# ---------------------------------------------------------------------------
# 自检
# ---------------------------------------------------------------------------
def _selftest():
    cfg = TINY
    p = init_params(cfg, seed=0)
    mc = cfg.to_model_config()
    assert count_params(p) == mc.n_params(), (count_params(p), mc.n_params())

    rng = np.random.default_rng(1)
    prompt = rng.integers(0, cfg.vocab, size=(2, 12))
    t1, l1 = generate(p, cfg, prompt, 8, use_cache=True)
    t2, l2 = generate(p, cfg, prompt, 8, use_cache=False)
    assert np.array_equal(t1, t2)
    diff = max(float(np.max(np.abs(a - b))) for a, b in zip(l1, l2))
    assert diff < 1e-9, diff

    # 记账 = 解析公式（prefill 和 decode 各一次）
    for B, T, pos0 in ((2, 12, 0), (3, 1, 17)):
        c = Counter()
        cache = KVCache(cfg, B)
        if pos0:
            forward(p, cfg, rng.integers(0, cfg.vocab, size=(B, pos0)), cache=cache)
        forward(p, cfg, rng.integers(0, cfg.vocab, size=(B, T)), cache=cache, counter=c)
        rows = analytic_ops(cfg, B, T, pos0)
        for op, r in rows.items():
            for key in ("flops", "useful", "w", "kv", "act"):
                assert abs(c.rows[op][key] - r[key]) < 1e-6, (op, key, c.rows[op][key], r[key])
    # 公式版的 FLOPs/token（llm_math）= 数出来的矩阵乘法 FLOPs / token
    c = Counter()
    forward(p, cfg, rng.integers(0, cfg.vocab, size=(1, 5)), counter=c)
    assert c.total("flops", MATMUL_OPS) == 5 * mc.linear_flops_per_token()
    assert einsum_flops("btd,df->btf", (2, 8, 64), (64, 256)) == 2 * 2 * 8 * 64 * 256
    print(f"transformer_lab 自检通过：参数量 {count_params(p):,} = 公式；带 / 不带 KV cache 的 logits 最大差 {diff:.1e}；"
          f"记账 = 解析公式")


if __name__ == "__main__":
    _selftest()