两个问题,一条路线。① 这两块硬件在硬件层面到底哪里不一样; ② 把同一个算子从头到尾跑一遍,看那些不一样分别在哪一步冒出来。 —— 全课的结论是:两边处处都不同,但这些不同全是同一个决定的后果。
上一课我们拆了 DeepSeek V3:总参数 671B,
但每个 token 只激活 37B —— 这就是 MoE。
于是很自然会想:既然只激活 37B,那它是不是就跟一个 37B 的稠密模型一样好养?
问的就是这个:MoE 到底省下了什么?
省算力,不省显存。
算的时候只走 37B 那条路,算力是真省了;
可 671B 的权重一个字节都不能少放 ——
你不知道下一个 token 会挑中哪几个专家,所以全都得在显存里待着。
「激活 37B」跟「像 37B 一样好养」是两回事。
⭐ 而这道题真正想让大家听见的,是它背后那件事:
显存和算力是两样东西,会各走各的。
上一课整整一课都在分开算这两笔账。
这一课要加的是第三样 —— 带宽。
而且下面马上会看到:这三样里,真正决定一个算子快不快的,
常常不是前两样。第 1 节只做一件事,就是把「算」和「搬」的兑换比算出来。
这门课后面会把两颗芯片拆开,一层一层往里看。
在开讲之前先猜一下 —— 它们各自把算力切成了多少块?
B200 上那个单元叫 Tensor Core,v7 上那个叫 MXU,
干的是同一件事:矩阵乘加。
问的是整颗芯片上各有几个。
两边各选各的 —— 都选完才出答案。
B200:592 个。v7:4 个。
592 = 148 个 SM × 每个 SM 4 个 Tensor Core;
4 = 2 个 TensorCore × 每个核 2 个 MXU。
份数差 148 倍。
⭐ 但真正该惊讶的是下一句:
这两颗芯片一拍能做的乘加总数只差 15%
—— B200 是 606,208,v7 是 524,288。
同样一份算力,一边切成 592 小块,一边切成 4 大块。
⭐ 数字给你了,「为什么」先欠着 —— 揣着这个问号往下听。 为什么份数能差 148 倍、总量却几乎一样?切得碎和切得整, 各自要付什么代价?第 3 节整整一节在回答这两个问题。 —— 到那儿的时候,你会自己认出这道题的答案长在哪一张图上。
正文从下面这张图开始,全课只有这一个前提 —— 看了它,后面每一节你都能自己预判「这里两边应该会不一样」。 看完就往下走,第 1 节只做一道除法,第 2 节就开始拆硬件。
这一版的节号跟完整版 L300 对齐,不重排。 好处是听完想往下挖的人,可以拿着同一个号直接翻 L300 的同一节; 代价是这一版会跳号:0 · 1 · 2 · 3 · 4 · 5 · 6–7 · 9 —— 第 8 节整节只在 L300 里,6 和 7 在这一版合成了一节。
⚠️ 还有一处只在图上:有一批图是整块从 L300 搬来的,图内文字用的是 L300 的号。 大节的对照关系是这三条:
小节号同理,而且更容易撞上。下面这些号在这一版里不存在, 看到就直接去 L300 翻,不用往回找: §1.4(那个 312 的拐点)、§2.6(约束能传多远)、 §2.8(把讲过的算子全放回一根轴)、§4.1(互联的边界)、 §5.5(推理引擎的移植难点)、§5.6(AOT 不占卡先算 OOM)。
—— 图不改是故意的:图是 L300 的产物,这一版一个字都不改它, 否则同一张图会有两个版本,改了一边另一边就悄悄漂掉。
这一节只做一件事:把一个数算出来。 那个数是这台机器上「算」和「搬」的兑换比 —— 搬一个字节要花一段时间,那段时间里这台机器能做多少次运算。 有了它,后面每遇到一个算子就拿自己的比值对一眼。低于它,带宽说了算;高于它,才轮到算力。 这一节不给任何结论。
两边的官方规格表都给了两个数:每秒能算多少次,每秒能搬多少字节。除一下。 这个比值的意思很实在 —— 搬一个字节要花多少时间,那段时间里这块硬件能算几次。 所以它是一条最低要求:算法每搬一个字节,配套得算够这么多次,才刚好把算力喂饱。
顺带答一个每次都会被问到的问题:两边同样是 8 颗 24 GB 的 HBM3E,
带宽为什么不一样?
—— 因为「HBM3E」只说了代次,没说档位。
每颗 stack 都是 1024 根数据线,带宽 = 1024 × 每根引脚的速率;
而这一代的引脚速率从 8.0 到 9.6 Gbps 都算 HBM3E。
反推一下:v7 跑的是 7.2 Gbps(7.37 TB/s ÷ 8 颗,算术上分毫不差 ——
但要注意 7.2 本身也在上面那个 8.0–9.6 的区间之外,见下面的折叠),
GB200 跑的是 7.8 上下。
用哪一档是买方挑的 —— 按功耗、散热、可靠性、供应挑,
不是芯片能力的上限。
(为什么不顶格跑、以及这件事怎么把本课自己的数审出一个洞,看下面。)
两边的 HBM 配置确实一模一样:都是 8 颗 HBM3E、每颗 24 GB、 合计 192 GB。(v7 那个「八颗」出自 Google Cloud 的 Ironwood 技术博客; B200 那个出自 NVIDIA 的 Blackwell 架构页。)
那带宽为什么不一样?—— 因为「HBM3E」不是一个速度,是一个速度区间。 每颗 stack 都是 1024 根数据线,带宽 = 1024 × 每根引脚的速率。 而 HBM3E 的引脚速率从 8.0 一路到 9.6 Gbps 都算 HBM3E (Micron 标「>9.2 Gbps、>1.2 TB/s」,Rambus 标「9.6 Gb/s、1,229 GB/s」)。 所以「同样的 HBM3E」推不出「同样的带宽」—— 它只说了代次,没说档位。
反推一下两边各跑在哪一档:
1024 bit × 7.2 Gbps ÷ 8 = 921.6。
算术上正好是 7.2 Gbps,分毫不差(1024 × 7.2 ÷ 8 = 921.6,
× 8 颗 = 7,372.8 GB/s)。
8 × 1024 GB/s = 8,192 GB/s,对外取整写成「8」
2,500 ÷ 8.000 = 312.5,而 2,500 ÷ 8.192 = 305.2。
再往下滚一屏,图还没完。这个除法每一层都能各做一次 —— 而这门课的主线干的正是把数据从一层挪到另一层。 现在只要留个印象:越靠近计算,这条线越低。本节最后讲融合时它才落地。
这张图把这条线用了三次。两个例子特意避开注意力 —— 那一笔第 2 节第一张图就要完整推一遍,在这儿先算等于剧透。
上面那两个例子里的矩阵乘是方阵,三边都是 N,所以约出来是 N/3。 但真实的矩阵乘是 M×K · K×N,三个维度不一样长。 那强度到底跟哪一个有关?
还是同一道除法,只是这次别急着代数:算的次数是 2·M·N·K,
搬的字节是 2·(MK + KN + MN)。约完是一个很干净的形式 ——
1/强度 = 1/M + 1/N + 1/K。
倒数相加,意味着谁最小谁说了算。
所以「K 越长强度越高」是对的 —— 但它有天花板:
K 拉到无穷,强度也只到 MN/(M+N)。
K 救不了 M 或 N 太小。
⛔ 最狠的一个特例:batch = 1 的 decode,M = 1,
于是 1/强度 ≥ 1 —— 强度上限就是 1,K 和 N 再大都没用。
这就是 LLM 解码怎么都卡在带宽上的那一行算术。
把注意力代进去:QKT 的两个外维都是序列长 S,
而收缩维就是 head_dim。多数模型 head_dim = 128,于是强度
≈ 128 —— 低于 v7 那条 312,也低于 v6e 那条 560。
—— 所以「注意力是矩阵乘,当然算力受限」是个错觉。
不做融合的话,它落在带宽那一侧,而卡住它的正是 head_dim。
那 Flash 改的是什么?下一节第一张图会给出完整那一笔;
这里先记住方向:它一个 FLOP 都没省,改的是分母。
⭐ 而同一个 128 还会再咬一口:MXU 是 256 见方,
只喂进 128 —— 那是 §3.4 / §3.6 的事。
强度问的是「要不要等数据」,喂满问的是「算的时候算得满不满」。
—— 这一节到此结束,你手上多了一把尺子。
它只回答一个问题:这个算子卡在带宽还是卡在算力。
下一节把它用在一个真算子上 ——
而那个算子会给出全课第一个大数:融合前后,强度差一千倍。
这一节这把尺子不只能判「算子卡在哪」。分子还是算力, 分母换成卡间网络带宽,同一道除法出来的就是另一个问题的答案: 每张卡每步至少要处理多少 token,跨卡通信才追得上计算。
算给 TPU v7 看(两个数都在本课 T-1 那张图上):
2,307 TFLOP/s ÷ 400 GB/s(单轴双向)≈ 5,768
三个轴一起用,带宽三倍,门槛降到 ≈ 1,922。
铺满一个 9,216 颗的域,全局批次得到 约 1,770 万 token
才不被通信拖累 —— 这就是大模型训练用超大批次的原因,
不是为了收敛,是为了让网络别拖后腿。
⚠️ 分母的口径是这里最容易错的一处。上面这个数取的是 单轴双向 400 GB/s(六条链路 1,200 ÷ 3 个轴), 不是整颗芯片的 1,200。换个口径,答案差三倍 —— 而两种写法在纸面上都读得通。
⭐ 所以这一节真正的产出是那道除法,不是 312。 分母换成 HBM 带宽,判的是算子;换成卡间带宽,判的是并行策略。 后一半是专题五的事,这里只演示一次「同一把尺子,量另一样东西」。
被塞进去的地方有名字:GPU 上是那块最多 227 KiB 的共享内存, TPU 上是 VMEM。HBM 那一层确实少搬了,但那批数据总得在片上落脚 —— 而那儿的兑换比是 64 不是 312(推导链在 1.1 那张图的下半段)。
⛔ 那为什么不干脆留在寄存器里?一个 SM 的寄存器堆总共 256 KiB, 跟共享内存一边大;而且你多占一份,能待命的 warp 就少一批。 大多数时候真做不到 —— 这笔账在 3.2 那张 SM 拆解图上。
问得对 —— 312 是「矩阵那条 roofline」的屋脊点,分子取的是 Tensor Core/MXU 的峰值,向量单元确实没算进去。
本课这么取,因为要判的算子 FLOP 几乎全在矩阵乘上 (注意力的两次矩阵乘、MLP 的大矩阵乘)。softmax、归一化那些向量运算 占的 FLOP 是零头。
向量单元到底小多少 —— 算一下:
GB200:18,944 个 CUDA Core × 2 FLOP/周期 × ≈2.06 GHz ≈ 78 TFLOPS
对上矩阵那边的 2,500,差 32 倍。
⚠️ 那个时钟是推出来的:NVIDIA 没公布 GB200 的 boost 频率,
这里用 2,500 TFLOPS ÷ (606,208 次乘加/拍 × 2 FLOP) ≈ 2.06 GHz
(606,208 = 148 SM × 4 个 Tensor Core × 1,024 次乘加/拍,
这条账 3.7 整节拆开 —— 这里先用结论。2026-09-05 补:
原来这里只有一个光秃秃的 606,208,而它的出处在两千多行之后。) ——
本课统一用这个数(此前另有一处按 1.83 外推得 2.03,已废;
1.83 是 HGX B200 的第三方数,它自己对 2,250 就差 1.4%)。
TPU 侧结构上同理:一个 TensorCore 的两个 MXU 合起来一拍 262,144 次乘加(每个 MXU 131,072),
VPU 一拍处理 1,024 个元素 —— 两个数不是同一个单位,
但量级差在两位数以上。
所以严格讲,判 element-wise 算子该用向量那条 roofline: 屋脊点是 77 ÷ 8 ≈ 9.6,不是 312。
但结论不翻。向量加的强度 0.17,比 9.6 还低 58 倍,照样卡在带宽上。 换个说法:那 1.3 TFLOPS 是矩阵峰值的 0.05%、向量峰值的 1.7% —— 换哪个分母都惨,只是后者才是公平的比法。
⭐ 真正该带走的是这件事本身:向量单元弱不是疏忽,是两边共同的设计选择。
一旦你的算子落到向量单元上,你输了两次 ——
一次是带宽受限,一次是就算不受带宽限制,可用算力也只剩三十分之一。
—— 这就是「必须做算子融合」的根本原因:
融合不是为了少算,是为了别让数据在这条又窄又弱的通路上来回跑。
按总量是 GPU 231 MiB 对 TPU 134 MiB,多 73% —— 数字在一个量级,看着就该归进「不值得记」那一堆。但不能这么归。 GPU 领先的那部分几乎全是 L2,而 L2 正是 TPU 整层没有的东西: 差的不是容量,是这块地由谁说了算。完整的三个口径在 3.2b。
* † ‡ § 四条脚注的出处说明)| 层 | TPU v7 | GB200 / B200 | 差别的本质 |
|---|---|---|---|
| 矩阵单元 | MXU 256 × 256,每核 2 个 | Tensor Core(第 5 代),每 SM 4 个 | 粒度:一块大的 vs 很多小的 |
| 标量 / 地址 | 每核只有一个标量单元 | 每 SM 128 个 CUDA core,148 个 SM* | 不规则访问谁来发地址(第 3 节) |
| 片上 · 显式 | VMEM 64 MiB / 核 · 编译器管 | shared memory 最多 228 KB / SM · 人管 | 谁决定块开多大(第 2 节) |
| 片上 · 自动 | CMEM = 0(没有) | L1 + L2 126 MB | ⭐ 整整一层只有一边有 |
| 可编程协处理器 | SparseCore × 4 / chip 跑自己的程序:前缀和 · 排序 · 计数 · gather / scatter |
(没有可编程的) TMA 是固定功能搬运,不是这一类 |
⭐ 另一层只有另一边有 |
| 零碎活的小单元 | XLU 跨 lane 单元 × 2† 转置 · 跨 lane 归约 · shuffle。超越函数( exp / tanh)是单独一类资源§ —— Pallas 的成本模型里 transcendentals 跟 flops 平级单开一栏;但部件叫什么、几个,公开资料没有 |
SFU × 4 / 处理块‡ exp · rcp · rsqrt 等超越函数。跨 lane 靠 warp shuffle 指令,官方没单列成部件 |
⚠️ 两边各点名了一个,又各自不说另一个 |
| HBM | 96 GiB / device(= 192 GiB ÷ 2);整 chip 7,372.8 GB/s(= 7.37 TB/s) | 186 GB / GPU(软件可见;物理 192 GB);8,000 GB/s | 同一代 HBM,量级相同 |
| 算力(每 chip / 每 GPU) | BF16 2,307 | FP8 4,614 TFLOPS | BF16 2,500 | FP8 5,000 TFLOPS(dense,NVL72 里那颗) | 同一量级 |
| 算力 ÷ 带宽 | 312.9 FLOP/byte | 312.5 FLOP/byte | ⭐⭐ 几乎完全相同 |
| 卡间 | ICI 1,200 GB/s 双向 · 3D torus · 最大 9,216 chip | NVLink 1.8 TB/s / GPU · NVL72 域 · 72 卡(本课编排口径按 64) | 出了域换不换协议(第 4 节) |
* 148 个 SM 来自第三方拆解,NVIDIA 官方规格表
没有公布这一项。表里其余数字都出自两家的官方文档。
† XLU 的个数出自 Google 公开的 How to Scale Your Model
(jax-ml.github.io/scaling-book)—— 原文是「一个标量核管着一个 VPU、
4 个 MXU、2 个 XLU 和多个 DMA 引擎」。⚠️ 同一句里的「4 个 MXU」对不上 v7
(v7 是 2 个 256×256),说明那句没有标代次,所以这个 2 只当「有不止一个」用。
‡ SFU 存在是 NVIDIA 官方明写的(Blackwell SM 构成里列了
Special Function Units),但个数官方框图没标,这里的 ×4 是沿用 Hopper 的画法。
§ 这一格值得单独讲,因为它是全课少见的「证据分三级」的例子。
「一颗专算神经网络的芯片,总得有块硬件负责 exp / tanh /
sin」—— 这个直觉是对的,但能公开证明到哪一步,跟直觉走得多远,是两回事:
① 曾经有过,白纸黑字。初代 TPU 的框图里就有一个独立的
Activation 单元,2017 年那篇论文的正文也提到芯片上有一块
「nonlinear function logic」。所以这类硬件在 TPU 上不是新鲜事。
② 到今天它仍然是「单独一类资源」,这条有正面证据。开源 JAX 里
Pallas 给 TPU kernel 报成本的接口是
CostEstimate(flops, transcendentals, bytes_accessed, …) ——
transcendentals 跟 flops 平级,单开一栏。
这一点很说明问题:如果超越函数只是「用普通 ALU 多跑几条指令」,
它就该被算进 flops 里,没有理由单开一栏计数。
③ 但部件叫什么、有几个、怎么实现,公开资料确实没有。
TPU v4 的 ISCA 论文对一个 TensorCore 的完整描述只有「四个 128×128 MXU,
一个带 128 条 lane、每 lane 16 个 ALU 的 VPU,以及 16 MiB VMEM」——
那 16 个 ALU 里有没有专算超越函数的,论文没拆。Google Cloud 的架构文档、
JAX 官方的 Pallas 硬件参考页、JAX 里那张机器可读的芯片规格表,
三处都只列 MXU / 向量单元 / 标量单元 / SparseCore。
至于 NVIDIA SFU 那套「查表 + 插值」的实现方式,TPU 侧也找不到公开说明。
⚠️ 最后一句最重要:规格表「没有列出」,不等于芯片上「没有」。
现成的反例就在同一批文档里 —— XLU 也没被任何官方规格表列过,
而它确实存在(它是从另一份公开材料里查到的)。规格表只列它想让你据以选型的东西。
所以本课这一格写的是「是单独一类资源,但部件未公开」,
既不写「没有」,也不写我拿不出公开出处的细节。
上一节说两边胃口一样大。这一节说吃法完全不同。
而「不同」不是指少了哪个零件 ——
是同一条取数的路上,一边有四个「跑起来才知道」的决定,另一边一个都没有。
cache 只是那四个里的两个;第一个决定是 warp 那三十二个地址能不能合并,
最后一个是没命中该换出谁。这门课后面所有的差别,都是从这一处长出来的。
不看参数,只跟着一次取数走一遍:从 HBM 里拿一块数进片上, 这一路要做几个决定,每个决定是谁做的、什么时候做的。
先把第一个决策点说清楚,它是四个里唯一跟 cache 无关的那个。
coalescing = 合并访存。一条 load 指令发出去,
warp 里那 32 个线程各自算出一个地址 ——
硬件要把这 32 个地址,合并成尽可能少的内存事务。
合得好不好,差 8 倍带宽,而且只有跑起来才知道 ——
因为编译器不知道你的地址是怎么算出来的。
TPU 那边没有这件事:数组按 8 × 128 分块、编译期定死,
搬的是一整块,根本不存在「32 个地址要不要合并」这个问题。
关键不是 cache line,是 sector。NVIDIA 的 cache line 是 128 字节, 但它切成 4 个 32 字节的 sector,而真正的搬运单位是 sector。
最好的情况:32 个线程读连续的 4 字节 = 128 字节,正好一条 cache line
—— 4 个 sector,一个字节都不浪费。
最坏的情况:32 个线程的地址跨度很大(比如按列去读一个行主序的矩阵),
每个线程落在不同的 sector —— 32 个 sector,
每个 sector 搬了 32 字节回来却只用了 4 字节。
—— 32 ÷ 4 = 8 倍。带宽照付,八分之七扔掉。
⭐ 所以「访问连续」是 GPU 优化的第一课,而它优化的不是延迟, 是你为这次访存实际付了多少带宽。 出处:cache line 128 B / sector 32 B 是 NVIDIA 的 cache 组织方式,Nsight Compute 直接按 sectors per request 报这个数 —— 完美访问是 4,多出来的每一个都是浪费。
⚠️ 那 TPU 凭什么敢说「形状已经对」? 因为它做了一件更狠的事:把向量寄存器的形状和数组在内存里的布局,做成了同一个形状。 —— 所以搬进来就能直接算,一次重排都不用。 代价不是「一行不再连续」(那恰恰是好处:一段正好铺满 128 条 lane), 而是这个形状被钉死了 —— 换最内维要真搬数据、不整除要补零、 跟外部交接要重排。这笔钱 3.3 算给你看。
判据有了,坐标有了。现在挑一个算子把它跑到底。 挑 FlashAttention 的理由很实际:每个大模型都要跑它。 而且它正好卡在 312 那条线的错误一侧 —— 两边都必须动手。
关键在第 ③ 栏 —— 前两栏图上已经说完了, 而「谁来安排搬运」正是上一张图上那个决定。 现在把它放回 312 那条线上,看看这笔搬运账到底有多离谱:
—— 所以必须分块,两边都得改写。剩下的整门课, 就是看这个「改写」在两边分别长什么样。
⭐ 请把这两个数并排看一眼:同一个注意力,FLOP 一模一样,
一次乘加都不多不少,算术强度差一千零二十四倍。
—— 所以算术强度不只由数学决定,也由你怎么实现决定。
这句话是这门课后面所有优化动作的许可证:
如果强度是算子的固有属性,那唯一能做的就是换硬件;
正因为它由实现决定,改写才有意义。
Br、列数 Bc,head_dim 是 d。Bc × d;
Q 那一块一直待在寄存器里不动。4 × Br × Bc × d 次。4 × Bc × d 字节(bf16)。Br = 128:128 > 64,过关。Br = 32:32 < 64,在片上重新变成带宽受限。问得对,而且这是个有标准答案的问题。先把「多算的」分成三种 —— 它们的性质完全不一样:
① 重算。真的把同样的运算再做一遍。Flash 的前向没有这一项 —— 上面那张图第 ①② 栏乘加次数一模一样,这是它少见的地方。 但反向有:既然中间那个矩阵没存下来,反向就得拿 Q、K 重新算一次。 省下来的存储,是用重算换的。
② 补零。形状对不齐时补出来的那些零 —— 它们照样占片上空间、照样过矩阵单元,但什么都没产出。 这是三种里最亏的一种,因为连「换来了什么」都说不上。第 3 节会看到它的量。
③ 额外的向量活儿。online softmax 每处理一块,都要把已经累加的 结果重新缩放一次。次数上不多,但它落在向量单元上 —— 而向量单元的算力只有矩阵单元的三十几分之一(第 1 节末尾那个折叠里算过)。 数 FLOP 会低估它,数时间才看得见。
那这笔买卖到底划不划算?—— 第 1 节那条线在这里第二次派上用场,但要先把一件事说准。
⭐ 算术强度的分子,是硬件实际执行的 FLOP,不是模型需要的那些。 —— 硬件不认识哪些是重算。所以判你落在 312 的哪一侧, 用的是含重算、含补零的那个强度。
但你真正要的产出是有用的 FLOP ÷ 时间。把两件事合起来, 天花板就有两块,谁低听谁的:
拿一个具体的数走一遍:含重算的强度 600,其中一半是重算。
有用比例 = 1/2 → 算力这一块 = 50% 峰值;
有用强度 = 300 → 带宽这一块 = 300 ÷ 312 = 96% 峰值。
取小的 —— 你只跑到屋顶线的一半。
⚠️ 顺便纠一个很容易顺口说错的:这时候带宽并没有被用满。 既然卡在算力上,搬运就得等计算,实际用掉的带宽只有 312 ÷ 600 ≈ 52%。矩阵单元 100% 忙着,一半忙在垃圾上, 而带宽在旁边闲着 —— 这就是「换过头了」的典型指纹。
最后回到 Flash:它离那个坏情况远得很。 前向两次矩阵乘,反向四次;Flash 因为没存中间矩阵,反向要多算一次 —— 一共 7 次对 6 次。 也就是 HFU ÷ MFU = 1.167,有用比例 85.7%。
换来的是搬运少了一千零二十四分之一。付 17% 的多算,换掉三个数量级的搬运 —— 所以它不是「优化了 kernel」,是把这个算子挪到了另一个受限区。
第 2 节数的是决策点,4 个对 0 个。上半节沿着一条路往里拆:
整颗芯片 → 一个核 → 核里那块地 → 并行层级 → 一条指令吃多大一块。
每一层都问同一个问题:这一层两边少了什么、多了什么,是谁的决定造成的。
拆完之后,下半节把零件装回去 —— 两边并排把 FlashAttention 走完一遍。
两颗芯片都是双 die 封装 —— 先看这一层,是因为有个差别 只在封装上才看得见,跳进核内部就漏掉了。
它数的是格子:两个 MXU 一共 2 × 256² = 131,072 个。
而每个格子一拍做 2 次乘加 —— 这一步图上没有。
⚠️ 「2 次乘加」不是「一个乘 + 一个加,一共两次」。
一次乘加本身就是一乘一加 —— a × b + c,一个周期里一起做完,记 2 FLOP。
所以「一格一拍 2 次乘加」= 4 FLOP,不是 2。这一步错了,最后就差一倍。
常见的错法正是把一格当成一拍只出一次乘加:
131,072 × 2 core × 2 FLOP × 2.2 GHz = 1,153 TFLOP/s,
刚好是图头 2,307 的一半。少的就是「每格 2 次」那一步。
完整那条链在 3.4,到那儿一步步走一遍。
往里一层。先看 GPU 这边有什么,再看 TPU 那边没有什么 —— 顺序不能反,不然会以为 TPU 是个「简化版」。
省掉了哪些东西,论文里是一句话列完的(Jouppi et al., ISCA 2017, In-Datacenter Performance Analysis of a Tensor Processing Unit, arXiv:1704.04760,第 8 页):
the single-threaded TPU has none of the sophisticated microarchitectural features that consume transistors and energy to improve the average case but not the 99th-percentile case: no caches, branch prediction, out-of-order execution, multiprocessing, speculative prefetching, address coalescing, multithreading, context switching, and so forth. Minimalism is a virtue of domain-specific processors.
省下来变成了什么,摘要里直说:
The lack of such features helps explain why, despite having myriad MACs and a big memory, the TPU is relatively small and low power.
⭐ 最硬的一个数在 Figure 2 那张 floor plan 上 —— 整颗 die 的面积分配:数据缓冲 37% · 计算 30% · I/O 10% · 控制逻辑只有 2%。论文紧接着那句是:
Control is much larger (and much more difficult to design) in a CPU or GPU.
—— 这就是「简单」在版图上的样子:2%。
「稳定性」那一条,官方讲的其实是另一个词:不是良率,是 确定性执行(deterministic execution)。论文的论点是 —— CPU / GPU 那些随时间变化的优化(cache、乱序、多线程) 提高的是平均吞吐,不是尾延迟;而 TPU 的确定性模型 更匹配这类应用 99 分位响应时间的要求。 ⚠️ 良率/可靠性我没有找到出处,所以这门课不说这一条。
① warp 调度器:部件在 GPU 显微镜 第 2 节,
它工作起来什么样在显微镜第 7 节(「延迟怎么被藏起来 —— 一边靠换人,一边靠排班」)。
② 线程 / warp 到底钉在哪块硅上:显微镜第 3 节,一层抽象一层硅地对。
③ SM 里的寄存器怎么用:显微镜第 2 节。一句话预告 ——
256 KiB / SM 切成四份,每个处理块 64 KiB,单线程最多 255 个;
它这么大,主要不是为了算得快,是为了同时装下几十份上下文。
④ L2 怎么用:容量与分区在显微镜第 1 节,一个数从 HBM 走到乘加单元中间几站在第 6 节。
⑤ SM 之间怎么协同:显微镜第 3 节讲 cluster / 分布式共享内存 ——
它之所以存在,正是因为 148 个 SM 之间除了 L2 再没有别的快捷通道;
另外上面那张 SM 图里的 cta_group::2 是更紧的一档:
两个 SM 的 Tensor Core 合起来做同一次矩阵乘。
148 个 SM × 228 KiB = 33,744 KiB,约 33 MiB。拆完部件,下一个必然要问的是片上那块地 —— 3.6 那个融合的全部收益就来自「中间结果留在片上」。 这个问题只有两半:多大、多快,下面两小节各答一半。
这是个必须正面回答的反问。把两颗芯片的片上 SRAM 全列出来, 一格都不漏:
| 这块 SRAM 干什么用 | 谁说了算 | GPU B200(整颗) | TPU v7(整颗 chip) |
|---|---|---|---|
| L2 缓存 | 硬件 | 126 MB ≈ 120 MiB | 没有这一层 |
| 线程上下文(寄存器堆) | 硬件 / 编译期分配 | 37 MiB(256 KiB × 148) | 不驻留上下文,没有 |
| 矩阵操作数暂存 | 指令显式搬 | TMEM 37 MiB | 累加器 4 MiB |
| 软件显式管的暂存 | 编译器 / 你 | L1+共享 37 MiB 单个线程块最多 227 KiB |
VMEM 128 MiB+SMEM(标量内存)2 MiB 每核一整块 64 MiB |
| 合计 | — | ≈ 231 MiB | ≈ 134 MiB 向量寄存器官方未公开,未计入 |
GPU 侧四项与 TPU 侧容量均取自两份《显微镜》的来源标注
(L2 与共享内存为第三方实测/规格,VMEM 64 MiB / core、SMEM 1 MiB / core 出自 JAX 开源代码)。
—— ⚠️ TPU 的 SMEM 是 scalar memory,存标量与 DMA 描述符,跟 CUDA 里的 shared memory 不是一回事;本表里对得上 CUDA shared memory 的是 VMEM 那一格。
合计为本文相加所得:37+37+37+126 MB→120 MiB = 231 MiB;
64×2+1×2+1×2×2 = 134 MiB。
64 MiB ÷ 227 KiB ≈ 289。
GPU 那 37 MiB 不是一整块,是 148 个互相看不见的 256 KiB 小岛
—— 而这 256 KiB 里 L1 还占着一部分,一个线程块能显式管到的最多 227 KiB。xla_tpu_scoped_vmem_limit_kib,遇到 VMEM 放不下时调它是常规操作)。
编译器切块的时候,除的是这个额度,不是那 64 MiB。
xla_tpu_scoped_vmem_limit_kib)。
⚠️ 默认值官方没公布 —— 但官方教程里常见的写法是把它设到
65536(=64 MiB),反过来说明默认低不少。
3.2b 回答的是「有多大、归谁管」。还差一半:跑多快。 补上之后会撞出一个几乎所有人都猜错的结果。
| 这一站 | 带宽 | 相对 HBM | 数从哪来 |
|---|---|---|---|
| HBM3e(片外) | 8.0 TB/s | 1× | 官方 |
| L2 本分区 | 21 TB/s | 2.6× | 第三方实测 |
| L2 跨到对面 die | 16.8 TB/s | 2.1× | 第三方实测 |
| L1 + 共享内存(全片合计) | ≈ 35 TB/s | ≈ 4.3× | 每 SM 128 B/周期 有实测,见折叠 |
| TPU VMEM | 官方未公开 —— 公开资料只说得到「比 HBM 高约一个数量级」 | 只给量级,原因见下 | |
⚠️ ≈ 35 TB/s 是推导值,128 B/周期是上限不是保证 —— B200 时钟本身也不确定,所以写 ≈ 35 而不是一个精确数。 推导链、第三方实测背书与两个限定,折在下面。
推导链是 32 个 bank × 每 bank 每周期 4 B = 128 B/周期/SM,
× 148 SM × 1.83 GHz = 34.7 TB/s。
128 B/周期这个数有第三方实测背书:Hopper 微基准(arXiv 2402.13499 表 V)
在 H800 上量到 共享内存 127.9、L1 125.8 B/周期/SM ——
两者撞同一个上限,说明它们确实共用一条数据通路。
另有两份公开材料独立给出同一个 128 B/周期/SM。
跨代锚点:按同一公式算 H100 是 132 × 128 × 1.755 GHz ≈ 29.7 TB/s,
第三方对 H100 报「约 33 TB/s」(反推时钟约 1.95 GHz)—— 量级对得上,公式没跑偏。
⚠️ 两个必须带上的限定:① B200 时钟本身不确定(1.83 是第三方 boost 值,
另一家实测反推约 1.99),所以写 ≈ 35 TB/s 而不是一个精确数;
② 128 B/周期是上限不是保证 —— 同一张表里 RTX 4090 标量 FP32 只跑到 63.7,
换成 FP32.v4 才上到 121.2;H800 标量就能到 125.8。
访问宽度不够,这条通路吃不满。
那 VMEM 相当于 GPU 的哪一级?—— 哪一级都不是。
按性格它是共享内存(软件显式搬、没有 tag、不会 miss),
按容量它是 L2 那一级(128 MiB vs 126 MB,几乎一样大,而共享内存全片才 37 MiB)。
—— 最准的说法是:TPU 把「共享内存那种性格」的东西做到了「L2 那么大」,
然后把 L2 整层删掉了。
① L2 不是「每 die 切一半」,是切成 4 个分区、每 die 两个(Hopper 的两倍)。 对面 die 的能用,而且不用你做任何事 —— 两个 die 由 NV-HBI(10 TB/s)连成单一 CUDA 设备,L2 全局一致。 代价就在上面那张表里:21 掉到 16.8,延迟也变高。 这正是 3.1 说的「你付了钱,但不知道自己付了」的具体形态。 TPU 反过来 —— 两半直接暴露成 2 个 device,想用对面得自己写一步。
② SM 之间换数据分三档,档与档之间差别很大:
同一个 cluster 内(≤ 8 个线程块、H100 起可 opt-in 到 16,须在同一个 GPC)——
分布式共享内存,直接读写别人的共享内存,不走 L2;
Blackwell 还多一档更紧的 cta_group::2,两个 SM 的 Tensor Core 配对干同一次矩阵乘。
出了 cluster,就只剩 L2 和 HBM —— 片上再没有快捷通道,同步只能靠原子操作、协作组或者拆 kernel。
⭐ 值得记的是因果方向:cluster 之所以被发明出来,正是因为 148 个 SM 之间除了 L2 再没有别的路。
它是在「一个 SM」和「整颗 GPU」之间硬插的一层。
—— TPU 完全不需要这一层:一个 chip 只有 2 个 TensorCore,
本来就不存在「一组核怎么协同」的问题。核少反而省掉一整层抽象,
这是规模带来的差别,不是设计水平的差别。
展开讲在 GPU 显微镜 第 3 节。
3.2b 问「有多大」,3.2c 问「跑多快」。还差第三个轴:等多久。 而这个轴,是三个里面唯一一个 GPU 自己也答不上来的。
先把提问摆清楚:一个 warp 要一块数据,L1 里没有 → 去 L2 查 → L2 也没有 → 再去显存搬。这一串下来要多久?
先补一对坐标词 —— 下面这一节和 3.4 全靠它们说话, 而这两个词的字面会骗人:
刚才少掉的是硅片上的部件。往上一层,编程模型里也会少掉东西 —— 同一个原因,第二次露面。
这一节先归到一个问题上。
既然不知道,你就没法点名。写不出「让 3 号 SM 干这个」。
你能做的只剩一件事:描述关系。 哪些活必须待在一块儿,哪些落哪儿都行。
「必须待在一块儿」的那一批,就是一个 block。
为什么是「必须」?因为这批活之间要互相传中间结果。
而中间结果放在 SM 自己身上那块小内存里,不在外面。
所以一批活一旦落到某个 SM 上,就不许再搬走 —— 搬走了,中间结果就找不着了。
反过来,互相不用说话的活,落哪儿都行。全部加起来,叫 grid。
这两层是你说了算的。 你写代码,其实就是在说两句话:这些必须在一起,总共这么多。
硬件那边还有它自己的分组,理由完全不一样 —— 是为了省电路。
取一条指令很贵。那就让 32 份数据共用同一条指令,一起算。
这 32 个绑成一捆,就是 warp。这个数你改不了,硬件焊死的。
到这儿四层齐了:底下两层硬件定死,上面两层你定。
剩下两层是后来补的,各被一个很具体的麻烦逼出来:
warp group —— 新一代的矩阵指令,要四捆一起发才凑得出形状。
cluster —— 相邻几个 SM 想直接够到对方那块小内存,不想绕远路。
六层就齐了。每一层各自被什么逼出来,横着摆一遍。
而 TPU 那边,编译的时候就已经把一切定死了 —— 不存在「不知道」。
所以这一整套语言,它一层都不需要。
两边的层根本不在回答同一个问题 ——
GPU 那些层是用来描述「谁必须在一起」的,而 TPU 压根不需要描述这件事。
所以下面这张按问题排,不按名词排。
CUDA 那套六层套娃
thread → warp →(warp group)→ block → cluster → grid
的完整包含关系折在下面,回查用 ——
真正要记的是它后面那三处「不对齐」。
| 层 | 有多大 | 谁定的 | 落在哪块硬件上 |
|---|---|---|---|
| thread 线程 | 1 个 | —— | 一个处理块里的 1 条 lane |
| warp | 32 个 thread | 硬件写死,不是你选的 | 一个 SM 四个处理块(sub-core)中的一个 |
| warp group | 4 个 warp = 128 thread | 指令集写死 | 一个 SM 的四个处理块各出一个 warp |
| block 线程块 (= CTA) |
≤ 1,024 thread = ≤ 32 warp | 你写代码时定 | 整块钉死在一个 SM 上,落下就不迁走 |
| cluster | ≤ 8 个 block H100 起可 opt-in 到 16 |
你定,可选(Hopper 才有) | 同一个 GPC 内的若干 SM |
| grid | 一次 kernel 的全部 block | 你定 | 整颗 GPU(B200 是跨两 die 的 148 个 SM) |
wgmma/TMEM 这类指令要求四个连续 warp 一起发才凑得出的形状。
这就是 3.4 那五代矩阵指令「动员多少线程」越来越多的那条线。
warp 出自织布:织布机上纵向绷紧、被一起拉动的那排线叫经线, 英文就是 warp。取的是「一排并排绷着、一起动」的意象。 中文标准译法线程束,「束」这个字抓到了重点 —— 它们是被捆在一起的。
为什么是 32?越宽,取指译码摊得越薄, 但分支分歧的惩罚越重、访存合并的粒度越粗。 AMD 早年用 64(叫 wavefront),新架构也退回了 32。
CTA 全称 Cooperative Thread Array,它就是 thread block —— CUDA 那一层叫 block,到了 PTX 和硬件那一层叫 CTA。看到两个词别当成两层。
warpgroup 除了发 wgmma,还带出一种写法叫
warp 分工(warp specialization):一个 warpgroup 专管从 HBM 往
共享内存搬数据,另几个专管发矩阵指令 —— 搬的和算的分开,流水线自然叠起来。
FlashAttention-3 就建在这套上 —— 3.6 并排走那一遍,GPU 侧就是这个形状。
再往里一层。「一个 SM」到此为止一直是个黑盒 —— 拆开它,上面那些数才不是背下来的。
一个寄存器就是 32 位、4 个字节,编译器报的「每线程多少个寄存器」数的就是它。 这条能自己验:65,536 × 4 B = 256 KiB,正好是图里那个寄存器堆。
每线程有上限:现在的卡都是 255 个(历史上变过,早期只有 63)——
就是上面那张图里「每线程 255 个(上限)」那一行。
⚠️ 这个数 NVIDIA 自己两处对不上:编程指南写 255,
而 CUDA 自带的 cuda_occupancy.h 从 Volta 起写 256。
我没能确认哪个是硬件真值(大概是「能分配 256、能用 255」,但那只是解释)——
按 255 算。
而且不是一个一个分的,是按 warp 打包:粒度是 256 个一包,
Kepler 到 Blackwell 一路没变 —— 先算「每线程个数 × 32」,再取整到 256 的倍数。
—— 换算到线程头上就是向上取整到 8 的倍数。
所以挤掉一两个寄存器占用率常常纹丝不动,要动得跨过 8 这道坎。
出处不是文档是代码:CUDA 自带的 cuda_occupancy.h
里按计算能力分档写着这几个数。
64 位的东西占两格 —— double 和指针。 多留几个指针,占用率掉得比你预期快。
最后回到最底下那一格。
一捆 warp 是 32 个线程。每个线程在这一捆里有个编号,0 到 31 —— 那个位置,就是一条 lane。
所以 GPU 上也有 lane。硬件上它就是一条 SIMD 通道, 跟 TPU 那边的 lane 是同一样东西。
差别只在一件事:GPU 的 lane 有名字,TPU 的没有。
GPU 那条有编号能写出来,你可以对它单独写 if、写循环,看着就像在写一个真线程。
TPU 那条你碰不到 —— 写不出「第 37 号 lane 单独干点别的」。
你只能让整条向量一起动,哪个元素落在哪条 lane 上,是编译器的事。
所以 CUDA 说的「线程」,跟你平时写的那个线程不是一回事 —— 它究竟是什么,得把左右两边都摆上才说得清。
「一个线程就是一条 lane」这句话还太抽象。拿一次矩阵乘当场落地 —— 而这里必须分两个时代讲,混在一起一定说不清。
① 没有 Tensor Core 的经典写法。 一个线程 = 输出矩阵里一格(或一小块)的所有者。 它自己跑那个 K 循环:从共享内存把 A 的一段、B 的一段读进自己的寄存器, 乘加,累加在自己手上,最后写回。 —— 搬运和计算,它都干。这就是大多数人脑子里那个画面。
② 用 Tensor Core 的现代写法 —— 那个画面直接作废。
矩阵乘指令是整个 warp 一起发的一条指令,
操作数是一个 fragment:同一块矩阵的元素,
按 PTX 规定的固定布局散在 32 个线程的寄存器里。
拿一条典型的 m16n8k16 举例:每个线程交出 A 的 8 个 fp16、
B 的 4 个 fp16、累加器 C 的 4 个 fp32。
—— ⭐ 单个线程自己算不出任何一个完整的输出元素。
结果也是按 fragment 散回各线程手里。
所以「一个线程负责哪一格」这个问题,
到 Tensor Core 上就失去意义了。
线程退化成两样东西:一份寄存器的持有者,和发那条指令的三十二分之一。
—— 这才是「一个线程其实是一条 lane」的完整含义。
⭐ 顺带回答「它管不管搬运」—— 管,但这一半正在退场:
Ampere 之前线程自己搬(全局 → 寄存器 → 共享),搬的时候它就被占住;
Ampere 起 cp.async 让数据绕过寄存器直通共享内存;
Hopper 起 TMA 一个线程发一条描述符就能搬一整块,另外 31 个不参与。
三代下来,线程从「搬运工 + 计算工」被剥成了「寄存器的容器」。
这条通路 3.3b 会拆开画(那张三代对照图),
到 Blackwell 更进一步 —— 累加器搬进 TMEM,连寄存器都不占了。
上一节的落点还只是编程模型那一层的说法 —— 它说的是你写代码时看到的形状。
往硬件里再走一步。这笔成本落地成三套电路,每一套只回答一个很窄的问题:
① 谁现在能发指令?
② 谁跟谁怎么对齐?
③ 数据怎么在不占用线程的前提下就位?
第一个问题这门课其实欠了一笔账:「记分板」这个词前面已经点过好几次名 —— 拆 SM 的时候、说 GPU 把 lane 暴露给你要付什么代价的时候, 却从没说过它是什么。先把这笔还上。
一句话还这笔账:记分板就是一张表 ——
表上记着每条指令要的操作数到齐了没有。
调度器每一拍扫一遍这张表,谁的操作数齐了就让谁发射。
—— 这就是「运行时才知道谁能跑」在硅片上的样子:
它得有个地方把「还没齐」这件事记下来,而 TPU 那边这张表整个不存在。
SB0–SB5);
生产者发出时加一、写回时减一;消费者带一个六位掩码,写明要等哪几号归零。
先说清哪句有据、哪句是推的。有据的是:这个词出自 CDC 6600, 那个部件的正式名字就是 Scoreboard。设计者本人为什么挑这个字,我没找到出处 —— 下面这段是推测,别当史料引。
英文 keep score 的本义不只是记比分,是把当下的状态记着: 球场边那块牌子的用处,是让你抬一次头就知道现在什么局面。这张表干的正是这件事。 所以中文译成「记分板」其实译窄了,叫「状态板」更贴。 —— 不过到了 GPU 这一代反倒名副其实:它真的在数数,发出去加一、写回来减一; 老那块牌子只是挂个「忙」的旗子。
更要紧的是别把两块板当成同一个东西。CDC 那块是一张集中的大表,
管全部寄存器的读写冲突,连「后写的别把先写的盖了」这类反向冲突也归它仲裁,
依赖判断全在运行时做。
NVIDIA 这块只有每个 warp 六个计数器,只回答「我等的那件事回来了没有」;
定长指令的依赖压根不经过它,那部分是编译器提前写进控制位排掉的。
—— 名字继承下来了,机制小了一大圈。
缩掉的那一圈正好是本课主线的一个小注脚:能在编译期定的就别留到运行时
—— 连 GPU 自己也在往这个方向挪。
第二个问题的答案是 SM 里一小组硬件 barrier:
warp 陆续到达,硬件数人头,数满全放行。
机制折在本节末尾,这里只带走最实用的那一条 ——
在 barrier 上等着的 warp 不腾位子。 所以同步密集的 kernel,占用率不会因为「大家都在等」而变好看。
第三个问题值一张图 —— 带着「哪一列在变、哪几列自始至终没变」去看。
cp.async 砍掉了寄存器那一道,地址仍要自己算;
到 TMA 才把地址生成整个接管过去 —— 一个线程填张描述符发出去就走。
⭐ 走到这儿,它和 TPU 那条编译器发的 DMA 已经是同一类东西;
剩下的分歧只有一个 —— 这条 DMA 是运行时某个线程发的,还是编译期就排好的一步。一、block 里怎么同步。
__syncthreads() 背后是 SM 里一小组硬件 barrier 资源:
warp 陆续到达,硬件数人头,数满全放行。三个容易忽略的点 ——
① 按 warp 到达,不按线程,因为一个 warp 本来就一起走;
② 它同时兼一个内存栅栏,barrier 之前写进共享内存的东西之后别人一定读得到,
所以不用另写 fence;
③ 等在 barrier 上的 warp 仍占着它的槽 ——
这一条已经提到主线上去了,不在这儿重复。
Ampere 起多了异步 barrier:把「到达」和「等待」拆开。 干完先喊一声「我到了」,不阻塞,接着去做别的独立的活,真需要数据时才等。 早到的人不用干站着。Hopper 又让等待的线程能真的睡过去 —— 以前是在共享内存上自旋。
还有一处很实际:两种 barrier 住在不同地方。
老的 bar / barrier 是 SM 里固定的一小组资源,
每个 CTA 能用的个数有限;mbarrier 是放在共享内存里的对象,
要多少个建多少个,只受共享内存容量限制。
—— 这句是 PTX 文档自己写的。
二、cluster 怎么读到别的 SM 的共享内存。三件事凑齐才成立 ——
① 同一个 cluster 的 block 被硬件保证同时调度上去;
② GPC 内部有一条专用的 SM 到 SM 网络,
cluster 必须整个落在一个 GPC 里,正是因为这条线只铺到 GPC 边界;
③ 地址靠 cluster.map_shared_rank(指针, 目标块编号)
换来一个指向对方那块内存的地址(编译成 PTX 的 mapa),
有了地址普通读写指令就发得出去,网络负责路由。
到这里,两边的部件都拆完了。但有一件更基本的事,前面一直没说: 一次矩阵乘,在硬件上究竟是按什么顺序算的?
这一节大概是全课唯一一处「几乎所有人都想错了、而且错得毫无察觉」的地方。 因为错的那个版本太自然了 —— 它就是你在纸上算矩阵乘的方法。
先把你脑子里那个说出来:拿左矩阵的第一行、右矩阵的第一列, 对应相乘再加起来,得到结果矩阵左上角那一个数;这个格子就算完了, 换下一个。—— 数学上完全正确。 而真实硬件,无论 GPU 还是 TPU,一个都不这么算。
| 累加器形状 | 复用率 | 读一个数干几次活 |
|---|---|---|
| 1 × N(只留一行) | N ÷(1+N) | ≈ 1,N 开到一百万也一样 |
| 128 × 128 | 16,384 ÷ 256 | 64 |
| 128 × 256 | 32,768 ÷ 384 | ≈ 85 |
⭐ 把这一节收成一句:同一个数学,区别只在「哪个维度摊在空间上,哪个维度走时间」。 这一个选择往下决定了所有事 —— 累加器要多大、tile 为什么是最核心的那个旋钮、 TPU 为什么非得让权重坐着不动。 接下来 3.4 要比的「一条指令吃多大一块」,比的就是这个选择的具体取值。
上一节讲的是谁来发这条指令。这一节换一个问题 —— 一条指令发出去,一次能吃下多大一块?
这个问题听起来很内部,但它会直接决定一件很外部的事:
你的 head_dim 选多少不吃亏。
16 —— 记住这个数,下一张图要用它。 —— 出自《GPU 显微镜》head_dim 上会怎么样,正文接着走。 —— 出自《TPU 显微镜》两张图里最该看的是同一个位置:那条指令在「求和的那一维」上有多深。
GPU 是 16,TPU 是 256。差 16 倍。
这一维有个名字叫收缩维 —— 就是矩阵乘里被加掉的那一维,一条指令一次要沿着它累加多少项。
这三个字母(M/K/N)后面到处都是,
先记一条就够:K 是唯一出现两次、也是唯一消失的那个
—— 它被加掉了,所以叫收缩维。
剩下两个就是答案的形状。
⭐ 而那个 128 的麻烦,一句话是:两头不靠
—— 当阵列的边嫌它小(只占 256 的一半),当流过去的行数又嫌它短
(摊不动装权重的钱)。
N 为什么也废一半、
以及转置到底等不等价(被问到再展开)M/K/N 一次说清 —— 后面到处都是它们
矩阵乘写成 A[M, K] × B[K, N] = C[M, N]。
三个字母只需要记一条规律:K 是唯一出现两次的,也是唯一消失的那个。M = 有多少条结果,N = 每条结果里有多少个数。
而 K = 每个数要加多少项。
m16n8k64 —— 输出 16×8,每个数由 64 项加出来。
(顺序也有规律:BLAS 先报输出形状 M、N,
最后才报 K。)i、j、k:
对每个 i(第几条结果)→ 对每个 j(这条里的第几个数)→
对每个 k(把第 k 项加进来)
M = i 跑多少次,N = j 跑多少次,
K = k 跑多少次。
k 是最内层那一层 —— 最内层干的正是「累加」。
所以「K 是收缩维」不是约定,是它本来就是那个累加循环的名字。
M K · K N → M N ——
挨在一起的那两个必须相等,然后它俩一起消失。
剩下的就是答案。
M 是时间,K 和 N 是空间。
在脉动阵列上,K×N 那一块是装进阵列驻留的
—— K 是阵列的高,N 是阵列的宽,
两个都被硬件尺寸卡死。而 M 只是有多少行数据流过去,
想多长有多长。
K=128 浪费一半、N=128 浪费一半,
却从来没提过 M —— 因为 M 根本不占地方,它只占时间。N=128 为什么也只剩一半?—— 同一个算术,完全不同的原因
先把阵列想成一块 256 × 256 的格子地。
要做 A[M,K] × B[K,N],硬件干的事是:
把 B 那一块 K×N 铺进这块地(K 是高,N 是宽),
然后让 A 的每一行从左边流进来。
K=128(高只占一半)—— 下半边 128 行头顶没有输入。
能不能塞点别的进去?不能 —— 同一列的格子会把结果加进同一个数,
塞什么都会污染那个和。
N=128(宽只占一半)—— 右半边 128 列脚下没人要。
它们照样能算,只是算出来的是没人需要的输出通道
—— 而你的输出就只有 128 个通道,没有更多的了。
利用率 = (K ÷ 256) × (N ÷ 256)。
所以 128 那一边占一半,另一边满,就是 50%。
⚠️ 这个式子是从阵列结构推的,不是官方公式 ——
但它跟官方那句「head_dim 是 128 时利用率无法超过 50%」对得上。
N 短能救 ——
只要有另一个矩阵乘用的是同一个输入 A,
就可以把两边的 B 在宽度方向拼起来,一次算完。
Q、K、V 三个投影融成一个大矩阵乘,干的正是这件事
—— 它们吃的都是同一个 X,所以把三个 N 拼成一个大的。
K 短救不了 ——
K 方向是求和,拼进去的东西会被加进同一个结果里。
这是它跟 N 的根本区别。
PV 那一步的 N=128
恰恰救不了:要拼就得有别的矩阵乘共享同一个 P,
而每个头的 P 都不一样。N 需要的是「同一个 A、不同的 B」;
而共享 KV 给你的是「同一个 B、不同的 A」。
后者只能往 M 方向摞 ——
而 M 是时间不是空间,摞多少都填不满那半块阵列。
head_dim。
它要是 256,两头的浪费一起消失。
但这是模型架构的决定,不是调参 ——
也就是说,模型选的那个 head_dim,直接决定了注意力在这块硬件上能拿到几成峰值。
PV 会撞 N=128;
但改成算它的转置 Vᵀ · Pᵀ,三个字母全换了位置:
M 变成 128,而 K 和 N 都变成了长的那两维。
(P · V)ᵀ = Vᵀ · Pᵀ。
这是精确相等,不是近似。
Vᵀ · Pᵀ 凭什么等于 (P · V)ᵀ?
—— 拆开走一遍(形状 → 累加 → 数字 → 计算量)第一步:形状先对上。
P 是 [q, kv],V 是 [kv, d]
→ O = P·V 是 [q, d]Vᵀ 是 [d, kv],Pᵀ 是 [kv, q]
→ Vᵀ·Pᵀ 是 [d, q][q, d] 转置过来正好是 [d, q]
—— 形状先对上了,剩下的是里面的数对不对。
第二步:一格一格看,被点乘的是同样那两条向量。
原来那格:O[i, j] =
P 的第 i 行(横的,长 kv)
点乘 V 的第 j 列(竖的,长 kv)。
转置那格:(Vᵀ·Pᵀ)[j, i] =
Vᵀ 的第 j 行 点乘 Pᵀ 的第 i 列。
而按转置的定义,Vᵀ 的第 j 行就是 V 的第 j 列,
Pᵀ 的第 i 列就是 P 的第 i 行。
⭐ 所以两边点乘的是一模一样的两条向量,只是先后换了个个儿。
第三步:为什么「换个个儿」不影响结果 —— 这是全部的关键。
矩阵乘之所以不能随便交换(A·B ≠ B·A),
是因为交换之后「谁的行去配谁的列」变了。
但这里我们同时把两个矩阵都转置了 ——
行列也跟着换了一遍。
⭐ 两次交换互相抵消,配对关系一点没变。 剩下的只是每一格里那两个数谁写前面 —— 而它们是标量,标量的乘法本来就可交换。
写成式子就一行:
O[i,j] = Σ P[i,k]·V[k,j] 与 (Vᵀ·Pᵀ)[j,i] = Σ V[k,j]·P[i,k]
同一串乘积、同一个求和,只是每一项里两个因子调了位置。
第四步:拿数字走一遍,看它逐项对上。
P = [ 1 2 3 ] V = [ 7 8 ]
[ 4 5 6 ] [ 9 10 ]
[ 11 12 ]
O = P·V 的左上角 = 1×7 + 2×9 + 3×11 = 7 + 18 + 33 = 58
Vᵀ = [ 7 9 11 ] Pᵀ = [ 1 4 ]
[ 8 10 12 ] [ 2 5 ]
[ 3 6 ]
Vᵀ·Pᵀ 的左上角 = 7×1 + 9×2 + 11×3 = 7 + 18 + 33 = 58
⭐ 注意不是「结果碰巧相等」—— 是 7+18+33
这三项一项一项都一样。
整体是 O = [[58,64],[139,154]],
而 Vᵀ·Pᵀ = [[58,139],[64,154]]
—— 正好是 O 的转置。
第五步:计算量一模一样。
原来:q × d 个输出,每个要 kv 次乘加
→ q·d·kv。
转置:d × q 个输出,每个还是 kv 次乘加
→ d·q·kv。
一次不多一次不少。
⛔ 那硬件上凭什么会快?—— 因为变的不是乘加的数量, 是它们怎么摆进那块 256×256 的地里。
M=q、K=kv、N=d=128
→ 阵列的宽只用了一半M=d=128、K=kv、N=q
→ 高和宽都能占满,那个 128 跑到 M 上去了⭐ 同样多的乘加,摆法不同,空转的格子数就不同。 —— 这就是「把 128 从空间挪到时间」那句话的全部含义。
Vᵀ,被装载的反而是 Pᵀ。
而不同 head 的 Vᵀ 必须配自己那份 Pᵀ
—— M 方向照样拼不了,跟前面 N 方向拼不了是同一个理由。
① 权重是移位进去的,装满一块要 256 个周期。
它不是「一下子拍进去」—— 权重从一个 FIFO 里一行一行往下推, 像卷帘一样铺满整块阵列。 TPU 初代那篇论文里写得很直白:装一块权重要 256 个周期。
② 那是不是要等它装完,左矩阵才能开始流?—— 靠双缓冲,通常不用等。
每个格子里有两套权重寄存器:一套在参与计算,另一套同时在装下一块。 所以那 256 个周期是藏在「上一块还在算」的时间里的 —— 论文里那句话的原话就是「用双缓冲来藏掉这 256 个周期」。
M 行过去 ≈ M 个周期。
M 行就是 M 拍,跟这一块里有多少次乘加无关。
乘加是在 65,536 个格子里同时发生的,它不占额外的时间。
M = 128 的时候:算 128 个周期,装要 256 个周期
—— 装的比算的还慢,双缓冲藏不住了。
你会卡在等权重上。
③ 装进去之后怎么「固定」? —— 每个格子里就有一个权重寄存器,值放进去就不动了, 整块 tile 算完之前一直是它。 这套做法有名字,叫权重驻留(weight stationary)。
④ 有没有「两个矩阵都在流」的算法?—— 有,而且是三大数据流之一。
脉动阵列的经典分类就三种,区别只在「谁留在格子里不动」:
B 的一块
—— 阵列装的是 K × N,流过去的是 M。
TPU 走的是这条。M × N,流过去的是 K。
这就是你问的那种。A 的一块
—— 阵列装 M × K,流 N。输出驻留的好处正是你想到的那个:没有「装权重」这一步, 两个操作数一起流进来。 代价是结果攒在格子里,最后要单独排出去一趟。
K 和 N,时间是 M
M 和 N,时间是 K
M 和 K,时间是 N
head_dim = 128 的麻烦在于:它两头不靠。
这一整串画成了一张图 —— 三种驻留方式、真尺寸怎么切块、乒乓怎么藏装载、段数怎么换重搬次数, 四条带一次走完。
Q 投影是怎么在 MXU 上跑完的
(M=128K,K=7168,N=1536)先把三个数摆好。
左矩阵 X 是 [131072, 7168],
右矩阵 W 是 [7168, 1536],
结果 Q 是 [131072, 1536]。
① 一次能放多大一块?—— 这个没得选:256 × 256。
阵列就这么大,右矩阵只能按这个尺寸切。
② 一共切成多少块?
K 方向:7168 ÷ 256 = 28N 方向:1536 ÷ 256 = 628 × 6 = 168 块。W 只有 7168 × 1536 × 2 B = 21 MiB
—— 一次就全塞进 VMEM 了(那儿有 64 MB)。这一步只发生一次。
③ 时间线长什么样。假设我们把左矩阵按 2048 行一段来喂:
装 B(k=0, n=0) ← 256 个周期。第一块没得藏,只能等
↓
流 X 的 2048 行 ← 2048 个周期
(同时在装 B(k=1,n=0)) ← 256 周期,藏在上面那 2048 里
↓
流 X 的 2048 行 ← 用 B(k=1,n=0)
(同时在装 B(k=2,n=0))
↓
… 28 次之后 …
↓
这一段 2048 行、这 256 个输出列,累加完了 → 写出去
↓
换 n=1,再走 28 次 … 六个 n 走完 → 这 2048 行彻底做完
↓
换下一段 2048 行 → 168 块 B 全部重来一遍
④ 「装完才开始流」吗?—— 只有第一块是。 之后靠双缓冲:装下一块和算这一块同时进行。 藏得住的条件是「流的行数 > 256」 —— 这里流 2048 行,8 倍富余,稳稳藏住。
⑤ 每块 B 只用一次吗?—— 不是。 它被重复搬的次数,正好等于左矩阵切成了几段。
| 左矩阵一段多少行 | 片上累加器要多大 | B 进阵列的总次数 | 装载藏得住吗 |
|---|---|---|---|
| 256 | 0.25 MiB | 86,016 | 藏不住(256 不大于 256) |
| 512 | 0.5 MiB | 43,008 | 能 |
| 2048 | 2 MiB | 10,752 | 能 |
| 8192 | 8 MiB | 2,688 | 能 |
⭐ 看这张表就明白左矩阵为什么不能一段太短: 段越长,B 重搬得越少;但累加器要一直待在片上,它跟着变大。 所以「一段多少行」是被片上容量顶住的 —— 又绕回那口灶台。
⑥ 那能不能让 M 一口气流完、B 只搬 168 次?
—— 算一下就知道不行:那样累加器要
131072 × 256 × 4 B = 128 MiB,
比整块 VMEM 还大。
所以 M 必须切段,而切段就必然重搬 B。这不是实现不好,是账本身如此。
M 有 13 万行,随便切一段都远超 256,装载稳稳藏住。
—— 投影这类矩阵乘根本不担心这个问题。
M = 128 ——
连一段的门槛(256)都不到。
同一个机制,一个稳稳藏住,一个连门槛都够不着。光看这两个数没有感觉。拿一个真实的形状套进去就有了。
注意力的 head_dim 常取 128。
GPU:128 ÷ 16 = 8 —— 切成八条指令,一点不浪费。
TPU:128 对 256 —— 只喂满一半,另一半空着。
那空着的那一半,找点别的活填进去不行吗?
不行。收缩维是一条累加链 —— 塞进去的任何别的数据,都会被加进同一个部分和。谁也顶不上来。
反过来问:那 GPU 为什么不吃这个亏?
不是因为它有调度器。
调度器换的是 warp,填的是时间上的气泡 —— 它填不了收缩维里空着的那一半。
是因为它的 K 只有 16,几乎什么形状都是 16 的整数倍。
而它敢把 K 做这么窄、把对齐留到运行时收拾,代价就是 3.2 里 TPU 少掉的那五样。
⚠️ 这里最容易翻车的一点先说在前面:
16 和 256 是一条指令的形状,不是一个周期做多少事
—— 把这两件事除在一起,算出来的峰值必然荒谬。
4,096 × 2 FLOP × 148 SM × 1.83 GHz = 2,218 TFLOPS,
官方 2,250,差 1.4% —— 那点缺口来自时钟不确定,不是公式错。
262,144 FLOP/周期/MXU × 2.2 GHz × 4 个 MXU = 2,307 TFLOP/s262,144 FLOP ÷ 65,536 格 = 每格每拍 4 FLOP,
也就是 2 次乘加 —— 一次乘加 = 一乘一加 = 2 FLOP,两次正好 4。mma.sync 的 m16n8k16
一共是 16 × 8 × 16 = 2,048 次乘加;一个 Tensor Core 每周期做 1,024 次
—— 所以这条指令占它 2 个周期。它是一包活,不是一拍的事。
最大那条 tcgen05.mma(128×256×16)是 524,288 次乘加,
一个 SM 每周期 4,096 次 → 要 128 个周期;两个 SM 配对就是 64 个。head_dim = 128、bf16、K/V 双缓冲、分数矩阵按 fp32 估;
真实 kernel 会把 Q 和累加器留在寄存器里,所以这是量级估算,不是逐字节账。)
| Br × Bc | K+V(双缓冲) | Q | 分数矩阵 | 合计 |
|---|---|---|---|---|
| 64 × 64 | 64 KiB | 16 KiB | 16 KiB | 96 KiB |
| 128 × 64 | 64 KiB | 32 KiB | 32 KiB | 128 KiB |
| 128 × 128 | 128 KiB | 32 KiB | 64 KiB | 224 KiB ← 顶到天花板 |
| 256 × 256 | 256 KiB | 64 KiB | 256 KiB | 576 KiB ← 爆了 |
| 512 × 1024 | 1.0 MiB | 128 KiB | 2.0 MiB | 3.1 MiB |
| 2048 × 2048 | 2.0 MiB | 512 KiB | 16.0 MiB | 18.5 MiB |
128 × 128 那 224 KiB
几乎是顶着天花板放进去的,再往上一档直接爆两倍半。
2048 × 2048 那 18.5 MiB 装得下,还余三分之二。
上半节到这儿为止,数的都是「一边有、另一边没有」。还有两块,是各自额外长出来的。
但这一节要问的不是「它们是什么」,是这个 —— 宣传页上写着「支持」的东西,你到底用不用得上?
这一节就两个例子,两边各一个。
第一个例子在 GPU 这边,叫块量化。 宣传口径是「B200 支持 FP4」。
这句话是真的。门槛在于:带块量化的指令有两条通路, 而它们不在同一颗 die 上 —— B200 只有其中一条。 (哪条对哪颗 die,图上写了。)
⭐ 块量化在硬件上到底怎么跑的(NVFP4 与 MXFP4、 Blackwell 加了哪三样、训练侧那个 16×16)—— 整套机制在 专题八 · 精度与量化。
第二个例子在 TPU 这边,叫 SparseCore —— 一颗专门查嵌入表的核。 宣传口径是「TPU 有个专门查 embedding 表的核」。
这句话也是真的。门槛在于:它只对「同一批行被反复查」的负载划算。 推荐系统那边少量热行被反复命中,重复度很高 —— 那是它的主场。
而语言模型的取行重复度不够高 —— 注意不是「没有重复」:一条十几万 token 的序列里,不同的词只有几千个, 而且绝大部分命中的是最高频那一两千行。 重复确实有,只是跟搜广推差着量级 —— SparseCore 的收益门槛是按那种集中度定的,这点重复度不划算。 编译器根本不往那儿派,生产任务里它一次表都没查过。
SparseCore All Gather。
那它到底怎么处理不规则访存?—— 把这颗核拆开看一眼就清楚了。
pltpu.get_tpu_info().sparse_core)。
⚠️ 「操作 VMEM 还是 HBM」是个陷阱二选一 —— 两块都碰,
但它碰的那块 VMEM 是它自己的(每个子核一份 + 一份共享,
OpenXLA 文档里叫 SPMEM),不是 TensorCore 那块。
⭐ 它扛延迟靠的不是每次更快,是十六路各追各的地址、同时欠着很多次取数。知道了它不是什么,接着该知道它是什么 —— 掀开盖子看一眼。
图上那条「集合通信」的线,落到 XLA 里就是一组开关。 列出来是为了让「卸载」这两个字有个能上手的抓手 —— 全部出自公开材料:
| flag | 它开的是什么 | 该知道的那一条 |
|---|---|---|
| ① 主开关 —— 哪些集合通信交给 SparseCore | ||
--xla_tpu_enable_sparse_core_collective_offload_all_gather |
All-Gather 交给 SparseCore | Ironwood(v7)上这三个默认就是 true。 也就是说:你不设任何 flag,它已经在替你干活了 |
…_offload_reduce_scatter | Reduce-Scatter | |
…_offload_all_reduce | All-Reduce | |
…_offload_2d_all_gather |
两层网格上的 All-Gather | 要显式开 |
| ② 地基 —— 主开关生效需要跟着一起设的 | ||
--xla_tpu_use_tc_device_shape_on_sc |
让 SC 沿用 TensorCore 那套 device shape | MaxText 把这四个打成一组叫
ENABLE_SPARSECORE_OFFLOADING_BASE_FLAGS,取值 true / false / false / true。
单开主开关往往看不到效果,坑就在这儿 |
--xla_sc_enable_instruction_fusion | SC 侧的指令融合 | |
--xla_sc_disjoint_spmem |
共享 SPMEM 是否按 tile 切成互不重叠的块 | |
--xla_sc_disable_megacore_partitioning |
关掉 megacore 划分 | |
| ③ 看得见 —— 不开这几个,profile 上那段是黑的 | ||
--xla_tpu_enable_all_gather_offload_tracing…_reduce_scatter_offload_tracing…_all_reduce_offload_tracing |
把卸载出去的那段打进 trace | 调优时才开。「看 trace 不看宣传页」这条规矩, 在这里需要你先把灯打开 |
| ④ 顺带一个不属于 SparseCore、但每次都要一起调的 | ||
--xla_tpu_scoped_vmem_limit_kib |
单个 op 最多能占多少 VMEM,剩下的留给下一个 op 预取 | MaxText 公开配方:稠密模型 98304、MoE 81920 |
⚠️ 两条必须一起记住的约束。
第一,Continuation Fusion 与 SparseCore 卸载只能二选一——
MaxText 的注释原话是「Either one of CF or SC can be enabled at a time」。
两者都是「让通信和计算重叠」的手段,走的是两条不同的路,不能叠加;
以为「都打开更快」是这一组 flag 上最常见的误用。
第二,98304 KiB = 96 MiB,比 v7 单核 64 MiB 的 VMEM 还大——
这个数在公开配方里是跨代通用的一个上限申请,实际生效值受目标机器约束。
这里只如实转述配方,不替它解释。
| NVIDIA 的 Tensor Core | TPU 的 TensorCore | |
|---|---|---|
| 是什么 | SM 里面的一个执行单元 | 整颗 device 的主计算核 |
| 里面装着 | 就是矩阵乘阵列本身 | MXU + VPU + 标量单元 + VMEM |
| 一颗芯片上几个 | 一百多个 SM,每个里面还有若干 | v7 上 2 个 |
D = A×B + C。下面三块都是现场追问逼出来的,主线上跳过不影响听懂, 但它们各自都堵住了一个很容易含混过去的地方。
src/maxtext/configs/README.md 与
benchmarks/xla_flags_library.py;
「continuation state 从一个 async collective fusion 传到下一个」
的描述见公开技术博客,NVIDIA/Google 均未发布 CF 的正式设计文档。
D = A×B + C。
| 代 | 谁来发这条指令 | 操作数放在哪 |
|---|---|---|
| Volta ~ Ampere | 一个 warp(32 线程)同步执行 | 各线程私有寄存器里各拿一片 |
| Hopper | 一个 warpgroup(128 线程),异步 | A 可以直接来自共享内存 |
| Blackwell | 一个线程发出去就不管了 | 共享内存 + Tensor Memory,累加器也在里面 |
Blackwell 那一行的「必须由单个线程发起」是官方文档的措辞, 理由很干脆:这条指令要用的数据一个字节都不在私有寄存器里, 全在 CTA 共享的内存空间中 —— 所以矩阵单元跟 warp 调度器解耦了。
那 TPU 有没有 MMA?—— 数学上有,形式上没有,而这个差别正好是本讲的主线。⭐ 最后一个观察,也是这一讲想让你带走的东西: Blackwell 这一代的动作 —— 累加器搬出寄存器、单线程发射、 操作数放进专用内存 —— 方向上是在朝 TPU 靠。 它在矩阵乘这一块主动放弃了「每个线程各管一片」的动态性, 换来更大的块和更少的调度开销。 而在别的地方,GPU 依然是那台什么都能干的动态机器。 两条路没有合并,但在最赚钱的那个部件上,它们靠近了一步。
出处:NVIDIA PTX ISA 文档(mma / wgmma /
tcgen05.mma 三代指令)、NVIDIA CUTLASS 的 Blackwell 功能文档、
Colfax Research 的 Blackwell Tensor Memory 教程。
「Blackwell 在朝 TPU 靠」是我的判断,不是任何一方的官方说法。
上半节拆出来的零件,现在全部要同时上场。
这半节做两件事:先把算子走完一遍(3.6),
再回头兑现前面立过的两个说法 ——
3.6 末尾兑现第 0 节那句「处处不同、只有一个成因」,
3.7 兑现第 1 节那个没解释过的巧合:两边峰值为什么在同一量级。
这一节只回答两句话:这一锅能做多大? 以及两边分别是被什么卡住的?
上半节拆出来的全是零件,这一节把零件装回去 —— 第 2 节挑的那个算子,现在真的在两块硬件上各跑一遍。
比方铺完,现在换成真数字。先回答一个还没正面回答的问题:为什么非分块不可? 第 2 节说「需求比供给多 5 倍」是个比值,这里给绝对量。
要删掉哪一项清楚了,接下来看两边各自怎么删。先各走各的,走完再并排 —— 这个顺序有意为之:直接看对照表会以为两边处处不同, 各走一遍才看得出绝大多数站是一样的。
两条路各走完了,现在把它们并排放在一起 —— 这是本课的主图。
128×128 的 bf16 tile(Q/K/V)=
3 × 128 × 128 × 2 B = 96 KiB,再留出双缓冲,基本就到顶了。head_dim=128:在 GPU 上是 16 的整 8 倍,怎么切都齐;
在 TPU 上只填满 256 那条边的一半。
head_dim = 128 撞上 256 的 MXU,
而且两个矩阵乘各撞一次 ——
QKᵀ 那步 128 落在收缩维,PV 那步 128 落在输出维。
于是这个算子的 MXU 利用率封顶 50%,调 kernel 参数救不了。
—— 这里说的是 kernel 参数那一层。真出路有两条,都在 3.4:
改 head_dim(模型架构的决定,不是调参 —— 32 头 × 128 换成 16 头 × 256,
实测 MFU 在 8K/16K/32K 上分别 +21% / +32% / +46%),或者全转置流水线。
head_dim = 128 卡在中间 ——
当空间嫌它大不起来,当时间又嫌它短。
GPU 那边阵列小得多,128 怎么摆都齐,所以它撞的是容量;
TPU 阵列大到 128 填不满,所以它撞的是形状。
(这条判据怎么来的、三种数据流为什么都救不了 ——
在 3.4 那两个折叠里。)⭐ 再往下就是 FlashAttention 自己的事了 —— 那张内外循环的经典图、两个版本为什么把循环掉了个个儿、 块大小的实测扫描、以及三层效率天花板的拆解, 都在 专题三 · 注意力演进 第三节。 这一讲只借它把两块硬件并排走一遍。
最后一张收账。不看站点,只看「你要写多少东西」 —— 这是这趟旅程留给写代码的人的那一面。
cp.async/TMA、双缓冲、warp 分工、对齐、要不要占 L1。BlockSpec 与 grid。右栏比左栏短,这件事本身就是结论。红带按住了最容易讲过头的那句:L2 绕不开,绕开的是对缓存的依赖。第 1 节开场把两个峰值并排摆出来(2,307 对 2,500),只说了「同一量级」, 没说为什么。走完全程之后,这个「为什么」现在算得出来了 —— 而答案恰好是上半节那两个差得最离谱的数,自己把自己抵消了。
看图的时候盯住右边那张账。同一把刀有两个尺寸,图上都给了,别混成一个 —— 刀有多宽(一次吃多深)就是 3.4 那个收缩维,16 对 256; 一下能切下多少(每拍多少次乘加)是 1,024 对 131,072,差 128 倍。
148 ÷ 128
它不是约等于,是精确等于:148 ÷ 128 = 1.156,
而 606,208 ÷ 524,288 = 1.156。
这一节是读第 6 节那组实测的前置。
那组数是 64 对 64 —— 而「64 颗怎么连起来」这件事,
两边的答案在结构上根本不是一回事。
这一节只画硬件:几条链路、每条多宽、连成什么形状、最远几跳。
哪种通信模式吃亏、EP 为什么对拓扑最挑剔 —— 那是「切法」的事,在专题五。
把这张图压成一句:NVL72 是一个「交换式的域」,ICI 是一张「直连的网」。
—— 交换式的域有上限(72 颗),但域内任意两点一跳、位置无关;
直连的网没有交换机、只连邻居,但同一套协议能一路铺到 9,216 颗。
而 64 这个数,恰好落在两边最不一样的地方:
它在 NVL72 的上限之内(整个域一跳可达),
而在 ICI 上它是一个 4×4×4、最远 6 跳的环面。
但拓扑能延展多远,跟你实际拿得到多少是两件事。 下面是我们自己跑的数 —— 不是厂商标称,是实测掉了多少。
⚠️ 先把这组数的身份说清楚,它跟上面两张图不是一回事。 上面画的是结构(两边怎么连),这组是一侧的实测 —— 只有 GPU 侧,而且是同一个 NVL72 域内从 8 卡扩到 32 卡, 不是两边对照。TPU 侧的对应实测这门课没跑过(环面上的集合通信一次都没测)。 放在这里只为说明一件事:「域内没有距离」不等于「域内没有代价」。 至于哪种通信模式在哪种拓扑上吃亏 —— 那还是专题五的事。
这对词你已经见过四次了 ——
2.1 那个「4 对 0」、3.3「哪几层是运行时才定的」、
3.3b 那三套电路、3.6 六站里两边分岔的那两站。
前面每一节都停在同一个决定上,只是每次换了个部件。这一节把那件事本身拿出来讲
—— 「谁来安排」这个决定,从硬件一路流到你写代码的那只手上。
最底下这一层最能说明问题。取一次数要等几百个周期,这段时间干什么? 两边的答案不是「快慢」之分,是「谁负责」之分。
硬件上的这个差别,会原样出现在你敲的字里。
threadIdx,右边 JAX 一个下标都没有 —— 差别写在代码里,不在形容词里。下半两栏刻意一样长:这是取舍,不是优劣。再往上一层,连集合通信也是同一个模式:
GPU 那边 all-reduce 是你在训练循环里亲手写的第二步,
TPU 那边你从头到尾没写过它 —— 它是编译器按你声明的并行方式自己插进去的。
(这一对代码在 L300 5.3。)
说到这儿会有一个自然的反应:「那什么都交给编译器不就行了?」 —— 不行,而且边界画得相当清楚。
O(n·d),和 O(n²·d) 比可以忽略。
「不改数学」说的是结果等价,不是「一条指令都没多」。这一节讲到这里,最强的一个反驳一定会出现,而且它有一半完全成立: 整个 step 的 kernel 图提前录好、之后一条道跑到黑,收益是真的, 很多负载上开与不开是数量级的差别。
但要判断它推翻了什么,得先讲准它究竟为什么快 —— 最常见的解释是错的:省掉的不是「分支」,是 CPU 一条一条发 kernel 的开销。 —— 所以问题就落在位置上了。
cudaGraphInstantiate,
instantiate 不是 compile。这个用词是准确的。
这一问比 CUDA Graph 那句更值得答,因为它说的方向是对的。
把 3.3b 那张四代图竖着读一遍,趋势藏不住:
cp.async —— 砍掉寄存器那一道,
地址仍然自己算。tcgen05 —— 这一步最狠。
矩阵指令由单个线程代表整个 CTA 发射,
而操作数与累加器搬进一块 256 KB/SM 的专用矩阵暂存,
连结果都不回寄存器堆。四代下来,那条通路的终点长这样:整块搬、按描述符搬、 不占线程、落在一块软件说了算的专用暂存里。 —— 而这正是 TPU 第一天的样子。
3.6 那六站走完了 —— 但下面这两个数,不是那六站跑出来的。
FlashAttention 那一遍全程是结构对照与理论上界
(2.3 那张图自己就标了「真实 Flash 会重复读 K/V,强度到不了 65,536」);
而两边同口径的实测只有一组,跑的是 Hunyuan3-295B 整网训练。
—— 所以它兑现的是第 1 节那条线,不是 3.6 那次并排。
这一节先说这两个数该被当成什么,再把它们摆出来。
最常见的两种比法都不对:比峰值(那是芯片能干什么,不是你的活跑多快), 比 MFU(分母自己选,换个分母就换个结论)。 —— 剩下唯一站得住的,是图上第 ③ 档。
⚠️ 图上那个 ④(每美元/每瓦)才是采购真正要的答案, 但它买不到捷径:④ 的分子就是 ③ —— 没有 ③ 的实测, ④ 只是把两个不确定的数相除。
854.0 ÷ 2,700 = 31.6%
(2026-08-05 写进仓库,挂了 26 天)
854.0 ÷ 2,500 = 34.16%
(2026-08-31 更正,现在的数)
上面那张入场券问的全是分母该拿哪个峰值。但分子也有两种取法, 而且这两种都有名字、都在被人用 —— 不知道对方报的是哪一种,两个数根本不能放一起。
区别只有一句话:分子是「模型数学上需要多少次运算」,还是 「硬件实际执行了多少次运算」。
图上那个 4/3 是上界(整层重算)。
本课这个算子落在哪儿 —— 当场算:FlashAttention 反向要重算一次
前向的 S,整层 7 次矩阵乘里有 1 次白做(前向 2 + 反向 4 + 重算 1):
7 ÷ 6 ≈ 1.167×,有用的占 85.7%。
1.167 和 1.33 都是从定义推出来的,不是测出来的 ——
而它们差了一倍,正说明「开了重算」这四个字本身说明不了任何事。
312.5 × 70% = 218.8。
有用强度才刚到 219,含重算的强度就已经顶在 312 上了。
你以为还在带宽那一侧留着一大截余量,其实已经越线。—— 记住这一条,6.3 末尾那张「这个倍数还缺什么」图的第 ③ 条才读得懂: 它说的正是两边的重算策略根本不是一回事。
入场券拿了,摆数。看这张图有个诀窍:先别看那两个大数,先看下面那两段红色。
最难的一步在后面:不把这个数说过头。
这个 1.31 倍还缺什么?—— 先自己列一遍,再看下面这张对答案。
⚠️ 先对一下口径:上图的 1.29 是每芯片
TFLOP/s(854.0 ÷ 662.3),下文一律用每芯片 token/s 的
1.31(6,242 ÷ 4,770)—— 差的就是图里那 1.5%。
走完全程,最后回答开课时那个没答的问题 —— 而答案只有一句,剩下的都是它的限定。
从封装拆到实测,两块硬件走完了一整趟。开场说「处处都不同,但只有一个成因」, 3.6 把全程摊开数过,六行里四行两边一样 —— 现在可以说那个成因到底是什么了。
讲完了。它没有告诉你该买哪个 —— 它想让你下次看到一份 「A 比 B 快 N 倍」的材料时,知道该问哪四句话,以及翻到第几页去找它的分母和分子。 —— 6.2 那两个都叫「利用率」的数,分母一样,差的全在分子。