Transformer架构深度解构：自注意力、几何缩放与显存IO约束

现代大语言模型在长上下文扩展中遭遇的物理瓶颈，从来不是浮点计算吞吐量（FLOPs），而是高带宽显存（HBM）与片上静态随机存取存储器（SRAM）之间的访存墙。Vaswani et al. (2017) 确立的标准缩放点积注意力算子（Scaled Dot-Product Attention），在序列长度 $N$ 膨胀时，其中间注意力矩阵产生 $O(N^2)$ 的显存搬运开销，使得自注意力算子本质上沦为内存带宽受限（Memory-Bound）任务。

理解 Transformer 的第一步，是剥离那些将注意力机制拟人化的认知比喻，回到矩阵投影的几何空间与现代 GPU 硬件架构的物理约束中。

自注意力机制的物理现实：计算受限还是内存受限？

#

在 GPU 体系结构中，算子的性能瓶颈由算术强度（Arithmetic Intensity，定义为每字节显存传输所完成的浮点运算次数 FLOPs/Byte）决定。

现代计算卡（以 NVIDIA A100 SXM4 40GB 为例）拥有 312 TFLOPS 的 FP16 Tensor Core 峰值计算能力，但其 HBM2 显存带宽仅为 1555 GB/s。这意味着硬件的平衡算术强度临界点约为：

当一个算子的算术强度低于该临界值时，计算单元必然处于饥饿等待状态，算子处于内存带宽受限状态；反之则为计算受限（Compute-Bound）状态。

标准自注意力的经典实现由三组离散的张量操作构成：

1. 计算注意力分数矩阵：$S = Q K^T$
2. 缩放与概率归一化：$P = \text{softmax}(S / \sqrt{d_k})$
3. 聚合上下文向量：$O = P V$

设批大小（Batch Size）为 $B$、头数（Head Number）为 $h$、序列长度为 $N$、头维度为 $d_k$。在步骤 2 中，Softmax 操作读取大小为 $B \times h \times N \times N$ 的中间矩阵 $S$，完成指数与除法运算后再写回大小相同的矩阵 $P$。

该过程浮点运算量约为 $O(B \cdot h \cdot N^2)$，而内存读写量同样为 $O(B \cdot h \cdot N^2)$ 个浮点数。其算术强度仅为几个 FLOPs/Byte，远远落后于硬件平衡点 200 FLOPs/Byte。

这意味着：在长序列场景下，自注意力机制绝大部分执行耗时都浪费在将 $N \times N$ 的庞大矩阵反复写入 GPU HBM 又读出，GPU 的 Tensor Core 处于严重的闲置等待状态。

Q/K/V 投影与点积寻址的几何本质

#

自注意力机制通过三个可学习的线性变换矩阵，将输入 Token 的高维嵌入投影到不同的表征子空间。

给定输入序列矩阵 $X \in \mathbb{R}^{N \times d_{\text{model}}}$，三个线性投影定义如下：

其中 $W_Q, W_K \in \mathbb{R}^{d_{\text{model}} \times d_k}$，$W_V \in \mathbb{R}^{d_{\text{model}} \times d_v}$。

从加法注意力到矩阵点积

#

在 Vaswani et al. (2017) 之前，Seq2Seq 领域主要采用 Bahdanau et al. (2014) 的加法注意力（Additive Attention）：

加法注意力使用单隐层多层感知机（MLP）评估相关性。虽然加法注意力和点积注意力在理论复杂度上类似，但在现代硬件上存在根本性差距：点积注意力可以通过两次高度优化的通用矩阵乘法（GEMM，General Matrix Multiply）实现，直接调用硬件层面的 Tensor Core 或脉动阵列（Systolic Array），其并行流水线执行效率比加法注意力的非线性激活分支高出一个数量级。

内积空间的动态寻址机制

#

内积 $q_i \cdot k_j^T = |q_i| |k_j| \cos \theta_{ij}$ 计算的是两个投影向量在几何子空间中的方向对齐程度与模长乘积。

• $Q$（Query，查询向量）：当前位置的语义探针，编码了当前 Token 寻求上下文信息的需求方向。
• $K$（Key，键向量）：被检索位置的特征索引，编码了该位置能够对外提供的特征契约。
• $V$（Value，值向量）：被检索位置的实质内容载荷，一旦键值匹配，该载荷将按匹配权重融入当前位置。

这种内积寻址的本质是可微分的软寻址（Soft Addressing）：系统不通过离散索引读取存储单元，而是将全序列的 Value 向量根据内积相似度投影形成的概率分布，进行加权线性组合。

Softmax 缩放因子 $\sqrt{d_k}$ 的数学推导与数值稳定

#

Vaswani et al. (2017) 在定义注意力算子时引入了缩放因子 $\frac{1}{\sqrt{d_k}}$：

为什么必须除以 $\sqrt{d_k}$ 而不是 $d_k$ 或其他常数？其核心根源在于高维内积的方差扩散与 Softmax 导数的饱和区特性。

方差随维度 $d_k$ 线性膨胀的推导

#

假设 Query 向量 $q$ 与 Key 向量 $k$ 的各分量 $q_m, k_m$（$m = 1, \dots, d_k$）为相互独立的随机变量，且满足均值为 0、方差为 1 的标准分布：

计算两者的点积 $z = q \cdot k = \sum_{m=1}^{d_k} q_m k_m$。

根据独立随机变量的期望乘法性质：

单个分量乘积的方差为：

由于各个分量相互独立，根据方差的可加性，点积和的期望与方差分别为：

这意味着，点积 $z$ 的标准差为 $\sigma = \sqrt{d_k}$。当向量维度 $d_k = 64$ 或 $128$ 时，$z$ 的取值范围会扩展到很大区间，产生极大正值或负值。

梯度消失与 Softmax 饱和区

#

考察 Softmax 函数对输入分量 $z_i$ 的偏导数：

当某些分量 $z_i \gg z_j$ 时，经指数映射后，$S_i \to 1$，而其余分量 $S_j \to 0$。

带入导数公式：

• 对于极大分量 $i$：$S_i (1 - S_i) \approx 1 \times (1 - 1) = 0$
• 对于其余分量 $j$：$S_j (0 - S_j) \approx 0$

整个 Softmax 输出进入严重的梯度饱和区（Saturation Region），其雅可比矩阵元素全部趋近于 0。在误差反向传播中，连乘项骤降为零，网络权重陷入停滞。

除以 $\sqrt{d_k}$ 严格使点积分布的标准差重新归一化为 1：

这确保了输入维持在 Softmax 梯度最陡峭的敏感工作区间内，保证训练梯度的顺畅反向传播。

多头注意力的正交表征与秩坍塌问题

#

凸组合限制与单头注意力的表征瓶颈

#

单头自注意力的聚合公式本质上是一个凸组合（Convex Combination）：

这意味着聚合后的向量 $y_i$ 必然落在输入集合 ${v_1, \dots, v_N}$ 的凸包（Convex Hull）内部。

如果序列中同时存在多组相互独立的依赖关系（例如代词指代、句法主谓对齐、时态修饰约束），单头 Softmax 概率分布会被其中幅值最大的一组特征强行主导，其他维度的微弱信号被指数放大后的主导项完全淹没。

深层网络中的注意力秩退化（Rank Collapse）

#

Dong et al. (2021) 在 Attention is Not All You Need: Pure Attention Loses Rank Doubly Exponentially with Depth 中给出了严格证明：

如果不引入残差连接（Residual Connection）和多层感知机（MLP），多层自注意力堆叠将以双指数速度发生秩坍塌（Rank Collapse）：

输出矩阵的各行迅速收敛到相同的向量，输出矩阵退化为秩为 1 的平凡矩阵。所有 Token 的表征丧失区分度，系统彻底丧失建模复杂序列的能力。

多头机制构建正交低维子空间

#

Vaswani et al. (2017) 提出的多头自注意力（Multi-Head Attention, MHA）通过参数正交分解抑制了这种退化：

将原始 $d_{\text{model}}$ 拆分为 $h$ 个低维子空间，每个子空间维度 $d_k = d_{\text{model}} / h$。

每个头在几何上建立了独立的投影超平面：

1. 头 1 可以将注意力集中在局部相邻的形态学搭配；
2. 头 2 可以将注意力投影到长距离的核心名词指代；
3. 输出投影矩阵 $W^O \in \mathbb{R}^{d_{\text{model}} \times d_{\text{model}}}$ 负责将各个正交子空间的特征重新混合融合。

在算力开销相当的前提下，多头注意力大幅拓宽了网络表达能力的几何自由度。

硬件瓶颈与 FlashAttention 的 SRAM 重分块重构

#

即使数学形式优雅，标准自注意力在工业落地时依然被 $O(N^2)$ 的物理显存占用死死卡住。

标准自注意力的 HBM 显存读写灾难

#

在标准 PyTorch 算子执行链路中，GPU 显存层次结构如下：

• SRAM（片上高速缓存）：单 SM 约 192KB–228KB，带宽超 19 TB/s，容量极小；
• HBM（片外高带宽显存）：容量 40GB–80GB，带宽 1.5–3.3 TB/s，但延迟比 SRAM 高一个数量级。

标准注意力的计算流程必须多次往返 HBM：

1
2
3

[HBM] Q, K  ---> [SRAM] 计算 QK^T      ---> [HBM] 写入 S (O(N^2) 显存)
[HBM] S     ---> [SRAM] 计算 Softmax   ---> [HBM] 写入 P (O(N^2) 显存)
[HBM] P, V  ---> [SRAM] 计算 PV        ---> [HBM] 写入 O

当上下文长度达到 $N = 32768$、批大小 $B = 2$、头数 $h = 32$ 时，单个中间矩阵 $S$ 采用 FP16 占用的显存为：

这单次中间变量的内存就足以撑爆两张 A100 80GB 显卡。

在线 Softmax 与 SRAM 局部递推计算

#

Tri Dao et al. (2022) 提出的 FlashAttention 从底层重写了注意力算子。其核心思想是：绝不将 $N \times N$ 的注意力矩阵实例化到 HBM 中，而是将计算完全局限在 SRAM 内部，利用分块（Tiling）与在线递推（Online Softmax）一次性完成整个计算。

标准 Softmax 要求预先获知全行向量的最大值以防指数溢出：

这导致计算必须有跨整行的全局数据依赖。Milakov & Gimelshteyn (2018) 与 Tri Dao et al. (2022) 采用分块在线递推算法打破了这一全局依赖。

设当前行划分为两个数据块 $x^{(1)}$ 与 $x^{(2)}$。已知第一块的局部统计量：

当接入第二块数据 $x^{(2)}$ 时，更新全局最大值与配分函数：

对于输出累加向量 $O$：

利用这一数学恒等变换，GPU 可以将 $Q$ 按行切块载入 SRAM，将 $K, V$ 按列切块流式载入 SRAM。每个数据块在 SRAM 内部完成局部的点积、局部 Softmax 统计量校正与 Value 向量相乘累加。

全过程完全不需要向 HBM 写出任何中间 $N \times N$ 矩阵。

反向传播的重计算策略与 IO 复杂度优化

#

传统深度学习反向传播要求前向过程缓存所有激活值。如果缓存 $P \in \mathbb{R}^{N \times N}$，显存复杂度依然为 $O(N^2)$。

FlashAttention 的反向传播做出了激进的工程权衡：前向传播只保存分块的局部统计量（每行的标量最大值 $m$ 与归一化因子 $d$，占用空间仅 $O(N)$），彻底丢弃 $P$。

在反向传播计算梯度时，利用保存在 HBM 中的原始 $Q, K, V$ 和统计量标量，在 SRAM 内部现场重新计算分块注意力权重。

这一设计增加了约 15% 的浮点运算量，但省去了巨大的 HBM 显存读写带宽消耗，使得前向和反向的端到端吞吐提升了 2–4 倍，将序列长度的内存开销直接从 $O(N^2)$ 压降至 $O(N)$。

生产实现考据：PyTorch 核心数学验证

#

以下代码使用标准 PyTorch 精确演示缩放点积的数学本质，以及模拟分块在线 Softmax 的递推数值稳定性：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48

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

def scaled_dot_product_attention_reference(Q: torch.Tensor, K: torch.Tensor, V: torch.Tensor) -> torch.Tensor:
    """
    标准点积注意力参考实现 (显存 O(N^2))
    Q, K, V 形状: [Batch, Heads, SeqLen, HeadDim]
    """
    d_k = Q.size(-1)
    
    # 1. 计算未缩放内积并强制缩放 sqrt(d_k)
    # scores: [B, H, N, N]
    scores = torch.matmul(Q, K.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 2. Softmax 归一化 (沿最后一个维度)
    attn_weights = F.softmax(scores, dim=-1)
    
    # 3. 聚合 Value 载荷
    # output: [B, H, N, HeadDim]
    output = torch.matmul(attn_weights, V)
    return output

def online_softmax_tiling_step(q_block: torch.Tensor, k_block: torch.Tensor, v_block: torch.Tensor,
                               prev_m: torch.Tensor, prev_d: torch.Tensor, prev_acc: torch.Tensor):
    """
    单分块在线递推逻辑 (FlashAttention 核心数学单步简化)
    """
    d_k = q_block.size(-1)
    # 局部点积
    local_scores = torch.matmul(q_block, k_block.transpose(-2, -1)) / math.sqrt(d_k)
    
    # 局部最大值更新
    curr_max = torch.max(local_scores, dim=-1, keepdim=True).values
    new_m = torch.maximum(prev_m, curr_max)
    
    # 局部缩放因子
    alpha = torch.exp(prev_m - new_m)
    beta = torch.exp(curr_max - new_m)
    
    # 更新局部配分函数
    exp_scores = torch.exp(local_scores - new_m)
    new_d = prev_d * alpha + torch.sum(exp_scores, dim=-1, keepdim=True)
    
    # 更新累加向量
    new_acc = prev_acc * alpha + torch.matmul(exp_scores, v_block)
    
    return new_m, new_d, new_acc

生产环境的工程权衡与变体演进

#

FlashAttention 解决了训练和推理 Prefill 阶段的长序列显存膨胀，但在自回归解码（Autoregressive Decoding）阶段，自注意力的物理约束演化出全新的形态。

推理阶段的 KV Cache 内存墙

#

在文本流式生成阶段，每一步只输入一个新 Token，产生单行 $Q \in \mathbb{R}^{1 \times d_k}$。此时已不存在并行矩阵乘矩阵（GEMM），而是退化为矩阵乘向量（GEMV）。

为了避免历史 Token 的 Key 与 Value 重复计算，系统将所有历史 $K$ 与 $V$ 缓存在显存中，这即是 KV Cache。

KV Cache 的单并发内存消耗公式为：

对于 LLaMA-3-70B（80 层，64 个头，$d_k = 128$），在并发序列长度达到 8192 时，单个请求的 KV Cache 占用高达：

单并发仅上下文缓存就占据了一张 80GB 显卡的四分之一以上容量，使并发吞吐急剧恶化。

MQA 与 GQA 的结构妥协

#

为了在工业部署中挽救推理并发吞吐，业界对原始注意力架构做出了妥协：

延伸阅读与技术索引

#

• 关于底层 GPU 显存层次体系与 CUDA 访存带宽的量化调优，参见 GPU加速与CUDA编程实践指南。
• 关于工业级长上下文管理中 KV 缓存的显存规避与状态截断设计，参见 Claude Code 上下文压缩与状态管理。
• 关于高维语义空间与向量投影的基础知识，参见 AI 大模型 Token 与向量基础。