【导读】：本文是LLM知识点第四篇，介绍LLM中的最重要的Attention机制，具体有Attention机制，Self-Attention，Multi-head Attention的原理和实现细节，接着会介绍围绕KV Cache进行的推理优化而提出的MQA，GQA和MLA。
Attention机制
【1】简要介绍Attention机制
针对问题改进：提升模型的并行运算能力，序列中的每个token能无损地捕获序列里的其他tokens信息。改进办法就是Attention。如蓝色方框为attention模型。在每个位置，例如在a2处产生b2时，attention将会同时看过a1到a4的每个token。此外，每个token生成其对应的输出的过程是同时进行的，计算不需要等待。
【推荐阅读】Transformer学习笔记二：Self-Attention（自注意力机制）-猛猿【202111】
Attention的优点：1.相对于传统的CNN和RNN,Attention模型的参数数量更少。2.使得Transformer模型在计算Query时实现并行计算。3.使得模型能够更好地捕捉长距离的依赖关系，模型效果更好。
【2】Attention机制的计算流程
【带权求和】
Attention机制计算流程：
1.计算Query和各个Key的相似度(点乘)，得到每个Key对应Value 的权重系数s；
【3】以机器翻译为例理解Attention过程
以李宏毅深度学习授课资料：Attention-based Model所讲述的机器翻译为例。注：传统Attention模块能捕获source端和target端的token间的依赖关系。
目标：将‘机器学习’翻译为‘machine learning’。attention其实就是一个当前的输入与输出的匹配度，即为h1和z0的匹配度。h1为当前时刻RNN的隐层输出向量，而不是原始输入的词向量，z0初始化向量，如rnn中的初始记忆。
第一步：求z与h的相似性得到a；第二步：softmax归一化处理得到概率值a^；第三步：对h加权求和得c；
【推荐阅读】https://www.jianshu.com/p/d7f50cc5560e
【4】Attention的机制的代码实现
注意力机制的代码实现
import torch
import torch.nn as nn
from typing import List
def get_input_embeddings(words: List[str], embeddings_dim: int):
# we are creating random vector of embeddings_dim size for each words
# normally we train a tokenizer to get the embeddings.
# check the blog on tokenizer to learn about this part
embeddings = [torch.randn(embeddings_dim) for word in words]
return embeddings
text = "I should sleep now"
words = text.split(" ")
len(words) # 4
embeddings_dim = 512 # 512 dim because the original paper uses it. we can use other dim also
embeddings = get_input_embeddings(words, embeddings_dim=embeddings_dim)
embeddings[0].shape # torch.Size([512])
# initialize the query, key and value metrices
query_matrix = nn.Linear(embeddings_dim, embeddings_dim)
key_matrix = nn.Linear(embeddings_dim, embeddings_dim)
value_matrix = nn.Linear(embeddings_dim, embeddings_dim)
query_matrix.weight.shape, key_matrix.weight.shape, value_matrix.weight.shape # torch.Size([512, 512]), torch.Size([512, 512]), torch.Size([512, 512])
# query, key and value vectors computation for each words embeddings
query_vectors = torch.stack([query_matrix(embedding) for embedding in embeddings])
key_vectors = torch.stack([key_matrix(embedding) for embedding in embeddings])
value_vectors = torch.stack([value_matrix(embedding) for embedding in embeddings])
query_vectors.shape, key_vectors.shape, value_vectors.shape # torch.Size([4, 512]), torch.Size([4, 512]), torch.Size([4, 512])
# compute the score
scores = torch.matmul(query_vectors, key_vectors.transpose(-2, -1)) / torch.sqrt(torch.tensor(embeddings_dim, dtype=torch.float32))
scores.shape # torch.Size([4, 4])
# compute the attention weights for each of the words with the other words
softmax = nn.Softmax(dim=-1)
attention_weights = softmax(scores)
attention_weights.shape # torch.Size([4, 4])
# attention output
output = torch.matmul(attention_weights, value_vectors)
output.shape # torch.Size([4, 512])
【5】Attention的应用场景
图片来自https://mp.weixin.qq.com/s/OabP1Apl5ullK4oxiyfqJA
【1】简要介绍Self-Attention
在一般任务的Encoder-Decoder框架中，输入 Source 和输出 Target 内容是不一样的，比如对于英 - 中机器翻译来说，Source是英文句子，Target是对应的翻译出的中文句子。
Self-Attention的计算方式:
Self-attention的意思是，给Attention的输入都来自同一个序列，其计算方式如下：
self-attention计算框架
【2】Self-Attention的计算公式, 参数量
缩放点积注意力：Scaled Dot-Product Attention
注：
X的维度(seq_length，d_model)=(单词个数,词向量维度);
WQ=WK维度(hidden_size，hidden_size);
WV维度(hidden_size，hidden_size);可设置
Q, K, V 维度(seq_length,hidden_size);
QKᐪ维度(seq_length,seq_length);
hidden_size=d_model//num_head,这里num_head=1,有hidden_size=d_model
Self-Attention的参数量：
上面是一个简化的算法，multi-head attention也符合这个参数量，但实际上需要分head计算再组合。
附Attention(Q,K,V)计算图解：
注意：此处得到矩阵Z,第1行代表x1的注意力得分。
【3】Self-Attention的计算过程和图解
【词向量矩阵X】Self-Attention的计算过程图解
a.Query矩阵，Key矩阵，Value矩阵的计算
b.最后计算得分和注意力输出
说明：
a.这里矩阵X计算self-attention，Q,K,V计算同源都使用词向量矩阵X;
b.对于每个词向量xi，计算得到Q,K,V后，都有Qi,Ki,Vi与之对应。
Self-Attention的计算过程详细步骤：
1.将词向量矩阵X与权重矩阵
2.将Q和K相乘得到权重Score矩阵,除以根号dk并softmax得到新的权重score矩阵。
3.将权重score矩阵和该序列的V矩阵相乘，得到该序列对应的Attention得分矩阵。
【3】Self-Attention的计算过程和图解2
【词向量矩阵X】Self-Attention的计算过程图解2
a.Query矩阵，Key矩阵，Value矩阵的计算
b.最后计算得分和注意力输出
【引用链接】Transformer学习笔记二：Self-Attention（自注意力机制）-猛猿【202111】
【4】Self-Attention的计算过程和图解
【词向量x】Self-Attention的计算过程图解：
说明：
a.对于每个词向量xi，计算得到Q,K,V后，都有qi,ki,vi与之对应。
b.每个词向量xi,其Attention-score计算结果为zi。zi的维度为(1,seq_length)。
c.对于单个词向量x1计算步骤，q1要和 K矩阵 相乘是 q1Kᐪ，经scaled和softmax得到 权重矩阵 a ，
将权重矩阵a1和 V矩阵相乘是 a1·V= z1。
某个单词Zi 的计算方法：softmax后的score矩阵的第1行表示单词x1与其他所有单词的attention系数。最终单词 x1的输出 z1 等于所有单词i 的值 ?? 根据 attention 系数的比例相加得到，如下图所示：
d.注意看QKᐪ,其实是个word2word的attention map！(加了softmax之后就是一个和为1的权重了)。
比如输入序列是 "i have a dream" 总共4个单词，这里就会形成4x4的注意力机制的图：
注意encoder里面是叫self-attention，decoder里面是叫masked self-attention。
mask就是沿着对角线把灰色的区域用0覆盖掉，不给模型看到未来的信息。
masked self-attention就是只能看见当前位置和之前token的信息。如a作为第三个单词，有和i,have,a 三个单词的attention。
【5】计算Self-Attention是为什么要除以d_k的开方？
Transformer计算 self-attention的计算QKᐪ后为什么要除以 d_k的开方？【重点】
计算QKᐪ后除以 d_k的开方是为了压缩softmax的输入值，使得输入值分布变得更好，能够进入softmax的敏感区间，这样在训练的时候不至于梯度值更新的太小，能够防止梯度消失问题。如果不进行scaleing，模型训练时将很难收敛。
【dk就是词向量维度/隐藏层维度】
论文中作者的解释是发现当维度dk值很大时，输入softmax的值QKᐪ就越大，会导致后面的softmax计算会有极小的梯度，不利于更新学习，因此除以dk，防止梯度消失。(softmax值过大,其偏导数趋于0)
We suspect that for large values of dk, the dot products grow large in magnitude,pushing the softmax function into regions where it has extremely small gradients.
计算QKᐪ后为什么scaled参数选择除根号d_k？
选择根号d_k是因为可以使得QKᐪ的结果满足期望为0，方差为1的分布，类似于归一化。于此同时，除以根号d_k使得输入值进入softmax的敏感区间，能够防止梯度消失问题。
计算QKᐪ后有其他方法不用除根号dk吗？
有，只要能做到每层参数的梯度保持在训练敏感的范围内，使得网络比较好训练。能缓解梯度消失的问题就可以。那么这个网络就比较好训练。不用除根号dk方式有，详情可以了解Google T5的Xavier初始化。
【6】为什么Q和K要使用不同的权重矩阵进行线性变换投影？
Self-Attention计算公式：
为什么要计算Q和K的点乘？
简单的说，Q和K的点乘是为了计算序列中每个token与其他token的相似度,最终得到attention score 矩阵，用来对V进行提纯。假设一个句子"Hello, how are you?"长度是6，embedding维度是300，那么Q，K，V都是(6, 300)的矩阵。比如说 "Hello, how are you?"这句话，当前token为”you"的时候，可以知道”you“对于"Hello", ” , “, "how", "are", "?"这几个token对应的关注度是多少。有了这个attention score，可以知道处理到”you“的时候，模型在关注句子中的哪些token。
如果WQ和WK一样，则QKᐪ结果是对称矩阵，这样就减弱了模型的表达能力。同时在对称矩阵中，对角线的值会比较大，导致每个token过分关注自己。使用不同的投影矩阵，参数增多，可以增强模型表达能力。
为什么要使用Q,K,V，仅仅使用QV，KV或者V为什么不行？
使用QKV主要为了增强网络的容量和表达能力。self-attention 使用 Q,K,V这三个参数独立，模型的表达能力和灵活性显然会比只用QV或者只用 V 要好些，
【7】Self-Attention实现的细节问题
Self-Attention的计算公式（QKV计算不加偏置项）：
缩放点积注意力：Scaled Dot-Product Attention
Self-Attention的计算公式（QKV计算加偏置项）：
Self-Attention计算时乘上WQ.WK.WV的好处？
1.增加了参数量，增加模型的表达能力。
2.加入了不同的线性变换相当于对x 做了不同的投影，将向量x 投影到不同空间，增加模型的泛化能力。
3.允许某个token对其他位置token的注意力大于对自己的注意力，才能更好的捕捉全局位置的注意力。
【引用链接】：https://zhuanlan.zhihu.com/p/626820422
Self-Attention 的时间复杂度是怎么计算的？
Self-Attention时间复杂度：O(n²⋅d) ，这里，n是序列的长度seq_length，d是embedding的维度d_model。
Self-Attention包括三个步骤：相似度计算，softmax和加权平均，时间复杂度分别是：
相似度计算可看作为(n,d)和(d,n)的矩阵相乘: (n,d)∗(d,n)=O(n²⋅d) ，
softmax就是直接计算了，时间复杂度为 O(n²)
加权平均可看作为(n,n)和(n,d)的矩阵相乘: (n,n)∗(n,d)=O(n²⋅d)
Self-Attention的时间复杂度是 O(n²⋅d) 。
【引用链接】：https://zhuanlan.zhihu.com/p/132554155
Self-Attention和Multi-head Attention的参数量怎么计算？
self-attention块的模型参数有Q,K,V的权重矩阵WQ,WK,WV和偏置,输出权重矩阵Wo和偏置。
4个权重矩阵的形状为【h,h】,4个偏置的形状为【h】。总参数量为4h²+4h.
Multi-head Attention也符合这个参数量，但实际上需要分head计算再组合。
Self-Attention 是如何 解决 长距离依赖问题的呢?
解决方式： 利用注意力机制来“动态”地生成不同连接的权重，从而处理变长的信息序列。
具体介绍： 对于当前query，需要与句子中所有 key 进行点乘后再 Softmax ，以获得句子中所有 key 对于当前query的score(可以理解为权重)，然后与所有词 的value向量进行加权融合之后，就能使当前token学习到句子中其他词的信息；
Transformer 中self-attention是如何并行化的？
Transformer 的并行化主要体现在 self-attention 模块，在Encoder端 Transformer可以并行处理整个序列，并得到整个输入序列经过 Encoder 端的输出，在 self-attention 模块，对于某个序列(x1,x2,...xn)，self-attention 模块可以直接计算xi,xj的点乘结果，而RNN系列的模型就必须按照顺序从x1计算到xn。
在Self-Attention能够并行的计算句子中不同的query，每个query之间并不存在先后依赖关系，使得transformer能够并行化；
Self-Attention 在计算的过程中，如何对padding位做mask？
【参考链接】：https://zhuanlan.zhihu.com/p/149634836
Attention与全连接层的区别何在？既然Attention是为了关注某些局部信息，那些不就相当于连上一层在关注的部分权重更大的全连接层吗，二者的区别何在？正如你所说的，Attention的最终输出可以看成是一个“在关注部分权重更大的全连接层”。但是它与全连接层的区别在于，注意力机制可以利用输入的特征信息来确定哪些部分更重要。
Self Attention的一些反思【重要】
深入思考，会发现它真的是一个很神奇的存在，它是BERT乃至整个预训练语言模型的基石，是接棒CNN/RNN，成为特征抽取的新利器。Attention is all you need !
0、深度学习中Attention与全连接层的区别何在？[15]
注：这是一个检验你是否真正理解Attention的问题
1、self-attention 的本质是什么？包括哪几个步骤？和普通 Attention 的差别在哪里？[4]
2、不考虑多头的原因，self-attention中词向量不乘QKV参数矩阵，会有什么问题？[4]
3、在普通 attention 中，一般有 k=v，那 self-attention 可以嘛？[4]
4、self-attention 在计算的过程中，如何对padding位做mask？[2]
5、bert的mask为何不学习transformer在attention处进行屏蔽score的技巧？[11]
6、XLNet为什么不直接在attention掩码矩阵中只把当前的单词掩盖住来获取上下文的信息呢？直接mask住左上到右下的对角线构建双向语言模型不行吗？[3]
【1】MHA提出的背景
MHA提出论文：《Attention Is All You Need》
MHA实现思路：简单来说，MHA在d_model维度进行拆分，将原来单个self-attention的计算拆分成h个小份self-attention分别计算，最后将计算结果合并和变换回原来的维度。h是head的个数。
所谓Multi-Head Attention其实是把QKV的计算并行化，原始attention计算d_model维的向量，而Multi-Head Attention则是将d_model维向量先经过一个Linear Layer，再分解为h个Head计算（d_model/h维）attention，最终将这些attention计算结果连在一起后再经过一层Linear Layer输出。[配合图解1理解]
MHA简要介绍：为了提高模型的表达能力，Transformer引入了多头注意力机制，允许模型学习多组不同的注意力权重。每个注意力头都产生一个输出，最后通过线性变换和拼接得到最终的多头注意力输出。
在MHA计算整个过程中需要4个输入和输出维度都是d_model的Linear Layer，而整个Model的输入是(batch_size, seq_length, d_model)，输出也是(batch_size, seq_length, d_model)。
Multi-head attention allows the model to jointly attend to information from different representation subspaces at different positions. MHA将隐状态向量分成多个头，形成多个子语义空间，可以让模型去关注不同维度语义空间的信息。
【2】MHA的计算公式
Multi-Head-Attention的公式：Multi-Head-Attention就是在单头Self-Attention的基础上，在隐状态维度的方向将其切分成H个头，再分别计算得到headi 结果后进行concat【隐状态维度=hidden_size=d_model】
公式如下所示：
假设H=3,d_model=dv=dq=dk=6, seq_length=3图解过程如下所示：
【推荐阅读】https://zhuanlan.zhihu.com/p/626820422
在《Attention Is All You Need》论文原文中解释了多头的作用：将隐状态向量分成多个头，形成多个子语义空间，可以让模型去关注不同维度语义空间的信息（或者说让模型去关注不同方面的信息）。
常见的设置是dk=dv= d/h,
对于LLAMA2-7b 有 d=4096,h=32,dk=dv=128
对于LLAMA2-70b有 d=8192,h=64,dk=dv=128
【3】MHA的计算流程图解
MHA多头注意力计算过程图解1
该图MHA的计算流程如下：
1.输入句子, 这里是Thinking Machines; [Input sentence]
2.对句子进行tokenize和embedding并得到X, 后续计算需使用encoder输出R代替X。[Embed each word]。X维度为(seq_length,d_model)=(2,4);
3.将WQ,WK,WV划分多头h=8,,构造多头注意力权重矩阵
4. 将X分别和 8 组权重矩阵
5.将8组注意力结果
Z'维度为(seq_length,d*h)=(2,24)。
【此处seq_length=2,d_model=4,dq=dv=dk=3,h=8】
【推荐阅读】图解Transformer最关键模块MHA【20240214】
https://mp.weixin.qq.com/s/-jeuRsrOehG5ydba0txLug
MHA多头注意力计算过程图解2
https://mp.weixin.qq.com/s/QjK_gL4pHCLHNi6y4vsM6g
MHA多头注意力计算过程图解3
https://mp.weixin.qq.com/s/yhEa1olOjyQ4oaDhHocB6A
【4】MHA的相关细节解析
3.MHA的多头使用几个头比较好？
已有论文证明，头数不是越多越好，论文实验结果如下：
由A组实验，可以看到多头h=8/16时，PPL/BLEU最好，h=4/32次之，h=1最差。
头的数量不是越多越好，头越多，每个qkv分到的维度就会降低，各个子空间的维度越低，表达能力也就变差，也未必能更好的捕捉到语法/句法/词法信息。具体多头的数量要视模型规模，任务而定。
目前可以看到的趋势是，模型越大（hidden size越大），头数的增多越能带来平均效果上的收益（或者说允许注意力头增大而不影响子空间的学习能力）。目前LLM主流的头数视乎模型结构和规模大致有12、16、24、48、96这样一些主流设置。
总结一下就是，Bert类的attention矩阵是稀疏的（这个是双向信息导致的），因此attention矩阵可以被低秩分解，也就是说虽然矩阵很大但是信息却很少。D远大于d的情况下，T*D 的向量得到的attention矩阵和 T*d 的向量得到的attention矩阵信息量差不多，那不如把D切成若干个d，搞出若干个头来，相当于融合了，因此有了多头。苏神虽然和本文得出的结论是一致的，都是多头融合理论，不过苏神给出了一个全新的视角，非常有意思。
【引用链接】https://zhuanlan.zhihu.com/p/626820422
先说结论: 区别在于模型容量增加，带来模型表现力的提升。
https://www.zhihu.com/question/446385446/answer/1982483918
【5】MHA的代码实现
#The Annotated Transformer中的MHA代码实现
def attention(query, key, value, mask=None, dropout=None):
"Compute 'Scaled Dot Product Attention'"
d_k = query.size(-1)
scores = torch.matmul(query, key.transpose(-2, -1)) \
/ math.sqrt(d_k)
if mask is not None:
scores = scores.masked_fill(mask == 0, -1e9)
p_attn = F.softmax(scores, dim = -1)
if dropout is not None:
p_attn = dropout(p_attn)
return torch.matmul(p_attn, value), p_attn
class MultiHeadedAttention(nn.Module):
def __init__(self, h, d_model, dropout=0.1):
'''
h: head number
'''
super(MultiHeadedAttention, self).__init__()
assert d_model % h == 0
# We assume d_v always equals d
self.d = d_model // h
self.h = h
self.linears = clones(nn.Linear(d_model, d_model), 4)
self.attn = None
self.dropout = nn.Dropout(p=dropout)
def forward(self, query, key, value, mask=None):
if mask is not None:
# Same mask applied to all h heads.
mask = mask.unsqueeze(1)
nbatches = query.size(0)
# 1) Do all the linear projections in batch from d_model => h x d
query, key, value = \
[l(x).view(nbatches, -1, self.h, self.d).transpose(1, 2)
for l, x in zip(self.linears, (query, key, value))]
# 2) Apply attention on all the projected vectors in batch.
x, self.attn = attention(query, key, value, mask=mask,
dropout=self.dropout)
# 3) "Concat" using a view and apply a final linear.
x = x.transpose(1, 2).contiguous() \
.view(nbatches, -1, self.h * self.d)
return self.linears[-1](x))# 1) Do all the linear projections in batch from d_model => h x d_k
query, key, value = \
[l(x).view(nbatches, -1, self.h, self.d_k).transpose(1, 2)
for l, x in zip(self.linears, (query, key, value))]
#重写这段代码
#第一行把QKV分别经过一层Linear变换，tensor size不变，第二行将QKV的d_model维向量分解为h * d_k。
query, key, value = [l(x) for l, x in zip(self.linears,(query, key, value))]
query, key, value = [x.view(nbatches, -1, self.h,self.d_k).transpose(1, 2)
for x in (query, key, value)]
#跑一个self-attention的实例，作为输入，query/key/value的shape为(batch_size, seq_lengh, d_model):
h = 8
d_model = 512
batch_size = 1
seq_length = 10
model = MultiHeadAttention(h, d_model)
query = torch.randn([batch_size, seq_length, d_model])
key = query
value = query
print ('Input size: ' + str(query.size()))https://github.com/zzzichen277/LLM_CodingSet/blob/main/10.multiheadattention%E5%8F%98%E6%8D%A2%E8%BF%87%E7%A8%8B%E4%BB%A3%E7%A0%81.ipynb
【推荐阅读】https://mp.weixin.qq.com/s/lMXf6zRWrH4SV8YbQ2Jzcw
【6】Self-Attention代码LLama实现
使用hf的代码实现（删掉RoPE）Multi-headed attention
#https://zhuanlan.zhihu.com/p/643829565
from torch import nn
class LlamaAttention(nn.Module):
"""Multi-headed attention from 'Attention Is All You Need' paper"""
def __init__(self, config: LlamaConfig): super().__init__()
self.config = config
self.hidden_size = config.hidden_size # 隐藏层大小
self.num_heads = config.num_attention_heads# 注意力头的数量
self.head_dim = self.hidden_size // self.num_heads # 每个注意力头的维度
self.max_position_embeddings = config.max_position_embeddings # 最大位置编码长度
# 线性变换层，用于产生查询（q）、键（k）、值（v）向量
self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.k_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.v_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)
self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)
def forward(
self,
hidden_states: torch.Tensor, # 输入的隐藏状态张量
attention_mask: Optional[torch.Tensor] = None, # 注意力掩码，可选
position_ids: Optional[torch.LongTensor] = None, # 位置编码，可选
past_key_value: Optional[Tuple[torch.Tensor]] = None, # 过去的键值对，可选
output_attentions: bool = False, # 是否输出注意力权重
use_cache: bool = False, # 是否使用缓存
) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
bsz, q_len, _ = hidden_states.size()
# 获得qkv向量
query_states = self.q_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
key_states = self.k_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
value_states = self.v_proj(hidden_states).view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)
# 拼接kvcache
kv_seq_len = key_states.shape[-2]
if past_key_value is not None:
kv_seq_len += past_key_value[0].shape[-2]
if past_key_value is not None:
# reuse k, v, self_attention
key_states = torch.cat([past_key_value[0], key_states], dim=2)
value_states = torch.cat([past_key_value[1], value_states], dim=2)
past_key_value = (key_states, value_states) if use_cache else None
# 计算attention权重
attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)
# 加入mask矩阵，decoder-only为下三角
if attention_mask is not None:
attn_weights = attn_weights + attention_mask
dtype_min = torch.tensor(
torch.finfo(attn_weights.dtype).min, device=attn_weights.device, dtype=attn_weights.dtype
)
attn_weights = torch.max(attn_weights, dtype_min)
# 计算softmax，这里需要从fp16升为fp32
# upcast attention to fp32
attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)
attn_output = torch.matmul(attn_weights, value_states)
attn_output = attn_output.transpose(1, 2)
attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)
attn_output = self.o_proj(attn_output)
if not output_attentions:
attn_weights = None
return attn_output, attn_weights, past_key_value
LLaMA的huggingface实现：https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py#L232
【7】MHA.MQA.GQA对比
GQA与MHA、MQA的对比如下图所示：
当前有如下 3 种主流的 Attention 计算方式：
a.MHA拥有H份查询头Q和H份KV对。查询头Q使用各自对应K,V对进行计算。所有头的K和V权重不共享。MHA需要的缓存量很大。（Multi-head Attention）
b.MQA拥有H份查询头和1份KV对。所有的查询头Q共享同1份K.V对进行计算。MQA从减少K和V矩阵的参数量,加快解码器推断的速度，但生成质量随着降低。（Multi-Query Attention）
c.GQA拥有H份查询头和G份KV对。将查询头Q分成G组且组内共享同份KV对进行计算。GQA是MHA效果和MQA速度的折中方案。（Grouped-Query Attention）
GQA-G是指具有G组的Grouped-Query Attention。当G=1时，就是MQA。当G=H时，就是MHA。
