大模型原理--多头注意力机制的代码实现（MHA引入KV cache）

import torch
import torch.nn as nn

# 在推理阶段将历史 token 的 Key 和 Value 缓存
past_key_value = {}

# 自注意力机制
class DecoderSelfAttentionHead(nn.Module):
    def __init__(self, d_model, head_size):
        super().__init__()
        self.head_size = head_size
        self.query = nn.Linear(d_model, head_size)
        self.key = nn.Linear(d_model, head_size)
        self.value = nn.Linear(d_model, head_size)

    def forward(self, x, use_cache, masked, layer_num, head_index):
        print("layer_num: ",layer_num)
        print("head_index: ",head_index)
        # 计算Q, K, V
        Q = self.query(x)
        K = self.key(x)
        V = self.value(x)
        # 如果use_cache是True,则代表当前是推理阶段,需要进行KV的缓存
        if use_cache:
            if layer_num in past_key_value and head_index in past_key_value[layer_num]:
                # 拼接KV
                K = torch.cat([past_key_value[layer_num][head_index][0], K], dim=1)
                V = torch.cat([past_key_value[layer_num][head_index][1], V], dim=1)
                # 更新KV缓存
                past_key_value[layer_num][head_index] = (K, V)
            elif layer_num not in past_key_value:
                # 每层的初始KV缓存
                past_key_value[layer_num] = {head_index:(K,V)}
            else:
                past_key_value[layer_num][head_index] = (K, V)
            print(past_key_value)
        # Q和K的点积
        attention = Q @ K.transpose(-2, -1)
        # 缩放
        attention = attention / (self.head_size ** 0.5)
        # 训练的时候使用
        if masked:
            mask = torch.triu(torch.ones_like(attention), diagonal=1).bool()
            attention.masked_fill(mask, float('-inf'))
        # 打分
        attention = torch.softmax(attention, dim=-1)
        # 加权汇总
        attention = attention @ V
        return attention

#
class MaskedMultiHeadAttention(nn.Module):
    def __init__(self, d_model, n_heads):
        super().__init__()
        # 注意力头的个数,这块也可以看出d_model和注意力头数目由整倍数关系
        self.head_size = d_model // n_heads
        self.heads = nn.ModuleList([
            DecoderSelfAttentionHead(d_model, self.head_size) for _ in range(n_heads)
        ])
        self.multi_header = nn.Linear(d_model, d_model)
    # 引入layer_num,表示模型层数. 需要为每一层的每一个自注意力头缓存KV
    def forward(self, x, use_cache=False, masked=False, layer_num=0):
        # 拼接多个注意力头
        multi_header = torch.cat(
            [head(x, use_cache, masked, layer_num, head_index) for head_index, head in enumerate(self.heads)], dim=-1)
        out = self.multi_header(multi_header)
        return out

if __name__ == '__main__':
    x1 = torch.randn((1, 5, 12))
    x2 = torch.randn((1, 1, 12))
    model = MaskedMultiHeadAttention(12, 4)
    y1 = model(x1, use_cache=True)
    y2 = model(x2, use_cache=True)