【大语言模型】—— 注意力机制及其变体的代码实现

摘要

本文介绍了注意力机制的几种变体及其PyTorch代码实现。主要包括：

  1. Self-Attention：基础自注意力机制，通过Q、K、V计算注意力权重，适用于序列内部建模。
  2. CrossAttention：让一个序列关注另一个序列，典型应用于Transformer解码器-编码器交互和多模态任务。
  3. CausalAttention：通过三角掩码实现因果性，确保只能关注当前位置之前的token，适用于自回归生成任务。
  4. MultiHeadSelfAttention：多头注意力机制，将输入分割到多个子空间并行计算注意力，最后合并结果。

代码实现中详细展示了各注意力的关键操作，包括线性变换、注意力分数计算、softmax归一化和掩码处理等。特别解释了dim=-1的作用、unsqueeze(0)的广播机制等实现细节。

Self-Attention


class SelfAttention(nn.module):
# torch.matmul     专用于批量矩阵乘法，适用于形状为 (batch_size, n, m) 和 (batch_size, m, p) 的 3D 张量。
# torch.matmul  支持更灵活的张量乘法运算
	def __init__(self, input_dim, dim_k,dim_v):
		super().__init__()
		self.q = nn.Linear(input_dim, dim_k)
		self.k = nn.Linear(input_dim, dim_k)
		selv.v = nn.Linear(input_dim, dim_v)
		self.scale = np.sqrt(dim_k)
		
	def forward(self, x):
		Q = self.q(x)
		K = self.k(x)
		V = self.v(x)
	
		atten = torch.softmax(torch.matmul(Q, K.permute(0,2,1))/self.scale, dim=-1)
		return torch.matmul(atten, V)

为什么 dim=-1

在自注意力机制中，nn.Softmax(dim=-1)的作用是对 ​​注意力分数矩阵​​ 进行归一化，使得每一行的权重之和为 1。这里 dim=-1表示在最后一个维度（即 seq_len维度）上进行 Softmax 计算。

在 Q K T QK^T QKT计算后，得到的注意力分数矩阵的形状是 [batch_size, seq_len, seq_len]，其中：
​​第 1 个 seq_len（dim=1）​​：代表 Q的序列长度（即当前 token 的位置）。
​​第 2 个 seq_len（dim=2）​​：代表 K的序列长度（即被计算注意力的 token 的位置）。
dim=-1（即 dim=2）表示 ​​对每个 token 计算它对所有 token 的注意力权重​​，即 ​​对每一行进行 Softmax​​，使得：
每一行的所有值加起来等于 1（概率分布）。
这样，每个 token 的注意力权重是独立计算的。
假设 Q K T QK^T QKT的结果是：

[
  [[1.0, 0.5, 0.2],  # Token 0 对所有 token 的注意力分数
   [0.3, 1.2, 0.7],  # Token 1 对所有 token 的注意力分数
   [0.1, 0.4, 1.5]]  # Token 2 对所有 token 的注意力分数
]

应用 nn.Softmax(dim=-1)后：

[
  [[0.55, 0.27, 0.18],  # Token 0 的注意力权重（总和=1）
   [0.16, 0.58, 0.26],  # Token 1 的注意力权重（总和=1）
   [0.07, 0.20, 0.73]]  # Token 2 的注意力权重（总和=1）
]

这样，每个 token 的注意力权重是独立的，且所有 token 对它的影响权重之和为 1。

CrossAttention

# 查询通常来自解码器，键和值通常来自编码器
import torch
import torch.nn as nn
import numpy as np

class CrossAttention(nn.Module):
    def __init__(self, input_dim, dim_k, dim_v):
        super().__init__()  # 必须调用父类初始化
        self.q = nn.Linear(input_dim, dim_k)
        self.k = nn.Linear(input_dim, dim_k)
        self.v = nn.Linear(input_dim, dim_v)
        self.scale = np.sqrt(dim_k)
    
    def forward(self, x1, x2):
        Q1 = self.q(x1)  # [batch_size, seq_len1, dim_k]
        K2 = self.k(x2)  # [batch_size, seq_len2, dim_k]
        V2 = self.v(x2)  # [batch_size, seq_len2, dim_v]
        
        # 计算注意力分数
        atten = torch.softmax(torch.matmul(Q1, K2.permute(0, 2, 1)) / self.scale, dim=-1)  # [batch_size, seq_len1, seq_len2]
        
        # 加权求和
        return torch.matmul(atten, V2)  # [batch_size, seq_len1, dim_v]

交叉注意力的作用

交叉注意力用于 让一个序列 x 1 x1 x1关注另一个序列 x 2 x2 x2，典型应用包括：

  1. ​​Transformer 解码器​​：
    • x1= 解码器的输入（当前生成的 token）
    • x2= 编码器的输出（源序列的表示）
    • 解码器通过交叉注意力关注编码器的信息。
  2. ​​多模态任务​​（如视觉-语言模型）：
    • x1= 文本序列
    • x2= 图像特征
    • 文本通过交叉注意力关注图像的关键区域。

CausalAttention

class CausalAttention(nn.Module):
	def __init__(self,input_dim, dim_k,dim_v):
		super().__init__()
		self.q = nn.Linear(input_dim, dim_k)
		self.k = nn.Linear(input_dim, dim_k)
		self.v = nn.Linear(input_dim, dim_v)
		self.scale = np.sqrt(dim_k)
    def forward(self, x):
        # x: [batch_size, seq_len, input_dim]
        Q = self.q(x)  # [batch_size, seq_len, dim_k]
        K = self.k(x)  # [batch_size, seq_len, dim_k]
        V = self.v(x)  # [batch_size, seq_len, dim_v]

        # 注意力分数
        atten = torch.matmul(Q, K.permute(0, 2, 1)) / self.scale  # [batch, seq, seq]

        # 下三角 mask，确保因果性（只能看到之前的token）
        seq_len = atten.size(-1)
        mask = torch.tril(torch.ones(seq_len, seq_len)).unsqueeze(0)
        atten = atten.masked_fill(mask == 0, float('-inf'))

        # softmax 归一化
        atten = nn.Softmax(dim=-1)(atten)

        # 输出加权和
        return torch.matmul(atten, V)  # [batch, seq, dim_v]

unsqueeze(0)的作用​​

假设我们有一个 2D 张量 mask，形状是 [seq_len, seq_len]

mask = torch.tril(torch.ones(seq_len, seq_len))  # 形状 [seq_len, seq_len]

如果我们想让它变成 [1, seq_len, seq_len]（即增加一个 batch 维度），可以使用：

mask = mask.unsqueeze(0)  # 形状变为 [1, seq_len, seq_len]

这样做的目的是：

  • 匹配注意力分数矩阵的形状​​（atten的形状是 [batch_size, seq_len, seq_len]）。
  • ​​支持批量计算​​，因为 mask需要广播到所有 batch 样本。

这样：
mask的形状变成 [1, seq_len, seq_len]。
PyTorch 会自动广播 mask到 [batch_size, seq_len, seq_len]，使其与 atten形状匹配。

MultiHeadSelfAttention

class MultiHeadAttention(nn.Module):
	def __init__(self, heads, input_dim, dim_k, dim_v):
		super().__init__()
		self.heads = heads
		
		self.dim_k_per_head = dim_k // heads
		self.dim_v_per_head = dim_v // heads
		
		self.q = nn.Linear(input_dim, dim_k)
		self.k = nn.Linear(input_dim, dim_k)
		self.v = nn.Linear(input_dim, dim_v)
		
		self.scale = np.sqrt(self.dim_k_per_head)
		
		self.out = nn.Linear(dim_v, input_dim)
	def forward(self, x):
		batch_size = x.size(0)
		
		Q = self.q(x)#[batch_size, seq_len, dim_k]
		K = self.k(x)
		V = self.v(x)
		
		#[batch_size, seq_len, heads, dim_k_per_head] 
		#		--> [batch_size, heads, seq_len, dim_k_per_head]
		Q = Q.view(batch_size, -1, self.heads, self.dim_k_per_head).permute(0,2,1,3)
		#[batch_size, seq_len, heads, dim_k_per_head] 
		#		--> [batch_size, heads, seq_len, dim_k_per_head]
		K = K.view(batch_size, -1, self.heads, self.dim_k_per_head).permute(0,2,1,3)
		#[batch_size, seq_len, heads, dim_v_per_head] 
		#		--> [batch_size, heads, seq_len, dim_v_per_head]
		V = V.view(batch_size, -1, self.heads, self.dim_v_per_head).permute(0,2,1,3)
		
		#转置[batch_size, heads, seq_len, dim_k_per_head] 
		#		--> [batch_size, heads, dim_k_per_head, seq_len]
		K = K.permute(0, 1, 3, 2)
		
		# [batch_size, heads, seq_len, seq_len]
		atten = torch.softmax(torch.matmul(Q,K) / self.scale, dim = -1)
		# [batch_size, heads, seq_len, dim_v_per_head]
		out = torch.matmul(atten, V)
		# [batch_size, seq_len, heads, dim_v_per_head]
		out = out.permute(0, 2, 1, 3).contiguous()
		# [batch_size, seq_len, heads* dim_v_per_head]
		out = out.view(batch_size, -1, self.heads * self.dim_v_per_head)
		return self.out(out) # [batch_size, seq_len, input_dim]

多头自注意力（Multi-Head Attention）的核心思想

多头自注意力（Multi-Head Attention）的核心思想是将输入向量分别映射为查询 Q、键 K、值 V，再按照头数切分到多个子空间中；每个头独立计算注意力分数并得到加权表示，最后拼接各头的结果，通过线性层 out 映射回输入维度，从而捕捉序列中多角度的相关性。

在使用时需要注意以下几点：

  1. 维度整除：要确保 dim_k 和 dim_v 能被 heads 整除，否则 view 时会报错。
  2. 缩放因子：缩放应该基于每个头的维度 sqrt(dim_k_per_head)，而不是整体的 dim_k。
  3. Softmax 顺序：正确做法是 Softmax(QK^T / scale)，不要写成 Softmax(QK^T) / scale。
  4. 张量连续性：permute 之后用 .contiguous().view()，否则可能报错；或者用 reshape 自动处理。
  5. 输出层作用：self.out 的作用是把多头拼接后的结果重新映射回输入维度，保持层间维度一致。

Multi-head Latent Attention

多头潜在注意力。它最常见于 DeepSeek-V2 / V3 系列。核心思想是：不要直接缓存每一层完整的 K/V，而是先把 hidden state 压缩成低维 latent 向量，再由 latent 还原出 K/V。这样推理时 KV cache 可以显著变小。DeepSeek-V3 技术报告也明确说其采用 MLA 来提高推理效率；后续分析论文也指出 MLA 通过把 K/V 压缩到 latent 空间来降低 KV-cache 和带宽开销。

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


class MultiHeadLatentAttention(nn.Module):
    def __init__(self, heads, input_dim, dim_k, dim_v, kv_lora_rank):
        super().__init__()

        self.heads = heads
        self.dim_k_per_head = dim_k // heads
        self.dim_v_per_head = dim_v // heads
        self.kv_lora_rank = kv_lora_rank

        # Q 仍然正常投影
        self.q = nn.Linear(input_dim, dim_k)

        # MLA 的核心：先把 x 压缩成低维 latent
        # C_KV: [batch_size, seq_len, kv_lora_rank]
        self.kv_down = nn.Linear(input_dim, kv_lora_rank)

        # 再从低维 latent 恢复出 K 和 V
        self.k_up = nn.Linear(kv_lora_rank, dim_k)
        self.v_up = nn.Linear(kv_lora_rank, dim_v)

        self.scale = np.sqrt(self.dim_k_per_head)

        self.out = nn.Linear(dim_v, input_dim)

    def forward(self, x):
        batch_size = x.size(0)

        # x: [batch_size, seq_len, input_dim]

        # Q: [batch_size, seq_len, dim_k]
        Q = self.q(x)

        # 压缩后的 KV latent:
        # C_KV: [batch_size, seq_len, kv_lora_rank]
        C_KV = self.kv_down(x)

        # 从 latent 还原 K / V
        # K: [batch_size, seq_len, dim_k]
        # V: [batch_size, seq_len, dim_v]
        K = self.k_up(C_KV)
        V = self.v_up(C_KV)

        # [batch_size, seq_len, heads, dim_k_per_head]
        #   -> [batch_size, heads, seq_len, dim_k_per_head]
        Q = Q.view(batch_size, -1, self.heads, self.dim_k_per_head).permute(0, 2, 1, 3)

        # [batch_size, seq_len, heads, dim_k_per_head]
        #   -> [batch_size, heads, seq_len, dim_k_per_head]
        K = K.view(batch_size, -1, self.heads, self.dim_k_per_head).permute(0, 2, 1, 3)

        # [batch_size, seq_len, heads, dim_v_per_head]
        #   -> [batch_size, heads, seq_len, dim_v_per_head]
        V = V.view(batch_size, -1, self.heads, self.dim_v_per_head).permute(0, 2, 1, 3)

        # K 转置:
        # [batch_size, heads, seq_len, dim_k_per_head]
        #   -> [batch_size, heads, dim_k_per_head, seq_len]
        K = K.permute(0, 1, 3, 2)

        # attention score:
        # [batch_size, heads, seq_len, seq_len]
        atten = torch.softmax(torch.matmul(Q, K) / self.scale, dim=-1)

        # attention output:
        # [batch_size, heads, seq_len, dim_v_per_head]
        out = torch.matmul(atten, V)

        # [batch_size, seq_len, heads, dim_v_per_head]
        out = out.permute(0, 2, 1, 3).contiguous()

        # [batch_size, seq_len, dim_v]
        out = out.view(batch_size, -1, self.heads * self.dim_v_per_head)

        # [batch_size, seq_len, input_dim]
        return self.out(out)

关键区别就是这几行：

self.kv_down = nn.Linear(input_dim, kv_lora_rank)

self.k_up = nn.Linear(kv_lora_rank, dim_k)
self.v_up = nn.Linear(kv_lora_rank, dim_v)

普通 MHA 是：

K = self.k(x)
V = self.v(x)

MLA 是：

C_KV = self.kv_down(x)

K = self.k_up(C_KV)
V = self.v_up(C_KV)

推理框架和 kernel 支持还不如 GQA 成熟。GQA 已经被 vLLM、TensorRT-LLM、FlashAttention 等主流推理栈广泛优化；MLA 需要专门处理 latent cache、RoPE 解耦、cache absorption、tensor parallel 切分等问题。

Grouped-Query Attention

它的核心思想是：

Q 仍然保留很多个 head，但 K/V 的 head 数量减少；多个 Q head 共享同一组 K/V head。

所以它介于 MHA 和 MQA 之间：

机制 Q heads K/V heads 特点
MHA 多个 多个 每个 Q head 对应自己的 K/V head
MQA 多个 1 个 所有 Q head 共享一组 K/V
GQA 多个 少量几个 每组 Q head 共享一组 K/V

例如：

heads = 8
kv_heads = 2

那么：

Q 有 8 个 head
K/V 只有 2 个 head
每 4 个 Q head 共享 1 个 K/V head

GQA 代码示例

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


class GroupedQueryAttention(nn.Module):
    def __init__(self, heads, kv_heads, input_dim, dim_k, dim_v):
        super().__init__()

        assert heads % kv_heads == 0, "heads must be divisible by kv_heads"
        assert dim_k % heads == 0, "dim_k must be divisible by heads"
        assert dim_v % heads == 0, "dim_v must be divisible by heads"

        self.heads = heads
        self.kv_heads = kv_heads
        self.group_size = heads // kv_heads

        self.dim_k_per_head = dim_k // heads
        self.dim_v_per_head = dim_v // heads

        # Q 仍然有 heads 个 head
        self.q = nn.Linear(input_dim, dim_k)

        # K/V 只有 kv_heads 个 head
        self.k = nn.Linear(input_dim, kv_heads * self.dim_k_per_head)
        self.v = nn.Linear(input_dim, kv_heads * self.dim_v_per_head)

        self.scale = np.sqrt(self.dim_k_per_head)

        self.out = nn.Linear(dim_v, input_dim)

    def forward(self, x):
        batch_size = x.size(0)
        seq_len = x.size(1)

        # x: [batch_size, seq_len, input_dim]

        # Q: [batch_size, seq_len, dim_k]
        Q = self.q(x)

        # K: [batch_size, seq_len, kv_heads * dim_k_per_head]
        # V: [batch_size, seq_len, kv_heads * dim_v_per_head]
        K = self.k(x)
        V = self.v(x)

        # Q:
        # [batch_size, seq_len, heads, dim_k_per_head]
        # -> [batch_size, heads, seq_len, dim_k_per_head]
        Q = Q.view(
            batch_size,
            seq_len,
            self.heads,
            self.dim_k_per_head
        ).permute(0, 2, 1, 3)

        # K:
        # [batch_size, seq_len, kv_heads, dim_k_per_head]
        # -> [batch_size, kv_heads, seq_len, dim_k_per_head]
        K = K.view(
            batch_size,
            seq_len,
            self.kv_heads,
            self.dim_k_per_head
        ).permute(0, 2, 1, 3)

        # V:
        # [batch_size, seq_len, kv_heads, dim_v_per_head]
        # -> [batch_size, kv_heads, seq_len, dim_v_per_head]
        V = V.view(
            batch_size,
            seq_len,
            self.kv_heads,
            self.dim_v_per_head
        ).permute(0, 2, 1, 3)

        # GQA 核心：
        # 把 K/V 从 kv_heads 扩展到 heads
        #
        # 例如 heads=8, kv_heads=2, group_size=4
        # K/V: [batch_size, 2, seq_len, head_dim]
        # ->   [batch_size, 8, seq_len, head_dim]
        #
        # 每 4 个 Q head 共享同一个 K/V head
        K = K.repeat_interleave(self.group_size, dim=1)
        V = V.repeat_interleave(self.group_size, dim=1)

        # K 转置:
        # [batch_size, heads, seq_len, dim_k_per_head]
        # -> [batch_size, heads, dim_k_per_head, seq_len]
        K = K.permute(0, 1, 3, 2)

        # attention score:
        # Q: [batch_size, heads, seq_len, dim_k_per_head]
        # K: [batch_size, heads, dim_k_per_head, seq_len]
        # atten: [batch_size, heads, seq_len, seq_len]
        atten = torch.softmax(torch.matmul(Q, K) / self.scale, dim=-1)

        # attention output:
        # atten: [batch_size, heads, seq_len, seq_len]
        # V:     [batch_size, heads, seq_len, dim_v_per_head]
        # out:   [batch_size, heads, seq_len, dim_v_per_head]
        out = torch.matmul(atten, V)

        # [batch_size, heads, seq_len, dim_v_per_head]
        # -> [batch_size, seq_len, heads, dim_v_per_head]
        out = out.permute(0, 2, 1, 3).contiguous()

        # [batch_size, seq_len, heads * dim_v_per_head]
        out = out.view(batch_size, seq_len, self.heads * self.dim_v_per_head)

        # [batch_size, seq_len, input_dim]
        return self.out(out)

GQA 的关键代码

最核心的是这两行：

K = K.repeat_interleave(self.group_size, dim=1)
V = V.repeat_interleave(self.group_size, dim=1)

假设：

heads = 8
kv_heads = 2
group_size = 4

原始 K/V 的形状是：

K: [batch_size, 2, seq_len, head_dim]
V: [batch_size, 2, seq_len, head_dim]

扩展后变成：

K: [batch_size, 8, seq_len, head_dim]
V: [batch_size, 8, seq_len, head_dim]

但是注意，这里的 8 个 K/V head 不是独立学习出来的，而是由 2 个 K/V head 复制出来的。

也就是：

Q head 0, 1, 2, 3 共享 K/V head 0
Q head 4, 5, 6, 7 共享 K/V head 1
注意力机制 当前应用状态
MHA 经典结构，但大模型推理显存开销大
MQA / GQA 目前主流大模型更常用，LLaMA/Qwen/Mistral 等很多模型采用或接近这种路线
MLA DeepSeek 系列强力验证，研究热度上升，但还不是行业默认配置
Linear Attention / Sparse Attention 长上下文方向活跃，但路线分散
# 语言模型 # 人工智能 # 自然语言处理
Logo
DAMO开发者矩阵

DAMO开发者矩阵，由阿里巴巴达摩院和中国互联网协会联合发起，致力于探讨最前沿的技术趋势与应用成果，搭建高质量的交流与分享平台，推动技术创新与产业应用链接，围绕“人工智能与新型计算”构建开放共享的开发者生态。

更多推荐

写给具身智能创始人：别让故事跑在产品前面太远

具身智能是一个需要乐观主义者的行业。如果只看今天的失败率、成本、交付难度和场景复杂度，很容易得出悲观结论。但真正重要的技术产业，往往都不是在完美时才开始。新能源汽车早期也被质疑续航、成本、补能和安全；智能手机早期也经历过性能、生态和网络的瓶颈；云计算、AI、大模型也都经历过从概念到基础设施的漫长周期。机器人也会经历这个过程。所以，创始人不该因为产品还不完美就停止讲未来。但也不该因为未来足够大，就忽

avatar DAMO开发者矩阵

具身智能开始需要新的讲述方式

2026 年 8 月，北京的机器人行业看起来很热。世界机器人大会期间，人形机器人在展馆里走路、搬箱、握手、演示灵巧手。几天后，第二届世界人形机器人运动会在北京举行，机器人跑步、踢球、格斗、接力。这个行业仍然有足够强的画面感，足够容易被传播，也足够适合被写成“未来已来”。但在展馆外，另一场更小的聚会持续到了晚上十点。WRC 期间，一个叫“聚深会”的小聚，从晚上六点半开始。来的不是单一角色，有具身智能

avatar DAMO开发者矩阵

具身智能进入深水区：机器人公司真正要补的六堂课

这轮争论的背景，是全行业从"拼融资"转向"比落地"。回看 2026 年 8 月的这一周，具身智能完成了某种成人礼：它有了自己的第一股，有了创纪录的融资，有了能打水漂一样激起全场欢呼的乒乓球对局，也有了上市公司董事长当众泼下的冷水。具身智能的产出是物理世界里的一次抓取、一次行走、一次递杯子——错了，碎的是杯子，磨损的是价值 6000 美元一只的灵巧手。而对站在河对岸的资本来说，8 月的这一周已经足够

avatar DAMO开发者矩阵