企业官网建设流程全解析

1. 为什么大模型面试总爱考Attention变体

面试官让你手撕Attention，从来不是想看你默写softmax(QK^T/sqrt(d))V。这个公式谁都会背，真正拉开差距的是：为什么现在的模型不用标准MHA了？MLA到底省了什么？GQA和MQA的取舍在哪里？这几个问题答不上来，基本就暴露了你只跑过transformers的from_pretrained，没真正理解推理时的显存账本。

我面过不少人，也被人面过。一个很深的感受是：Attention变体这个题，表面考代码，实际考的是你对推理成本结构的理解。MHA、MQA、GQA、MLA这四个东西，本质上是同一件事在不同约束下的四种解法——约束就是KV Cache的显存和带宽。

先把结论摆出来，方便你建立全局观：

这张表你先记住，后面每一节我都会把里面的数字拆开算给你看。面试时如果你能主动把这张表画出来，再逐行解释，基本就稳了一半。

需要提前说明的是，MLA是DeepSeek提出的结构，网上公开的技术报告讲得比较清楚，我下面关于MLA的推导基于公开资料和我的理解，具体实现细节以官方代码为准。MHA、MQA、GQA则是业界共识，可以放心讲。

2. 从KV Cache这笔账算起

2.1 推理时显存到底花在哪

很多人一上来就讲结构，其实应该先算账。大模型推理分两个阶段：prefill和decode。prefill阶段把整个prompt一次性算完，是计算密集型的；decode阶段一个token一个token往外吐，是访存密集型的。KV Cache就是decode阶段的命根子。

为什么？因为自回归生成时，每生成一个新token，都要和前面所有token做Attention。如果每次都重算前面token的K和V，计算量会爆炸。所以业界做法是把历史token的K、V缓存下来，这就是KV Cache。

KV Cache的大小公式：

KV Cache = 2 × num_layers × num_kv_heads × head_dim × seq_len × batch_size × dtype_bytes

注意那个2，是因为K和V各存一份。num_kv_heads就是关键变量——MHA里它等于num_heads，MQA里它等于1，GQA里它是中间值。这就是所有变体的核心分歧点。

2.2 一个具体的数字感受一下

拿Llama 2 70B举例。它有80层，64个Q头，head_dim是128，用fp16（2字节）。假设seq_len=4096，batch_size=1。

如果是MHA（KV头数=64）：

2 × 80 × 64 × 128 × 4096 × 1 × 2 = 10.7 GB

光是KV Cache就10.7GB，而模型权重本身fp16也就140GB左右。batch_size一上去，或者seq_len拉到32K，这个数字会线性膨胀到让你怀疑人生。

如果是GQA（Llama 2 70B实际用8个KV头）：

2 × 80 × 8 × 128 × 4096 × 1 × 2 = 1.34 GB

直接降到1/8。这就是为什么Llama 2 70B要用GQA——不是为了效果更好，是为了能跑得起来。

提示：面试时如果被问到"GQA省了多少"，不要只答"省了KV Cache"，要能报出这个公式和具体倍数。KV头数从64降到8，就是8倍，这是硬账。

2.3 为什么不能无脑用MQA

既然MQA把KV头数压到1，省得最狠，为什么不全用MQA？

因为表达能力会掉。MHA里每个Q头有自己专属的K、V，可以关注不同的模式；MQA里所有Q头共享同一组K、V，相当于强行让64个注意力头看同一份"记忆"，多样性被砍掉了。实测中MQA在长文本、复杂推理任务上掉点比较明显。

GQA就是折中：分组，组内共享KV，组间不共享。Llama 2 70B用8组，相当于64个Q头分成8组，每组8个Q头共享一份KV。既省了显存，又保留了组间的多样性。

MLA走的是另一条路——不减少KV头数，而是把KV压缩成低秩表示。这个后面单独讲。

3. MHA：一切的原点，也是显存杀手

3.1 标准MHA的结构拆解

MHA（Multi-Head Attention）是《Attention Is All You Need》里的原始设计。核心思想是：把d_model维度的输入投影到多个子空间，每个子空间独立做Attention，最后拼接。

具体来说，假设d_model=512，num_heads=8，那么每个头的head_dim=64。输入X（seq_len × 512）分别乘以W_q、W_k、W_v（都是512×512），得到Q、K、V，然后reshape成(seq_len, 8, 64)，再transpose成(8, seq_len, 64)。

每个头独立计算：

head_i = softmax(Q_i @ K_i^T / sqrt(64)) @ V_i

最后8个头concat起来，再过一个W_o投影回512维。

用PyTorch写出来大概是这样：

import torch import torch.nn as nn import math class MHA(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.d_model = d_model self.num_heads = num_heads self.head_dim = d_model // num_heads assert self.head_dim * num_heads == d_model self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, d_model) self.W_v = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, T, C = x.shape q = self.W_q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = self.W_k(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v = self.W_v(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: attn = attn.masked_fill(mask == 0, float('-inf')) attn = torch.softmax(attn, dim=-1) out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) return self.W_o(out)

这段代码面试时能默写出来是基本要求。但光写出来不够，面试官会追问：KV Cache怎么加？

3.2 MHA的KV Cache实现要点

推理时，K和V要缓存。上面代码里k和v的shape是(B, num_heads, T, head_dim)，缓存的就是这个。每来一个新token，只算新token的q、k、v，然后把新的k、v拼到缓存的k、v后面。

# 推理时的简化逻辑 def forward_with_cache(self, x, kv_cache=None): B, T, C = x.shape # T=1 in decode q = self.W_q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = self.W_k(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v = self.W_v(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) if kv_cache is not None: k = torch.cat([kv_cache[0], k], dim=2) v = torch.cat([kv_cache[1], v], dim=2) new_cache = (k, v) attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) attn = torch.softmax(attn, dim=-1) out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) return self.W_o(out), new_cache

注意torch.cat这个操作，每次decode都要重新分配内存并拷贝，这是效率杀手。实际工程里会用预分配的cache buffer，直接往固定位置写。面试时如果你能提到这一点，会加分。

注意：手撕代码时，面试官通常不要求你写完整的cache管理，但你要能说清楚"cache的shape是(B, num_heads, max_seq_len, head_dim)，用position索引写入"。这个细节能体现你真正部署过。

3.3 MHA为什么被淘汰

MHA的问题就一个字：贵。KV Cache随num_heads线性增长，而num_heads通常等于d_model/head_dim，d_model动辄几千，num_heads动辄几十上百。在长上下文场景下，KV Cache甚至能超过模型权重本身。

但MHA的表达能力是最强的，所以它没有完全消失——训练时用MHA，推理时转成GQA，或者小模型继续用MHA。理解MHA是理解其他变体的前提，这个地基必须打牢。

4. MQA：把KV头砍到1的激进方案

4.1 MQA的核心改动

MQA（Multi-Query Attention）的思路极其简单粗暴：所有Q头共享同一组K、V。也就是说，num_kv_heads=1。

原来MHA里W_k和W_v的输出维度是d_model（num_heads × head_dim），MQA里改成head_dim（1 × head_dim）。Q还是num_heads份，K、V只有1份。

class MQA(nn.Module): def __init__(self, d_model, num_heads): super().__init__() self.d_model = d_model self.num_heads = num_heads self.head_dim = d_model // num_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, self.head_dim) # 只有1个头 self.W_v = nn.Linear(d_model, self.head_dim) self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, T, C = x.shape q = self.W_q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = self.W_k(x).view(B, T, 1, self.head_dim).transpose(1, 2) v = self.W_v(x).view(B, T, 1, self.head_dim).transpose(1, 2) # k, v broadcast到num_heads attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: attn = attn.masked_fill(mask == 0, float('-inf')) attn = torch.softmax(attn, dim=-1) out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) return self.W_o(out)

关键在k和v的shape是(B, 1, T, head_dim)，而q是(B, num_heads, T, head_dim)。q @ k.transpose(-2,-1)时，PyTorch会自动broadcast，k被复制num_heads份参与计算。计算时还是num_heads份，但存储时只有1份——这就是省显存的本质。

4.2 MQA省了多少，代价是什么

还是Llama 2 70B的例子，MQA的KV Cache：

2 × 80 × 1 × 128 × 4096 × 1 × 2 = 0.168 GB

对比MHA的10.7GB，省了64倍。这个数字非常夸张，意味着同样的显存可以支持64倍的batch_size或seq_len。

代价是质量下降。原论文《Fast Transformer Decoding: One Write-Head is All You Need》里提到，MQA在训练时可能不稳定，需要调学习率。实测中，MQA在翻译、摘要等任务上掉点不明显，但在需要精细区分的任务（如长文档问答、代码生成）上掉点比较明显。

提示：面试时如果被问"MQA为什么快"，不要只说"KV Cache小"。还要提内存带宽——decode阶段是访存瓶颈，KV Cache小了，每次读取的数据量就小，GPU的HBM带宽压力就小，这才是真正的加速来源。

4.3 MQA的适用场景

MQA适合推理成本极度敏感、质量要求相对宽松的场景。比如：

• 大规模在线服务的草稿模型
• 边缘设备部署
• 对延迟要求极高的实时应用

但主流大模型现在很少纯用MQA了，因为GQA在几乎不增加成本的情况下能拿回大部分质量。MQA更多是作为一个"极端点"存在，理解它是理解GQA的跳板。

5. GQA：工业界最爱的折中方案

5.1 GQA的分组逻辑

GQA（Grouped-Query Attention）是MQA和MHA的插值。把num_heads个Q头分成G组，每组共享一份K、V。G=1就是MQA，G=num_heads就是MHA。

Llama 2 70B用G=8，num_heads=64，所以每组8个Q头共享一份KV。Llama 3 8B用G=8，num_heads=32，每组4个Q头共享一份KV。

class GQA(nn.Module): def __init__(self, d_model, num_heads, num_kv_heads): super().__init__() self.d_model = d_model self.num_heads = num_heads self.num_kv_heads = num_kv_heads self.head_dim = d_model // num_heads self.num_groups = num_heads // num_kv_heads self.W_q = nn.Linear(d_model, d_model) self.W_k = nn.Linear(d_model, num_kv_heads * self.head_dim) self.W_v = nn.Linear(d_model, num_kv_heads * self.head_dim) self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, T, C = x.shape q = self.W_q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) k = self.W_k(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.W_v(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2) # 把k, v在头维度上repeat，使每个Q头对应到组内的KV k = k.repeat_interleave(self.num_groups, dim=1) v = v.repeat_interleave(self.num_groups, dim=1) attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: attn = attn.masked_fill(mask == 0, float('-inf')) attn = torch.softmax(attn, dim=-1) out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) return self.W_o(out)

repeat_interleave是GQA实现的关键。它把(B, num_kv_heads, T, head_dim)的k扩展成(B, num_heads, T, head_dim)，扩展方式是每个KV头连续重复num_groups次。这样第i个Q头就对应到第i//num_groups个KV头。

5.2 GQA的显存账和效果账

Llama 2 70B的GQA KV Cache是1.34GB，是MHA的1/8，MQA的8倍。这个位置很微妙——显存省了8倍，但质量几乎不掉。

原论文《GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints》里做了大量实验，结论是：GQA在质量上接近MHA，在速度上接近MQA。而且论文还提出了一种从MHA checkpoint转换到GQA的方法——对K、V的投影矩阵做mean pooling，把num_heads个头的参数平均成num_kv_heads份。这样不用从头训练，省了大量算力。

提示：面试时如果被问到"GQA怎么从MHA初始化"，答"mean pooling"是加分项。具体做法是把W_k和W_v的权重reshape成(num_kv_heads, num_groups, head_dim, d_model)，然后在num_groups维度上取平均。

5.3 GQA的工程实现细节

实际部署GQA时，有几个坑：

第一个坑是repeat_interleave的效率。每次forward都repeat一遍k、v，虽然比MHA省了存储，但计算时还是num_heads份，repeat操作本身有开销。优化做法是在kernel层面处理，比如FlashAttention就支持GQA的原生输入，不需要显式repeat。

第二个坑是KV Cache的布局。缓存的是repeat前的(B, num_kv_heads, T, head_dim)，不是repeat后的。这样存储才省。读取时在kernel里做broadcast。

第三个坑是num_kv_heads必须整除num_heads。Llama 2 70B的64/8=8，Llama 3 8B的32/8=4，都是整除的。如果不整除，分组就不均匀，实现会复杂很多。

# 用FlashAttention的GQA支持（伪代码） from flash_attn import flash_attn_func # q: (B, T, num_heads, head_dim) # k: (B, T, num_kv_heads, head_dim) # v: (B, T, num_kv_heads, head_dim) out = flash_attn_func(q, k, v, causal=True) # FlashAttention内部处理GQA的broadcast，不需要手动repeat

这个细节能体现你真正用过FlashAttention，而不是只会写naive实现。

6. MLA：DeepSeek的低秩压缩思路

6.1 MLA到底在压缩什么

MLA（Multi-head Latent Attention）是DeepSeek-V2提出的结构。它和MQA、GQA的思路完全不同——不减少KV头数，而是把KV压缩到一个低维的latent向量。

具体来说，MHA里K、V是从d_model投影到num_heads × head_dim。MLA里先投影到一个更小的维度d_c（比如512），这个d_c远小于num_heads × head_dim（比如64 × 128 = 8192）。推理时只缓存这个d_c维的latent向量，需要K、V时再从latent上投影出来。

这就好比：MHA是给每个头存一份完整的K、V；MLA是存一份压缩包，用的时候解压。压缩包比原始数据小得多，但解压需要额外计算。

6.2 MLA的数学推导

设输入为h（d_model维），MLA的KV计算：

c_kv = W_dkv @ h # 压缩到d_c维 k = W_uk @ c_kv # 从d_c投影到num_heads × head_dim v = W_uv @ c_kv # 从d_c投影到num_heads × head_dim

推理时缓存的是c_kv，大小是d_c × seq_len × batch_size × dtype_bytes。对比MHA的2 × num_heads × head_dim × seq_len × ...，d_c通常远小于2 × num_heads × head_dim。

DeepSeek-V2里d_c=512，num_heads=128，head_dim=128，所以2 × 128 × 128 = 32768，而d_c只有512，压缩了64倍。这个压缩比比GQA还狠。

但MLA有个问题：位置编码。RoPE是作用在K上的，如果K是从latent投影出来的，RoPE怎么加？DeepSeek的解法是解耦RoPE——额外算一份带RoPE的K（叫k_rope），和从latent投影出的k_nope拼接。这样既保留了RoPE的位置信息，又享受了latent压缩。

class MLA(nn.Module): def __init__(self, d_model, num_heads, head_dim, d_c, d_rope): super().__init__() self.num_heads = num_heads self.head_dim = head_dim self.d_c = d_c self.d_rope = d_rope self.W_dkv = nn.Linear(d_model, d_c) self.W_uk = nn.Linear(d_c, num_heads * head_dim) self.W_uv = nn.Linear(d_c, num_heads * head_dim) self.W_kr = nn.Linear(d_model, d_rope) # RoPE部分 self.W_q = nn.Linear(d_model, num_heads * head_dim) self.W_o = nn.Linear(num_heads * head_dim, d_model) def forward(self, x, kv_cache=None): B, T, C = x.shape c_kv = self.W_dkv(x) # 压缩 k_nope = self.W_uk(c_kv).view(B, T, self.num_heads, self.head_dim) v = self.W_uv(c_kv).view(B, T, self.num_heads, self.head_dim) k_rope = self.W_kr(x).view(B, T, 1, self.d_rope) # k = concat(k_nope, k_rope broadcast) # ... 后续Attention计算

这段是简化版，实际MLA的RoPE处理更复杂，涉及位置编码的旋转操作。面试时如果被问到MLA，能说清楚"低秩压缩+解耦RoPE"这两个核心点就够了，不需要手推全部公式。

6.3 MLA的取舍

MLA的优势是压缩率极高，KV Cache比MQA还小。DeepSeek-V2的KV Cache只有MHA的几十分之一，同时质量接近MHA。

代价是计算量增加。每次都要从latent投影出K、V，这是额外的矩阵乘法。在prefill阶段，这个开销可以接受；在decode阶段，因为要反复投影，计算量比GQA大。但decode阶段是访存瓶颈，多算一点换少读很多，总体是划算的。

另一个代价是实现复杂。解耦RoPE、latent投影、cache管理，都比GQA复杂。这也是为什么MLA目前主要是DeepSeek在用，其他家还在观望。

提示：面试时如果被问"MLA和GQA哪个好"，不要给绝对答案。要说"取决于场景"——如果显存极度紧张、能接受实现复杂度，MLA更优；如果要快速落地、生态支持好，GQA更稳。这种有条件的回答比站队更专业。

7. 手撕代码：四个变体的统一实现

7.1 用一套代码覆盖四种变体

面试时最怕的是让你分别写四遍。其实可以用一套代码，通过参数控制。核心就是num_kv_heads和use_mla两个开关。

class UnifiedAttention(nn.Module): def __init__(self, d_model, num_heads, num_kv_heads=None, use_mla=False, d_c=512, d_rope=64): super().__init__() self.d_model = d_model self.num_heads = num_heads self.head_dim = d_model // num_heads self.use_mla = use_mla if use_mla: self.d_c = d_c self.d_rope = d_rope self.W_dkv = nn.Linear(d_model, d_c) self.W_uk = nn.Linear(d_c, num_heads * self.head_dim) self.W_uv = nn.Linear(d_c, num_heads * self.head_dim) self.W_kr = nn.Linear(d_model, d_rope) else: # num_kv_heads=None -> MHA, =1 -> MQA, 其他 -> GQA self.num_kv_heads = num_kv_heads if num_kv_heads else num_heads self.num_groups = num_heads // self.num_kv_heads self.W_k = nn.Linear(d_model, self.num_kv_heads * self.head_dim) self.W_v = nn.Linear(d_model, self.num_kv_heads * self.head_dim) self.W_q = nn.Linear(d_model, d_model) self.W_o = nn.Linear(d_model, d_model) def forward(self, x, mask=None): B, T, C = x.shape q = self.W_q(x).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) if self.use_mla: c_kv = self.W_dkv(x) k = self.W_uk(c_kv).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) v = self.W_uv(c_kv).view(B, T, self.num_heads, self.head_dim).transpose(1, 2) else: k = self.W_k(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2) v = self.W_v(x).view(B, T, self.num_kv_heads, self.head_dim).transpose(1, 2) if self.num_groups > 1: k = k.repeat_interleave(self.num_groups, dim=1) v = v.repeat_interleave(self.num_groups, dim=1) attn = (q @ k.transpose(-2, -1)) / math.sqrt(self.head_dim) if mask is not None: attn = attn.masked_fill(mask == 0, float('-inf')) attn = torch.softmax(attn, dim=-1) out = (attn @ v).transpose(1, 2).contiguous().view(B, T, C) return self.W_o(out)

用这个类，num_kv_heads=num_heads就是MHA，=1就是MQA，中间值就是GQA，use_mla=True就是MLA。面试时写出这个，再解释每个分支，比写四遍强得多。

7.2 面试现场的时间分配

手撕Attention通常给15-20分钟。我的建议是：

• 前3分钟：写MHA的骨架，把Q、K、V投影和Attention计算写出来。
• 接下来5分钟：加KV Cache逻辑，说明cache的shape和更新方式。
• 再5分钟：改造成GQA，加num_kv_heads参数和repeat_interleave。
• 最后5分钟：口述MQA和MLA的改动点，不需要全写。

这样既展示了代码能力，又展示了知识广度。如果面试官只让写一个，优先写GQA——因为它覆盖了MHA和MQA的极端情况，最能体现理解深度。

注意：写代码时变量命名要清晰，num_kv_heads、num_groups、head_dim这些名字要统一。面试官看代码首先看命名，命名乱会扣印象分。

7.3 数值稳定性：面试官爱追问的细节

softmax(QK^T/sqrt(d))里的sqrt(d)不是随便加的。如果d=128，QK^T的点积结果方差大约是128，数值可能到几十甚至上百。exp(100)直接溢出。除以sqrt(128)≈11.3，把方差拉回1附近，exp就不会溢出。

但即使除了sqrt(d)，极端情况下还是可能溢出。工程上更稳的做法是减去最大值：

attn = q @ k.transpose(-2, -1) / math.sqrt(self.head_dim) attn = attn - attn.max(dim=-1, keepdim=True).values # 减最大值 attn = torch.softmax(attn, dim=-1)

PyTorch的softmax内部已经做了这个，所以直接调torch.softmax是安全的。但如果你手写softmax，一定要减最大值。面试时如果被问"softmax怎么防溢出"，答"减最大值"是标准答案。

8. 常见问题与排查技巧实录

8.1 面试高频追问速查表

这张表建议背下来，面试时被问到能脱口而出。

8.2 实操中踩过的坑

坑一：repeat_interleave的内存布局。repeat_interleave返回的是新tensor，不是view。如果num_groups很大，这个操作会分配大量内存。优化做法是用expand+reshape，或者直接用FlashAttention。

坑二：KV Cache的预分配。推理时不要用torch.cat动态拼接，要预分配一个(B, num_kv_heads, max_seq_len, head_dim)的buffer，用position索引写入。torch.cat每次都要重新分配和拷贝，在长序列下开销巨大。

坑三：GQA的num_kv_heads不整除num_heads。如果64个Q头配3个KV头，分组就不均匀。实际模型都会选整除的值，比如8、4、2。自己设计时要注意。

坑四：MLA的cache管理。MLA缓存的是latent向量c_kv，不是K、V。如果按MHA的思路去缓存K、V，就完全失去了MLA的意义。这个点面试时容易被追问。

坑五：训练和推理的不一致。训练时可以用MHA，推理时转GQA，但转换后要重新校准（比如跑一遍验证集）。直接转换不校准，可能掉点。

8.3 性能对比实测数据

我在单卡A100上做过一组对比，模型配置是d_model=4096，num_heads=32，head_dim=128，num_layers=32，seq_len=8192，batch_size=8，fp16。

这组数据是特定配置下的，不同模型、不同硬件会有差异，但趋势是一致的：KV Cache越小，decode吞吐越高。MQA最快但质量最差，GQA和MLA在质量和速度之间取得了不错的平衡。

提示：面试时如果被问"你实测过吗"，能报出这样的数据会非常有说服力。但要注意说明测试条件，不要给绝对数字。

9. 从面试题到工程落地的延伸

9.1 怎么选：一张决策图

实际做项目时，选哪个变体不是拍脑袋，要看约束：

• 显存极度紧张、质量要求高：MLA，但要做好实现复杂的准备。
• 显存紧张、要快速落地：GQA，生态支持最好，Llama系列验证过。
• 显存不紧张、质量优先：MHA，或者GQA配更多KV头。
• 延迟极度敏感、质量可妥协：MQA，或者GQA配更少KV头。

我的经验是，GQA是默认选择。除非有特殊需求，否则GQA的性价比最高。MLA适合愿意啃硬骨头的团队，MQA适合特定场景。

9.2 训练时的注意事项

训练GQA时，有个细节容易忽略：KV头的初始化。如果从零训练，K、V的投影矩阵初始化要正常；如果从MHA转换，要做mean pooling。mean pooling后，KV头的输出尺度会变小（因为平均了num_groups份），可能需要调整学习率或加个缩放。

另外，GQA的训练和推理要一致。有些实现训练时用MHA，推理时转GQA，中间没有fine-tune，这样会有train-inference gap。稳妥做法是训练时就用GQA，或者转换后做少量fine-tune。

9.3 和FlashAttention的配合

FlashAttention 2开始原生支持GQA和MQA。用的时候直接把(B, T, num_kv_heads, head_dim)的k、v传进去，不需要手动repeat。这样既省显存又省计算。

from flash_attn import flash_attn_func # q: (B, T, num_heads, head_dim) # k: (B, T, num_kv_heads, head_dim) # 不需要repeat # v: (B, T, num_kv_heads, head_dim) out = flash_attn_func(q, k, v, causal=True)

MLA目前FlashAttention支持还不完善，DeepSeek有自己的kernel实现。如果要用MLA，可能需要自己写kernel或者用官方实现。

9.4 后续可以扩展的方向

Attention变体这个领域还在快速演进。除了MHA、MQA、GQA、MLA，还有NSA（Native Sparse Attention）、MoBA（Mixture of Block Attention）等新结构。核心思路都是在质量、显存、计算之间找更好的平衡点。

如果你想深入，建议从两个方向入手：一是读FlashAttention的论文和代码，理解kernel层面的优化；二是读DeepSeek的技术报告，理解MLA的设计动机。这两个方向能让你从"会背公式"进阶到"理解工程取舍"。

我个人在实际项目中的体会是，不要迷信某个变体。MHA、GQA、MLA都是工具，关键看你的约束是什么。面试时展示的应该是"我能根据约束选工具"的能力，而不是"我会背某个工具"的能力。这个思维转变，比多背几个公式重要得多。

最后分享一个小技巧：面试前把Llama 2、Llama 3、DeepSeek-V2的config.json找出来，看看它们的num_heads、num_kv_heads、head_dim分别是多少，自己算一遍KV Cache。算过一遍，这些数字就长在脑子里了，面试时随手就能报出来。