"""
ssm_lab.py —— 第 21 讲的共享模块：状态空间模型的三种写法、HiPPO、选择性扫描、线性注意力及其变体、混合模型的账
=====================================================================

【这个文件要回答的问题】
    Transformer 训练时并行，推理时要一个随长度增长的 KV 缓存；RNN 推理时只要一个固定大小的状态，训练时却只能一步一步算。
    状态空间模型（SSM）两样都要：同一个模型可以写成**连续**的微分方程、**递推**（像 RNN）和**卷积**（像 CNN，可以并行）。
    Mamba 让参数随输入变化（"选择性"），卷积写法没了，就用**并行扫描**来并行；线性注意力去掉 softmax 之后，
    也能写成一个"状态 = 外积之和"的递推——它和 SSM 是一回事。

【本文件提供】
    1. 连续与离散：discretize（Euler / 双线性 / 零阶保持 ZOH）、ssm_recurrent、ssm_kernel、causal_conv、fft_conv
    2. HiPPO：hippo_legs（S4 用的 HiPPO-LegS 矩阵）、legs_online（在线地把整段历史压进 N 个系数）、
       legendre_project（直接投影，用来核对）、legendre_reconstruct（用系数还原历史）
    3. 扫描：sequential_scan、parallel_scan（Hillis–Steele 式的对数深度扫描）——h_t = a_t ∘ h_{t−1} + b_t
    4. 线性注意力族：linear_attention_parallel / _recurrent（含固定衰减 γ、随输入变化的门 α_t）、delta_rule_recurrent
       （DeltaNet / Gated DeltaNet）、store_and_recall（不训练：固定大小的状态能存多少、能不能改写）
    5. 可训练的层与模型：S4DLayer（对角、LTI，用 FFT 卷积）、SelectiveSSM（Mamba 的 S6 核心，用并行扫描）、
       LinearAttnLayer（线性 / 门控线性 / DeltaNet，递推实现）、SoftmaxAttention、SeqModel（按层的类型列表搭混合模型）
    6. 任务：mqar_batch（多查询联想记忆，Zoology 的 MQAR）
    7. 混合模型的账：MODELS（Llama-2 7B、Mistral 7B、Mixtral、Jamba、Nemotron-H-8B、Qwen3-Next-80B 的注意力配置）、
       kv_bytes_per_token

【边界】
    - 为了能逐位核对，1–4 用 float64；5 的可训练模型用 float32。
    - S4 的真正实现（DPLR 分解、Cauchy 核）、Mamba 的硬件感知扫描、DeltaNet 的分块并行算法都没有实现，
      这里只保留"数学上等价"的最简单写法。
"""

import math

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from numpy.polynomial import legendre


# =====================================================================
# 1. 连续与离散：三种写法
# =====================================================================

def discretize(A, B, dt, method="zoh"):
    """
    把连续的 h'(t) = A h(t) + B x(t) 离散成 h_k = Ā h_{k−1} + B̄ x_k。

    【三种方法】
        euler    ：Ā = I + Δ·A，B̄ = Δ·B（最简单，一阶，Δ 大时不稳定）
        bilinear ：Ā = (I − Δ/2·A)⁻¹(I + Δ/2·A)，B̄ = (I − Δ/2·A)⁻¹·Δ·B（S4 用的；一阶 Padé 近似，二阶精度）
        zoh      ：Ā = exp(Δ·A)，B̄ = A⁻¹(exp(Δ·A) − I)·B（零阶保持：假设每个采样区间内输入不变，此时是**精确**解；Mamba 用的）
    【参数】A：(N, N)；B：(N,) 或 (N, 1)；dt：步长 Δ（标量）
    【返回】(Ā, B̄)，B̄ 形状 (N,)
    """
    N = A.shape[0]
    I = torch.eye(N, dtype=A.dtype)
    B = B.reshape(N)
    if method == "euler":
        return I + dt * A, dt * B
    if method == "bilinear":
        inv = torch.linalg.inv(I - dt / 2 * A)
        return inv @ (I + dt / 2 * A), inv @ (dt * B)
    if method == "zoh":
        Ad = torch.linalg.matrix_exp(dt * A)
        return Ad, torch.linalg.solve(A, (Ad - I) @ B)
    raise ValueError(method)


def ssm_recurrent(Ab, Bb, C, D, x):
    """
    递推写法：h_k = Ā h_{k−1} + B̄ x_k，y_k = C h_k + D x_k，h_{−1} = 0。
    【参数】x：(L,) 一维输入；C：(N,)；D：标量
    【返回】y：(L,)
    【直觉】像 RNN：一次只看一个输入、只保留 N 维的状态——生成时不需要 KV 缓存。
    """
    h = torch.zeros(Ab.shape[0], dtype=Ab.dtype)
    ys = []
    for xk in x:
        h = Ab @ h + Bb * xk
        ys.append(C @ h + D * xk)
    return torch.stack(ys)


def ssm_kernel(Ab, Bb, C, L):
    """
    卷积核 K = (C B̄, C Ā B̄, C Ā² B̄, …, C Ā^{L−1} B̄)。
    【推导】把递推展开：y_k = Σ_{j≤k} C Ā^{k−j} B̄ x_j（+ D x_k）——这正是一个长度为 L 的因果卷积 y = K ∗ x。
    """
    ks, v = [], Bb.clone()
    for _ in range(L):
        ks.append(C @ v)
        v = Ab @ v
    return torch.stack(ks)


def causal_conv(K, x):
    """直接按定义算因果卷积：y_k = Σ_{j=0}^{k} K_{k−j} x_j。O(L²)。"""
    L = len(x)
    return torch.stack([sum(K[k - j] * x[j] for j in range(k + 1)) for k in range(L)])


def fft_conv(K, x):
    """
    用 FFT 算因果卷积，O(L log L)：两边补零到 2L（避免循环卷积"绕回来"），逐频率相乘，再逆变换取前 L 个。
    【GenAI 联系】S4 训练时就是这样把一整段序列一次算完的——卷积写法让 SSM 能像 Transformer 一样并行训练。
    """
    L = len(x)
    n = 2 * L
    return torch.fft.irfft(torch.fft.rfft(K, n) * torch.fft.rfft(x, n), n)[:L]


# =====================================================================
# 2. HiPPO：把整段历史压进 N 个数
# =====================================================================

def hippo_legs(N):
    """
    HiPPO-LegS 矩阵（Gu et al., 2020；S4 用它初始化 A）：
        A_nk = −√(2n+1)·√(2k+1)（n > k），−(n+1)（n = k），0（n < k）；B_n = √(2n+1)
    【意义】状态 c(t) 是"整段历史 f(0…t) 在勒让德多项式上的展开系数"——用 N 个数概括任意长的历史，
        越早的部分和越近的部分一视同仁（均匀的测度）。
    【返回】(A (N, N), B (N,))，float64
    """
    n = torch.arange(N, dtype=torch.float64)
    r = torch.sqrt(2 * n + 1)
    A = -torch.outer(r, r)
    A = torch.tril(A, -1) - torch.diag(n + 1)
    return A, r.clone()


def legs_online(f, N):
    """
    在线地计算 HiPPO-LegS 的系数：读到第 k 个样本时（k = 1, 2, …），
        c_k = (I − A/(2k))⁻¹ [ (I + A/(2k)) c_{k−1} + (B/k)·f_k ]
    （连续形式 c'(t) = A c(t)/t + B f(t)/t 的双线性离散，A 用 hippo_legs 的负号约定；
    时间尺度随 t 变化（1/t），所以叫"scaled Legendre"：不管历史多长，都均匀地概括整段。）
    【参数】f：(L,) 的历史；N：系数个数
    【返回】(N,) 的系数（读完全部 L 个样本之后）
    """
    A, B = hippo_legs(N)
    I = torch.eye(N, dtype=torch.float64)
    c = torch.zeros(N, dtype=torch.float64)
    for k in range(1, len(f) + 1):
        M = torch.linalg.inv(I - A / (2 * k))
        c = M @ ((I + A / (2 * k)) @ c + B / k * f[k - 1])
    return c


def legendre_project(f, N):
    """直接投影（中点求积）：c_n = (1/L) Σ_s f(s)·√(2n+1)·P_n(2(s+½)/L − 1)。用来核对 legs_online。"""
    L = len(f)
    s = (np.arange(L) + 0.5) / L * 2 - 1
    cs = []
    for n in range(N):
        coef = np.zeros(n + 1)
        coef[n] = 1
        cs.append(float(np.mean(f.numpy() * math.sqrt(2 * n + 1) * legendre.legval(s, coef))))
    return torch.tensor(cs, dtype=torch.float64)


def legendre_reconstruct(c, L):
    """由系数还原历史：f(s) ≈ Σ_n c_n·√(2n+1)·P_n(2(s+½)/L − 1)，s = 0…L−1。"""
    s = (np.arange(L) + 0.5) / L * 2 - 1
    scaled = c.numpy() * np.sqrt(2 * np.arange(len(c)) + 1)
    return torch.tensor(legendre.legval(s, scaled), dtype=torch.float64)


# =====================================================================
# 3. 扫描：h_t = a_t ∘ h_{t−1} + b_t
# =====================================================================

def sequential_scan(a, b):
    """
    逐步扫描。【参数】a、b：形状 (..., L, ...)，沿第 1 维（时间）扫描；这里约定时间是 dim=1。
    【返回】h：和 b 同形状，h_t = a_t ∘ h_{t−1} + b_t，h_{−1} = 0
    """
    hs, h = [], torch.zeros_like(b[:, 0])
    for t in range(b.shape[1]):
        h = a[:, t] * h + b[:, t]
        hs.append(h)
    return torch.stack(hs, dim=1)


def parallel_scan(a, b):
    """
    对数深度的并行扫描（Hillis–Steele）。

    【思路】把一步更新看成一对 (a, b)：h ↦ a·h + b。两步合起来还是这种形式：
            先 (a₁, b₁) 再 (a₂, b₂) = (a₂a₁, a₂b₁ + b₂)——这个"合并"满足结合律。
        于是可以像前缀和一样：第 1 轮每个位置和左边相邻 1 个合并，第 2 轮和左边相邻 2 个合并，再 4 个……
        log₂L 轮之后，每个位置都合并了它左边的全部，(a, b) 中的 b 就是 h_t。每一轮都是整个张量一起算，可以并行。
    【返回】h，和 sequential_scan 的结果相同（只差浮点舍入）
    【GenAI 联系】Mamba 的参数随输入变化，没有固定的卷积核，不能用 FFT；它靠的就是这样的并行扫描（在 GPU 的 SRAM 里做）。
    """
    a, b = a.clone(), b.clone()
    L, off = b.shape[1], 1
    while off < L:
        a_prev = torch.cat([torch.ones_like(a[:, :off]), a[:, :-off]], dim=1)
        b_prev = torch.cat([torch.zeros_like(b[:, :off]), b[:, :-off]], dim=1)
        b = a * b_prev + b
        a = a * a_prev
        off *= 2
    return b


# =====================================================================
# 4. 线性注意力族
# =====================================================================

def linear_attention_parallel(Q, K, V, decay=None):
    """
    并行写法：O = (QKᵀ ⊙ M) V，M 是因果掩码（slides 的公式：去掉 softmax 的注意力）。
    decay：None（M_ts = 1，t ≥ s）；一个数 γ（RetNet：M_ts = γ^{t−s}）；或长度 L 的向量 α（随输入变化的门：
        M_ts = α_{s+1}·…·α_t，GLA / Mamba-2 的标量衰减）。
    【参数】Q、K：(L, d_k)；V：(L, d_v)
    """
    L = Q.shape[0]
    S = Q @ K.T
    t = torch.arange(L).unsqueeze(1)
    s = torch.arange(L).unsqueeze(0)
    if decay is None:
        M = (s <= t).to(Q.dtype)
    elif isinstance(decay, float):
        M = torch.where(s <= t, decay ** (t - s).to(Q.dtype), torch.zeros((), dtype=Q.dtype))
    else:
        logc = torch.cumsum(torch.log(decay), 0)               # log(α_1⋯α_t)
        M = torch.where(s <= t, torch.exp(logc[t] - logc[s]), torch.zeros((), dtype=Q.dtype))
    return (S * M) @ V


def linear_attention_recurrent(Q, K, V, decay=None):
    """
    递推写法：S_t = α_t·S_{t−1} + k_t v_tᵀ（d_k × d_v 的状态），o_t = S_tᵀ q_t。
    和 linear_attention_parallel 的结果相同——这就是"线性注意力是一个 RNN"（Katharopoulos et al., 2020），
    也是"线性注意力 = SSM"的含义：状态 S_t 的大小固定，和已经读过多少 token 无关。
    """
    L, dk = K.shape
    S = torch.zeros(dk, V.shape[1], dtype=Q.dtype)
    out = []
    for t in range(L):
        a = 1.0 if decay is None else (decay if isinstance(decay, float) else decay[t])
        S = a * S + torch.outer(K[t], V[t])
        out.append(S.T @ Q[t])
    return torch.stack(out)


def delta_rule_recurrent(Q, K, V, beta, alpha=None):
    """
    DeltaNet（Schlag et al., 2021；Yang et al., 2024）：写入前先"擦掉"这个键原来对应的值：
        S_t = S_{t−1} + β_t (v_t − S_{t−1}ᵀ k_t) k_tᵀ        （S 是 d_k × d_v，k_t 归一化）
            = (I − β_t k_t k_tᵀ) S_{t−1} + β_t k_t v_tᵀ
    Gated DeltaNet（Qwen3-Next 用的）再乘一个衰减门：S_t = α_t (I − β_t k_t k_tᵀ) S_{t−1} + β_t k_t v_tᵀ。
    【参数】beta：(L,)，写入强度；alpha：None 或 (L,)
    【返回】(L, d_v)
    【直觉】普通线性注意力只会"往状态里加"，同一个键写两次就把两个值混在一起；delta 规则像一个联想记忆的"覆盖写"。
    """
    L, dk = K.shape
    S = torch.zeros(dk, V.shape[1], dtype=Q.dtype)
    out = []
    for t in range(L):
        k = K[t]
        a = 1.0 if alpha is None else alpha[t]
        S = a * (S - beta[t] * torch.outer(k, S.T @ k)) + beta[t] * torch.outer(k, V[t])
        out.append(S.T @ Q[t])
    return torch.stack(out)


def store_and_recall(n, d, rule, seed=0, passes=1, rewrite=0):
    """
    不训练：把 n 对随机的（单位长度的）键 k_i 和值 v_i 写进一个 d × d 的状态 S，再用每个键去读 S k_i。
    【参数】rule："sum"（线性注意力：S = Σ k_i v_iᵀ）、"delta"（delta 规则，β = 1；passes：把全部键值重复写几遍）；
        rewrite：写完之后，再把前 rewrite 个键各写一个**新**的值（"这个键的值改了"）
    【返回】字典：acc —— 读出来的向量和哪个候选值最像（余弦最大），选对的比例（rewrite > 0 时，被改过的键要选中新值）；
        cos —— 读出来的向量和正确的值的平均余弦
    【对照】softmax 注意力存的是全部键值（KV 缓存），总能读回——代价是存储随 n 增长；这里的状态大小固定是 d × d。
    """
    g = torch.Generator().manual_seed(seed)
    K = F.normalize(torch.randn(n, d, generator=g, dtype=torch.float64), dim=1)
    V = F.normalize(torch.randn(n, d, generator=g, dtype=torch.float64), dim=1)
    V_new = F.normalize(torch.randn(max(rewrite, 1), d, generator=g, dtype=torch.float64), dim=1)[:rewrite]

    def write(S, k, v):
        return S + torch.outer(k, v) if rule == "sum" else S + torch.outer(k, v - S.T @ k)

    S = torch.zeros(d, d, dtype=torch.float64)
    for _ in range(passes if rule == "delta" else 1):
        for i in range(n):
            S = write(S, K[i], V[i])
    for i in range(rewrite):
        S = write(S, K[i], V_new[i])
    truth = V.clone()
    truth[:rewrite] = V_new
    cands = torch.cat([V, V_new])                              # 候选：全部旧值和新值
    R = F.normalize(K @ S, dim=1)
    pick = (R @ cands.T).argmax(1)
    want = torch.arange(n)
    want[:rewrite] = n + torch.arange(rewrite)
    return {"acc": float(pick.eq(want).double().mean()), "cos": float((R * truth).sum(1).mean())}


# =====================================================================
# 5. 可训练的层与模型
# =====================================================================

class S4DLayer(nn.Module):
    """
    对角的线性时不变 SSM（S4D，Gu et al., 2022）：每个通道一个独立的 SSM（slides："像注意力的多个头、卷积的多个通道"）。
        Λ_n = −½ + iπn（S4D-Lin 初始化），B = 1，C 复数可学，步长 Δ = exp(log_dt) 可学；ZOH：Λ̄ = exp(ΔΛ)，B̄ = (Λ̄ − 1)/Λ
        卷积核 K_l = 2·Re(Σ_n C_n B̄_n Λ̄_n^l)（共轭对合起来是实数）；y = K ∗ x + D·x，用 FFT 一次算完整段。
    【参数】H：通道数；N：状态维度（复数模式 N/2 个）
    【边界】参数和输入无关（LTI）——同一个卷积核作用在每个位置上，看不到"现在读到的是什么"。
    """

    def __init__(self, H, N=16):
        super().__init__()
        n = N // 2
        self.log_dt = nn.Parameter(torch.empty(H).uniform_(math.log(1e-3), math.log(1e-1)))
        self.A_re = nn.Parameter(torch.full((H, n), math.log(0.5)))          # Re(Λ) = −exp(A_re)
        self.register_buffer("A_im", math.pi * torch.arange(n).float().expand(H, n).clone())
        self.C = nn.Parameter(torch.randn(H, n, 2) * 0.5)
        self.D = nn.Parameter(torch.ones(H))

    def kernel(self, L):
        lam = torch.complex(-torch.exp(self.A_re), self.A_im)                 # (H, n)
        dt = torch.exp(self.log_dt).unsqueeze(-1)
        lam_bar = torch.exp(dt * lam)
        Bbar = (lam_bar - 1) / lam
        C = torch.view_as_complex(self.C)
        powers = lam_bar.unsqueeze(-1) ** torch.arange(L)                    # (H, n, L)
        return 2 * torch.einsum("hn,hnl->hl", C * Bbar, powers).real

    def forward(self, x):                                                     # x: (B, L, H)
        L = x.shape[1]
        K = self.kernel(L)
        n = 2 * L
        y = torch.fft.irfft(torch.fft.rfft(x.transpose(1, 2), n) * torch.fft.rfft(K, n), n)[..., :L]
        return y.transpose(1, 2) + x * self.D


class SelectiveSSM(nn.Module):
    """
    选择性 SSM（Mamba 的 S6 核心，Gu & Dao, 2023）：Δ、B、C 都由当前输入算出来：
        Δ_t = softplus(W_Δ x_t + b)（每个通道一个，初始约 0.05），B_t = W_B x_t，C_t = W_C x_t（每个 N 维），
        A = −exp(A_log)（对角，初始 −0.1·(1…N)）
        h_t = exp(Δ_t A) ∘ h_{t−1} + Δ_t B_t x_t，y_t = C_t · h_t + D x_t
    【为什么叫"选择性"】Δ_t 大：几乎忘掉旧状态、记住当前输入；Δ_t 小：几乎不写入、保持旧状态。
        模型可以根据**内容**决定记什么、忘什么——LTI 的 S4 做不到（它对每个位置一视同仁）。
    【实现】用 parallel_scan（对数深度）一次算完整段。
    """

    def __init__(self, H, N=16):
        super().__init__()
        self.N = N
        self.dt_proj = nn.Linear(H, H)
        nn.init.constant_(self.dt_proj.bias, math.log(math.expm1(0.05)))
        self.B_proj = nn.Linear(H, N, bias=False)
        self.C_proj = nn.Linear(H, N, bias=False)
        # A = −0.1·(1…N)：比 Mamba / S4D-Real 的 −(1…N) 慢 10 倍。用原来的初始化（以及 Mamba 的 Δ ~ 对数均匀 [0.001, 0.1]），
        # 我们的小模型在 MQAR 上 2400 步都学不会：Δ ≈ 0.05 时大部分状态维度每步衰减很快，20 个 token 之前写入的键值早就忘了（README）。
        self.A_log = nn.Parameter(torch.log(0.1 * torch.arange(1, N + 1).float()).expand(H, N).clone())
        self.D = nn.Parameter(torch.ones(H))

    def forward(self, x):                                                     # x: (B, L, H)
        dt = F.softplus(self.dt_proj(x))                                       # (B, L, H)
        A = -torch.exp(self.A_log)                                             # (H, N)
        a = torch.exp(dt.unsqueeze(-1) * A)                                    # (B, L, H, N)
        b = dt.unsqueeze(-1) * self.B_proj(x).unsqueeze(2) * x.unsqueeze(-1)  # (B, L, H, N)
        h = parallel_scan(a, b)
        return (h * self.C_proj(x).unsqueeze(2)).sum(-1) + x * self.D


class LinearAttnLayer(nn.Module):
    """
    线性注意力族的一层（多头）：kind =
        "linear" —— S_t = S_{t−1} + φ(k_t) v_tᵀ，φ = elu + 1（Katharopoulos et al., 2020），输出再除以 φ(q)ᵀΣφ(k) 归一化；
        "gated"  —— S_t = α_t S_{t−1} + k_t v_tᵀ，α_t = sigmoid(·) 由输入决定（GLA / Mamba-2 式的标量门）；
        "delta"  —— Gated DeltaNet：S_t = α_t (I − β_t k_t k_tᵀ) S_{t−1} + β_t k_t v_tᵀ，k 归一化，β_t = sigmoid(·)
    【参数】d_model、n_heads
    【实现】linear、gated 用并行写法（掩码矩阵）训练；delta 只用了最简单的逐步递推（真实实现有分块并行的算法）。
    """

    def __init__(self, d_model, n_heads, kind="linear"):
        super().__init__()
        self.h, self.dk, self.kind = n_heads, d_model // n_heads, kind
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.gate = nn.Linear(d_model, 2 * n_heads)
        nn.init.constant_(self.gate.bias, 3.0)                                 # α 一开始接近 1（几乎不遗忘）
        self.out = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, L, D = x.shape
        q, k, v = self.qkv(x).view(B, L, 3, self.h, self.dk).unbind(2)        # (B, L, h, dk)
        g = self.gate(x).view(B, L, 2, self.h)
        alpha, beta = torch.sigmoid(g[:, :, 0]), torch.sigmoid(g[:, :, 1])     # (B, L, h)
        if self.kind == "linear":
            q, k = F.elu(q) + 1, F.elu(k) + 1
        else:
            q, k = F.normalize(q, dim=-1), F.normalize(k, dim=-1)
        if self.kind in ("linear", "gated"):
            # 并行写法：O = (Q Kᵀ ⊙ M) V，M_ts = α_{s+1}⋯α_t（linear 时 M 就是因果掩码）——训练时比逐步递推快得多
            q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)   # (B, h, L, dk)
            causal = torch.tril(torch.ones(L, L, dtype=torch.bool))
            if self.kind == "gated":
                logc = torch.cumsum(torch.log(alpha.transpose(1, 2)), -1)       # (B, h, L)
                M = torch.exp((logc.unsqueeze(-1) - logc.unsqueeze(-2)).masked_fill(~causal, -1e9))
            else:
                M = causal.to(x.dtype)
            W = (q @ k.transpose(-2, -1)) * M
            o = W @ v
            if self.kind == "linear":
                o = o / (W.sum(-1, keepdim=True) + 1e-6)
            return self.out(o.transpose(1, 2).reshape(B, L, D))
        S = x.new_zeros(B, self.h, self.dk, self.dk)                           # DeltaNet：逐步递推
        outs = []
        for t in range(L):
            kt, vt, qt = k[:, t], v[:, t], q[:, t]                             # (B, h, dk)
            a, bt = alpha[:, t, :, None, None], beta[:, t, :, None, None]
            pred = torch.einsum("bhk,bhkv->bhv", kt, S)
            S = a * (S - bt * kt.unsqueeze(-1) * pred.unsqueeze(-2)) + bt * kt.unsqueeze(-1) * vt.unsqueeze(-2)
            outs.append(torch.einsum("bhk,bhkv->bhv", qt, S))
        return self.out(torch.stack(outs, 1).reshape(B, L, D))


class SoftmaxAttention(nn.Module):
    """普通的因果多头注意力（可学习的位置嵌入在模型里加）。"""

    def __init__(self, d_model, n_heads):
        super().__init__()
        self.h, self.dk = n_heads, d_model // n_heads
        self.qkv = nn.Linear(d_model, 3 * d_model, bias=False)
        self.out = nn.Linear(d_model, d_model, bias=False)

    def forward(self, x):
        B, L, D = x.shape
        q, k, v = self.qkv(x).view(B, L, 3, self.h, self.dk).permute(2, 0, 3, 1, 4)
        o = F.scaled_dot_product_attention(q, k, v, is_causal=True)
        return self.out(o.transpose(1, 2).reshape(B, L, D))


class ShortConv(nn.Module):
    """长度为 4 的因果深度卷积（Mamba、H3 等都在 SSM 前面放一个）：让每个位置先看到前几个 token。"""

    def __init__(self, H, width=4):
        super().__init__()
        self.conv = nn.Conv1d(H, H, width, groups=H, padding=width - 1)

    def forward(self, x):
        return self.conv(x.transpose(1, 2))[..., :x.shape[1]].transpose(1, 2)


class SeqModel(nn.Module):
    """
    通用的序列模型：词嵌入（+ 可学习的位置嵌入）→ 若干层 [x += 混合器(LN(x))；x += MLP(LN(x))] → LN → 输出层。
    【参数】layers：每层的混合器类型列表，"attn" / "s4d" / "mamba" / "linear" / "gated" / "delta"
        （线性注意力类的层可以加后缀 "-noconv" 去掉前面的短卷积，做对照）；
        例如 ["mamba", "attn"] 就是一个混合模型（Jamba、Nemotron-H 的思路）
        d_model、n_heads、N（SSM 的状态维度）、max_len
    【细节】SSM 类和线性注意力类的层前面都加一个短卷积（ShortConv），和 Mamba、DeltaNet 的常见实现一样：
        它让每个位置直接看到前 3 个 token（例如"值"的位置看到它前面的"键"）。不加时，普通线性注意力和 DeltaNet
        在联想记忆上学不会（README）；注意力不需要它（自己就能看任意位置）。
    """

    def __init__(self, vocab, layers, d_model=64, n_heads=2, N=16, max_len=128):
        super().__init__()
        self.tok = nn.Embedding(vocab, d_model)
        self.pos = nn.Embedding(max_len, d_model)
        self.kinds = list(layers)
        mixers = []
        for kind in layers:
            if kind == "attn":
                mixers.append(SoftmaxAttention(d_model, n_heads))
            elif kind == "s4d":
                mixers.append(nn.Sequential(ShortConv(d_model), S4DLayer(d_model, N)))
            elif kind == "mamba":
                mixers.append(nn.Sequential(ShortConv(d_model), SelectiveSSM(d_model, N)))
            elif kind.endswith("-noconv"):                                    # 对照：不加短卷积
                mixers.append(LinearAttnLayer(d_model, n_heads, kind[:-len("-noconv")]))
            else:
                mixers.append(nn.Sequential(ShortConv(d_model), LinearAttnLayer(d_model, n_heads, kind)))
        self.mixers = nn.ModuleList(mixers)
        self.ln1 = nn.ModuleList(nn.LayerNorm(d_model) for _ in layers)
        self.ln2 = nn.ModuleList(nn.LayerNorm(d_model) for _ in layers)
        self.mlps = nn.ModuleList(nn.Sequential(nn.Linear(d_model, 4 * d_model), nn.GELU(), nn.Linear(4 * d_model, d_model))
                                  for _ in layers)
        self.ln_f = nn.LayerNorm(d_model)
        self.head = nn.Linear(d_model, vocab, bias=False)

    def forward(self, ids):
        x = self.tok(ids) + self.pos(torch.arange(ids.shape[1]))
        for mix, ln1, ln2, mlp in zip(self.mixers, self.ln1, self.ln2, self.mlps):
            x = x + mix(ln1(x))
            x = x + mlp(ln2(x))
        return self.head(self.ln_f(x))


# =====================================================================
# 6. 任务：多查询联想记忆（MQAR）
# =====================================================================

MQAR_KEYS, MQAR_VALUES = 32, 32
MQAR_VOCAB = 2 + MQAR_KEYS + MQAR_VALUES                                     # <pad>=0、<s>=1、键、值


def mqar_batch(seed, B, n_pairs, n_queries, L):
    """
    多查询联想记忆（Arora et al., 2023 的 MQAR）：
        <s> k₁ v₁ k₂ v₂ … k_n v_n  k_a v_a  k_b v_b …    （先给 n 对键值，再问若干个键，键是有放回地抽的）
    只在"问的键"之后预测值的位置计损失。键值配对每个序列都不同，只能从上下文里查。
    【返回】(ids (B, L), targets (B, L))；序列不够 L 时后面补 <pad>
    """
    import random
    rng = random.Random(seed)
    ids = torch.zeros(B, L, dtype=torch.long)
    tgt = torch.full((B, L), -100, dtype=torch.long)
    for b in range(B):
        keys = rng.sample(range(MQAR_KEYS), n_pairs)
        vals = [rng.randrange(MQAR_VALUES) for _ in keys]
        seq = [1]
        for k, v in zip(keys, vals):
            seq += [2 + k, 2 + MQAR_KEYS + v]
        for _ in range(n_queries):
            i = rng.randrange(n_pairs)
            seq.append(2 + keys[i])
            tgt[b, len(seq) - 1] = 2 + MQAR_KEYS + vals[i]
            seq.append(2 + MQAR_KEYS + vals[i])
        ids[b, :len(seq)] = torch.tensor(seq[:L])
    return ids, tgt


# =====================================================================
# 7. 混合模型的账
# =====================================================================

# (名字, 层数, 其中注意力层数, KV 头数, 头维度, 全部参数（十亿）, 激活参数（十亿）)
MODELS = [
    ("Llama-2 7B", 32, 32, 32, 128, 6.7, 6.7),
    ("Mistral 7B", 32, 32, 8, 128, 7.2, 7.2),
    ("Mixtral 8x7B", 32, 32, 8, 128, 46.7, 12.9),
    ("Jamba", 32, 4, 8, 128, 52.0, 12.0),
    ("Nemotron-H-8B", 52, 4, 8, 128, 8.0, 8.0),
    ("Qwen3-Next-80B-A3B", 48, 12, 2, 256, 80.0, 3.0),
]
"""
注意力配置来自各自的论文或模型卡：Llama-2 7B 是普通多头注意力（32 个 KV 头）；Mistral / Mixtral 用 GQA（8 个 KV 头）；
Jamba 每 8 层里 1 层注意力（共 4 层）；Nemotron-H-8B 52 层里 4 层是自注意力；Qwen3-Next 48 层里每 4 层 1 层门控注意力
（16 个查询头、2 个 KV 头、头维度 256）。
"""


def kv_bytes_per_token(n_attn_layers, n_kv_heads, head_dim, bytes_per=2):
    """每个 token 的 KV 缓存 = 2（K 和 V）× 注意力层数 × KV 头数 × 头维度 × 每个数的字节数。线性层 / SSM 层不随长度增长。"""
    return 2 * n_attn_layers * n_kv_heads * head_dim * bytes_per