DeepSeek MLA + MoE

DeepSeek 的核心思想：用更少的计算资源达到大模型的效果。传统大模型（如 LLaMA）在推理时需要激活所有参数，计算成本极高。DeepSeek 通过两大创新大幅降低了计算开销。

生活类比：想象一家公司遇到一个复杂问题。传统做法是请一个全能专家（工资很高），而 DeepSeek 的做法是请一组各有所长的专家团队，每次只激活最相关的 2 位专家来解决问题。这样既保证了专业性，又大幅降低了成本。

两大核心创新

标准 MHA vs MLA 的 KV Cache 对比

假设：模型维度 d=128，注意力头数 h=8，每头维度 d_k=16，序列长度 n=1024，压缩维度 d_c=32

输入隐藏状态 h

假设输入向量 h 维度 d=4（简化演示），压缩维度 d_c=2：

Step 3a：下投影 — 压缩 c_KV = W_DKV · h

W_DKV 将 d=4 维压缩到 d_c=2 维：

Step 3b：上投影 — 恢复 K 和 V

从压缩向量 c_KV 恢复出完整的 K 和 V：

Step 4a：路由器计算

路由器是一个简单的线性层 + Softmax，输出每个专家的"被选中概率"。假设 4 个专家，输入 x 维度=2：

Step 4b：专家计算与加权输出

只有被选中的 Top-2 专家参与计算，其他专家不激活（稀疏激活）：

输入特征

专家激活权重

柱状图显示每个专家的路由器权重，绿色为被选中的 Top-2 专家：

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadLatentAttention(nn.Module):
    """Multi-head Latent Attention (MLA) — DeepSeek-V2 核心创新"""

    def __init__(self, d_model=512, n_heads=8, d_compress=64):
        super().__init__()
        self.d_model = d_model
        self.n_heads = n_heads
        self.d_k = d_model // n_heads
        self.d_compress = d_compress

        # Q 投影（与标准 MHA 相同）
        self.W_q = nn.Linear(d_model, d_model)

        # MLA 核心：下投影 + 上投影（替代标准的 K/V 投影）
        self.W_DKV = nn.Linear(d_model, d_compress)    # 下投影：压缩
        self.W_UK = nn.Linear(d_compress, d_model)      # 上投影：恢复 K
        self.W_UV = nn.Linear(d_compress, d_model)      # 上投影：恢复 V

        self.W_o = nn.Linear(d_model, d_model)

    def forward(self, h, kv_cache=None):
        batch, seq_len, _ = h.shape

        # Q：标准投影
        Q = self.W_q(h)

        # MLA：压缩 → 缓存 → 恢复
        c_KV = self.W_DKV(h)             # [batch, seq, d_compress]
        if kv_cache is not None:
            c_KV = torch.cat([kv_cache, c_KV], dim=1)

        K = self.W_UK(c_KV)              # 恢复完整 K
        V = self.W_UV(c_KV)              # 恢复完整 V

        # 多头拆分 + 注意力计算
        Q = Q.view(batch, seq_len, self.n_heads, self.d_k).transpose(1, 2)
        K = K.view(batch, -1, self.n_heads, self.d_k).transpose(1, 2)
        V = V.view(batch, -1, self.n_heads, self.d_k).transpose(1, 2)

        scores = torch.matmul(Q, K.transpose(-2, -1)) / (self.d_k ** 0.5)
        attn = F.softmax(scores, dim=-1)
        out = torch.matmul(attn, V)

        out = out.transpose(1, 2).contiguous().view(batch, seq_len, self.d_model)
        return self.W_o(out), c_KV  # 返回压缩缓存供下次使用

class MixtureOfExperts(nn.Module):
    """Mixture of Experts (MoE) — 稀疏专家混合"""

    def __init__(self, d_model=512, n_experts=8, top_k=2, d_ff=1024):
        super().__init__()
        self.n_experts = n_experts
        self.top_k = top_k

        # 路由器：决定激活哪些专家
        self.router = nn.Linear(d_model, n_experts)

        # 专家列表：每个专家是一个简单的 FFN
        self.experts = nn.ModuleList([
            nn.Sequential(
                nn.Linear(d_model, d_ff),
                nn.GELU(),
                nn.Linear(d_ff, d_model),
            )
            for _ in range(n_experts)
        ])

    def forward(self, x):
        batch, seq_len, d = x.shape
        x_flat = x.view(-1, d)           # [batch*seq, d]

        # 路由器计算
        router_logits = self.router(x_flat)  # [batch*seq, n_experts]
        router_weights = F.softmax(router_logits, dim=-1)

        # Top-K 选择
        topk_weights, topk_indices = torch.topk(router_weights, self.top_k, dim=-1)
        topk_weights = topk_weights / topk_weights.sum(dim=-1, keepdim=True)  # 归一化

        # 只计算被选中的专家
        output = torch.zeros_like(x_flat)
        for k in range(self.top_k):
            expert_idx = topk_indices[:, k]    # 第 k 个被选中的专家
            weight = topk_weights[:, k:k+1]  # 对应权重

            for e in range(self.n_experts):
                mask = (expert_idx == e)
                if mask.any():
                    expert_out = self.experts[e](x_flat[mask])
                    output[mask] += weight[mask] * expert_out

        return output.view(batch, seq_len, d)