transformer的kv缓存技术学习笔记

传统transformer的kv缓存：

从提示语（输入文本）到生成文本，一个LLM模型大体经过以下几步：

1、将模型加载到GPU上

2、在CPU上对提示语进行分词，把token张量（词表中的单词id序列）传输到GPU

3、将张量输入神经网络，生成扩展的第一个令牌。

上述过程称为启动阶段，也称预填充阶段（pre-fillphase），与模型响应效率有着十分密切的关系

4、将新生成的令牌附加到输入的令牌序列中，并将其初始化(清空梯度，padding操作，使得setence_length始终一致）

5、生成扩展文本中第二个令牌的输入，重复上述过程，直到生成停止token，如EOS，或者达到模型配置的最大序列长度。

***上述的所有阶段称之为生成阶段，解码阶段，自回归阶段，甚至是增量阶段。***

5、将完成的tokens从GPU获取到CPU，将它们进行逆tokenize，即detokenize(从token序列还原成文本、符号，最终生成一个句子的过程。)

投机误差（speculative Sampleing[3]）或前向解码（lookahead解码[4]）却不遵守上述过程——————目前仅做了解

综上所述，文本生成经历了两个阶段。

1、接受新的token序列，单个token生成的启动阶段

2、后续token生成的多步生成阶段

在第二个阶段，每一次的token生成都将使用整个序列作为输入，引入了不必要的计算。

encoder部分和decoder部分都会产生计算冗余。

先说encoder。

我们知道，在encoder中，分别有多头注意力机制，embedding向量化和mask的操作。

encoder的mask本质是一个padding后的置零操作，其处理在模型输入之前，因此对模型中的计算冗余影响不大。

当token序列输入encoder的多头注意力层时，会计算一次词向量的embedding和positional  embedding，它们的形状分别为

[batch_size,sequence_length,hidden_size] 和 [batch_size,sequence_length,sequence_length]

positional embedding :[max_position_embeddings,hidden_size] 直接作用于token_embedding，希望它间接作用于q、k、v的计算，从而使输入的序列即使单词相同，注意力分数也会因为它们位置之间的变化而产生不同，从而反应序列中词与词相对位置关系对句子语义的影响，positional embedding的输出不依赖于input_ids，而是 torch.arange(seq_length, dtype=torch.long).unsqueeze(0)，生成的向量为0到seq_length-1的整数序列，这代表着序列中的位置。

在推理过程中，由于输入序列的位置每个时间步增加一个，因此positional embedding的向量每次也增加一行，同一个位置的embedding结果相同。

词向量的embedding也是由经过一次nn.embedding层计算而得，同个词embedding结果也相同。

positional embedding+token embedding，得到最终的embedding输出，由于是加法运算，所以每个词最终的embedding向量都是一样的。

q、k、v是输入序列的线性变换，所以encoder中已有序列的q、k、v是不变的。

decoder有些特殊，它的输入是一个从开始标志符[sos]的新序列，由于每一个时间步生成的token随机，所以encoder-decoder层得到的q输入也随机，而k、v值由于是encoder的输出线性变换而来，所以

这里先对负责文本生成的decoder模型进行大概理解。

decoder每一步只解码一个词，起始输入一般是起始token(sos),经过第一个多头注意力层，获得第二个多头注意力层(也叫encoder-decoder层）q的输入。

,然后将encoder的输出作为第二个多头注意力层的k、v输入。

针对encoder-decoder层的输入（不是真实的q、k、v，真实的q、k、v要对输入做一次线性变换）

q一定是变化的，因为无法预测decoder模型输出的下一个token是什么，它会随着时间步变化而变化

k、v来源于encoder的输出，而输出与encoder的q、k、v挂钩。

掩码的张量形状是[batch_size,sequence,sequence]，新增一个词的时候变成[batch_size,sequence+1,sequence+1]，它本身是一个下三角矩阵，下三角矩阵只需要一步计算，因此计算复杂低，不需要额外储存。

decoder中掩码的作用是屏蔽未知token的信息。

例如，经过embedding层（token_embedding+positional embedding）之后得到的序列数据为 [batch_size,sequence_length, hidden_dim]

经过q,k计算注意力得分后得到score [batch_size,sequence_length,sequence_length]

mask的形状为 [batch_size,sequence_length,sequence_length]，两者相加得到，[batch_size,sequence_length,sequence_length]，mask中- inf 所对应的位置将会使scores的对应位置也成为-inf，softmax操作之后，注意力权重无限趋近于零。

我们用向量运算来模拟这个过程：

在训练的时候，mask常用于防止模型内部在生成多头自注意力以及后续的操作中提前知晓当前词与之后的词关系，保证了语言序列生成的合理性。

因为日常生活中，我们不会根据说出的最后一个词来判断我们当前应该说什么，而是仅依赖之前说过的话，来判断我们现在该说什么话（有些照猫画虎的感觉）

既然在训练中，decoder的mask是为了防止模型查看未来的信息，那在推理中，模型的任务就是生成之后的单词，不存在查看未来单词的可能，理论上mask是不需要的。

但实际上呢。

在训练阶段，由于mask的作用，模型在依据现有的序列信息，推测下一个单词的任务上经过了大量的训练，它更擅长从有限的信息中掌握序列中每一个单词的特性，如果将每个单词的注意力得分都暴露出来，违背了训练时给模型设定的目标。

其次在训练阶段，模型的注意力层输出由 soft（(q*kT/sqrt(d_k))*mask）*v计算获得获得，mask影响了模型在注意力层对每个词的输出结果，从而对整个模型的输出产生了难以估计的影响。

如非特殊情况，最好的做法是将mask保留。

————————————————————————

在encoder中，多头自注意力的最终scores得分不需要mask。

所以最终scores的计算公式为 (q*kT)/sqrt(d_k)

对于其中位于（1，1）的元素，实际上是q的第一条embedding向量与k的第一条向量矩阵运算，也就是第一个单词自己的q和自己的k点积

推广到（1，3)元素，就是q的第一条embedding向量与k的第3条向量矩阵运算，也就是第一个单词的q和第三个单词的k点积

因此，对于第一行，只需要知道1索引位置token对应的q、k，和其他token对应的k，我们就能计算出scores中第一行的所有元素。

以此类推。

最终scores的计算流程便可以拆解为如下的小步骤：

对于每一个token，计算其q，与全部k值做点积。

这个过程十分繁琐，要把q、k、v全部计算一遍。

在启动阶段，这一过程是无法避免的，而encoder的多头自注意力层也在启动阶段计算结束，后续时间步都不会参与计算，所以提升encoder层的多头自注意力层十分有限。

启动阶段结束，正式进入decoder的生成阶段。

decoder的生成阶段分为两种。

训练时的decoder生成。

推理时的decoder生成。

训练时的文本序列都是正确且已知的，每一个时间步的输入都是正确的序列，而不依赖模型输出，所以输入形式是一串序列。（教师强迫）。

训练阶段因为应用了掩码，因此可以并行计算预测每一个位置的token，最终输出也是一串序列。

如图，第一行只有一个sos编码，它没有任何意义，只代表句子的开始，第一行输入到后续过程前向传播，生成第1个token。

第二行同理，生成第二个token。

.......

直到第五行，生成预测的第五个token。

当然，实际的token生成策略要更加复杂，这里不展开。

这些token都是并行计算的，所以说transformer的训练是并行的。

而推理时，每一个token的生成依赖于上一个token的生成，输入是一个token一个token进行的。

后面的token必须依赖已有的token才能生成，每一时间步生成的scores得分只有一行。

如图。

同理，每一个时间步的v值依赖于之前的V值加上当前的v值。

例如

在该例子中，

scores的形状为[1,1,5],v的形状为[1,5,4]，scores softmax之后与v矩阵相乘，得到[1,1,4]的注意力层输出。

经过后续的层，它将生成第5个预测的token。

下一个时间步，只需要生成形状为[1,1,6]的scores和[1,6,4]的v,就能同理获得第6个token。

本质上，上述过程等同于训练时并行预测的最后一行运算。

我们已经知晓，最后一行的计算应为最后一个token的q值与之前所有token已生成的k(包括自己）的点积运算。

那么，只要将上一个时间步生成的新token输入模型，得到其q和k，把之前的时间步的k值保存下来，等到计算score时，再注入到点积运算中即可。

这便是kv缓存的第一个作用，将decoder的第一个多头自注意力层每一步的k值保存下来，后续计算依赖前一步的k值，所以无需重新计算，直接加载。

再看第二层多头自注意力，也就是encoder-decoder自注意力层。

由于每个时间步模型生成的token是随机的，所以第一层关于token的输入是随机的。

第一层的输出成为第二层q权重的输入，导致第二层得到的q也发生变化，每一时间步也不一样，所以需要重新计算。

k和v来源于encoder的输出，是个定值，每一个时间步都要用到，可以缓存，所以这也是kv缓存的一部分。