🔍 2026 五大注意力 · 从 MHA 到 MLA 精读版

class MHA(nn.Module):   # 对应 图MHA-1 的蓝色区域
    def __init__(self, d_model=4096, n_heads=32):
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads          # 128
        self.W_q = nn.Linear(d_model, d_model)      # 图MHA-2 的 Q 投影
        self.W_k = nn.Linear(d_model, d_model)      # 图MHA-2 的 K 投影（同样 32 头，浪费点 1）
        self.W_v = nn.Linear(d_model, d_model)      # 图MHA-2 的 V 投影（同样 32 头，浪费点 2）
        self.W_o = nn.Linear(d_model, d_model)      # 图MHA-1 的 O 投影

    def forward(self, x, past_kv=None):
        B, T, _ = x.shape
        # 投影：每个 token 生成自己的一套 H=32 组 Q/K/V
        Q = self.W_q(x).view(B,T,self.n_heads,self.head_dim).transpose(1,2)  # (B,32,T,128)
        K = self.W_k(x).view(B,T,self.n_heads,self.head_dim).transpose(1,2)  # (B,32,T,128)
        V = self.W_v(x).view(B,T,self.n_heads,self.head_dim).transpose(1,2)  # (B,32,T,128)
        # 图MHA-3 推理：把新 token 的 K/V 追加进缓存（缓存长度 = L，随 L 线性增长）
        if past_kv:
            K = torch.cat([past_kv[0], K], dim=2)   # (B,32,L,128)
            V = torch.cat([past_kv[1], V], dim=2)
        # 图MHA-2 注意力：Q·K^T/√d → softmax → ·V（Q 扫过全部 L 个历史位置）
        scores = Q @ K.transpose(-2,-1) / math.sqrt(self.head_dim)  # (B,32,T,L)
        attn = F.softmax(scores, dim=-1)
        out = attn @ V                                            # (B,32,T,128)
        return self.W_o(out.transpose(1,2).reshape(B,T,-1)), (K, V)  # O 投影 + 回传缓存

class MQA(nn.Module):   # 对应 图MQA-1：K/V 改为 1 个头
    def __init__(self, d_model=4096, n_heads=32):
        self.n_heads = n_heads
        self.head_dim = d_model // n_heads          # 128
        self.W_q = nn.Linear(d_model, d_model)      # 图MQA-2 的 Q 投影（仍 32 头）
        self.W_k = nn.Linear(d_model, self.head_dim) # ⚡ 图MQA-2 的 K 投影改为只 1 个头
        self.W_v = nn.Linear(d_model, self.head_dim) # ⚡ 图MQA-2 的 V 投影改为只 1 个头
        self.W_o = nn.Linear(d_model, d_model)      # O 投影

    def forward(self, x, past_kv=None):
        B, T, _ = x.shape
        Q = self.W_q(x).view(B,T,self.n_heads,self.head_dim).transpose(1,2)   # (B,32,T,128)
        K = self.W_k(x).unsqueeze(1)                                          # (B,1,T,128) 只 1 个头
        V = self.W_v(x).unsqueeze(1)
        if past_kv:
            K = torch.cat([past_kv[0], K], dim=2)   # (B,1,L,128) 缓存极瘦
            V = torch.cat([past_kv[1], V], dim=2)
        # 图MQA-2 广播：把 1 个头复制成 32 个头（等价 repeat_interleave）
        K = K.repeat(1, self.n_heads, 1, 1)         # (B,32,L,128)
        V = V.repeat(1, self.n_heads, 1, 1)
        scores = Q @ K.transpose(-2,-1) / math.sqrt(self.head_dim)
        attn = F.softmax(scores, dim=-1)
        out = attn @ V
        return self.W_o(out.transpose(1,2).reshape(B,T,-1)), (K[:, :1], V[:, :1])  # 缓存只存 1 份

class GQA(nn.Module):   # 对应 图GQA-1：K/V 改为 G 个头
    def __init__(self, d_model=4096, n_heads=32, n_groups=4):
        self.n_heads, self.n_groups = n_heads, n_groups
        self.head_dim = d_model // n_heads               # 128
        self.W_q = nn.Linear(d_model, d_model)           # 图GQA-2：Q 仍出 32 个头
        self.W_k = nn.Linear(d_model, n_groups * self.head_dim)  # ⚡ 图GQA-2：K 只出 G=4 个头
        self.W_v = nn.Linear(d_model, n_groups * self.head_dim)  # ⚡ 图GQA-2：V 只出 G=4 个头
        self.W_o = nn.Linear(d_model, d_model)           # O 投影

    def forward(self, x, past_kv=None):
        B, T, _ = x.shape
        Q = self.W_q(x).view(B,T,self.n_heads,self.head_dim).transpose(1,2)    # (B,32,T,128)
        K = self.W_k(x).view(B,T,self.n_groups,self.head_dim).transpose(1,2)  # (B,4,T,128)
        V = self.W_v(x).view(B,T,self.n_groups,self.head_dim).transpose(1,2)  # (B,4,T,128)
        # 图GQA-1 缓存：只追加 G=4 组，永远比 MHA 小 8 倍
        if past_kv:
            K = torch.cat([past_kv[0], K], dim=2)        # (B,4,L,128)
            V = torch.cat([past_kv[1], V], dim=2)
        # 图GQA-3 repeat_interleave：沿头维把 4 组复制成 32 组，与 Q 头对齐
        K = K.repeat_interleave(self.n_heads // self.n_groups, dim=1)  # (B,32,L,128)
        V = V.repeat_interleave(self.n_heads // self.n_groups, dim=1)  # 视图操作，缓存仍是 4 组
        # 之后的注意力与 MHA 完全一致
        scores = Q @ K.transpose(-2,-1) / math.sqrt(self.head_dim)
        attn = F.softmax(scores, dim=-1)
        out = attn @ V
        # 回传缓存时只保留 G 组（复制前的原始 4 组）
        return self.W_o(out.transpose(1,2).reshape(B,T,-1)), (K[:, :self.n_groups], V[:, :self.n_groups])

class MLA(nn.Module):   # 对应 图MLA-1：K/V 改为 latent 压缩
    def __init__(self, d_model=4096, n_heads=32, head_dim=128, d_latent=256):
        self.n_heads, self.head_dim, self.d_latent = n_heads, head_dim, d_latent
        self.W_q   = nn.Linear(d_model, n_heads * head_dim)       # 图MLA：Q 投影（同 MHA）
        self.W_kv  = nn.Linear(d_model, 2 * d_latent)             # ⚡ 图MLA-2 压缩层：x → (c_K, c_V)
        # 解压权重（训练用；推理被吸收进 Q/O，见 forward_fused）
        self.W_k_up = nn.Linear(d_latent, n_heads * head_dim)     # 图MLA-2：c_K → 32×128 的 K
        self.W_v_up = nn.Linear(d_latent, n_heads * head_dim)     # 图MLA-2：c_V → 32×128 的 V
        self.W_o   = nn.Linear(n_heads * head_dim, d_model)       # O 投影

    def forward_naive(self, x, past_latent=None):   # 训练/理解版：解压再算
        B, T, _ = x.shape
        Q = self.W_q(x).view(B,T,self.n_heads,self.head_dim).transpose(1,2)  # (B,32,T,128)
        c = self.W_kv(x)                                               # (B,T,512)
        c_K, c_V = c[..., :self.d_latent], c[..., self.d_latent:]      # 各 (B,T,256)
        if past_latent:                                        # 图MLA-1 缓存：只存 latent
            c_K = torch.cat([past_latent[0], c_K], dim=1)              # (B,L,256)
            c_V = torch.cat([past_latent[1], c_V], dim=1)
        K = self.W_k_up(c_K).view(B,-1,self.n_heads,self.head_dim).transpose(1,2)  # (B,32,L,128)
        V = self.W_v_up(c_V).view(B,-1,self.n_heads,self.head_dim).transpose(1,2)
        scores = Q @ K.transpose(-2,-1) / math.sqrt(self.head_dim)     # 与 MHA 相同
        attn = F.softmax(scores, dim=-1)
        out = (attn @ V).transpose(1,2).reshape(B,T,-1)
        return self.W_o(out), (c_K, c_V)

    def forward_fused(self, x, past_latent=None):    # 推理版：解压被吸收（图MLA-3）
        B, T, _ = x.shape
        Q = self.W_q(x).view(B,T,self.n_heads,self.head_dim)           # (B,T,32,128)
        c = self.W_kv(x)
        c_K, c_V = c[..., :self.d_latent], c[..., self.d_latent:]
        if past_latent:
            c_K = torch.cat([past_latent[0], c_K], dim=1)              # (B,L,256)
            c_V = torch.cat([past_latent[1], c_V], dim=1)
        # 吸收：q' = W_k_up^T · q（按头分块，W_k_up 每头占 (128,256)）
        Wk_up_T = self.W_k_up.weight.view(self.n_heads, self.head_dim, self.d_latent)
        q_abs = torch.einsum('bthd,hdL->bthL', Q, Wk_up_T)             # (B,T,32,256)
        scores = torch.einsum('bthL,blL->bthl', q_abs, c_K) / math.sqrt(self.head_dim)
        attn = torch.softmax(scores, dim=-1)
        o_lat = torch.einsum('bthl,blL->bthL', attn, c_V)              # 仍在 latent 空间 (B,T,32,256)
        # V 的解压也吸收进 O：out = W_o · (W_v_up^T · o_lat)
        Wv_up_T = self.W_v_up.weight.view(self.n_heads, self.head_dim, self.d_latent)
        o = torch.einsum('bthL,hdl->bthd', o_lat, Wv_up_T)             # (B,T,32,128)
        return self.W_o(o.reshape(B,T,-1)), (c_K, c_V)             # 缓存仍是 latent

class KDA(nn.Module):   # 对应 图KDA-1：完全替换图B蓝色框（线性注意力）
    def __init__(self, d_model=4096, s_dim=2048):
        self.s_dim = s_dim
        self.W_q   = nn.Linear(d_model, s_dim)     # 图KDA-2：q —— 读指针
        self.W_k   = nn.Linear(d_model, s_dim)     # 图KDA-2：k —— 写地址
        self.W_v   = nn.Linear(d_model, s_dim)     # 图KDA-2：v —— 写内容
        self.W_beta = nn.Linear(d_model, s_dim)    # 图KDA-2：δ —— 遗忘门（0~1）
        self.W_read = nn.Linear(s_dim, d_model)    # 图KDA-2：读状态 → 输出

    def forward(self, x, state=None):
        B = x.shape[0]
        # —— QKV 在哪里：Q=读、K=写地址、V=写内容，都在状态机内部 ——
        q = self.W_q(x).squeeze(1)                         # (B, s) 读取用
        k = self.W_k(x).squeeze(1)                         # (B, s) 写地址
        v = self.W_v(x).squeeze(1)                         # (B, s) 写内容
        decay = torch.sigmoid(self.W_beta(x)).squeeze(1)   # (B, s) 遗忘率 0~1
        if state is None:
            state = torch.zeros(B, self.s_dim, self.s_dim, device=x.device)
        # 图KDA-1 状态更新：S ← δ·S + v·k^T（外积累加，历史信息写进固定状态）
        state = decay.unsqueeze(1) * state + torch.einsum('bs,bt->bst', v, k)  # (B,s,s)
        # 图KDA-1 读取：out = W_read(q·S)，从状态读回压缩后的全部历史
        out = self.W_read(torch.einsum('bs,bst->bt', q, state)).unsqueeze(1)   # (B,1,d_model)
        return out, state   # 状态大小恒定，与 L 无关 → O(1) 显存、O(1) 算力