大模型系列（掩码注意力，KV Cache，GQA）

文章目录一. 掩码注意力二. KV Cache三. GQA (Grouped-Query Attention分组查询注意力)一. 掩码注意力假设我们正在训练一个语言模型比如GPT当前抓取到的一条训练数据是一句话bos 我 爱吃 苹果bos 代表句子开头。在传统的RNN循环神经网络时代GPU只能这样按顺序“排队”干活1看到 bos去预测“我”。2等第1步算完把 和“我”拼起来再去预测“爱吃”。3接着排队……但是在Transformer的GPU现场情况是完全不同的。所有词是“并排”站着同时开始干活的。GPU不会一个一个读取词它是一次性把bos , 我, 爱吃, 苹果 这4个词的Embedding向量全部装进显存拼成一个4 × d 4 \times d4×d的矩阵X XX。紧接着GPU一声令下执行一次矩阵乘法。一瞬间4个词同时变身各自生成了自己的Q QQ查询、K KK键、V VV值向量。此时显存里有4个Q QQ、4个K KK、4个V VV。现在每个词作为Q QQ都要去和其他所有词作为K KK计算相关性得分。如果没有Mask4个词就会互相看光光。Mask 就像是给每个词发了一个不同形状的“眼罩”。 在GPU进行Q × K T Q \times K^TQ×KT这个巨大的矩阵运算并加上 Mask 后我们得到了一个4 × 4 4 \times 44×4的注意力权重矩阵。让我们一行一行地看这个矩阵这正是“拆解为无数个已知上文”的真相1第 1 行主角是 它的眼罩遮挡了后面所有的词。它的注意力权重是 [1.0, 0, 0, 0]。这代表的任务 “已知上文只有 bos”。它拿着这个权重去乘V VV矩阵提取出的特征V 新1 V_{\text{新1}}V新1​只包含 的信息。2第 2 行主角是 我它的眼罩遮挡了它后面的词。它的注意力权重可能是 [0.4, 0.6, 0, 0]。这代表的任务 “已知上文是 bos, 我”。它去乘V VV矩阵提取出的特征V 新2 V_{\text{新2}}V新2​融合了 和 我 的信息绝不包含未来的词。3第 3 行主角是 爱吃注意力权重可能是 [0.2, 0.3, 0.5, 0]。这代表的任务 “已知上文是 bos, 我, 爱吃”。它提取出的特征V 新3 V_{\text{新3}}V新3​融合了前三个词的信息。4第 4 行主角是 苹果它的眼罩没有遮挡任何词。注意力权重没有0 00。这代表的任务 “已知上文是完整的 bos, 我, 爱吃, 苹果”。提取出的特征V 新4 V_{\text{新4}}V新4​融合了整句话的信息。极其关键的一点这4行的计算在GPU中是绝对同时发生的 第3行的计算完全不需要等待第1行或第2行算完。Mask 矩阵通过数学手段把特定位置强行置为− ∞ -\infty−∞在一次并行的矩阵运算中硬生生地切分出了4个不同长度的历史上下文。现在GPU的手里握着4个全新的特征向量V 新1 , V 新2 , V 新3 , V 新4 V_{\text{新1}}, V_{\text{新2}}, V_{\text{新3}}, V_{\text{新4}}V新1​,V新2​,V新3​,V新4​。这4个向量会同时通过模型后面的前馈神经网络FFN最后输出4个词表概率分布。在训练阶段的最后模型会同时进行4次“对答案”拿V 新1 V_{\text{新1}}V新1​的预测结果去和标签 “我” 对答案算出 Loss 1。拿V 新2 V_{\text{新2}}V新2​的预测结果去和标签 “爱吃” 对答案算出 Loss 2。拿V 新3 V_{\text{新3}}V新3​的预测结果去和标签 “苹果” 对答案算出 Loss 3。拿V 新4 V_{\text{新4}}V新4​的预测结果去和标签 结束符 对答案算出 Loss 4。最后把这4个Loss加起来求平均执行一次反向传播更新所有的参数。总结模型输入的虽然是一整句话但它在内部通过 Mask 机制并行地把它拆解成了无数个“已知上文预测下一个词”的小任务。二. KV Cache为什么只存 K 和 V不存 Q在生成新词时当前的 QueryQ只代表最新生成的那个词它要去和历史所有词的 KeyK匹配。过去的词已经成为了“历史背景”K 和 V它们不会再主动去查询别人了所以过去的 Q 毫无用处直接丢弃即可。带有 KV Cache 的工作流程时刻t − 1 t-1t−1历史已经缓存假设模型已经处理了前面的文本显存中已经保存了过去的 K 和 V 矩阵记作K p a s t K_{past}Kpast​和V p a s t V_{past}Vpast​。时刻t tt当前步生成现在输入最新生成的一个 Token例如x t x_txt​。计算当前步的特征 仅仅计算这个新 Token 的q t q_tqt​、k t k_tkt​和v t v_tvt​。q t x t W q q_t x_t W_qqt​xt​Wq​k t x t W k k_t x_t W_kkt​xt​Wk​v t x t W v v_t x_t W_vvt​xt​Wv​更新缓存 将新算出来的k t k_tkt​和v t v_tvt​拼接到之前的缓存中。K c a c h e [ K p a s t ; k t ] K_{cache} [K_{past}; k_t]Kcache​[Kpast​;kt​]V c a c h e [ V p a s t ; v t ] V_{cache} [V_{past}; v_t]Vcache​[Vpast​;vt​]计算注意力 用当前唯一的q t q_tqt​去乘以拼接好的长串K c a c h e K_{cache}Kcache​得到注意力得分再乘以V c a c h e V_{cache}Vcache​。O u t p u t t s o f t m a x ( q t K c a c h e T d k ) V c a c h e Output_t softmax(\frac{q_t K_{cache}^T}{\sqrt{d_k}}) V_{cache}Outputt​softmax(dk​​qt​KcacheT​​)Vcache​通过这种方式每生成一个新的 Token我们只需要做一次针对该 Token 的特征提取运算极大地节省了算力。三. GQA (Grouped-Query Attention分组查询注意力)阶段一MHA (Multi-Head Attention - 多头注意力)在多头注意力机制MHA的初始阶段模型中 Query、Key 和 Value 的头数是完全相等的。这种设计的核心优势在于它赋予了模型极强的特征表达能力因为每一个 Query 头都能在一个完整且独立的 Key 和 Value 特征子空间中进行信息提取从而充分捕捉极高维度的复杂语义特征。然而这种设计的代价也是极其高昂的。在模型进行自回归生成的推理阶段必须缓存的所有历史 Key 和 Value 数据即 KV Cache会随着序列长度呈线性暴涨。这种极其庞大的显存占用直接导致显存带宽成为致命瓶颈严重限制了模型的生成速度以及能够同时处理的并发任务量。阶段二MQA (Multi-Query Attention - 多查询注意力)为了彻底解决 MHA 带来的显存危机多查询注意力机制MQA走向了另一个极端。在 MQA 的架构中虽然 Query 依然保持着原有的多头状态但 Key 和 Value 被强制压缩到仅剩一个全局头。这就意味着所有的 Query 头在计算注意力分数和聚合特征时只能被迫共享这同一组 Key 和 Value 矩阵。这种极端的压缩方案带来了立竿见影的工程收益KV Cache 的体积被骤减了数十倍从根本上释放了显存带宽的压力让推理速度得到了极大的提升。但其弊端也随之暴露由于丧失了独立的特征检索空间模型在处理复杂逻辑或长文本时的特征多样性严重流失导致生成质量出现明显的“容量衰减”。阶段三GQA (Grouped-Query Attention - 分组查询注意力)GQA 依然保留了多头的 Query但不再将 Key 和 Value 缩减为单一头而是将其头数设定为一个大于 1 却远小于 Query 头数的参数。在具体运行机制上模型会将所有的 Query 头按固定数量划分为若干个独立的分组让每一组内的 Query 头共享专属的一个 Key 头和一个 Value 头。