因果注意力（or 掩码注意力）代码实现

import torch
import torch.nn as nn
import math
import numpy as np


class CausalSelfAttention(nn.Module):
    def __init__(self, hidden_dim, n_heads, dropout=0.0):
        super().__init__()
        self.hidden_dim = hidden_dim
        self.n_heads = n_heads
        self.head_dim = hidden_dim // n_heads
        assert hidden_dim % n_heads == 0, "hidden_dim 必须能被 n_heads 整除"
        self.query = nn.Linear(hidden_dim, hidden_dim)
        self.key = nn.Linear(hidden_dim, hidden_dim)
        self.value = nn.Linear(hidden_dim, hidden_dim)
        self.output = nn.Linear(hidden_dim, hidden_dim)
        self.dropout = nn.Dropout(p=dropout)
    

    def forward(self, x, mask=None):
        q = self.query(x)
        k = self.key(x)
        v = self.value(x)
        return self.get_attention_scores(q, k, v, mask=mask)

    def get_attention_scores(self, q, k, v, mask=None):
        B, L, H = q.shape
        # 重塑张量以适应多头注意力
        q = q.view(B, L, self.n_heads, self.head_dim)
        k = k.view(B, L, self.n_heads, self.head_dim)
        v = v.view(B, L, self.n_heads, self.head_dim)
        # 调整维度顺序以进行批量矩阵乘法
        q = q.permute(0, 2, 1, 3) # (batch_size, num_heads, seq_len, head_dim)
        k = k.permute(0, 2, 1, 3)
        v = v.permute(0, 2, 1, 3)
        # 计算注意力分数
        attention_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(self.head_dim)  # (batch_size, num_heads, seq_len, seq_len)
        if mask is not None:
            attention_scores = attention_scores.masked_fill_(mask==0, float('-inf'))
        # 计算注意力权重
        attention_weights = torch.softmax(attention_scores, dim=-1)
        attention_weights = self.dropout(attention_weights)
        # 计算输出
        out = torch.matmul(attention_weights, v) # (batch_size, num_heads, seq_len, head_dim)
        # 调整维度顺序并重塑回原始形状
        out = out.permute(0, 2, 1, 3).contiguous().view(B, L, H)
        out = self.output(out)
        return out

def generate_mask(x, causal=False):
    B, L, _ = x.shape
    original = torch.ones((L, L), dtype=torch.bool)
    if causal:
        mask = torch.tril(torch.ones((L, L), dtype=torch.bool)).unsqueeze(0).unsqueeze(0) # [1, 1, seq_len, seq_len]
    else:
        mask = torch.ones((L, L), dtype=torch.bool).unsqueeze(0).unsqueeze(0)  # [1, 1, seq_len, seq_len]
    return mask

if __name__ == "__main__":
    batch_size = 8
    seq_len = 16
    hidden_dim = 512
    num_heads = 8
    causal = True
    dropout = 0.1
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    # 创建模型实例并移动到目标设备
    model = CausalSelfAttention(hidden_dim=hidden_dim, n_heads=num_heads, dropout=dropout).to(device)
    # 生成输入张量并移动到目标设备
    x = torch.randn(batch_size, seq_len, hidden_dim).to(device)
    # 生成掩码并移动到目标设备
    if causal:
        mask = generate_mask(x, causal=True).to(device)
    else:
        mask = generate_mask(x, causal=False).to(device)
    # 前向传播
    output = model(x, mask=mask)
    print(output.shape)