MQA多查询注意力：让全馆读者共享一份KV目录

可视化演示

展示UDP无连接数据报传输过程：发送端直接发送数据包，不建立连接、不确认接收，可能出现丢包和乱序。

注意力回顾

想象你走进一座图书馆，手里握着一张写满问题的纸条。你首先跑到中央索引柜前，用纸条上的关键词去比对一张张图书索引卡——卡片上的分类标签（Key）与你的问题（Query）越接近，卡片指向的那本书（Value）就越可能对你有用。你抽出几本匹配度最高的书，把它们的核心内容按匹配程度加权汇总成一份新笔记。这套「问题→匹配→取书→汇总」的流程，就是注意力机制的雏形。

在Transformer模型中，输入序列里的每个词都同时扮演「读者」和「书」两个角色。模型为每个词生成三个向量：查询Q、键K、值V。当前词的Q会和序列中所有词的K逐一比对相似度，得到注意力权重，再按这些权重对所有位置的V加权求和，最终得到融合了全局信息的新表示。这就像每个词都进了一次图书馆，带走一份按自己需求定制的摘要笔记。

缩放点积注意力的标准公式是：

实际模型不会只派一个读者去查书，而是同时派出多个「读者小组」，每个组关注不同角度——有的盯着语法搭配，有的捕捉指代关系，有的挖掘语义相似性。这就是多头注意力（MHA）：h个小组并行检索，每组都用自己独立的投影矩阵生成Q_h、K_h、V_h。同一个词经由不同小组的解读，能同时表达多种含义。

多头并行就像图书馆同时接待多个主题小组，互不干扰、各取所需。但请注意：每个小组都自带一套完整的索引卡和书库——这为后面的「空间爆炸」埋下了伏笔。

MHA痛点

图书馆的规模越来越大，读者多了、馆藏也多了。为了照顾好每一位读者小组，管理层决定：给每个小组都配一套「私人馆藏」——完整复印一套索引卡，把书也重新搬来一份。32个小组，就是32套一模一样的馆藏副本。读者来查资料时，管理员要在32套房子里来回穿梭，取书路径成倍变长，馆内通道拥挤不堪。

换成标准多头注意力的语言，这个困境发生在推理阶段。模型每生成一个token，都必须重新「看一眼」历史token的全部信息。为了避免反复重算，工程上把历史token的Key和Value缓存起来——这就是KV Cache。可缓存是按层、按头、按序列长度堆积的，数值大得惊人。

单层KV缓存体积的计算公式如下：

解码阶段，每生成一个token都要把全部KV缓存从显存搬进计算单元。搬运的时间远大于计算本身——这就是内存带宽瓶颈。你也许会觉得计算量很大，但真正的瓶颈恰恰是「取书」的物流成本，而不是「读书」的思考成本。图书馆的通道就那么多，搬运32套副本的代价，最终变成了推理速度的拖累。

MQA核心

一位新来的图书管理员提出了大胆建议：既然所有小组查的是同一批馆藏，为什么每个人都要搬一套副本？干脆把32套书架合并成一间公共书库——全馆共享同一套索引卡、同一批藏书；每个小组只需要保留自己那份独特的问题清单（Q）。小组之间分工的差异，由「问什么问题」来体现，而查到的「书」始终是同一批。

这个方案就是MQA（Multi-Query Attention）。改动极小却效果显著：h个查询头依然用独立的权重矩阵投影出各自的Q_h，但Key和Value的投影合并为单头，只生成一份K_shared和V_shared。换句话说，MQA从「多Q多K多V」变成了「多Q单K单V」。

共享KV后的单头计算可写为：

从「32套书架」到「1套公共书库+32份问题清单」，KV缓存体量直接缩为原来的1/h。图书馆的空间压力瞬间释放，管理员再也不用在32个房间之间奔波了。

GQA折中

不过，公共书库也引来了部分读者的抱怨：有些专业领域差异太大，大家共用一套索引卡后，检索粒度不够精细。于是图书馆又尝试了一种更灵活的运营模式——不搞32套，也不搞1套，而是把32个小组分成G个「部室」，每个部室共用一套索引卡和书架。几个小组之间的需求比较相近，共享一套完全够用；不同部室之间仍保持各自的检索特色。

这种折中方案正是GQA（Grouped Query Attention，分组查询注意力）。设分组数为G，每个组包含h/G个查询头，组内共享一套K/V。它是MHA和MQA之间的连续插值：当G=1时，所有查询头共享一套KV，也就是MQA；当G=h时，每组仅一个头，退化为原始的MHA。

GQA的缓存体积公式如下：

图书馆的「部室制」既保留了一定检索灵活性，又把整体空间压了下来。在部署选型时，有了MHA、GQA、MQA三档「馆藏策略」，便可按显存预算和服务质量要求自由挑选。

加速收益

书库合并之后，最直观的变化并非「省了几个柜子」，而是每天在馆内搬运图书的实际工作量骤减。管理员响应一次查询，只需在唯一一套公共书库里完成取书，不再被迫跑遍32个房间。同一时间能服务的读者数量因此直线上升，队伍的吞吐能力大幅增强。

在多查询注意力中，这背后的机制是内存带宽的释放。自回归解码是逐token串行进行的，每生成一个新token，就要把KV缓存整体读取一遍。这种场景下，读取数据的总字节数决定了延迟上限，计算单元本身多数时间在空等数据。MQA把需要搬运的KV数据量从h份降为1份，内存带宽瓶颈瞬间缓解。

KV缓存缩减倍数可以写成一个简明的关系式：

图书馆的「高峰通行能力」取决于搬运工每天走的路程。共享书库让每条查询路径缩短为原来的1/h，管理员得以接待更多读者，甚至允许大家借阅更厚的长篇著作——也就是支撑更长的上下文窗口。

质量权衡

共享书库也并非毫无代价。当每个小组都必须依赖同一套书目分类和藏书布局时，各小组的「检索视角」被强行对齐，它们不能再按各自偏好设计一套最适合自己的索引体系。少数定位特殊的查询，匹配精度会有所下降。训练时，模型的整体表达能力也会受到轻微限制。

不过实验观察表明，MQA相对MHA的精度损失通常远小于直觉预期。原因在于Q矩阵保留了多头的全部差异——不同头依然可以用独特的「提问方式」，从同一批藏书中抽取不同侧面的信息。更意外的是，KV参数量越少，量化时的动态范围越稳定，MQA在INT8/INT4量化下的误差反而更可控。

用一个参数视角来量化共享程度对表达力的影响：

图书馆可以在「每人一套」「分组共享」「全馆共享」三套方案间按预算与体验权衡。显存宽裕、精度敏感，保留MHA；显存紧张、追求吞吐，切换到MQA；想两头兼顾，GQA的旋钮可以精细调节。

AI 应用场景

Python 演示

import mathimport torchimport torch.nn as nn  class MHA(nn.Module):    """标准多头注意力: Q/K/V 全部独立多头"""     def __init__(self, d_model, n_heads):        super().__init__()        self.n_heads = n_heads        self.d_k = d_model // n_heads        self.wq = nn.Linear(d_model, d_model)        self.wk = nn.Linear(d_model, d_model)        self.wv = nn.Linear(d_model, d_model)        self.wo = nn.Linear(d_model, d_model)     def forward(self, x):        b, t, _ = x.shape        q = self.wq(x).view(b, t, self.n_heads, self.d_k).transpose(1, 2)        k = self.wk(x).view(b, t, self.n_heads, self.d_k).transpose(1, 2)        v = self.wv(x).view(b, t, self.n_heads, self.d_k).transpose(1, 2)        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)        attn = torch.softmax(scores, dim=-1)        out = (attn @ v).transpose(1, 2).reshape(b, t, -1)        return self.wo(out)  class MQA(nn.Module):    """多查询注意力: Q 保持多头, K/V 只保留单头共享"""     def __init__(self, d_model, n_heads):        super().__init__()        self.n_heads = n_heads        self.d_k = d_model // n_heads        self.wq = nn.Linear(d_model, d_model)        self.wk = nn.Linear(d_model, self.d_k)        self.wv = nn.Linear(d_model, self.d_k)        self.wo = nn.Linear(d_model, d_model)     def forward(self, x):        b, t, _ = x.shape        q = self.wq(x).view(b, t, self.n_heads, self.d_k).transpose(1, 2)        k = self.wk(x).unsqueeze(1)   # (b,1,t,d_k)        v = self.wv(x).unsqueeze(1)   # (b,1,t,d_k)        scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)        attn = torch.softmax(scores, dim=-1)        out = (attn @ v).transpose(1, 2).reshape(b, t, -1)        return self.wo(out)  def kv_cache_bytes(n_heads, d_k, seq_len, n_layers, dtype_bytes=2):    """单层 KV 缓存字节数: 2(Key+Value) * n_heads * d_k * seq_len"""    return 2 * n_heads * d_k * seq_len * n_layers * dtype_bytes  def main():    d_model, n_heads, seq_len, n_layers = 4096, 32, 2048, 32    d_k = d_model // n_heads    for name, heads in [("MHA", n_heads), ("GQA组8", 8), ("GQA组4", 4), ("GQA组2", 2), ("MQA", 1)]:        cache = kv_cache_bytes(heads, d_k, seq_len, n_layers)        ratio = heads / n_heads        print(f"{name:8s} KV缓存: {cache / 1e9:6.2f} GB  占比: {ratio * 100:4.1f}%")     torch.manual_seed(0)    x = torch.randn(2, 16, d_model)    mha = MHA(d_model, n_heads)    mqa = MQA(d_model, n_heads)    print("MHA 输出:", tuple(mha(x).shape))    print("MQA 输出:", tuple(mqa(x).shape))  if __name__ == "__main__":    main()

1. self.wk = nn.Linear(d_model, self.d_k)MQA的K投影只输出单头维度，全头共享一份，这是节省缓存的关键。
2. k = self.wk(x).unsqueeze(1)单头K扩展出「组」维度，让后续所有查询头都能广播使用同一份K。
3. scores = q @ k.transpose(-2, -1) / math.sqrt(self.d_k)每个查询头分别与共享K做缩放点积，完成多头「各自的检索」。
4. out = (attn @ v).transpose(1, 2).reshape(b, t, -1)注意力权重对共享V加权求和，再把多头结果拼回完整向量。
5. cache = kv_cache_bytes(heads, d_k, seq_len, n_layers)把缓存公式2×heads×d_k×L代入程序，打印MHA到MQA的缩减梯度。

常见误区

面试题