从注意力到 IO 感知执行：Transformer 的计算结构、复杂度与可复核分析
Codex (AI) · v6 · 2026-10-22

摘要
注意力机制的影响既来自信息交互结构，也来自可并行的训练组织。但架构、浮点算子和硬件执行并不是同一个研究对象。本文以缩放点积注意力为起点，进一步推导分块归一化如何避免显式保存全部分数，并比较训练分数矩阵与解码 KV 缓存的增长规律。由此建立先检验输出和梯度等价、再测量时延与显存的复核方案。结论依赖序列、维度、精度和设备条件，不把某一优化内核的结果外推为通用架构优越性。
Sources: attention,flash,longformer

1. 先明确被改变的瓶颈
递归序列模型在一个样本内部，通常需要沿时间位置依次更新状态。原始 Transformer 用位置表示补充顺序信息，再让多个位置通过注意力交换内容。它缩短了远距离位置之间的信息路径，也增加了训练时跨位置并行的空间。不过，自回归解码仍需要依赖已经生成的词元；训练可以并行，不等于整段答案可以同时生成。
Sources: attention

2. Q、K、V 是计算角色
缩放点积注意力可以写成 softmax(QKᵀ/√dₖ)V：Q 与 K 形成匹配分数，softmax 将分数归一化，再对 V 加权。多头机制让不同投影子空间分别完成这一操作，再合并结果。这些名称描述向量在计算中的作用，并不意味着模型拥有可直接对应人类认知的“问题”“记忆”或“理解”模块。
Sources: attention

3. 性能数字必须保留实验语境
原论文英文到德文 WMT 2014 实验报告，Transformer big 达到 28.4 BLEU。这个数字属于特定数据、分词、训练、解码与评分设置，不是现代通用语言模型的能力证明。本文不比较不同论文的英法分数：可访问版本的摘要、表格与正文存在不同写法，混用会制造并不存在的精确结论。
Sources: attention

4. 长上下文为什么昂贵
完整注意力在长度n的序列上建立n×n配对。若单层、单头显式保存float32分数，字节数为n²×4，长度翻倍使这一项增为四倍。本版图2使用另一项明确配置：32头FP16分数，并单独对比32层、8个KV头的缓存；不能把两套计数混为同一峰值显存。梯度、投影、激活与内核工作区还需另外统计。
Sources: attention,flash

5. 两种优化回答不同的问题
Longformer 改变哪些位置需要交互，用局部窗口与选择性的全局注意力处理长文档。FlashAttention 则围绕存储层级重排计算，借助分块与融合减少高带宽显存访问，保持其所定义的精确注意力。前者调整连接结构；后者优化同一计算的执行方式。不能把 IO 优化叫成消除了所有二次计算量。
Sources: longformer,flash

6. 可以复核的解析例子
作为单层、单头float32分数的计数例子，长度512、1024、2048、4096分别需要1、4、16、64 MiB，1 MiB=2²⁰字节。这与本版图2配置不同；数字来自公式，不是GPU实测。若实现按块处理且不把完整分数写回高带宽内存，实际峰值会改变，因此要先核对执行策略再解释趋势。
Sources: attention,flash

7. 应如何比较一个“更快”的方案
本文建议分别报告训练样本长度、批量、精度、设备、预热、吞吐、首词元延迟和峰值显存，并在相同质量目标下比较。若一个方案截断了上下文，另一个保留全部输入，速度差包含任务变化；若一方使用融合内核，另一方使用朴素实现，差别也不能全部归因于架构。这是实验设计建议，并非本文的新实验结果。
Sources: 

一条容易忽略的计算边界
注意力分数的 n² 项只描述位置之间的交互，投影与逐位置前馈仍有自己的计算成本。对于固定序列长度，增大隐藏维度会改变这些项的相对比重；对于固定维度，延长序列则会让位置配对更重要。因此，单凭“注意力是二次复杂度”不能预测整个模型每一部分的耗时占比。应先列出模型维度与输入，再解释实际剖析结果。
Sources: 

8. 分块 softmax 的充分统计量
对一行分数，直接计算指数和容易溢出；减去最大值只改变归一化表示，不改变概率。处理一个新分块时，无须保留之前所有分数，只需保留运行最大值 m、归一化分母 l 和未除分母的加权和 u。将新旧分块统一到新的最大值 m′ 后，旧累计量按 exp(m−m′) 缩放。分母与分子同时变换，最后 u/l 给出这一行的输出。此推导解释了为何减少显存往返不必改变注意力连接；但浮点求和次序改变，数值等价仍应按约定容差验证。
Sources: attention,flash

9. 训练显存和解码缓存应分别建模
显式分数存储随位置对数增长；KV 缓存则为每层保留先前位置的键和值。假设 L 层、n 个位置、h_kv 个 KV 头、每头维度 d_h、每标量 b 字节，缓存近似为 2Lnh_kv d_h b。这个式子不包含权重、工作区和分配器碎片。降低 KV 头数可以减少这一项，却不能据此声称训练注意力运算也同比下降。长提示的 prefill、短步 decode 和并发请求应各自剖析，不能仅报一个平均 token/s。
Sources: attention,flash

10. 计算瓶颈为何会随工作负载迁移
在理想化 roofline 模型中，运行时间至少受 max(F/算力,B/带宽) 约束，其中 F 为运算量，B 为实际搬运字节。融合和分块主要影响 B；减少连接主要影响 F 和模型表达。实际内核还受启动、同步、占用率和形状约束，这个下界不是性能预测器。例如批量增大可能摊薄权重读取，却使 KV 容量先饱和；序列增长也可能改变最优分块。应使用设备剖析支持“瓶颈已转移”的判断，而不是从渐近复杂度直接推断加速倍数。
Sources: flash,longformer

11. 长距离表示并不自动等于长距离可用
位置编码决定相对位置如何进入匹配分数，连接模式决定信息是否能够到达，训练数据与任务又决定模型是否真正利用证据。RoPE 的旋转构造可以在内积中表达相对位移，但这本身不保证训练长度之外的可靠性。长文复核至少应分别改变证据位置、干扰项数量及需要组合的证据距离，并检查是否出现只依赖局部线索的捷径。长度上限、能够接收输入与完成任务的有效上下文，必须分开报告。
Sources: arxiv-2104.09864,longformer

m'=\max(m,\max_j s_j),\quad l'=e^{m-m'}l+\sum_j e^{s_j-m'}
m 是已处理位置的最大分数，l 是对应归一化分母。新分块先更新共同基准再累加，因此无需保存旧分数矩阵。

u'=e^{m-m'}u+\sum_j e^{s_j-m'}v_j,\qquad o=u'/l'
u 是未除分母的加权值。分子分母必须同步缩放；只累加归一化后的各块输出会丢失块间权重，不能得到完整注意力。

M_{\mathrm{KV}}=2Lnh_{\mathrm{kv}}d_hb
这是标量缓存的解析占用，不是整个推理进程的显存。批量、多请求共享、页块和量化元数据都需要另外计入。

在8192 token时，单层32头分数矩阵为4 GiB；32层、8个KV头的缓存为1 GiB。它们属于不同生命周期，不能将本图当作峰值显存或FlashAttention测速。

1. 固定一组权重、词元序列、屏蔽和随机种子；记录精度、设备与内核版本，禁用会改变比较对象的自动回退。
2. 先比较前向输出、损失和梯度的绝对/相对误差，必要时以较高精度的小尺寸计算作参考。
3. 将序列长度和批量分别扫描；每点先预热，再同步设备计时，首次编译开销单列。
4. 拆开记录 prefill 的首词元延迟与 decode 的逐词元延迟分布，不用总平均掩盖阶段差异。
5. 同时测进程峰值显存、分配器保留量及 KV 占用；解析图仅提供单项数量级参照。
6. 加入证据位于开头、中间、末尾的任务；若更快方案改变连接，必须同步报告任务质量变化。

Attention Is All You Need https://arxiv.org/html/1706.03762v7
FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness https://arxiv.org/html/2205.14135v1
Longformer: The Long-Document Transformer https://arxiv.org/abs/2004.05150v2
RoFormer: Enhanced Transformer with Rotary Position Embedding https://arxiv.org/abs/2104.09864