第 10 讲 · 在 GPU 上加速 Transformer（一）：算子融合与推理加速 · Accelerating Transformer on GPU (I): Kernel Fusion and LightSeq Inference

对应 Maps to：CMU 11-868 / 11-968 Large Language Model Systems（Spring 2026，Lei Li），Lecture 10 · Mon 2/16 · Accelerating Transformer on GPU Part 1。 阅读材料：Wang et al. (2021) LightSeq: A High Performance Inference Library for Transformers（NAACL）；Wang et al. (2022) LightSeq2: Accelerated Training for Transformer-based Models on GPUs（SC）；原课 HW4 说明（融合的 Softmax / LayerNorm CUDA 内核）。 依据说明：原课第 10、11 讲共用一份 slides（llmsys-10-transformer-acc.pdf，66 页）。本讲覆盖其中的：Transformer 作为通用架构、训练的代价与阶段、加速库的对比、LightSeq / LightSeq2 优化总览、技术 1：算子融合（发射开销、TriMatAdd、MHA 与 FFN 的结构、非 GEMM 运算、融合的嵌入前向 / 反向及代码、交叉熵的梯度与融合、层批处理的交叉注意力）、技术 2 的前向部分（LayerNorm 前向的归约改写、softmax 前向的两次归约与按形状选参数），以及推理部分（束搜索与代码、HARS 层级搜索及其例子、GPU 推理的实现细节、LightSeq 的软件架构与 API、GPU 占用、推理加速结果、总结与其他方法）。反向的归约改写、混合精度训练器、显存复用和训练实验留给第 11 讲。模拟器里的未融合 / 融合内核、纸面 GPU 的时间估算、C 语言的 CPU 实测、Welford 版 LayerNorm、原子加的不确定性实验、交叉熵的流量与中间张量核算、HARS 的正确性论证与大词表实验、束缓存重排与跨层复用的显存核算、float4 的访存计数，都是本讲义自己设计的，不代表课堂上讲了什么。 配套代码 Code：code/lecture10/（5 个脚本 + fused_kernels.py + csrc/fusion.c；都在 20 秒以内） 预计阅读时间：3–4 小时（含运行代码）

0. 本讲导航 · How to read this lecture

0.1 学习目标 Learning goals

学完这一讲，你应该能够：

1. 说出 Transformer 训练和推理里时间花在哪里：矩阵乘法（交给 cuBLAS）之外，大量非 GEMM 运算的 FLOPs 微不足道，却因为访存受限和内核发射开销占了一半以上的时间。
2. 用”内核数 × 发射开销 + 字节 ÷ 带宽”估算一串运算的时间，判断融合能省多少、省的是哪一笔。
3. 写出 LightSeq 的几个融合内核的结构：bias + dropout + 残差、注意力 softmax（带掩码）、LayerNorm、嵌入层前向 / 反向、带标签平滑的交叉熵（梯度 q − p）。
4. 解释”减少同步”的代数改写（方差 = E[x²] − E[x]²）省了什么、有什么数值风险，以及 Welford 怎样两全。
5. 说出推理专用的优化：束搜索的层级 top-k（HARS）为什么不会漏、束的 KV 缓存重排、跨层复用显存、FP16 与向量化访存。

0.2 先修与本讲范围 Prerequisites & scope

0.3 本讲地图 Map

§1  背景：Transformer 是通用架构、训练很贵；训练的阶段；加速库；LightSeq 的优化总览
§2  技术 1：算子融合——发射开销、TriMatAdd、bias + dropout + 残差                 ← 代码 01
§3  一层 Transformer 里非 GEMM 运算的比重；LightSeq 融合了哪些                       ← 代码 02
§4  融合的 softmax 与 LayerNorm（HW4）；归约的改写                                  ← 代码 03
§5  融合的嵌入层与交叉熵；层批处理的交叉注意力                                       ← 代码 04
§6  推理加速：束搜索与 HARS、缓存重排、显存复用、FP16 与向量化；LightSeq 的软件与结果  ← 代码 05
§7 小结   §8 常见误区   §9 练习   §10 延伸阅读   §11 下一讲预告

1. 背景 · Background

1.1 Transformer 是通用架构，训练很贵

slides 第 2 页：同一个 Transformer 结构，用在语言模型（BERT、T5、GPT）、视觉（ViT、Swin）、语音（wav2vec、HuBERT）、图文（CLIP）、文生图（Stable Diffusion）上。所以”让 Transformer 在 GPU 上跑得快”是一项回报极高的工作。slides 第 3 页引用的估算：训练 GPT-3 的碳排放相当于一辆车开 146 年。

1.2 训练的阶段与加速库

slides 第 5 页把一次训练迭代分成：数据分片 → 前向 → 反向 → 梯度同步（多卡）→ 参数更新。本讲和第 11 讲讨论单卡（模型能放进一块 GPU）上前向、反向和更新的加速；多卡同步是第 14–18 讲。

slides 第 7 页对比了几个加速库：TensorRT-LLM（原 FasterTransformer）只做推理；DeepSpeed 2022 年 10 月才加入 Transformer 内核；LightSeq（Wang et al. 2021，NAACL）和 LightSeq2（2022，SC）同时覆盖训练和推理，支持 PyTorch 和 TensorFlow。

1.3 LightSeq / LightSeq2 的优化总览

slides 第 8 页列了四类技术：

1. 算子融合（本讲 §2–§5）：把矩阵乘法之间的小运算合成少数几个内核；
2. 依赖感知的归约改写（本讲 §4 讲前向，第 11 讲讲反向）：用代数变换减少块内同步；
3. 混合精度训练器（第 11 讲）：参数和梯度放在连续的显存里，一个内核更新全部参数；
4. 显存管理（第 11 讲；推理部分见本讲 §6.3）：按生命周期复用显存。

2. 技术 1：算子融合 · Kernel fusion

2.1 两笔账：发射开销与中间结果

slides 第 9 页：每发射一个 CUDA 内核要 3–5 μs，约 4000 个时钟周期。01 号脚本第 1 部分对一个逐元素加法 C = A + B（float32）算两笔时间：读写 12n 字节 ÷ 1.555 TB/s，以及 4 μs 的发射开销：

两者在约 52 万个元素处相等。**小张量（例如解码时批大小 1 的隐藏状态，只有几千个数）的时间几乎全是发射开销；大张量（训练时几千个 token × 上千维）的时间由读写的字节决定。**两笔账都指向同一个办法：把多个小内核合成一个——少发射几次，中间结果也不用写回显存再读出来。

2.2 TriMatAdd

slides 第 10 页的例子：C = MatAdd(A, B)、E = MatAdd(C, D) 是两个内核、4 次读 2 次写；写一个 TriMatAdd(A, B, D) 是一个内核、3 次读 1 次写。01 号脚本第 2 部分在模拟器里数出来正好是”6 份数组流量 → 4 份、2 个内核 → 1 个”，两种写法结果逐位相同。中间结果 C 在寄存器里就被用掉了。代价是灵活性：TriMatAdd 只能算”三个数相加”，换一种组合就要另写一个内核——所以融合内核要么针对固定的模型结构手写（LightSeq、FasterTransformer），要么交给编译器自动生成（第 12–13 讲）。

2.3 bias + dropout + 残差

Transformer 里最常见的一串小运算是”线性层之后加偏置、做 dropout、再加残差”（slides 第 12 页）。01 号脚本第 3 部分：未融合 3 个内核要把两个中间张量写回显存，流量是融合版的 1.8 倍；按纸面 A100 放大到真实规模，训练形状（8 × 512 token、d = 1024）融合快 1.9 倍，解码形状（1 个 token、d = 4096）快 3.0 倍——前者靠少搬字节，后者几乎全靠少发射两个内核。

⏱ 第 4 部分用 C 在本机 CPU 上实测同样的计算：三遍循环读写 8 个数组流，一遍循环读写 4 个，融合版稳定地快约 2 倍，而两种写法的有效带宽一样（约 110 GB/s）——这类运算的时间就是”搬了多少字节”。（第一版 C 代码用 i % d 取偏置，取模比加法贵得多，融合版只快 1.3 倍：访存受限的前提是计算足够便宜。）

3. 一层里有多少非 GEMM 运算 · How much time goes to non-GEMM ops

slides 第 11–12 页把一个编码器层画成一串运算：Q/K/V 投影、QKᵀ、softmax、乘 V、输出投影、前馈两层是矩阵乘法（cuBLAS GEMM），它们之间夹着加偏置、拆头 / 拼头的 reshape、缩放、掩码、softmax、dropout、残差、LayerNorm、ReLU。02 号脚本按”PyTorch 不做融合时一个运算一个内核”列出全部 24 个运算（8 条序列 × 256 个 token、d = 1024、16 个头、fp16），用纸面 A100 的 roofline 估算：

• GEMM 占 FLOPs 的 99.73%，却只占时间的 39.9%；
• 非 GEMM 运算的算术强度都在 0.2–2 FLOP/字节，远低于 A100 约 200 的平衡点，完全是访存受限的；
• 最大的几个非 GEMM 运算（缩放、掩码、softmax、注意力 dropout）都作用在 B × h × L × L 的分数矩阵上——这也是第 21 讲 FlashAttention 干脆不把它写回显存的原因。

按 LightSeq 的融合方式（slides 第 12 页右栏：加偏置 + 拆头、缩放 + 掩码 + softmax、加偏置 + dropout + 残差、加偏置 + ReLU + dropout、融合的 LayerNorm）重算：

两种形状下整层前向都快约 1.44 倍（GEMM 部分不变）。LightSeq2 论文报告的完整训练加速是 1.4–3.5 倍（slides 第 31–32 页，还包括反向、优化器、显存管理），其中 LayerNorm 快 4 倍、softmax 2.5–3.4 倍、dropout 1.1–2.5 倍、参数更新 2.3 倍（slides 第 34 页）；第 33 页的时间线显示一次训练迭代从 Fairseq 的 457 ms 降到 214 ms。GPU 占用图（slides 第 60 页）里，非 GEMM 内核所占的比例明显下降。

4. 融合的 softmax 与 LayerNorm · Fused softmax and LayerNorm (HW4)

4.1 注意力 softmax

注意力的 softmax 作用在 B × h × L × L 的分数上，每一行（一个查询对所有键）独立做：缩放、加掩码、求最大、减最大取 exp、求和、相除。未融合是 5 个内核、3 个和分数矩阵一样大的中间张量。HW4 要写的 ker_attn_softmax 是融合版本：一行一个线程块，每个线程把自己负责的几个元素读进寄存器，在寄存器里做缩放和掩码，块内归约求最大，算 exp，块内归约求和，归一化后写回。03 号脚本第 1 部分（一个 64 × 64 的分数矩阵）：

结果都与 NumPy 一致到 float32 的舍入。融合后分数矩阵只读 1 次、写 1 次；有因果掩码时，被掩掉的那一半根本不用读。放大到批 2 × 16 头 × 256 × 256（fp16）的纸面估算快 4.6–5.2 倍（HW4 要求融合 softmax 平均快约 6.5 倍，同一量级）。

注意栅栏没有减少：求最大、求和两次块内归约照做，融合版还多一个”大家读完最大值再复用共享内存”的栅栏。slides 第 23 页说的”两次线程同步（两次归约）代价高”，融合本身解决不了。

块大小要按形状挑（slides 第 24–26 页）：块越大，每个线程负责的元素越少，但树形归约的轮数 log₂B 越多；块比行还长时线程空闲（03 号脚本第 2 部分：L = 256 时块大小从 16 到 256，每行栅栏从 11 增加到 19）。行很短（L ≤ 32）时一行交给一个线程束，用 __shfl_xor_sync 在寄存器之间归约，根本不需要块级栅栏——HW4 里现成的 ker_attn_softmax_lt32 就是这样。LightSeq 用 C++ 模板把块大小、每线程元素数做成编译期参数，按 L 的范围分派到不同的内核实例。

4.2 LayerNorm 与归约的改写

LayerNorm 前向要两次归约：先求均值 μ，再求方差 σ² = mean((x − μ)²)，每次归约一轮同步。slides 第 21 页（技术 2：减少同步）把方差改写成

$\sigma^2 = \mu(x^2) - \mu(x)^2$

于是 Σx 和 Σx² 可以在同一次归约里一起累加，两次线程同步变成一次。03 号脚本第 3 部分（8 行 × d = 64，块大小 32）：

代价是数值：均值远大于标准差时，E[x²] 和 E[x]² 都很大、差却很小，float32 里发生灾难性抵消（P2 §4）。输入每行 ~ N(均值, 1)：

均值 10000 时方差被算成 ≤ 0，结果完全错了。LightSeq 的做法是计算用 FP32、存储用 FP16（slides 第 21 页），在常见的激活范围里误差可以接受。两全的办法是 Welford / Chan 的并行算法：每个线程维护（个数, 均值, 平方偏差和），归约时两两合并（n = n_a + n_b，δ = μ_b − μ_a，μ = μ_a + δ·n_b/n，M2 = M2_a + M2_b + δ²·n_a·n_b/n）——栅栏和一段归约一样少，误差和两段归约一样小。很多生产级 LayerNorm 内核用的就是它。

5. 融合的嵌入层与交叉熵 · Fused embedding and criterion

5.1 嵌入层

slides 第 13 页：嵌入层前向 y = Dropout(s · E[w] + P[p])（s 是嵌入的缩放系数，第 6 讲 §4.4）。未融合要 5 个内核：查词表、查位置表、缩放、相加、dropout。04 号脚本第 1 部分：融合成 1 个内核后流量少了 2.8 倍（四个中间张量都不需要了），结果逐位相同。slides 第 15 页是 LightSeq 的 CUDA 代码：每个线程生成 dropout 掩码、查词表、查位置表、缩放、相加、写出，一气呵成。

反向（slides 第 14 页）：∇E = ReduceSum(s · ∇y ⊙ Mask)——同一个词在一批里出现几次，它那一行的梯度就要加几次（第 5 讲：scatter-add）。未融合是 dropout 反向、缩放、累加 3 个内核；融合成 1 个内核，用原子加（atomicAdd）直接加到嵌入表的梯度上。04 号脚本第 2 部分：流量减半以上，结果与 NumPy 的 np.add.at 一致。

原子加保证不丢任何一份，但不保证顺序；浮点加法不满足结合律，所以结果可能每次运行都差在最后几位。脚本里出现次数最多的词有 7 份梯度，按 200 种顺序相加得到 5 个不同的 float32 结果。需要逐位可复现时，要先按词排序再按固定顺序归约（torch.use_deterministic_algorithms 会换成这类实现，更慢）。LightSeq 的反向代码（slides 第 16 页）里还顺带做了梯度裁剪和梯度累加。

5.2 交叉熵：梯度就是 q − p

slides 第 17 页：带标签平滑的目标 p = (1 − α)·y + α/V（V 是词表大小；注意这和第 6 讲 slides 的 ε/(V−1) 约定不同，和 PyTorch 的 label_smoothing 相同），q = softmax(x)，损失 L = −Σᵢ pᵢ log qᵢ。求导：

$\frac{\partial L}{\partial x_j} = -\sum_i p_i(\delta_{ij} - q_j) = q_j\sum_i p_i - p_j = q_j - p_j$

（因为 Σ p = 1）。04 号脚本第 3 部分用 PyTorch autograd 验证：差 5.6e-17（float64）。梯度只是一个逐元素的减法，所以前向和反向可以融合进一个内核（slides 第 18 页）：一行一个块，先归约最大值 m，再一起归约 Σexp(x − m) 和 Σx，损失 = log Z + m − [(1 − α)x_y + (α/V)Σx]，梯度 q − p 直接写出。

词表越大，这一项越重要：第 8 讲 06 号脚本算过，128K 词表、几万个 token 的 logits 比模型权重还大。

5.3 层批处理的交叉注意力

slides 第 19 页（推理时的编码器–解码器模型）：解码器的每一层都有交叉注意力，要把编码器的输出分别投影成这一层的 K、V。编码器输出在整个解码过程中不变，而各层的投影矩阵也是固定的——所以可以把 L 层的投影矩阵拼成一个 [W_1, …, W_L]，一次大 GEMM 算出所有层的 K、V，再按层切开。一个大矩阵乘法比 L 个小的更能喂饱 GPU，也只发射一次。

6. 推理加速 · Fast inference

6.1 束搜索与 HARS

第 9 讲的束搜索每一步要做两件事（slides 第 43–45 页）：对每个束算词表上的 log softmax，再从 k × V 个累计得分里选出前 k 个。slides 第 44 页的 TensorFlow 代码先把 log 概率加到序列得分上、reshape 成 [batch, k·V]、调用 top_k，最后 refresh_cache 按选出的束重排解码器的 KV 缓存。词表几万时，“从 k·V 个数里选前 k 个”就是一次大排序，GPU 上需要多轮全局同步。

**HARS（层级自回归搜索，slides 第 45–54 页）**把它拆成两步：

1. 检索：每个束的 V 个分数切成 k 组，求每组的最大值 mᵢ；R = min mᵢ 叫”粗略的第 k 大”；把 ≥ R 的候选挑出来写进显存。
2. 重排：只在这些候选里排序，取前 k。

为什么不会漏？每个束有 k 个组最大值都 ≥ R，所以这个束的第 k 大 ≥ R，它的前 k 个都会被挑出来；全局前 k 一定在各束的前 k 里。05 号脚本第 1 部分复现了 slides 的例子（束宽 2、词表 8：两个束的 R 是 4 和 7，挑出 4、4、5 和 7、8 共 5 个候选，排出前 2 名 8 和 7），并在随机 logits 上试了 600 次：

粗筛只用”分组求最大""求最小""逐元素比较”这些完全并行的操作（slides 第 63 页：把排序变成并行的 max、filter、re-rank），之后只要在几十个候选上排序。⏱ 本机 CPU 上 HARS 0.28 ms、全排序 8.4 ms、argpartition（O(n) 的部分排序）0.82 ms——CPU 上 O(n) 的选择算法已经够快，HARS 的价值在 GPU 上避免全局排序。

6.2 束的 KV 缓存重排

每一步选出新的 k 个束后，新束的历史来自旧束（两个新束可能来自同一个旧束），KV 缓存要按”父束”重排（slides 的 refresh_cache）。05 号脚本第 2 部分：朴素地整份拷贝，7B 级模型（32 层、d = 4096）束宽 4、生成 100 个 token 要拷贝 9.86 GiB——和注意力本身读缓存的量一样多，等于把解码的显存流量翻倍，还随长度平方增长。第 24 讲的 PagedAttention 让多个束共享相同前缀的缓存块，只在分叉时写时复制。

6.3 跨层复用显存、FP16、向量化

slides 第 55 页列了 GPU 推理实现的几个细节：

• 跨层共享张量显存：推理没有反向，第 i 层的输出被第 i + 1 层读完就没用了，所有层可以轮流使用同几块缓冲区。05 号脚本第 3 部分：24 层、d = 1024、批 8 × 256，训练要为反向保留约 1.88 GiB 激活，推理只要约 0.02 GiB 的临时缓冲区。省下的显存给更大的批或更长的 KV 缓存。
• 不保留中间结果和梯度：相当于 torch.no_grad()。
• 混合精度，主要用 FP16 计算（slides 第 27 页）：同样的元素个数只有一半的字节；FP16 的张量核算力是 FP32 的数倍。
• 用 float4、half2 提高带宽利用率：05 号脚本第 4 部分在模拟器里拷贝 4096 个 float32，float 版本 4096 个线程、128 条读指令，float4 版本 1024 个线程、32 条读指令，读事务都是 512 个——字节一样，指令只有 1/4。

融合（少读写）、FP16（每个数更小）、向量化（每条指令搬更多）攻的是同一个瓶颈：显存带宽。

6.4 LightSeq 的软件架构与结果

slides 第 56 页：最上层是 Python API 和模型库（Transformer、BERT、GPT-2、ViT），中间是 C++ 算子（编码器、解码器、嵌入、softmax、训练器），最底层是 LightSeq2 的 CUDA 内核加上 cuBLAS、CUB 等第三方库。使用时可以直接替换 HuggingFace 的层（slides 第 57 页：LSTransformerEncoderLayer 替换 BERT 的第一层），或者在 Fairseq 里用 lightseq-train 和带 ls_ 前缀的模型、优化器、损失（slides 第 58–59 页），并兼容 Apex 和 DeepSpeed。

slides 报告的推理加速：机器翻译最高 14 倍、GPT-2 最高 6 倍（与 PyTorch / TensorFlow 相比，批越大优势越明显；slides 第 61–62 页）。slides 第 64 页还列了其他方向：更高效的模型结构（Linformer、Reformer）、训练策略（由浅到深、层 dropout、LAMB）、量化（第 19–20 讲）和硬件优化。

7. 小结 · Summary

1. Transformer 的 FLOPs 几乎全在矩阵乘法里，但时间有一半以上花在非 GEMM 运算上：它们访存受限、而且每个都要发射一次内核。
2. 融合同时省两笔：内核发射（小张量、解码时是大头）和中间结果的显存读写（大张量、训练时是大头）。TriMatAdd 6 份 → 4 份流量；bias + dropout + 残差 1.8 倍流量；一层前向融合后快约 1.44 倍（纸面）。
3. 融合 softmax：一行一个块、元素留在寄存器里，读写从 8.75 倍降到 2 倍；块大小按行长挑，短行用线程束级归约。融合不减少同步。
4. 归约改写：σ² = E[x²] − E[x]² 把 LayerNorm 的两次同步变一次，但在均值很大时有灾难性抵消；Welford 两全。
5. 嵌入层前向 5 → 1、反向 3 → 1 个内核（原子加带来浮点不确定性）；交叉熵梯度就是 q − p，前向 + 反向一个内核，大词表时省掉好几个 GB 的中间张量。
6. 推理：HARS 用”分组最大 → 粗阈值 → 筛选 → 小排序”代替全排序且不会漏；束缓存重排要搬和读缓存一样多的字节；推理可跨层复用显存；FP16 和向量化进一步省带宽。

8. 常见误区 · Common pitfalls

• “矩阵乘法占了 99% 的 FLOPs，所以只要优化矩阵乘法。” 时间不按 FLOPs 分配；非 GEMM 运算在训练形状下仍占一半以上的时间。
• “融合只是为了省内核发射。” 大张量时省的主要是中间结果的显存读写；小张量时才主要是发射。
• “融合以后同步也少了。” 融合 softmax 的块内归约一次都没少；减少同步要靠代数改写。
• “E[x²] − E[x]² 和两段归约在数值上等价。” 数学上等价，float32 里在均值大时差几个数量级。
• “原子加的结果是确定的。” 不丢份，但顺序不定；浮点加法顺序不同，结果就可能不同。
• “交叉熵的反向要单独算 softmax、log、乘法……” 梯度化简后就是 q − p，可以和前向合成一个内核。
• “HARS 是近似的 top-k。” 粗阈值 R 保证每个束的前 k 都在候选里，结果和全排序完全相同。
• “束搜索的代价只是 k 倍的前向。” 还有每一步按父束重排 KV 缓存的显存流量。
• “float4 能减少要搬的字节。” 字节不变，减少的是访存指令和线程数。
• “一个写好的融合内核在所有形状上都快。” 块大小、每线程元素数要按形状挑，LightSeq 用模板生成多个实例。

9. 练习 · Exercises

1. 用 TriMatAdd 计算 E = A + B + D，数组各有 100 万个 float32。未融合（两个 MatAdd）和融合各读写多少 MB？在 1.555 TB/s、每个内核 4 μs 的纸面 A100 上各要多久？
2. 一个逐元素运算读 a 个数组、写 1 个数组（float32），发射开销 4 μs、带宽 B。写出”发射开销 = 读写时间”时的元素个数 n*。a = 2、B = 1.555 TB/s 时 n* 是多少？
3. 03 号脚本里，未融合的 softmax 把分数矩阵读写了 8.75 倍。按 5 个内核逐个列出每个内核读写了哪些张量，解释这个数。
4. 两个数 1000 和 1001：用 σ² = E[x²] − E[x]² 计算方差，中间量各是多少？如果 E[x²] 只能精确到 0.0625（float32 在 10⁶ 附近的间距），相对误差可能有多大？
5. 证明带标签平滑的交叉熵对 logits 的梯度是 q − p；如果用第 6 讲 slides 的约定（正确类 1 − ε、其余 ε/(V−1)），梯度还是 q − p 吗？
6. 嵌入层反向为什么要用原子加？如果改成”每个词一个线程，按顺序把它出现的所有位置加起来”，有什么好处和坏处？
7. 用 HARS 处理一个束、k = 2、分数 (5, 1, 4, 2)：组最大值、R、候选和最终结果各是什么？
8. 7B 级模型每个 token 的 KV 缓存 512 KiB（fp16）。束宽 4、生成到第 50 步时，朴素的缓存重排这一步要拷贝多少？
9. 拷贝 100 万个 float32，每条线程束指令处理 32 个线程：用 float 和 float4 各需要多少条读指令？读事务（32 字节）各多少个？
10. 解释层批处理的交叉注意力（slides 第 19 页）为什么能加速；它为什么只适用于交叉注意力、不适用于解码器的自注意力？
11. 计算 fp16 的”加偏置”运算的算术强度（FLOP / 字节）。在 A100 上它是计算受限还是访存受限？
12. 行长 L = 64、块大小 B = 32 的融合 softmax，每行要几个栅栏？（按 03 号脚本的实现：两次树形归约各 log₂B 轮，外加归约前和复用前的栅栏。）

参考答案

1. 每个数组 4 MB。未融合：4 读 + 2 写 = 24 MB，两个内核：2 × 4 μs + 24 MB ÷ 1.555 TB/s ≈ 8 + 15.4 = 23.4 μs。融合：3 读 + 1 写 = 16 MB，4 + 10.3 = 14.3 μs。
2. 4 μs = 4(a + 1)n* ÷ B，n* = B × 4 μs ÷ (4(a + 1))。a = 2：1.555e12 × 4e-6 ÷ 12 ≈ 52 万个（01 号脚本：518,333）。
3. 缩放加掩码：读 X、写 T1（2 份）；求最大：读 T1（1 份），写每行一个数；减最大取 exp：读 T1、写 E（2 份，另读每行的最大值）；求和：读 E（1 份）；相除：读 E、写 Y（2 份，另读每行的和）。合计 8 份分数矩阵。另外，每行的最大值和和被每个元素各读一次（一个线程束的 32 个线程读同一个数，一个 32 字节事务，约合 0.25 份 × 2），每行写一次（一个事务分摊到 64 个元素，约 0.125 份 × 2），共约 0.75 份，合计 8.75。
4. E[x] = 1000.5，E[x]² = 1001000.25；E[x²] = (1000000 + 1002001)/2 = 1001000.5；方差 = 0.25。若 E[x²] 误差 ±0.0625（半个间距约 0.03），方差的相对误差可达 12%–25%——两个相近的大数相减，有效数字几乎全丢了。
5. 推导见 §5.2，只用到 Σᵢ pᵢ = 1。第 6 讲 slides 的约定下 p 的各项之和也是 (1 − ε) + (V − 1)·ε/(V − 1) = 1，所以梯度仍然是 q − p，只是 p 的数值不同（第 6 讲 05 号脚本：两种约定的梯度差 0.010）。
6. 多个线程可能同时往嵌入表的同一行加（同一个词出现多次），普通的读–改–写会互相覆盖、丢掉一部分梯度，原子加保证每一份都加上。“每个词一个线程顺序累加”是确定性的、没有原子操作的竞争，但要先把位置按词分组（排序），而且出现次数很不均匀时（高频词）有的线程很忙、有的很闲，负载不均衡。
7. 两组 (5, 1)、(4, 2)，组最大值 5、4，R = 4；候选 5、4（≥ 4）；最终前 2 名 5、4（与全排序相同）。
8. 4 个束 × 50 个 token × 512 KiB = 100 MiB（这一步）。
9. float：100 万 ÷ 32 = 31,250 条读指令；float4：25 万个线程 ÷ 32 = 7,813 条（约 1/4）。读事务都是 4 MB ÷ 32 字节 = 125,000 个。
10. 交叉注意力的 K、V 来自编码器的输出，它在整个解码过程中不变，L 层的投影矩阵也固定，所以可以拼成一个大矩阵、一次 GEMM 算出所有层的 K、V（只算一次，后面每一步直接用）。解码器自注意力的 K、V 来自这一层的输入，而这一层的输入要等前一层算完才知道，层与层之间有依赖，无法提前一起算。
11. 每个元素 1 次加法，读 2 字节、写 2 字节（偏置很小、在缓存里）：1 / 4 = 0.25 FLOP/字节，远低于 A100 约 200 的平衡点，访存受限。
12. 归约前 1 个 + 求最大 log₂32 = 5 个 + 复用前 1 个 + 归约前 1 个 + 求和 5 个 = 13 个（03 号脚本第 2 部分：L = 256、B = 32 时每行 13 个，与 L 无关）。

10. 延伸阅读 · Further reading

• Wang et al. (2021) LightSeq: A High Performance Inference Library for Transformers（NAACL Industry Track）。
• Wang et al. (2022) LightSeq2: Accelerated Training for Transformer-based Models on GPUs（SC）。
• NVIDIA FasterTransformer / TensorRT-LLM 的开源代码：融合内核与推理优化的工业实现。
• Ivanov et al. (2021) Data Movement Is All You Need: A Case Study on Optimizing Transformers（MLSys）：系统地量化 Transformer 训练中数据搬运的比重，并做算子融合。
• Welford (1962) 与 Chan, Golub & LeVeque (1979)：一遍计算方差与并行合并的算法。
• 原课 HW4 说明：ker_attn_softmax_lt32 / ker_attn_softmax 与 LayerNorm 前向 / 反向内核的结构。
• Dao et al. (2022) FlashAttention：把融合推到极致——整个注意力一个内核，分数矩阵不写回显存（第 21 讲）。

11. 下一讲预告 · Next lecture

这一讲主要看了前向和推理。训练还有一半：反向传播要把每个融合内核的梯度也写出来，而且 LayerNorm、softmax 的反向里同样有多次归约；优化器要更新几百个参数张量，逐个张量发射内核又是一笔发射开销；训练时的激活和梯度显存怎样规划。第 11 讲（LightSeq2）讲 slides 的后半部分：依赖感知的计算图改写（LayerNorm 反向的 α、β 系数把两次同步变一次、softmax 反向）、混合精度训练器（把所有参数和梯度放进一整块连续显存，一个内核完成全部参数的 FP32 更新）、按生命周期复用显存（slides 第 30 页自注意力反向的显存规划），以及 LightSeq2 在机器翻译、GPT-2、BERT、ViT 上的训练加速。对应原课 HW4 的后半：Softmax 和 LayerNorm 的反向内核。