前言

今天我们来聊聊注意力机制家族里三个绕不开的名字：MHA、GQA、MQA。它们既有关联，又有各自的应用场景，尤其在解码器推理效率上，差别不小。下面就从最基本的自注意力开始，一步步拆解它们的设计思路和代码实现。

self-attention

自注意力机制的输入是一个统一的矩阵，它通过三个不同的线性变换，分别生成查询向量（Q）、键向量（K）和值向量（V）。这三个变换对应的权重矩阵W_Q、W_K、W_V都是可学习的，并且维度由模型设计预先定义。之所以要分成三组，是为了让模型能在不同的表示空间里分别优化匹配（Q和K）和信息合成（V）——这样灵活性更高，表达能力也更强。

在注意力计算中，点积结果会除以 √d_k（d_k是键向量的维度），这是为了防止维度较高时数值不稳定。如果不做这个缩放，当d_k很大时，点积的值会变得非常大，softmax的指数运算就会产生极端结果（有的接近1，有的接近0），导致梯度消失，学习效率大打折扣。除以√d_k之后，softmax的输入控制在一个合适的范围内，保证了数值稳定性和梯度流的有效性。这一步看似简单，却是训练能顺利进行下去的关键。

代码实现

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

class SelfAttention(nn.Module):
    def __init__(self, seq_length):
        super(SelfAttention, self).__init__()
        self.input_size = seq_length
        self.Wq = nn.Linear(seq_length, seq_length)
        self.Wk = nn.Linear(seq_length, seq_length)
        self.Wv = nn.Linear(seq_length, seq_length)

    def forward(self, input):
        q = self.Wq(input)
        k = self.Wk(input)
        v = self.Wv(input)
        attention_scores = torch.matmul(q, k.transpose(-1, -2)) / torch.sqrt(torch.tensor(float(self.input_size)))
        attention_weight = F.softmax(attention_scores, dim=-1)
        output = torch.matmul(attention_weight, v)
        return output

x = torch.randn(2, 3, 4)
Self_Attention = SelfAttention(4)
output = Self_Attention(x)
print(output.shape)

MHA（多头注意力）

MHA的核心思路很简单：把Q、K、V拆成多个头，每个头独立计算自注意力，最后再拼接起来。假设原来Q的形状是(batch_size, seq_len, d_model)，拆成h个头后，每个头维度变成d_model/h。每个头各自计算注意力，得到的输出形状还是(batch_size, seq_len, d_model/h)，拼接后恢复成(batch_size, seq_len, d_model)，再经过一个线性层W_O整合。整个过程可以理解为：把原来的一次大注意力分解成多次小注意力并行执行，让模型从不同子空间学习不同的关系。

,每个头对进行变换后进行注意力机制

代码实现

import torch
import torch.nn as nn

class MultiHeadAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiHeadAttention, self).__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.wq = nn.Linear(embed_dim, embed_dim)
        self.wk = nn.Linear(embed_dim, embed_dim)
        self.wv = nn.Linear(embed_dim, embed_dim)
        self.wo = nn.Linear(embed_dim, embed_dim)

    def mh_split(self, hidden):
        batch_size = hidden.shape[0]
        x = hidden.view(batch_size, -1, self.num_heads, self.head_dim).transpose(1, 2)
        return x

    def forward(self, hidden_states, mask=None):
        batch_size = hidden_states.size(0)
        q, k, v = self.wq(hidden_states), self.wk(hidden_states), self.wv(hidden_states)
        q, k, v = self.mh_split(q), self.mh_split(k), self.mh_split(v)
        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attention = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention, v)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
        output = self.wo(output)
        return output

x = torch.rand(2, 3, 36)
print(x)
output = MultiHeadAttention(36, 6)
y = output(x)
print(y.shape)

MHA能捕捉输入中不同部分之间的复杂关系，但这种能力是有代价的——对内存带宽的需求非常大，尤其是在解码器推理阶段。自回归模型中，每个解码步骤都需要加载解码器权重以及所有历史键和值，这个过程既消耗算力又消耗带宽。随着模型规模增大，这种开销会变得越来越棘手。于是，多查询注意力（MQA）应运而生，它的理念很直接：保留多个查询头，但只用一个键头和一个值头，从而显著减少内存负载，提升推理速度。

MQA（多查询注意力）

MQA是MHA的一种变体，专为自回归解码设计。图1和图2很直观地展示了MHA和MQA的区别：MQA让所有头共享同一份K和V矩阵（意味着K和V的计算是唯一的），只有Q保留了多头的特性（每个头有不同的变换）。这样一来，K和V的参数量以及KV Cache的显存占用都大幅下降，推理速度自然就上去了。当然，精度上的损失是难免的，但胜在效率提升明显。目前很多LLM（比如ChatGLM2）都用了MQA。

一个很实际的问题：如何把现有的预训练多头注意力模型转换成MQA？通常需要两步走：先是模型结构转换，再是额外的预训练。

• 结构转换：将原始模型多个头的键和值投影矩阵（线性层）做均值池化，合并成单一的键和值投影矩阵。经验表明，这种均值池化方式比直接选其中一个头或随机初始化效果更好。转换后的结构就具备了MQA的特征。
• 预训练阶段：结构转换后，模型需要再训一小段时间（通常是原始训练步数的一个小比例α），让它适应新的简化注意力机制。训练方法和原始保持一致，确保学习动态的连续性。

代码实现

import torch
import torch.nn as nn

class MultiQuerySelfAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(MultiQuerySelfAttention, self).__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.wq = nn.Linear(embed_dim, embed_dim)
        # MQA: K和V的输出维度是一个头的维度
        self.wk = nn.Linear(embed_dim, self.head_dim)
        self.wv = nn.Linear(embed_dim, self.head_dim)
        self.wo = nn.Linear(embed_dim, embed_dim)

    def q_h_split(self, hidden, head_num=None):
        batch_size, seq_len = hidden.size()[:2]
        if head_num is None:
            x = hidden.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
            return x
        else:
            # MQA: K和V只需要拆成1个头
            return hidden.view(batch_size, seq_len, head_num, self.head_dim).transpose(1, 2)

    def forward(self, hidden_states, mask=None):
        batch_size = hidden_states.size(0)
        q, k, v = self.wq(hidden_states), self.wk(hidden_states), self.wv(hidden_states)
        q, k, v = self.q_h_split(q), self.q_h_split(k, 1), self.q_h_split(v, 1)
        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
        print("scores:", scores.shape)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attention = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention, v)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
        output = self.wo(output)
        return output

x = torch.rand(3, 12, 512)
atten = MultiQuerySelfAttention(512, 8)
y = atten(x)
print(y.shape)

GQA（分组查询注意力）

MQA虽然参数量大幅减少，推理加速明显，但性能损失和训练不稳定也是实实在在的毛病——复杂度降低往往伴随着质量下降。于是GQA被提了出来：把Query分成若干组，每组内部共享一组Key和Value。这一思路在LLaMA-2和Mistral 7B中已经得到了应用。

GQA的数学原理很简单：将传统多头模型中的查询头分成G个组，每组分配一个键头和值头。这种配置记作GQA-G，G就是组数。两个极端情况：G=1时就是MQA；G=H（头数）时退化为MHA。从MHA转换到GQA同样可以使用均值池化：对每组中原始头的键和值投影矩阵做平均，得到该组单一的键值和值投影。

通过GQA，模型能在MHA的质量和MQA的速度之间找到一个平衡点。键值对少了，内存带宽和数据加载需求自然降低。G的选择代表了一种权衡：更多组（接近MHA）质量更高但速度较慢，更少组（接近MQA）速度更快但可能牺牲质量。而且随着模型规模变大，GQA允许内存带宽和模型容量按比例减少，而MQA中只有一个键值头，对超大模型的限制可能过于严苛。

代码实现

import torch
import torch.nn as nn

class GroupedQueryAttention(nn.Module):
    def __init__(self, embed_dim, num_heads):
        super(GroupedQueryAttention, self).__init__()
        self.num_heads = num_heads
        self.head_dim = embed_dim // num_heads
        self.wq = nn.Linear(embed_dim, embed_dim)
        # GQA: 设置4个组，K和V的输出维度是 group_num * head_dim
        self.group_num = 4
        self.wk = nn.Linear(embed_dim, self.group_num * self.head_dim)
        self.wv = nn.Linear(embed_dim, self.group_num * self.head_dim)
        self.wo = nn.Linear(embed_dim, embed_dim)

    def split(self, hidden, group_num=None):
        batch_size, seq_len = hidden.size()[:2]
        if group_num is None:
            # Q拆成所有头
            x = hidden.view(batch_size, seq_len, self.num_heads, self.head_dim).transpose(1, 2)
            return x
        else:
            # K和V拆成 group_num 个头，然后扩展为多头
            x = hidden.view(batch_size, seq_len, group_num, self.head_dim).transpose(1, 2)
            x = x[:, :, None, :, :].expand(batch_size, group_num, self.num_heads // group_num, seq_len, self.head_dim).reshape(batch_size, self.num_heads, seq_len, self.head_dim)
            return x

    def forward(self, hidden_states, mask=None):
        batch_size = hidden_states.size(0)
        q, k, v = self.wq(hidden_states), self.wk(hidden_states), self.wv(hidden_states)
        q, k, v = self.split(q), self.split(k, self.group_num), self.split(v, self.group_num)
        scores = torch.matmul(q, k.transpose(-2, -1)) / torch.sqrt(torch.tensor(self.head_dim, dtype=torch.float32))
        print("scores:", scores.shape)
        if mask is not None:
            scores = scores.masked_fill(mask == 0, float('-inf'))
        attention = torch.softmax(scores, dim=-1)
        output = torch.matmul(attention, v)
        output = output.transpose(1, 2).contiguous().view(batch_size, -1, self.num_heads * self.head_dim)
        output = self.wo(output)
        return output

x = torch.ones(3, 12, 512)
atten = GroupedQueryAttention(512, 8)
y = atten(x)
print(y.shape)