Deepseek优化学习笔记——MLA详细流程与优化原理

    最近比较仔细的阅读了Mla相关的优化，下面从下面几个方面梳理一下该部分的一些内容：

1. deekseek的模型结构以及mla部分的计算流程简介；
2. MQA，GQA的kv cache优化思路以及网络结构对比；

1.1简介

    DeepSeek-V3的基本架构仍然基于Transformer框架，为了实现高效推理和经济高效的训练，DeepSeek-V3还采用了MLA（多头潜在注意力）。
    MHA（多头注意力）通过多个注意力头并行工作捕捉序列特征，但面临高计算成本和显存占用；MLA（多头潜在注意力）则通过低秩压缩优化键值矩阵，降低显存占用并提高推理效率。

1.2 多头注意力（MHA）

    多头注意力（MHA）是Transformer模型架构中的一个核心组件，它允许模型在处理输入序列时能够同时关注来自不同位置的不同表示子空间的信息。
    MHA通过将输入向量分割成多个并行的注意力“头”，每个头独立地计算注意力权重并产生输出，然后将这些输出通过拼接和线性变换进行合并以生成最终的注意力表示。

    多头注意力如何进行Q,K,V计算？多头注意力（MHA）通过线性变换将张量分别转换为查询(Q),键（K）和值（V）矩阵，每个矩阵再被分割成多个头进行并行处理。

1. 输入变换：输入序列首先通过三个不同的线性变换层，分别得到查询（Query）、键（key）和值（value）矩阵，这些变换通常是通过全连接层实现的。
2. 分头：将查询、键和值分成多个头（即多个子空间），每个头具有不同的线性变换参数。
3. 注意力计算：对于每个头，都执行一次缩放点积注意力（scaled Dot-Product Atention）运算，具体来说，计算QK的点积经过缩放，加上偏置后，使用softmax函数得到注意力权重。这些权重用于加权值矩阵，生成加权和作为每个头的输出。
4. 拼接与融合：将所有头的输出拼接在一起，形成一个长向量。然后，对拼接后的向量进行一个最终的线性变换，以整合来自不同头的信息，得到最终的多头注意力输出。
   
       MHA计算过程图：
   
       其中 $d_k$ 表示每个头的维度。

    多头注意力机制和注意力机制区别是什么？多头注意力机制通过引入多个并行的注意力头，提高了模型对输入数据的全面捕捉和处理能力，使其在处理大规模数据和复杂任务时更具有优势。

• 注意力机制：通过聚焦于关键信息，提高了模型对输入数据的理解和处理能力；
• 多头注意力机制：通过并行处理和集成多个注意力头的结果，从不同的角度捕捉数据的多样性，进一步增强了模型的学习能力和表达力。

1.3 MQA与GQA

• Multi-Query Attention(MQA):所有的查询头（Query Heads）共享相同的 key 和 value。即消减了 Flops,也降低了cache,并且压缩了频繁矩阵拼接的IO耗时。
• Group-Query Attenation(GQA):是 MQA 的改进版，他通过在多个查询头之间共享 key 和value，在 MHA 和 MQA 之间找到了一种折中方案。

    GQA旨在推理速度和模型质量之间取得更好的平衡，减少 MQA 带来的模型质量下降问题，同时保留比 MHA 更快的推理速度。DeepseekV1 67B,llama2 70B 和 llama3 系列全部都用了 GQA。

    下面来看一下 MQA 与 GQA 带来的 kv cache压缩效果：
    首先给出变量的含义：$d$ 代表输入维度(input dim)，$n_h$ 代表头数(head数)，$d_h$ 代表每个头的维度，$h_h$代表输入的第t个向量，$l$ 代表transformer的层数。
    对于标准的MHA而言，对于每一个token，KV cache占用的缓存的大小为 $2n_h{d}_{h}l$ 。其中2为fp16的字节数。
    对于 MQA 来说，缓存大小变为 $2d_{h}l$ ,相比于 MHA 的 $2n_h{d}_{h}l$ 而言大大减少，但是模型精度也会差一些；
    对于 GQA 来说，缓存大小变为 $2n_g{d}_{h}l$ ,其中 $n_g$ 代表head的分组数。

    MQA 与 GQA 如何进行多头 KV 共享 ，或许有的时候会有这个疑问，但是当再看看模型结构的时候，心里就没有这种疑问，这是一个标准的MHA网络的attenation 网络层结构：

llama2-7b网络层结构：

model: LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 4096)
    (layers): ModuleList(
      (0-5): 6 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (v_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        )
        ...
)

下面是一个GQA的llama3-8b的网络层结构：

model: LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(128256, 4096)
    (layers): ModuleList(
      (0-5): 6 x LlamaDecoderLayer(
        (self_attn): LlamaAttention(
          (q_proj): Linear(in_features=4096, out_features=4096, bias=False)
          (k_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (v_proj): Linear(in_features=4096, out_features=1024, bias=False)
          (o_proj): Linear(in_features=4096, out_features=4096, bias=False)
        )
        ...
  )
  (lm_head): Linear(in_features=4096, out_features=128256, bias=False)
)

GQA的llama2-72b的网络层结构：

model: LlamaForCausalLM(
  (model): LlamaModel(
    (embed_tokens): Embedding(32000, 8192, padding_idx=0)
    (layers): ModuleList(
      (0-5): 6 x LlamaDecoderLayer(
        (self_attn): LlamaSdpaAttention(
          (q_proj): Linear(in_features=8192, out_features=8192, bias=False)
          (k_proj): Linear(in_features=8192, out_features=1024, bias=False)
          (v_proj): Linear(in_features=8192, out_features=1024, bias=False)
          (o_proj): Linear(in_features=8192, out_features=8192, bias=False)
          (rotary_emb): LlamaRotaryEmbedding()
        )
       ...
)

    两种网络在 k_proj 和 v_proj 的输出维度是不同的，MHA 是与 q 是相同的，但是 GQA 的这部分并不是相同的，而是多对一的关系，在 llama3-8b 种有32个头，每个头的维度为4096/32=128。kv 1024/128=8组，在8B模型种32个头用8组，4个头共用一组，同理可以推算72B。对于MQA的模型，则只有一组kv,所有头共用，对应的out-feature就是头的维度。
对于deepseek v3的网络结构如下：

model: DeepseekV3ForCausalLM(
  (model): DeepseekV3Model(
    (embed_tokens): Embedding(129280, 7168)
    (layers): ModuleList(
      (0-2): 3 x DeepseekV3DecoderLayer(
        (self_attn): DeepseekV3Attention(
          (q_a_proj): Linear(in_features=7168, out_features=1536, bias=False)
          (q_a_layernorm): DeepseekV3RMSNorm()
          (q_b_proj): Linear(in_features=1536, out_features=24576, bias=False)
          (kv_a_proj_with_mqa): Linear(in_features=7168, out_features=576, bias=False)
          (kv_a_layernorm): DeepseekV3RMSNorm()
          (kv_b_proj): Linear(in_features=512, out_features=32768, bias=False)
          (o_proj): Linear