第 4 章 · 自注意力 › 显存带宽墙

显存带宽墙

深入阅读 · 第 4 章 自注意力——softmax(QKᵀ) 到底会有多大。

再看一眼 Q · Kᵀ：它是一个 [seq_len, seq_len] 矩阵——每一对 token 一个分数。对一条 4,096 token 的提示词，这就是每个头、每一层约 1,680 万个分数。8 个 query 头加起来，单个注意力层就是约 1.34 亿个数——在 bf16（每个 2 字节）下约 268 MB。模型从来不需要同时用到全部这些数，朴素做法却把整个矩阵都构造出来、放在一旁。

先认识两种 GPU 内存

要明白这为什么疼，你需要知道 GPU 构造上的一个事实。它有两个截然不同的地方存放数字：

  • SRAM——紧贴计算单元的一小条内存，真正的数学运算就在这里发生。它快得惊人，但极小（几个 MB）。把它想成你干活用的小桌面。
  • HBM——GPU 的主内存，权重和激活都住在这里。它很大（数 GB），但相对遥远、访问慢。把它想成走廊尽头的仓库。

桌面放不下的东西只能存进仓库，需要时再搬回来。N×N 分数矩阵远远超出 SRAM 的容量——所以它住在 HBM 里，每一步触碰它的操作都要为这趟搬运付费。

代价在带宽，不在存储

带宽指的是数据在 HBM 与计算单元之间能移动得有多快——而不是能装多少。朴素的注意力做法对带宽极其残忍：它把完整的 N×N 矩阵写出到 HBM，读回来做 softmax，把归一化后的分数再写回去，然后第三次读出来给 V 加权。矩阵在慢速内存总线上往返大约四趟，而中间的算术相比之下微不足道。注意力的大部分时间花在搬运那个矩阵上，而不是计算——所以我们说它受显存带宽限制（memory-bandwidth-bound）。

看它如何发生——拖动序列长度，感受 N² 的爆炸：

显存带宽墙——注意力为何受带宽限制
内存总线慢——瓶颈所在HBM · GPU 主内存巨大——但遥远而慢SRAM · 片上极小——但快约 10×计算N×N 分数矩阵

Q·Kᵀ 在快速 SRAM 中算出——但完整的 N×N 分数矩阵远远放不进去。

序列长度 N = 4,096 个 token
N×N 矩阵 = 33.6 MB 每头
至今搬过总线： 0 B 每头
≈ 0 B——该层全部 8 个头合计

诚实起见的一个限定：这只在 prefill（预填充——一口气消化整条长提示词）阶段、且只在 6 个全量注意力层上才会咬人。逐 token 的解码（每次只有一行新 query，也就是你在本章演示里看的部分）和 18 个 GatedDeltaNet 层从头到尾都不会构造 N×N 矩阵。

而这恰好就是下一个子章节绕开的问题：在完全不把 N×N 矩阵写进 HBM 的情况下，算出同样的答案。