技术博客

PyTorch 实现 multi-head attention 逐行讲解

第 63 / 147 篇 · 2026-06-02 · ai · 注意力机制动手实现PyTorch

上一篇用 NumPy 写了注意力的数学骨架，这一篇写它的 PyTorch 工业形态。目标是一个能直接训练的 nn.Module，逐行讲 shape 怎么变、每个设计对应什么工程动机。代码以 PyTorch 2.x（F.scaled_dot_product_attention 可用的版本）为准。

完整实现

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadAttention(nn.Module):
    def __init__(self, d_model: int, n_head: int):
        super().__init__()
        assert d_model % n_head == 0            # 头维必须整除，否则没法均分
        self.n_head = n_head
        self.d_head = d_model // n_head
        # 一次性建三个投影，后面 cat 成一个大矩阵用——这是 fused QKV
        self.q_proj = nn.Linear(d_model, d_model)
        self.k_proj = nn.Linear(d_model, d_model)
        self.v_proj = nn.Linear(d_model, d_model)
        self.out_proj = nn.Linear(d_model, d_model)

    def forward(self, x, causal=True):
        B, T, C = x.shape                        # batch, 序列长, 模型维
        H, D = self.n_head, self.d_head

        # (B, T, C) -> 三个 (B, T, C)。fused 版：一个 Linear(3C, C) 再 chunk
        q = self.q_proj(x).view(B, T, H, D).transpose(1, 2)   # (B, H, T, D)
        k = self.k_proj(x).view(B, T, H, D).transpose(1, 2)   # (B, H, T, D)
        v = self.v_proj(x).view(B, T, H, D).transpose(1, 2)   # (B, H, T, D)

        # SDPA 内部：scores = q @ k.transpose(-2,-1) / sqrt(D)，softmax，乘 v
        y = F.scaled_dot_product_attention(q, k, v, is_causal=causal)  # (B, H, T, D)

        y = y.transpose(1, 2).contiguous().view(B, T, C)       # 头维折叠回通道维
        return self.out_proj(y)                                # (B, T, C)

逐段拆解

投影与维度折叠（第 17–19 行）。 三行同构，以 q 为例：q_proj 输出 $(B,T,C)$，view(B, T, H, D) 把最后一维切成头数和头维，transpose(1, 2) 把头维挪到序列维之前，得到 $(B, H, T, D)$。为什么要挪？因为后面的 scaled_dot_product_attention 把倒数两个维度当作矩阵乘的双方，把 H 放在第 1 维后，batch 和 head 合并成一个广播维，$H$ 个头共享同一次调用——这正是 NumPy 版里 np.stack 循环的向量化形态，循环消失了，语义一字不差。

contiguous() 那一行不能省。 transpose 只是交换 stride 不复制数据，$(B,H,T,D)$ 转回 $(B,T,H,D)$ 后内存不连续，view 要求连续内存所以先 .contiguous()。漏掉会抛 RuntimeError，这是这个实现里最常见的报错，报错信息会直接提示你加 .contiguous() 或改用 reshape（reshape 会自动处理，但多一次隐式拷贝，显式写出来更好读）。

SDPA 的 backend 选择。 F.scaled_dot_product_attention 是个分发器，按输入条件在三种实现间自动选择：满足条件时用 FlashAttention（cuda，半精度，无自定义 mask 以外的复杂 attn_mask 时）、否则 memory-efficient（xformers 风格），都不可用退回纯 PyTorch 的 math 路径（物化 $n^2$ 矩阵）。也就是说这 30 行代码在 Ampere 以后的卡上、bf16 下，白拿前两篇讲的 FlashAttention 加速；而在 CPU 或 debug 模式下自动降级为可读的数学路径。手动调 backend 的接口也存在（torch.nn.attention.sdpa_kernel 上下文管理器），压测时用来强制某条路径对比耗时。

因果掩码的参数语义。 is_causal=True 等价于上三角置 $-\infty$ 的下三角掩码，要求 q 和 k 等长（自注意力场景成立）。cross attention 或 prefix 场景 q/k 不同长时用 attn_mask 传布尔张量（True = 保留），注意它是逐点加到 scores 上或逐位置掩码的接口，不要和 is_causal 同时传。

fused QKV 投影。 上面为了逐行清楚用了三个独立 Linear；工业实现（nn.MultiheadAttention、HF transformers）是单个 nn.Linear(d_model, 3*d_model) 再 chunk(3, dim=-1)。动机是 kernel 层面的：三个小 GEMM 合成一个大 GEMM，一次 launch、更好的内存局部性，在 decode 阶段 GEMM 本就很小时，launch 开销占比不可忽视。三份权重拼一个矩阵，数学完全等价。

和 nn.MultiheadAttention 的对应

手写版每行都能映射到 nn.MultiheadAttention 的内部：in_proj_weight 就是 fused QKV 的三倍宽矩阵（外加 in_proj_bias），out_proj 对应同名属性，core 计算走 SDPA。训练时直接用官方实现即可，它额外处理了 batch-first/batch-second 布局、增量解码的 kv 复用（key_value_states 参数）等边角。但有两个理由值得手写过一遍：一是 debug attention 相关问题时（比如检查某个头的注意力图），知道形状在哪一步变才能正确地 hook 或 register_forward_hook 抓中间量；二是做变体实验——GQA、滑动窗口、ALiBi 这些改动，在 30 行手写版上都是三五行的事，在官方模块上就要绕。

shape 流转：(B, T, C) 经过 MHA 的每一步 输入 x (B, T, C) fused QKV 投影 (B, T, 3C) → chunk×3 view 折叠 + transpose (B, H, T, D) SDPA flash/mem/math 自动三选一 out_proj (B, T, C) 逐行要点 · d_head = d_model // n_head，view 之前先整除断言 · transpose 后内存不连续，view 前必须 .contiguous()（或改用 reshape） · is_causal=True 仅适用于 q/k 等长；不等长走 attn_mask 布尔张量 · 变体实验（GQA / 滑窗 / ALiBi）在 30 行手写版上改，别绕官方模块

小结

这 30 行是 NumPy 骨架的工业翻译：三处投影对应上一篇的 $W^Q,W^K,W^V$，view + transpose 对应那次 reshape 折叠，SDPA 是 softmax 核的硬件加速分发器，out_proj 是把多头输出重组回模型维的 $W^O$。读懂了 shape 的每一步流转，nn.MultiheadAttention、HF 的 Attention 类、各种变体论文里的架构图，就都落在同一张地图上。