Lecture 27: Deep Neural Networks and ML Accelerators
目录 · ← l25 · appendix →
Lecture 27: Deep Neural Networks and ML Accelerators
1. 章节标题与概述
Lecture 27: Deep Neural Networks and ML Accelerators
- 本讲核心问题：现代深度神经网络（DNN）到底在算什么，为什么它既是并行计算最”友好”的负载，又是对硬件能效最”苛刻”的负载？本讲要回答两个层面的问题：(1) 软件层面——如何把一个卷积层、一行 softmax、一个 attention 块，变换成算术强度（arithmetic intensity）足够高的 GEMM（通用矩阵乘）与融合算子，使得程序从”带宽受限”（bandwidth bound）移动到”计算受限”（compute bound）；(2) 硬件层面——既然通用处理器（general-purpose processor）执行一条指令的固定开销（取指、译码、发射、寄存器读写）远远超过一次算术运算本身，那么专业化硬件（specialization：DSP / FPGA / 域专用加速器 / ASIC）能把每瓦性能提高多少，代价又是什么？两个层面的交汇点是一句话：GEMM 很便宜，搬数据很贵（GEMM computation is cheap, but data movement is expensive——它消耗硅面积、瓦特和纳秒）。
- 涉及的主要硬件/软件机制：
- 硬件侧：Tensor Core（张量核，单条指令做 4×8 × 8×4 甚至 16×16×16 的矩阵乘加，A/B 用 fp16、累加器用 fp32）、TMA（Tensor Memory Accelerator，异步张量搬运单元）、TMEM/共享内存层次的片上存储、SM（Streaming Multiprocessor）的 warp 调度结构、TPU 的 systolic array（脉动阵列，weight-stationary 数据流）、FPGA 的 LUT（查找表）与可重构互连、ASIC/dataflow 架构（PMU/PCU + 片上网络）、以及低精度数值格式（fp16 / bf16 / fp8 / fp4）。
- 软件侧：implicit GEMM 与 explicit GEMM（是否物化 im2col 矩阵）、多层 blocking / tiling（寄存器-共享内存-L2-DRAM 的层次化分块）、寄存器分块（register blocking）与 SIMD 向量化、”算子融合”（operator fusion）、online/chunked softmax 的代数重构（FlashAttention）、Halide 的”算法 / 调度”分离（algorithm-schedule separation）与自动调度搜索、以及 Tile 级编程系统（CUTLASS、Triton、Thunderkittens、CuTe）。
- 贯穿全讲的分析工具：Roofline 模型、算术强度、work-span 模型、Amdahl 定律、数据移动的能耗量级（pJ 级）。
- 在并行计算知识体系中的角色：这是全课程”工具集”的一次总集成。前面学过的 SIMD 与多核分块（第 3–7 讲）、CUDA 执行与存储层次（第 5 讲）、性能优化与 Roofline（第 7–8 讲）、通信与带宽（第 14 讲）在这里全部被同一个负载——矩阵乘——同时调用；而本讲新引入的第四条设计轴是”能效”（energy efficiency）：
Power = (Ops/second) × (Joules/Op)，当性能提升到一定程度后，约束不再是”快不快”，而是”每焦耳能做多少运算”。它也解释了为什么今天最贵的计算机是”为矩阵乘而生的”：AI 的算力需求指数增长，而数据中心在能耗与散热上被死死约束。 - 配套材料：
- Fall 2026 课程主页与日程表：https://www.cs.cmu.edu/~418/、https://www.cs.cmu.edu/~418/schedule.html —— 已公开（无需登录）。日程表 Oct 27 一行为 “Deep neural networks”，Oct 29 一行为 “ML accelerators at Amazon（guest lecture by Randy Huang and Ron Diamant）”。
- 本讲的 Fall 2026 讲义 PDF：在 https://www.cs.cmu.edu/~418/lectures/ 之下尚未发布（该目录当前公开的 PDF 覆盖 Why Parallelism、ILP、Basic Architecture、Programming Models、CUDA、Parallel Programming Basics、Performance Optimization I/II、Cache Coherence、Directory Coherence、Snoop Implementation、Virtual Memory、Interconnects、Consistency、Synchronization、Lock-Free、Heterogeneity、Specialization，以及
24-parallel_deep_learning_data_parallel.pdf、25-parallel_deep_learning_model_pipeline_parallel.pdf，不含本讲的 DNN / 加速器讲义）。 - 日程表中引用的讲义：Oct 27 一行的 slides 链接指向
https://www.cs.cmu.edu/afs/cs/academic/class/15418-f22/public/lectures/25_dnn.pdf，位于/afs/cs/academic/class/15418-f22/public/之下，访问需要 CMU 登录 —— 属未公开。该行同时给出一个标注为 “lecture 24 video” 的 YouTube 链接（历史学期录像）；Fall 2026 日程表中多数 Panopto 录像链接被 HTML 注释隐藏，属未发布。 - Oct 29 客座讲（ML accelerators at Amazon）：日程表该行没有任何 slides 链接，讲者也不提供公开讲义 —— 属未公开。本笔记的硬件加速器部分因此完全依据下述公开材料与公开成熟知识撰写。
- 需登录的资源：Ed 讨论区、Autolab、Canvas/Gradescope 均需 CMU 账号，属需登录未公开。
- 本笔记的事实基础（本地已读取的公开讲义抽取文本）：
cs149_supp/dnninference.txt（75 页，页眉标注 “Stanford CS149, Fall 2025 — Lecture 9: Efficiently Evaluating DNNs”）、cs149_supp/accelerators.txt（71 页，”Lecture 10: Hardware Specialization”）、cs149_supp/aiperfoptimization.txt（55 页，”Lecture 13: Domain-Specific Programming Systems and Automatic Performance Optimization”）。这三份是 Stanford CS149 的公开课程讲义逐页抽取文本，不是 CMU 15-418 的讲义，但覆盖了与本讲完全相同的主题（DNN 计算特征、GEMM/融合优化、硬件专业化、DSL 与自动调优），因此本笔记的全部术语、公式、示例与结论均以这三份文本为准；其中未涉及的部分（如具体芯片型号的公开规格、能耗 ballpark 数据的物理含义）依据公开成熟知识补充，并已尽量与讲义中的说法对齐。 - Fall 2026 授课教师：Brian Railing 与 Dimitrios Skarlatos；课程由 Kayvon Fatahalian 创建。讲义 PDF 首页可能写有历史学期字样（如 Fall 2025），属正常的讲义沿用。
2. 核心概念与硬件/软件架构图解
2.1 DNN 的计算本质：它只是一个电路，而且几乎全是矩阵乘
- 定义与目的：一个神经网络就是一个有向无环的计算图（circuit）。基本单元是”神经元（unit / neuron）”：给定 n 个输入 $x_0..x_{n-1}$ 与 n+1 个参数（n 个权重 $w_i$ 加 1 个偏置 $b$），计算 $f(\sum_i x_i w_i + b)$，其中 $f$ 是非线性激活函数（例如 ReLU：$f(x)=\max(0,x)$，或 sigmoid：$f(x)=1/(1+e^{-x})$）。把许多这样的单元按层连起来就是 DNN。本讲的第一个目的，是让”神经网络”褪去神秘感：把它当成电路看，它的性能问题就完全变成数据复用问题。
- 直观解释（”它是什么？”）：把 DNN 想象成一条巨大的流水线工厂，每一道工序都是一台”称重机”：每个神经元手上有 n 个砝码（权重），它把送进来的 n 种原料（激活值）按砝码加权求和，再决定是否放行（ReLU 相当于”负值一律当 0 处理”）。全连接层（fully connected layer）意味着”每台称重机都从前面所有工序取料”；卷积层（convolutional layer）意味着”每台称重机只看前面工序的一小块区域，而且同一层里所有位置共用同一套砝码“。这个”权重共享”（weight sharing）是卷积层节省参数量的关键，也是它能变成矩阵乘的关键。
- 架构/机制图解（软件模型：从神经元到层）：
一个神经元 (unit) 全连接层 = 矩阵-向量乘
x0 ──w0──┐ [w00 w01 w02] [x0] [b0] [y0]
x1 ──w1──┼──►( Σ )──► f(·) ──► y [w10 w11 w12] [x1] + [b1] = [y1]
x2 ──w2──┤ + b [w20 w21 w22] [x2] [b2] [y2]
... │ [w30 w31 w32] [b3] [y3]
xn ──wn──┘ \_____ A (M×K) ____/\_x_/ \_b_/ \_y_/
n+1 个参数 M×K 矩阵 × K 向量 = M 向量
(M 个输出单元, 每个 K 个权重)
卷积层 = "局部连接 + 权重共享" 深度网络 = 层叠的 DAG
输入平面 (W×H) 输入 ──►[Conv]──►[ReLU]──►[Pool]──►
┌───────────────────────┐ ──►[Conv]──►[ReLU]──►...
│ ┌───┐ │ 3×3 的窗口滑过整幅图, ──►[FC]──►[Softmax]──► 输出
│ │w w w│ │ 9 个权重被所有输出位置复用
│ │w w w│ ──► 1 个输出 │ (对比全连接: 每个位置独立权重)
│ │w w w│ │
│ └───┘ │ AlexNet / Inception-ResNet / MobileNet
└───────────────────────┘ 拓扑与规模差异极大 → "极端效率挑战"
(Extreme efficiency challenge)
- 关键操作与性能特征：这三类层最终都归约为同一个 kernel——稠密矩阵乘（dense GEMM）：
- 全连接层 = 矩阵 × 矩阵（把 batch 里的样本堆成矩阵的列）。
- 卷积层 = 隐式的矩阵乘（implicit GEMM，见 2.3）。
- Transformer 的 attention 块 = 两次矩阵乘（$S=QK^T$ 与 $O=PV$）。 因此讲义把 GEMM 称为”现代 AI 的 kernel”（The kernel for: fully-connected layers, convolutional layers, the attention block of a transformer）。性能特征由算术强度决定：如果每次算术运算需要从片上存储之外搬入大量数据，程序必然带宽受限；而矩阵乘的算术强度随分块尺寸线性增长（$OI \propto n$，OI = operational intensity / 算术强度），这正是它能吃掉海量算力的原因。
2.2 直接卷积实现：七重循环与”巨大但未被利用的数据复用”
- 定义与目的：最直接的卷积实现是一个七重循环嵌套（batch × 输出行 × 输出列 × 输出通道 × 输入通道 × 卷积核 Y × 卷积核 X）。它的目的是暴露数据复用机会：同一个滤波器权重在一次卷积内被反复使用（沿输出空间滑动），同一个输入激活值被不同滤波器反复使用（沿输出通道遍历）。
- 直观解释（”它是什么？”）：把它想象成在 64 张照片上用 64 副不同的”眼镜”看。direct 实现就是”对每个输出像素、每副眼镜，都把 3×3 邻域重新摸一遍”。数据其实就在手边（同一个输入像素会被 64 副眼镜各看一次），但如果循环顺序排错了，这 64 次都会变成从 DRAM 重新取一次。所以直接写的卷积通常是”算术上正确、带宽上灾难”。
- 架构/机制图解（硬件侧：卷积的 halo/tile 数据复用）：
输入特征图 (W+2)×(H+2)×C 的某个 channel 一个输出 tile 的输入"晕轮"(halo)
┌───────────────────────────────────┐
│ ┌─────────────────────────────┐ │ tile 尺寸 T×T (输出)
│ │ │ │ 需要的输入 (T+2)×(T+2)
│ │ T×T 输出 tile │ │ ── 重复计算的边界 = halo
│ │ (每点被 K 个滤波器复用)│ │
│ │ │ │ 3×3 卷积的 halo 开销:
│ └─────────────────────────────┘ │ 额外输入 = ((T+2)²−T²)/T² ≈ 4/T
│ ▲ 每个输出点需要 3×3 邻域 │ T=32 → 12.5% ; T=64 → 6.3%
└───────────────────────────────────┘
复用维度: (a) 权重沿输出空间复用 (b) 输入沿输出通道复用
(c) 权重/输入沿 batch 复用 (identical weights for all images)
- 关键操作与性能特征：direct 卷积每次输出要读 $C \times R \times S$ 个输入元素与权重，做 $C \times R \times S$ 次乘加，算术强度约为 2 FLOP / 每次 load（若完全无缓存复用）。把它变成计算受限，必须人为构造复用：循环分块（blocking）——把输出 tile 的输入（含 halo）留在寄存器/共享内存里，然后让 K 个滤波器反复访问它。性能特征上，halo 开销随 tile 增大而下降（$O(1/T)$），但 tile 越大寄存器压力越大，超过寄存器容量会 spill，反而更慢。
2.3 explicit GEMM（im2col）与 implicit GEMM：物化与否的代价
- 定义与目的：既然卷积和 GEMM 是同一件事，最省事的做法就是把卷积变成一次真正的矩阵乘：把输入特征图的每个 $R\times S\times C$ 邻域”拍平”成一个行向量，拼成”卷积矩阵”（im2col / explicit GEMM）。这样可以直接调用高度优化的 GEMM 库。implicit GEMM 则相反：不物化这个矩阵，而是在 GEMM 的分块内部按需索引激活张量，只把当前需要的 sub-block 搬到片上共享内存，再用 well-tuned 的 shared-memory GEMM 做子块乘法（讲义点名 CUTLASS）。
- 直观解释（”它是什么？”）：im2col 像把一叠照片的每个 3×3 小窗口都复印一份、装订成一本巨大的相册——好处是”查相册”非常简单（标准矩阵乘），坏处是相册比原照片厚 9 倍，而且复印本身要占仓库（DRAM）。implicit GEMM 是”不印相册，现场翻照片”，用手指按坐标去找——代码复杂，但省掉了整本相册的存储与搬运。
- 架构/机制图解（软件模型：im2col 展开）：
输入 3×3 窗口 + 零填充 卷积矩阵 (im2col) K 列 = R×S×C
┌───────────────┐ ┌──────────────────────────────┐ [w00 w01 ... w0N]
│ 0 0 0 0 │ 展开 │ 0 0 0 | 0 x00 x01 | 0 x10 x11 │ [w10 w11 ... w1N]
│ 0 x00 x01 x02 │ ────────► │ 0 0 0 | x00 x01 x02 | x10 ... │ × [ ... ]
│ 0 x10 x11 x12 │ (每行是一个 │ 0 0 0 | x01 x02 x03 | x11 ... │ [w80 w81 ... w8N]
│ 0 x20 x21 x22 │ 输出位置的 │ x00 x01 x02 | x10 x11 x12 |.. │ \____ K×N ______/
└───────────────┘ 感受野) └──────────────────────────────┘
(W×H) 行 num_filters = N 列
代价: 存储开销 O(滤波器元素数 N) 级别重复 → 实际是 R×S 倍膨胀
DRAM 流量增加 R×S 倍 (讲义: "Increases DRAM traffic by a factor of R x S")
implicit GEMM: 不建这张大矩阵, 只在片上共享内存里物化一个 sub-block
→ 不需要额外 off-chip 存储, 也不增加 DRAM 流量
- 关键操作与性能特征：explicit GEMM 的计算是最优的（可以吃到 GEMM 库的全部优化），但它的数据移动被放大了 $R\times S$ 倍（3×3 卷积就是 9 倍），且需要一块可能与激活张量同量级的临时存储。implicit GEMM 把这次放大消掉，代价是索引逻辑复杂、边界处理（padding）要写进 kernel 内部。CUTLASS 正是为这种”自定义 DNN 层的构建基元”而生：in-shared-memory 的快速 GEMM、warp 级 GEMM、用于快速块加载/张量索引的 iterator、tensor reduction 等。
2.4 GEMM 的层次化分块：把算术强度从 O(1) 抬到 O(n)
- 定义与目的：blocking（分块 / tiling）是把矩阵切成能装进各级缓存的小块，在块内完成复用，从而提高算术强度。单层 blocking 只能利用一级缓存；层次化 blocking 让每一级存储层次（寄存器 → L1 → L2 → DRAM）都做一次分块，逐级提高”每字节 DRAM 流量对应的浮点运算数”。
- 直观解释（”它是什么？”）：把 GEMM 想成做菜。DRAM 是超市（一次采购很贵、很慢），L2 是冰箱，L1 是案板，寄存器是手里的刀。笨做法是”每切一刀就跑一趟超市”；分块就是”一次买回一整块案板分量的原料（L2 块），在案板上再分成小份（L1 块），手里同时拿 4 行 × 16 列的原料（寄存器块）反复下刀”。块越大，去超市的次数越少——所以讲义的自检问题是：”你想要尽可能大的 BLOCKSIZE 吗？为什么？”答案是：在能被该级存储装下的前提下越大越好，但超过容量就会失效（thrash）。
- 架构/机制图解（软件模型：三级分块 + 寄存器微内核）：
DRAM (慢/大/贵) L2 (中) L1 / 共享内存 寄存器
┌──────────────┐ ┌───────────────┐ ┌──────────────┐ ┌──────────┐
│ A B │ │ A(kb) B(kb) │ │ A(jb,kb) │ │ a0..a3 │
│ ┌───┐ ┌───┐ │ ──► │ ┌───┐ ┌───┐ │ ──► │ ┌──────────┐ │ ──► │ × │
│ │jb │ │kb │ │ │ │ │ │ │ │ │ │ jb × kb │ │ │ b0..b15 │
│ └───┘ └───┘ │ │ └───┘ └───┘ │ │ └──────────┘ │ │ = c(4×16)│
│ ------- C │ │ B(kb,ib) │ │ B(kb,ib) │ └──────────┘
└──────────────┘ └───────────────┘ └──────────────┘
C 的一个块 每级都用"三次读、n³ 次算"把 OI 放大
OI(块大小 b, 三级块都在缓存里) = 2b³ FLOP / 3b² 元素 = 2b/3 FLOP/元素 = b/6 FLOP/Byte (fp32)
→ b=8(寄存器): 1.3 ; b=64(L1): 10.7 ; b=256(L2): 42.7 [FLOP/Byte]
三种把 SIMD 塞进分块的方式 (讲义给出三种):
(1) 向量化 i 循环: C_accum = vec_load(C); 循环 k 做 splat(A)·vec_load(B) → 好: 改善 B 的空间局部性
坏: 工作集 ×SIMD_WIDTH, B 仍跨大步长
(2) i 维太小时: 预转置 B 的一个块到 BT, 向量化最内层 k (点积)
(3) 双方都预转置 (A→AT, C→CT), 用 SIMD_WIDTH × SIMD_WIDTH 的累加器数组, 最内层 kk 无依赖
- 关键操作与性能特征：单层 blocking 后，DRAM 流量降为 $O(n^3/b)$ 而计算仍是 $O(n^3)$，算术强度 $OI = 2b/3$ FLOP/元素。层次化 blocking 让这个 $b$ 逐级变大（寄存器级只关心微内核的复用，L1/L2 级关心更大的块），从而让 DRAM 侧的有效算术强度逼近甚至超过 Roofline 的 ridge point。代价：代码复杂度爆炸——讲义专门提醒”不同层可能需要不同的调度策略（矩阵维度不同）”，例如 MobileNet 的 body 里既有
3×3×3×32、112×112×32这样的”通道少、空间大”的层，也有1×1×1024×1024、7×7×1024这样的”通道多、空间小”的层，二者对分块尺寸、向量化维度、转置策略的要求正好相反（讲义的原话是：”Ug for library implementers!”）。
2.5 Roofline 与算术强度：DNN 优化的全部”软件侧”答案
- 定义与目的：Roofline 模型用一条折线回答”这个程序最快能多快”：$P = \min(\pi_{peak},\; I \times \beta)$，其中 $\pi_{peak}$ 是机器的峰值算力（ops/sec），$\beta$ 是可用的通信/内存带宽（bytes/sec），$I$ 是程序的算术强度（ops/byte）。交点 $I^* = \pi_{peak}/\beta$ 称为 ridge point（屋脊点）：$I < I^$ 时程序带宽受限，$I > I^$ 时计算受限。
- 直观解释（”它是什么？”）：把机器想成一条生产线（算力）配一条进料传送带（带宽）。传送带每分钟能送 20 公斤原料，生产线每分钟能加工 768 份。如果 1 公斤原料能做 1 份产品（$I=1$），生产线永远吃不饱（带宽受限）；如果 1 公斤原料能做 100 份产品（$I=100$），传送带轻松喂饱，瓶颈在生产线上（计算受限）。提高算力（换更快的机器）只会让程序更容易掉进带宽受限区；提高算术强度（改程序：分块、融合）才能让它回到计算受限区——这是讲义用整整一页强调的结论。
- 架构/机制图解（软件模型：Roofline 曲线与”更强算力把屋脊点推向右”）：
吞吐量 (GFLOP/s, log) ← 计算受限区 (Compute bound): P = π_peak
^ 768 ┤─────────────────────────────●──────────────── π_peak = 16核×3.0GHz×8宽×2(FMA)
| │ ..--´ 768 GFLOP/s
| │ ..--´ (升级算力 → 水平线抬高、
| │ ..--´ 屋脊点右移 → 更容易带宽受限)
| │ ..--´ ← 带宽受限区 (BW bound): P = I × β
| │ ..--´ 斜率 = β (DRAM 带宽)
| │..--´
+──────┴──────────────────────────────────────────────► 算术强度 I (FLOP/Byte, log)
1 2 4 8 16 38.4 100 200 400
▲
ridge point I* = π/β = 768/20 = 38.4 FLOP/B
· 点 A (I=1, 带宽受限): P = 1 × 20 GB/s = 20 GFLOP/s (只用到 2.6% 算力)
· 点 B (I=64, 计算受限): P = 768 GFLOP/s (满载)
· ΔR (融合前→后): 把 I 提高 3/5 ÷ 1/3 = 1.8 倍 → 吞吐也涨 1.8 倍 (仍在带宽受限区)
- 关键操作与性能特征：Roofline 给出了 DNN 优化的完整”操作清单”：
- 提高 $I$：循环融合（loop fusion）、分块（blocking）、算子融合（fusion）——把中间结果留在片上。讲义的经典例子：
E = D + (A+B)*C分成三个循环写，算术强度是 $1/3$；融合成一个循环后是 $3/5$（每次遍历从 2 load + 1 store / 1 op，变成 4 load + 1 store / 3 op）。 - 抬高 $\beta$：用更高带宽的存储（HBM、更宽的接口）、更好的访问模式（合并访问、避免 cache line 浪费）。
- 降低 $\pi_{peak}$ 的需求：用更低精度（fp16/bf16/fp8）在同等硅面积下做更多运算——或者更激进，用专用硬件（第 2.8–2.12 节）。
- 性能特征的关键提醒：重叠通信与计算需要额外的片上缓冲（double buffering），而片上存储的面积本来可以用于计算——讲义把它列为”其余知识”的第一条：数据移动消耗能量、片上存储与计算争抢芯片资源，所以缓冲区要尽可能小。
- 提高 $I$：循环融合（loop fusion）、分块（blocking）、算子融合（fusion）——把中间结果留在片上。讲义的经典例子：
2.6 算子融合（fusion）与 FlashAttention：不改数学，只改数据流
- 定义与目的：融合是把相邻算子的循环体合并，使中间结果在寄存器/共享内存中被消费掉，而不是写回 DRAM 再读回。讲义给出三个由浅入深的例子：(1) Conv + Scale/Bias + MaxPool 融合；(2) 逐行 softmax 的融合；(3) Transformer attention 的 online（chunked）softmax 融合，即 FlashAttention 的核心思想。
- 直观解释（”它是什么？”）：想象搬家。不融合 = 每搬一样东西都要装箱、运到仓库、再取回来；融合 = 东西一手交一手，全程不进仓库。softmax 的例子特别直白：naive 实现要读矩阵 5MN+2M 个元素、写 3MN+2M 个元素（因为 max、exp、sum、除法各走一遍）；融合后”读一行 → 在片上算完这一行 → 写一行”，只需读 MN、写 MN。attention 的例子更极端：naive 实现要物化 $N\times N$ 的分数矩阵 $S=QK^T$（N 可以是几千，$N^2$ 空间就爆了），而 chunked 版本永远不物化这个矩阵。
- 架构/机制图解（软件模型：FlashAttention 的分块循环嵌套）：
Q (N×d) K^T (d×N) naive: S = QK^T ──► P = softmax(S) ──► O = PV
┌──────┐ ┌──────┐ N×N 矩阵被物化 2 次 (写+读) → 内存/带宽爆炸
│ Qi │ │ KTj │ chunked (FlashAttention):
│ │ │ │ for each j (key tile):
└──────┘ └──────┘ for each i (query tile):
┌──────┐ ┌──────┐ Load Qi, KTj, Vj, Oi (Oi 常驻片上)
│ Vj │ │ Oi │ Sij = Qi · KTj
└──────┘ └──────┘ mij = m(Sij) ; Pij = f(Sij) ; lij = l(Sij)
j 循环 (key tile) Oi += Pij·Vj (按上页公式重标定)
i 循环 (query tile) end
end
online softmax 的结合律 (关键数学):
把行向量 x 切成 x(1), x(2):
m(x) = max(m(x(1)), m(x(2))) ← max 可结合
f(x) = [ e^{m(x1)-m(x)} f(x(1)) , e^{m(x2)-m(x)} f(x(2)) ]
l(x) = e^{m(x1)-m(x)} l(x(1)) + e^{m(x2)-m(x)} l(x(2)) ← 归一化和可结合
⇒ "先算局部, 再用重标定因子 alpha = e^{m_old - m_new} 合并" 与"全量算完再除"完全等价
代价: 额外计算 (每次 i 循环都要重标定已有的 O 累加器), 但省下 N² 的存储与带宽
收益: 读 3 个块 (Q,K,V)、做两次矩阵乘 + 少量行求和、累加进 cache 常驻的 O 块
- 关键操作与性能特征：融合把”多次 DRAM 往返”压成”一次读 + 一次写”，代价是额外计算与片上缓冲。讲义明确点出这个 trade-off：”Note there is additional computation vs. the original version (must re-scale prior values of O each step of the i-loop)”。数值上重标定的额外乘法是每行每 tile 一次 $d$ 维缩放（约 $N/B_J \times d$ 次），相对主计算 $2Nd$ 的比例约 $1/(2B_J)$——当 key tile $B_J=64$ 时不到 1%，非常划算。另外，融合的历史演化是：早期由库作者手写几个固定的 fused op（如 TensorFlow），后来由编译器自动生成融合实现（cuDNN backend、
torch.compile），目标是”在一个节点里高效执行而不通过内存传递中间结果”。
2.7 低精度数值格式：用”更小的字节”换”更高的算术强度”
- 定义与目的：DNN 对精度的容忍度远高于科学计算。把权重与激活从 fp32 降到 fp16 / bf16（16 bit）、fp8、甚至 4 bit 或 1 bit，可以在不改变数据布局的前提下把同一份数据的字节数减半甚至更多，从而直接提高算术强度、降低带宽需求与能耗，并让同样面积的乘法器做更多运算。
- 直观解释（”它是什么？”）：把 fp32 想成”用 32 位写一个数字”，bf16 是”用 16 位写，但保留和 fp32 一样的指数范围（只是尾数更短）”，fp8 是”只有 8 位的草稿纸”。就像快递箱：fp32 是一个大箱子装一件小商品（浪费空间），bf16 是刚好合身的箱子，fp8 是信封——前提是商品（数值）不需要那么精细的保护。
- 架构/机制图解（数值格式的位分配）：
fp32 S│EEEEEEEE│MMMMMMMMMMMMMMMMMMMMMMM 1 + 8 + 23 位
-1^S × (1 + M×2^-23) × 2^(E-127)
bf16 S│EEEEEEEE│MMMMMMM 1 + 8 + 7 位
← 指数与 fp32 相同 ⇒ 动态范围相同; 尾数只有 7 位 ⇒ 精度下降
fp8 E4M3 S│EEEE│MMM 1+4+3 范围 0 ~ 448
fp8 E5M2 S│EEEEE│MM 1+5+2 范围 0 ~ 57344
← E5M2 范围更大、精度更差; E4M3 反之。训练中常"前向用 E4M3、反向用 E5M2"
Tensor Core 的混合精度约定: A, B 存为 fp16 → 乘法用 fp16 → 累加用 fp32
⇒ 用"低精度存储 + 高精度累加"避免误差累积
- 关键操作与性能特征：
表 1：低精度带来的收益与代价
格式 每元素字节 相对 fp32 的算术强度 主要风险 典型用途 fp32 4 1× 无 训练的主权重/累加、科学计算 tf32 4 1× 尾数 10 位 Tensor Core 上的 fp32 训练折中 fp16 2 2× 范围小（易溢出，需 loss scaling） 推理与训练激活 bf16 2 2× 精度低 7 位尾数 训练（范围与 fp32 同） fp8 (E4M3/E5M2) 1 4× 精度/范围都紧 大模型训练/推理（Transformer Engine） int8 1 4× 需要量化标定（scale/zero-point） 推理 fp4 / 1-bit 0.5 / 0.125 8× / 32× 需要专用算法与格式 研究前沿（讲义原话：1-bit 是 “In the extreme case: 1-bit ;-)”） 关键约束：累加必须在更高精度（Tensor Core 就是 A,B = fp16、D = fp32），否则长归约链的舍入误差会吞掉结果；低精度同时降低了带宽与能耗（搬 2 字节比搬 4 字节便宜），这是它比”仅提高峰值算力”更有价值的地方。
2.8 硬件侧的问题：为什么”通用处理器”如此低效？
- 定义与目的：本讲从软件转向硬件时，第一步是承认一个反直觉的事实：整个课程都在教如何高效使用多核 CPU/GPU，但这两个平台相对于专用硬件仍然是”低效”的。原因是通用处理器为”通用性”付出了巨大代价：执行每一条指令都要经过取指、译码、检查依赖/流水线冒险、选择执行单元、读寄存器堆、把数据搬到执行单元、运算、再搬回寄存器堆、写回 SRAM 等一长串固定开销；这些开销在能耗上甚至超过运算本身。
- 直观解释（”它是什么？”）：通用 CPU 像一个什么菜都会做的大厨：他每做一道菜都要先”读菜谱（取指）、理解菜谱（译码）、去仓库取料（load）、洗手（流水线控制）”。如果一天只做三道大菜（复杂、数据密集），这些准备成本可以忽略；但如果要做一百万份小菜（SIMD 位运算、视频编码里的整数运算），大厨 90% 的时间花在翻菜谱上。专业硬件（ASIC）就像一台只会切片的机器：它不会别的，但每秒切一万片，而且几乎不耗电在”理解要做什么”上。
- 架构/机制图解（硬件侧：一条指令的固定开销 vs 一次运算）：
执行一条指令的最简路径 (讲义逐条列出):
┌────────────┐ ┌────────────┐ ┌──────────────┐ ┌────────────┐
│ 取指 │→│ 译码 │→│ 依赖/冒险检查 │→│ 选执行资源 │
│ (i-cache, │ │ (→uops, │ │ (记分牌/ │ │ (哪个 ALU?) │
│ 地址翻译) │ │ uop cache)│ │ 调度窗口) │ │ │
└────────────┘ └────────────┘ └──────────────┘ └────────────┘
↓ ↓
┌────────────┐ ┌────────────┐ ┌──────────────┐ ┌────────────┐
│ 读寄存器堆 │→│ 数据搬到 │→│ 执行运算 │→│ 结果搬回 │
│ SRAM │ │ 执行单元 │ │ (这才是"工作")│ │ 寄存器堆 │
└────────────┘ └────────────┘ └──────────────┘ └────────────┘
H.264 视频编码的能耗分解 (即使已用 SIMD 指令实现):
┌───────────────────────────────────────────────────────────┐
│ int motion est. │ frac. motion est. │ intra pred/DCT │ arith │
│ FU ▓▓ │ FU ▓▓ │ FU ▓▓ │ FU ▓▓ │ ← 功能单元 FU
│ Ctrl ▓▓▓▓▓▓▓▓ │ Ctrl ▓▓▓▓▓▓▓ │ Ctrl ▓▓▓▓▓▓▓ │ ... │ ← 控制开销占大头
│ RF ▓▓▓▓▓▓ │ RF ▓▓▓▓▓ │ RF ▓▓▓▓▓ │ ... │ ← 寄存器堆
│ IF/D-$ ▓▓▓▓▓▓▓ │ ... │ ... │ ... │ ← 取指/数据缓存
└───────────────────────────────────────────────────────────┘
FU = functional units, RF = register fetch, Ctrl = 流水线控制,
Pip = 流水线寄存器, IF = 取指+i-cache, D-$ = 数据缓存
⇒ 功能单元（真正干活的部分）只占能耗的一小部分
- 关键操作与性能特征：两条由此推出的设计原则（贯穿后续所有硬件加速器）：
- 摊薄指令流开销（amortize cost of instruction stream processing）：用”更复杂的指令”让一次取指/译码服务更多运算。讲义给出量化对比（相对简单指令的可编程性开销）：半精度 FMA 为 2000%、半精度 DP4（vec4 点积）为 500%、半精度 4×4 MMA（矩阵乘加）为 27%——也就是说，用一条 MMA 指令代替 16 次 FMA，指令流开销被摊到几乎可以忽略。
- DSP 与 VLIW 的思路：Qualcomm Hexagon DSP 是”可编程但控制路径更简单、指令更复杂（SIMD/VLIW）”的处理器，单条 VLIW 指令同时指定多个不同操作（与 SIMD 的”同一操作多个数据”不同），在其 FFT 最内层循环中可以每周期执行 29 个”RISC 操作”。
2.9 GPU 的应对：Tensor Core、TMA 与”算力全在张量核里”
- 定义与目的：GPU 的做法是把”通用性开销”摊到极致：用大量 SIMD 通道共享控制，再把矩阵乘加做成专用指令（Tensor Core MMA）。A100 的一个 SM 有 64 个 fp32 ALU、32 个 int32 ALU、4 个 tensor core；每个 tensor core 执行 8×4 × 4×8 的矩阵乘加（$A \times B + D$，A/B 为 fp16，D 为 fp32）；GA100 有 108 个 SM，即 6912 个 fp32 ALU 与 432 个 tensor core，1.4 GHz 下 fp32 为 19.5 TFLOPS，而 fp16/32 混合的 tensor core 峰值是 312 TFLOPS——相差 16 倍。
- 直观解释（”它是什么？”）：Tensor Core 像一台”打包好的”乘法车间：普通 SIMD 是”一组工人各拧一颗螺丝”，MMA 是”一台机器一次压出一整块 8×8 的成品”。代价是这台机器只吃特定形状的原料（16×16×16 的 tile、特定数据类型），所以程序员必须把数据摆放成它喜欢的布局——这就是 tile 级编程模型的由来。
- 架构/机制图解（硬件侧：H100 的 SM 与存储层次）：
┌──────────────────────────── 一个 H100 SM (144 个 SM / GPU) ─────────────────────────────┐
│ Warp Scheduler 0..3 (每周期发射 1 个 warp, 最多 64 warps/SM) │
│ ┌──────────┐ ┌──────────┐ ┌──────────┐ ┌──────────┐ │
│ │ Fetch/ │ │ Fetch/ │ │ Fetch/ │ │ Fetch/ │ ← 控制被 16/32 条 lane 共享 │
│ │ Decode 0 │ │ Decode 1 │ │ Decode 2 │ │ Decode 3 │ (每周期 1 条 32 宽 SIMD 指令) │
│ └────┬─────┘ └────┬─────┘ └────┬─────┘ └────┬─────┘ │
│ ┌────▼────────────▼────────────▼────────────▼─────────────────────────────────────┐ │
│ │ SIMD fp32 FU (32 MUL-ADD/clock) │ SIMD int FU (16/clk) │ SIMD fp64 FU (16/clk) │ │
│ ├──────────────────────────────────┼──────────────────────┼───────────────────────┤ │
│ │ Load/Store 单元 │ Tensor Core ×4 │ 16×16×16 │ │
│ │ │ (systolic MMA) │ [fp16×fp16→fp32] │ │
│ └──────────────────────────────────┴──────────────────────┴───────────────────────┘ │
│ 寄存器: 64 KB / sub-core，256 KB / SM，分给 64 个 warp (每线程约 1 KB) │
│ 共享内存 / L1: 256 KB / SM │
└────────────────────────────────────────────────────────────────────────────────────────┘
▲ ▲
│ TMA (Tensor Memory Accelerator): 单线程发起的异步拷贝 │
│ 用 copy descriptor 描述区域 → global ↔ shared memory │
│ cuda::memcpy_async, 完成时通过 mbarrier 通知 │
H100 全芯片: 144 SM, Tensor Core 989 TFLOPS (fp16), SIMD 134 TFLOPS (fp16) / 67 TFLOPS (fp32)
历代 GPU 中 Tensor Core 贡献的 FLOPs 占比: 89% → 50% → 94% → 96% → 98% (逐代上升)
CUDA/计算/存储三级层次 (讲义给出的对应关系):
Grid ↕ GPU ↕ 80 GB HBM3 / 50 MB L2
Cluster ↕ CPC ↕ 256 KB shared / SM (Cluster = 最多 16 个线程块, 保证同 SM 同时运行)
Thread Block↕ SM ↕ 256 KB shared
Threads ↕ SIMD lanes ↕ 1 KB RF/thread
- 关键操作与性能特征：
- 异步性是让这些单元吃满的关键：异步计算（
mma_async）、异步内存访问（TMA + TMEM）、异步片间通信（overlap compute / memory / communication）。讲义的”理想 AI 加速器”特性表逐条对照 NVIDIA GPU 是否满足：tiled tensors ✅（16×16/32×32 tile 换最高 GEMM TFLOPS 与低指令开销）、asynchronous compute ✅、asynchronous memory access ✅（TMA+TMEM）、asynchronous chip-to-chip ❓、compute-unit-to-compute-unit comm ❓（靠 TB Cluster 部分满足）。 - 新一代的编程难度：B100 的 tensor core 数据放在 SMEM 与 TMEM 中，由单个线程执行 MMA（不再需要 warp！），编程接口变成
tcgen05.alloc/tcgen05.mma/tcgen05.commit/tcgen05.fence+cp.async.bulk.tensor+ mbarrier 协调——讲义的原话是 “Not your father’s CUDA”。这正是 DSL（Mosaic GPU、CuTe-DSL（Python 里的 CUTLASS））存在的理由。 - 填满机器需要”很多并行工作”：讲义用卷积的输出规模说明这一点——$N=1, P=Q=64$ 时只有 524K 个输出（2 MB），而 $N=32, P=Q=256$ 时有 256M 个输出（1 GB）。小 batch 的推理根本喂不饱 80 个 SM。
- 异步性是让这些单元吃满的关键：异步计算（
2.10 TPU 与 Systolic Array：数据驱动而非指令驱动
- 定义与目的：Google 的 TPU（Jouppi et al. 2017）是”域专用加速器”的样板：芯片面积中算术单元约占 30%，控制部分面积极小——因为它的指令集只有寥寥几条：读主机内存、写主机内存、读权重、matrix_multiply/convolve、activate。核心是一个 systolic array（脉动阵列）：权重从左侧 FIFO 逐列灌入并驻留在 PE（processing element）中（weight-stationary），激活值从上方沿对角线以”波前（wavefront）”方式推进，每个 PE 做一次乘加并把部分和传给下方邻居。
- 直观解释（”它是什么？”）：把 systolic array 想象成一排传菜的服务员（这也是 “systolic” 一词的来源——像心脏搏动一样有节奏地泵数据）。数据不像冯·诺依曼机器那样”存到仓库、用时再取”，而是像水流一样从邻居手里接过来、用一下、再传给下一个邻居。每份数据一旦进入阵列，就会被沿途所有 PE 使用；只有阵列的边界需要与外界通信。因此通信是局部的（只与邻居 PE），控制是分布式的（每个 PE 只需知道自己的节奏）。
- 架构/机制图解（硬件侧：矩阵-向量乘 y = Wx 的波前推进）：
权重从左侧 FIFO 逐行灌入 (weight-stationary, 停在 PE 里):
Weights FIFO: w00 w01 w02 w03 | w10 w11 w12 w13 | w20 ... | w30 ...
────────► 按行从左向右推进
x0 x1 x2 x3 ← 激活值从上方逐拍注入 (斜对角)
│ │ │ │
┌──────▼─┐ ┌───▼────┐ ┌───▼────┐ ┌───▼────┐
w0x │ PE │→│ PE │→│ PE │→│ PE │ 每一拍: PE 做 1 次乘加,
────►│ w00·x0 │ │ w01·x1 │ │ w02·x2 │ │ w03·x3 │ 并把部分和向"下"传
└───┬────┘ └───┬────┘ └───┬────┘ └───┬────┘ (竖直箭头 = 部分和累加)
w1x ┌───▼────┐ ┌───▼────┐ ┌───▼────┐ ┌───▼────┐
────►│ w10·x0 │ │ w11·x1 │ │ w12·x2 │ │ w13·x3 │
└───┬────┘ └───┬────┘ └───┬────┘ └───┬────┘
w2x ┌───▼────┐ ... 第 4 拍后, 每列底部的 32 位累加器
────►│ w20·x0 │ ... 得到 y0..y3 的一个完整点积
└────────┘
时序 (matrix-vector, y=Wx): tap 0: x0 进; tap 1: x1 进, PE0 出 x0·w00
tap 2: x2 进, PE0 出 x0·w00+x1·w01, PE1 出 x0·w10
tap 3: x3 进, PE0 得到 y0 = Σ xk·w0k, PE1 得 x0w10+x1w11 ...
矩阵-矩阵 (Y=WX): 沿斜对角灌入 X 的多列 → 需要多个 4×32bit 的累加器列来保存输出列
扩展到更大的矩阵: 例如 A=8×8, B=8×4096, C=8×4096 需要 4096 个累加器
(讲义用"把 B 的列块切给阵列、每拍换一块"的方式放大规模)
表 2：SIMD 与 Systolic Array 的对比（讲义原表）
特性 SIMD Systolic Array 数据流（Dataflow） 控制驱动（control-driven，靠指令） 数据驱动（data-driven，靠波前 wavefront） 局部性（数据复用） 有限（Limited） 时间局部性 + 空间局部性（temporal and spatial） 通信 全局（寄存器/内存） 局部（相邻 PE 之间） 控制 集中式（Centralized） 分布式（Distributed） 效率（perf/mm², perf/Watt） 中（Medium） 非常高（Very high） 关键操作与性能特征：systolic array 的代价是灵活性——它只擅长”形状整齐的稠密矩阵乘”，矩阵维度必须被 padding 到阵列尺寸的整数倍，否则 PE 空转；它也无法表达任意控制流。它带来的收益是能效：Jouppi 等人的 TPU 论文给出的性能/瓦特数据（讲义引用其图）显示专用加速器相对通用处理器的巨大优势。讲义还用 “硬件的彩票”（Hardware Lottery，Sara Hooker） 提醒我们：”当一个研究想法胜出，是因为它契合当时可用的软件与硬件，而不是因为它普遍优于其他方向”——TPU 稠密矩阵乘的算术强度随 $n$ 线性增长（$OI \propto n$），于是 Transformer 模型胜出，于是硬件进一步为矩阵乘特化，形成正反馈。
2.11 FPGA、ASIC 与 Dataflow 架构：从”可编程”到”可重构”
- 定义与目的：在通用处理器与全定制 ASIC 之间，有两类中间形态：
- FPGA（Field Programmable Gate Array）：芯片提供逻辑块阵列（可编程 LUT + 触发器）与可编程互连，逻辑由”配置”直接实现。讲义例子：Xilinx Virtex-7 的 LUT6 可理解为一个 64 项查找表（6 输入 1 输出），因此一个 6 输入 AND 只需 1 个 LUT6；40 输入 AND 可用 8 个 LUT6 链接而成，延迟为 3。现代 FPGA 还含大量硬核：SRAM 存储块、DSP 乘法块、ARM/RISC-V CPU 核，用硬件描述语言（Verilog）编程；Amazon EC2 的 F1/F2 实例提供云端 FPGA。
- ASIC / 域专用加速器：为某一领域定制的电路。讲义用两个经典案例量化收益：FFT 上 ASIC 用约 1/1000 的芯片面积达到与一个 CPU 核相同的性能，而 GPU 核相对 CPU 核的面积效率只有约 5–7 倍；ASIC 还能用约 1/100 的功耗达到一个 CPU 核的性能。另一个案例是 DE Shaw 的 Anton 分子动力学超级计算机（2008，机器内含 512 个粒子相互作用 ASIC + 面向 FFT 的吞吐型子系统 + 为 N-body 通信模式定制的低延迟网络），Anton 3（2025）比同期 GPU 快约 20 倍。
- 可重构 Dataflow 架构：如 Plasticine（Prabhakar、Zhang 等，ISCA 2017），由 PCU（Pattern Compute Unit） 与 PMU（Pattern Memory Unit） 通过片上 Switch 网络连接；把 AI 模型的 dataflow graph（GEMM 1 → Pool → GEMM 2 → SoftMax，外加 map/filter/reduce 等并行模式）空间映射到芯片上。它”没有指令，因此没有取指/译码开销”，并且是”极端异步的：不存在顺序指令执行”，从而同时具备异步计算、异步访存、异步片间通信、计算单元间直接通信（融合与流水）四种能力——正好满足”理想 AI 加速器”表格的全部条目。
- 直观解释（”它是什么？”）：把三种硬件想成三种交通方式：CPU 是万能出租车（哪都能去，但每趟都要绕路、计价）；FPGA 是可拆卸拼装的轨道（你可以按需铺一条专用线路，但铺路本身很费工）；ASIC 是一次性建好的高铁（建成后极快极省电，但只为某条线路服务，造价上千万到上亿美元）。dataflow 架构相当于“把程序图直接铺成铁路网”：算子变成车站（PCU），中间结果通过铁轨（Switch）直接开到下一站，不用回总站（DRAM）换乘。
- 架构/机制图解（硬件侧：Plasticine 式 dataflow 与 FlashAttention 的空间映射）：
┌──────────────────────── reconfigurable dataflow fabric ──────────────────────┐
│ ┌──────┐ ┌──────┐ ┌──────┐ ┌──────┐ │
│ │ PMU │ │ PCU │ │ PMU │ │ PCU │ PMU = Pattern Memory│
│ │(存储)│ │(计算)│ │ │ │ │ Unit (片上模式存储) │
│ └──┬───┘ └──┬───┘ └──┬───┘ └──┬───┘ PCU = Pattern Compute│
│ │ ╲ ╱ │ Switch (S) │ ╲ ╱ │ Unit (模式计算单元) │
│ ┌──▼──────▼───▼──┐ ┌──────────▼──────▼──────▼──┐ S = 可重构互连开关 │
│ │ S S │ │ S S S │ │
│ └──────┬─────────┘ └────────┬──────────────────┘ 没有指令流 ⇒ │
│ │ │ 没有取指/译码开销 │
│ 数据在 PCU/PMU 之间直接流动 (compute-unit → compute-unit) │
└───────────────────────────────────────────────────────────────────────────────┘
把 FlashAttention 映射到 16 个 tile 的 dataflow (讲义图):
Tile0..Tile15 ──► [QKT] ──► [Mask] ──► [Softmax] ──► [Dropout] ──► [×V] ──►
空间上摊开, 多级流水 (MetaPipeline): 不同 tile 同时处于不同阶段
⇒ "dataflow kernel fusion": 融合不是循环变换, 而是把算子直接连成硬件流水线
- 关键操作与性能特征：讲义给出的”专业化能效经验法则”（相对于 CPU 上的高质量 C 代码）：
- 吞吐型处理器架构（GPU 核）：约 10× perf/watt（前提：代码能映射为宽数据并行且计算受限）。
- 固定功能 ASIC：可达 100–1000× 或更高（前提：计算受限，且不是浮点数学）。
- 域专用加速器（DSL 可编程，如 Google TPU）：约 20×。
- FPGA / 可重构逻辑：约 50×？（讲义原话 “jury still out”——尚无定论；难编程，让它变简单是活跃的研究方向）。
- 而 ASIC”不可编程，且设计/验证/流片要花上千万到上亿美元”——这就是”效率 vs 可编程性”曲线（一组从”最易编程”到”最难编程”的连续谱）的核心权衡。
2.12 能量账本：为什么”搬数据”比”算数据”贵
- 定义与目的：现代系统设计的第一条经验法则：”永远设法减少计算机中的数据移动量“。原因是数量级差异：一次整数运算约 1 pJ，一次浮点运算约 20 pJ（这还只是”逻辑运算本身”的成本，不含译码、寄存器读写等开销）；而从 1 mm 外的小 SRAM 读 64 bit 约 26 pJ；从低功耗移动 DRAM（LPDDR）读 64 bit 约 1200 pJ。也就是说，从片外 DRAM 读 8 个字节的能量，相当于做 60 次浮点运算。
- 直观解释（”它是什么？”）：把运算想成”拧一颗螺丝”（很便宜），把访存想成”开车去五金店取螺丝”（很贵）。芯片里的 SRAM 是”桌上的零件盒”（26 pJ 一趟），DRAM 是”城里的五金店”（1200 pJ 一趟）。所以架构师的全部努力，就是把”去五金店的次数”降到最少——这就是 tiled tensor（tile 化张量）、异步搬运、计算单元间直接通信这三种机制的统一动机。
表 3：能效 / 可编程性阶梯（讲义 “Efficiency vs Programability” 图）
平台 相对能效 可编程性 代价 能效优化 CPU 1×（基准） 最容易 峰值算力低 吞吐型处理器（GPU） ~10× 需要数据并行 + 计算受限 编程模型复杂（CUDA/tile） 可编程 DSP（如 Hexagon，VLIW） 介于两者之间 较易（复杂指令摊薄控制） 领域受限（信号处理） FPGA / 可重构逻辑 ~50×？（尚无定论） 难（Verilog/HLS） 编译时间长、工具链复杂 域专用加速器（TPU，DSL 编程） ~20× 中等（受限于 DSL 表达力） 表达能力受限 固定功能 ASIC 100–1000×+ 不可编程 NRE 成本数千万~上亿美元；只对计算受限、非浮点负载有效 - 关键操作与性能特征：这些数字与本讲的三大机制一一对应：tiled tensors（16×16、32×32 的 tile 才能拿到 GEMM 的最高 TFLOPS 与最低指令开销）、异步执行（异步计算、异步访存、异步片间通信，三者叠加才能同时重叠 compute / memory / communication）、片上通信（计算单元到计算单元的直接通信，实现融合与流水、streaming dataflow）。TPU 的芯片面积分配（算术单元 ~30%、控制面积极小）和”关键指令只有 5 条”正是这一哲学的极端体现。
2.13 编程系统与自动优化：让”性能专家”这件事变得可复制
- 定义与目的：性能优化需要极高专业度、极其枯燥、且换一台机器就要重做——而公司每年在 AI 算力上花费数千万到数亿美元，所以”自动化性能优化”本身就是巨大的经济动机。讲义给出三条思路：(1) 提高抽象层次（DSL）；(2) 智能搜索（autotuning / autoscheduler）；(3) [新兴] 利用 LLM 的问题求解与代码生成能力。
- 直观解释（”它是什么？”）：Halide 的核心是把程序拆成两份描述：”算法”（算什么，声明式、无循环、无副作用）与”调度”（怎么算：循环顺序、分块、向量化、并行化）。这就像菜谱与厨房排班分开写：菜谱只写”3×3 均值模糊”，排班表才写”256×32 分块、内层向量化 8 宽、外层多线程、blurx 在 tile 内按需计算”。
表 4：Halide 的”算法 / 调度”分离（同一个算法的三种实现）
调度写法 产生的 C 等价结构 中间缓冲 性能特征 blurx.compute_root()先整幅算 blurx（读 1024+2 × 1024+2），再整幅算 out blurx：W × (H+2)，全尺寸 需一次完整 DRAM 往返；朴素 out.tile(x,y,xi,yi,256,32); blurx.compute_at(out, xi)tile 内对每个 xi 只算 3 个 blurx 元素 blurx：1 × 3 复用最省，但重复计算多（每个输出点重算 blurx） out.tile(...).vectorize(xi,8).parallel(y); blurx.compute_at(out, x).vectorize(x,8)tile 级分配 blurx(258,34)，逐 tile 计算并立即消费 blurx：258 × 34，装得进缓存 生产者-消费者局部性最优；8 宽 SIMD + 多线程 - 关键操作与性能特征：
- Halide 的语言约束是为了让编译器能提供服务：只支持规则 N 维域上的计算、只支持前馈流水线（外加 reduction 与固定深度递归的特殊支持）、所有依赖必须可被编译器推断。有了这些约束，编译器才能机械地把”调度”翻译成 pthread + AVX intrinsics（含边界条件的自动生成，例如 256+2 不能被 8 整除时的收尾处理）。
- 早期学术结果：相机 RAW 处理流水线（原 463 行手写 ARM NEON 汇编）用 Halide 写代码量少 2.75×、速度快 5%；bilateral filter（原 122 行 C++）用 34 行算法 + 6 行调度达到 CPU 快 5.9×、GPU 比手写 CUDA 快 2×。
- 自动调度：现实是”会写 Halide 的人多，会写 Halide 调度的人极少”（Google 有 80+ 人写 Halide，但被信任写调度的只有极少数）。于是把调度建模为一串离散选择（对 DAG 从末尾往前，逐个决定
compute_at的位置与 tile 尺寸），用贪心/beam search 搜索，用一个小 MLP 在”几十微秒”内估计代价（166 秒内评估 140 万个调度，实际输出 27 个系数代入手工代价模型）。结果：在图像处理 CPU 上，自动调度器与人类最佳调度相当。 - LLM 路线：用”生成 → 执行/剖析 → 反思 → 修改”的循环（剖析给出 SM 利用率 42%、DRAM 利用率 89%、L2 命中 68%、32 ms 等反馈），在 KernelBench（数百个 PyTorch kernel）这样的基准上自动产出正确且快速的 CUDA kernel；关键技巧是让 LLM 组装高层原语（Triton / CUTLASS-CuTe / TileLang）而不是写底层 CUDA，从而减少正确性错误与幻觉。进一步的想法包括：用经验微调模型、建立”优秀解法数据库”做检索增强（存储的不只是解法，还有优化决策序列）、用 prompt 优化器把优化轨迹总结成原则、以及把穷举搜索与 LLM agent 结合（代价极高但效果最好）。
3. 代码示例与性能分析
3.1 示例一：全融合的卷积 + Scale/Bias + ReLU + 2×2 MaxPool（OpenMP）
这个例子来自讲义中反复出现的”融合”主线：Conv → Scale/Bias → Max Pool 三步如果各自物化中间张量，会产生数倍于必要量的 DRAM 流量；把它们融合进一个循环嵌套，中间值就永远留在寄存器里。
// conv_fused.cpp — 直接卷积 + Scale/Bias + ReLU + 2x2 MaxPool 的单遍融合实现
// 编译: g++ -O3 -march=native -fopenmp conv_fused.cpp -o conv_fused
// 运行: OMP_NUM_THREADS=16 ./conv_fused
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <chrono>
#include <omp.h>
static double now_ms() {
using namespace std::chrono;
return duration<double, std::milli>(steady_clock::now().time_since_epoch()).count();
}
int main() {
// ---------------- 问题规模 ----------------
const int N = 8; // batch
const int H = 64, W = 64; // 输入空间尺寸 (均为偶数, 便于 2x2 pooling)
const int C = 32; // 输入通道 (INPUT_DEPTH)
const int K = 64; // 输出通道 (滤波器个数)
const int R = 3, S = 3; // 滤波器空间支撑 (stride = 1)
const int HH = H + 2, WW = W + 2; // 带 1 像素零填充的 halo
const int OH = H / 2, OW = W / 2;
// ---------------- 分配数据 ----------------
float* in = (float*)aligned_alloc(64, (size_t)N * HH * WW * C * sizeof(float));
float* w = (float*)aligned_alloc(64, (size_t)K * R * S * C * sizeof(float));
float* scale = (float*)aligned_alloc(64, (size_t)K * sizeof(float));
float* bias = (float*)aligned_alloc(64, (size_t)K * sizeof(float));
float* out = (float*)aligned_alloc(64, (size_t)N * OH * OW * K * sizeof(float));
srand(0);
auto rnd = []() { return (float)rand() / (float)RAND_MAX - 0.5f; };
// 输入: 内区填随机数, halo 保持 0 (calloc 语义)
for (size_t i = 0; i < (size_t)N * HH * WW * C; i++) in[i] = 0.f;
for (int n = 0; n < N; n++)
for (int j = 0; j < H; j++)
for (int i = 0; i < W; i++)
for (int c = 0; c < C; c++)
in[(((size_t)n * HH + (j + 1)) * WW + (i + 1)) * C + c] = rnd();
for (size_t i = 0; i < (size_t)K * R * S * C; i++) w[i] = rnd();
for (int f = 0; f < K; f++) { scale[f] = 1.0f + 0.01f * f; bias[f] = 0.5f; }
for (size_t i = 0; i < (size_t)N * OH * OW * K; i++) out[i] = 0.f;
// ---------------- 融合内核 ----------------
const double t0 = now_ms();
#pragma omp parallel for collapse(3) schedule(static)
for (int n = 0; n < N; n++)
for (int oh = 0; oh < H; oh += 2) // 每次产出 2x2 卷积输出 -> 1 个 pool 输出
for (int ow = 0; ow < W; ow += 2)
for (int f = 0; f < K; f++) { // 输出通道
float pool = -3.4e38f;
for (int dy = 0; dy < 2; dy++)
for (int dx = 0; dx < 2; dx++) {
const int j = oh + dy, i = ow + dx;
float acc = 0.0f;
for (int jj = 0; jj < R; jj++)
for (int ii = 0; ii < S; ii++)
for (int kk = 0; kk < C; kk++) // kk 连续 => 可向量化
acc += w[((f * R + jj) * S + ii) * C + kk]
* in[(((size_t)n * HH + (j + jj)) * WW + (i + ii)) * C + kk];
acc = acc * scale[f] + bias[f]; // ← 融合 scale/bias (无需物化)
if (acc < 0.0f) acc = 0.0f; // ← 融合 ReLU
pool = fmaxf(pool, acc); // ← 融合 2x2 MaxPool
}
out[(((size_t)n * OH + oh / 2) * OW + ow / 2) * K + f] = pool;
}
const double t1 = now_ms();
// ---------------- 统计与校验 ----------------
const double conv_flops = 2.0 * N * (double)H * W * K * R * S * C; // 1 FMA = 2 FLOP
double checksum = 0.0;
for (size_t i = 0; i < (size_t)N * OH * OW * K; i++) checksum += out[i];
printf("conv FLOPs = %.4f GFLOP\n", conv_flops / 1e9);
printf("fused kernel time = %.3f ms -> %.2f GFLOP/s\n",
t1 - t0, conv_flops / ((t1 - t0) * 1e6));
printf("checksum = %.6f (threads = %d)\n", checksum, omp_get_max_threads());
printf("DRAM 流量(理想) = in %.2f MB + w %.2f KB + out %.2f MB\n",
N * HH * WW * C * 4 / 1e6, K * R * S * C * 4 / 1e3, N * OH * OW * K * 4 / 1e6);
return 0;
}
【代码做什么？】
- 构造一批随机激活
in[N][HH][WW][C]（含 1 像素零填充 halo，所以空间尺寸比输出大 2）、随机权重w[K][R][S][C]、逐通道的scale/bias。 - 以
(n, oh, ow)三层循环做并行分解（collapse(3)），每个工作项负责一个 2×2 的卷积输出块——这个块正好对应一个 max-pool 输出，因此池化不需要额外遍历。 - 对块内 4 个空间位置、每个位置累加 $C \times R \times S = 288$ 次乘加，随后在同一段代码里依次施加
scale/bias、ReLU、以及 2×2 的fmaxf池化，最终只写一个结果。 - 最后统计 FLOPs、计时、打印 GFLOPS、校验和与理论 DRAM 流量。
【并行机制与性能解说】
- 并行机制：
#pragma omp parallel for把 $N \times (H/2) \times (W/2) = 8 \times 32 \times 32 = 8192$ 个工作项切分给线程（collapse(3)让 OpenMP 把三维迭代空间线性化后再切块）；每个线程独立执行内层的f、dy/dx、jj/ii/kk循环，共享数据是只读的in、w、scale、bias，输出写入互不重叠的地址（写集合不重叠 ⇒ 无需锁、无数据竞争）。向量化发生在最内层kk（输入与权重在该维都是连续的），但如果编译器不把acc归约拆开，会形成 288 长的串行依赖链——生产实现应当用 4 个独立部分和累加器或把f提到内层做向量化。 - Work（总工作量）：$W = N \cdot H \cdot W \cdot K \cdot R \cdot S \cdot C = 8 \times 64 \times 64 \times 64 \times 9 \times 32 = 6.04\times10^{8}$ 次 FMA $= 1.21$ GFLOP。池化部分只增加 $3 \times N\cdot OH\cdot OW\cdot K \approx 1.6\times10^{6}$ 次比较，可忽略。
- Span（关键路径）：单个输出点的归约链为 $C\cdot R\cdot S = 288$ 次 FMA，池化再叠加 3 次
fmax，所以 $L \approx 4\times(288 + 3) + 3 \approx 1167$ 次串行运算。 - 并行度 = Work / Span $= 6.04\times10^{8} / 1167 \approx 5.2\times10^{5}$。这远大于任何现有机器（16 核 × 8 lane × 2 issue ≈ 256 路并行），说明这个 kernel 的并行度不是问题，问题在内存与计算效率。
- 瓶颈分析：
- DRAM 流量（理想）：输入内区 4.19 MB（含 halo 为 4.46 MB）+ 权重 73.7 KB + 输出 2.10 MB ≈ 6.4 MB；对应算术强度 $1.21\times10^{9}/6.4\times10^{6} \approx 189$ FLOP/Byte，远高于 ridge point 38.4（见第 4 节）⇒ 理论上计算受限。
- 但实际上：这个循环顺序把输入 halo 的重读次数限制在每次输出块一次；若把
f循环移到jj/ii之外重新排列，输入会被重读 64 次，DRAM 流量变成 64 × 4 MB ≈ 268 MB，算术强度掉到 4.5 FLOP/Byte ⇒ 立刻变成严重带宽受限。这正是”循环顺序决定一切”的实例。 - 负载不均：每个工作项的工作量完全相同（8192 个同构任务），所以
schedule(static)足够；若把并行维度改成按行切分（H不整除线程数）才会出现不均。 - 同步开销：整个内核只有一次隐式的 fork-join 屏障（parallel for 结束），同步成本被 1.21 GFLOP 完全摊薄。
3.2 示例二：分层分块 + 寄存器分块 + SIMD 的 GEMM（OpenMP）
// gemm_blocked.cpp — 三级分块 (L2/L1) + 4x16 寄存器分块 + SIMD 的单精度 GEMM
// 编译: g++ -O3 -march=native -fopenmp gemm_blocked.cpp -o gemm_blocked
// 运行: OMP_NUM_THREADS=16 ./gemm_blocked 1024
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <chrono>
#include <omp.h>
#define MB 64 // j (行) 方向的块: 让 A 的一个条带在 L2 里常驻
#define NB 256 // i (列) 方向的块: B 的一个条带在 L2 里常驻
#define KB 256 // k (归约) 方向的块: 让 C 的累加块在缓存里完成
#define MR 4 // 微内核行数 (A 的复用度)
#define NR 16 // 微内核列数 (B 的复用度, 也是 SIMD 宽度)
static double now_ms() {
using namespace std::chrono;
return duration<double, std::milli>(steady_clock::now().time_since_epoch()).count();
}
int main(int argc, char** argv) {
const int n = (argc > 1) ? atoi(argv[1]) : 1024;
const int M = n, K = n, N = n; // 方阵, 且保证 n 是 64 的倍数
const size_t bytes = (size_t)M * N * sizeof(float);
float* A = (float*)aligned_alloc(64, bytes);
float* B = (float*)aligned_alloc(64, bytes);
float* C = (float*)aligned_alloc(64, bytes);
srand(1);
for (size_t i = 0; i < (size_t)M * K; i++) A[i] = (float)rand() / RAND_MAX - 0.5f;
for (size_t i = 0; i < (size_t)K * N; i++) B[i] = (float)rand() / RAND_MAX - 0.5f;
for (size_t i = 0; i < bytes / 4; i++) C[i] = 0.f;
const double t0 = now_ms();
// ---- 并行分解在最外层 (j 块), 避免跨线程写同一 cache line ----
#pragma omp parallel for schedule(static)
for (int jb = 0; jb < M; jb += MB)
for (int kb = 0; kb < K; kb += KB)
for (int ib = 0; ib < N; ib += NB)
for (int j = jb; j < jb + MB && j + MR <= M; j += MR) {
float* c0 = &C[(size_t)(j ) * N];
float* c1 = &C[(size_t)(j + 1) * N];
float* c2 = &C[(size_t)(j + 2) * N];
float* c3 = &C[(size_t)(j + 3) * N];
for (int k = kb; k < kb + KB && k < K; k++) {
const float a0 = A[(size_t)(j ) * K + k];
const float a1 = A[(size_t)(j + 1) * K + k];
const float a2 = A[(size_t)(j + 2) * K + k];
const float a3 = A[(size_t)(j + 3) * K + k];
const float* brow = &B[(size_t)k * N];
// 规范循环形式: 上界先算成变量 (OpenMP 不接受 i < ib+NB && i < N 这种谓词)
const int ib_hi = (ib + NB < N) ? (ib + NB) : N;
#pragma omp simd
for (int i = ib; i < ib_hi; i++) {
const float bv = brow[i];
c0[i] += a0 * bv; // 每个 b 值被 4 次复用 (寄存器分块)
c1[i] += a1 * bv;
c2[i] += a2 * bv;
c3[i] += a3 * bv;
}
}
}
const double t1 = now_ms();
double checksum = 0.0;
for (size_t i = 0; i < (size_t)M * N; i++) checksum += C[i];
const double flops = 2.0 * M * N * K;
printf("M=N=K=%d, threads=%d\n", n, omp_get_max_threads());
printf("time = %.3f ms\n", t1 - t0);
printf("%.2f GFLOP/s (checksum %.3f)\n", flops / ((t1 - t0) * 1e6), checksum);
return 0;
}
【代码做什么？】
- 生成两个 $n\times n$ 的随机矩阵（$n$ 取 64 的倍数），零初始化 $C$。
- 外层
parallel for只并行化最外的jb分块循环——每个线程拿若干行条带，互不共享输出 cache line（避免伪共享）。 - 每个线程串行执行三层分块循环
kb、ib，在最内层用 4 行 × 16 列的寄存器微内核：一次读入 4 个 A 元素，然后对 16 个连续的 B 元素做广播乘法（#pragma omp simd让编译器生成宽 SIMD 指令）。 - 打印时间、GFLOP/s 与校验和。
【并行机制与性能解说】
- 并行机制：线程数 $P$ 由
OMP_NUM_THREADS决定；每线程分到 $M/(P \cdot MB)$ 个jb块（静态分配）。SIMD 通道由#pragma omp simd在最内层i循环上启用（-march=native允许 AVX2/AVX-512）：同一时刻 8 或 16 个 lane 各自处理不同的输出列，共享同一条browcache line 的读请求——这就是讲义第一种 SIMD 方案（”向量化 i 循环，同时改善 B 的空间局部性”）。 - Work（总工作量）：$W = M \cdot N \cdot K$ 次 FMA $= 2 M N K$ FLOP。$n=1024$ 时为 $1.07\times10^{9}$ FMA $= 2.15$ GFLOP。
- Span（关键路径）：每个输出元素是 $K$ 长的归约链。若允许树形归约，深度为 $\log_2 K \approx 10$；若用 4 路部分和（本代码是 4 个独立行，不是 4 路拆分同一元素），单元素仍是 $K$ 深。取较保守的 $L = K = 1024$ 次串行 FMA。
- 并行度：$W/L = M N K / K = M\cdot N = 1.05\times10^{6}$。并行度与 $K$ 无关，只与输出规模成正比——这也是为什么 GEMM 能轻松填满 GPU（输出元素数量极大）。
- 瓶颈分析：
- 内存带宽：$n=1024$ 时三矩阵各 4 MB，共 12 MB（L2 若为 32 MB 级别可大部分驻留）。若假定 L2 装不下、所有 B 都从 DRAM 读：B 被读 $M/MB = 16$ 次 ⇒ 流量 $16\times4$ MB $=64$ MB；实测 16 核上这类实现通常在几十到一两百 GFLOP/s（取决于带宽与端口），远低于 768 GFLOP/s 的 SIMD 峰值 ⇒ 带宽受限。
- 寄存器压力：
MR×NR = 4×16 = 64个累加器 + 4 个 A 值 + 1 个广播值 ≈ 70 个向量寄存器需求，在 AVX2（16 个 ymm）上必须靠编译器把NR降级或用更小的NR；这是BLOCKSIZE不能无限增大的现实约束（讲义的自检问题）。 - 伪共享：由于并行维度是最外层
jb且MB=64（一行 256 字节 = 4 条 cache line），不会出现两个线程写同一条 line。 - 同步开销：只有一次 fork-join。
3.3 示例三：CUDA 上的融合 online-softmax 注意力（FlashAttention 风格）
// flash_attn.cu — FlashAttention 风格的融合注意力: 永不物化 N×N 的分数矩阵
// 编译: nvcc -O3 -arch=sm_80 -o flash_attn flash_attn.cu
// 运行: ./flash_attn 1024 64
#include <cstdio>
#include <cstdlib>
#include <cmath>
#include <vector>
#include <cuda_runtime.h>
#define WARP 32
#define BLOCK_J 64 // 每次处理 64 个 key (key tile)
#define DPT_MAX 8 // 每个 lane 最多负责 d/32 个 feature (支持 d <= 256)
__device__ __forceinline__ float warp_max(float v) {
#pragma unroll
for (int o = 16; o > 0; o >>= 1) v = fmaxf(v, __shfl_xor_sync(0xffffffffu, v, o));
return v;
}
__device__ __forceinline__ float warp_sum(float v) {
#pragma unroll
for (int o = 16; o > 0; o >>= 1) v += __shfl_xor_sync(0xffffffffu, v, o);
return v;
}
// 一个 warp (= 32 线程) 负责一行 query; d 必须是 32 的倍数
__global__ void flash_attn_kernel(const float* __restrict__ Q,
const float* __restrict__ K,
const float* __restrict__ V,
float* __restrict__ O,
int N, int d, float scale)
{
extern __shared__ float smem[];
float* kbuf = smem; // BLOCK_J × d 的 K tile
float* vbuf = kbuf + BLOCK_J * d; // BLOCK_J × d 的 V tile
const int row = blockIdx.x; // query 行号
const int lane = threadIdx.x; // 0..31
const int nd = d / WARP; // 每个 lane 负责的 feature 数
if (row >= N || nd > DPT_MAX) return;
float q[DPT_MAX], acc[DPT_MAX];
for (int p = 0; p < nd; p++) {
q[p] = Q[(size_t)row * d + lane + p * WARP];
acc[p] = 0.0f;
}
float m = -1.0e30f; // running max (online softmax 的状态)
float l = 0.0f; // running sum of exp
for (int jb = 0; jb < N; jb += BLOCK_J) {
const int nj = (N - jb < BLOCK_J) ? (N - jb) : BLOCK_J;
// 协同把 K/V 的 tile 搬进片上共享内存 (jb..jb+nj-1 行在全局内存中连续)
for (int t = threadIdx.x; t < nj * d; t += WARP) {
kbuf[t] = K[(size_t)jb * d + t];
vbuf[t] = V[(size_t)jb * d + t];
}
__syncwarp();
for (int jj = 0; jj < nj; jj++) {
// ---- 分数 s = q · k_jj (warp 内归约) ----
float s = 0.0f;
for (int p = 0; p < nd; p++) s += q[p] * kbuf[jj * d + lane + p * WARP];
s = warp_sum(s) * scale;
// ---- online softmax 更新: 重标定旧的累加器 ----
const float m_new = fmaxf(m, s);
const float alpha = __expf(m - m_new); // 旧值 (含 m、l、acc) 的缩放因子
const float pij = __expf(s - m_new);
l = l * alpha + pij;
for (int p = 0; p < nd; p++)
acc[p] = acc[p] * alpha + pij * vbuf[jj * d + lane + p * WARP];
m = m_new;
}
__syncwarp();
}
for (int p = 0; p < nd; p++) O[(size_t)row * d + lane + p * WARP] = acc[p] / l;
}
static void check(cudaError_t e, const char* msg) {
if (e != cudaSuccess) { printf("CUDA error (%s): %s\n", msg, cudaGetErrorString(e)); exit(1); }
}
// CPU 参考实现: 老老实实物化 N×N 矩阵, 三遍走完
static void naive_ref(const float* Q, const float* K, const float* V, float* O,
int N, int d, float scale) {
std::vector<float> Smat((size_t)N * N);
for (int i = 0; i < N; i++)
for (int j = 0; j < N; j++) {
float s = 0.0f;
for (int t = 0; t < d; t++) s += Q[(size_t)i * d + t] * K[(size_t)j * d + t];
Smat[(size_t)i * N + j] = s * scale;
}
for (int i = 0; i < N; i++) {
float mx = -1e30f;
for (int j = 0; j < N; j++) mx = fmaxf(mx, Smat[(size_t)i * N + j]);
float sum = 0.0f;
for (int j = 0; j < N; j++) {
Smat[(size_t)i * N + j] = expf(Smat[(size_t)i * N + j] - mx);
sum += Smat[(size_t)i * N + j];
}
for (int t = 0; t < d; t++) {
float a = 0.0f;
for (int j = 0; j < N; j++) a += Smat[(size_t)i * N + j] * V[(size_t)j * d + t];
O[(size_t)i * d + t] = a / sum;
}
}
}
int main(int argc, char** argv) {
const int N = (argc > 1) ? atoi(argv[1]) : 1024;
const int d = (argc > 2) ? atoi(argv[2]) : 64;
const float scale = 1.0f / sqrtf((float)d);
const size_t bytes = (size_t)N * d * sizeof(float);
float *hQ = (float*)malloc(bytes), *hK = (float*)malloc(bytes),
*hV = (float*)malloc(bytes), *hO = (float*)malloc(bytes),
*hRef = (float*)malloc(bytes);
srand(2);
for (size_t i = 0; i < bytes / 4; i++) {
hQ[i] = (float)rand() / RAND_MAX - 0.5f;
hK[i] = (float)rand() / RAND_MAX - 0.5f;
hV[i] = (float)rand() / RAND_MAX - 0.5f;
}
float *dQ, *dK, *dV, *dO;
check(cudaMalloc(&dQ, bytes), "malloc Q"); check(cudaMalloc(&dK, bytes), "malloc K");
check(cudaMalloc(&dV, bytes), "malloc V"); check(cudaMalloc(&dO, bytes), "malloc O");
check(cudaMemcpy(dQ, hQ, bytes, cudaMemcpyHostToDevice), "copy Q");
check(cudaMemcpy(dK, hK, bytes, cudaMemcpyHostToDevice), "copy K");
check(cudaMemcpy(dV, hV, bytes, cudaMemcpyHostToDevice), "copy V");
const size_t smem = 2 * (size_t)BLOCK_J * d * sizeof(float);
// 若 smem > 48 KB 需显式申请: cudaFuncSetAttribute(flash_attn_kernel,
// cudaFuncAttributeMaxDynamicSharedMemorySize, smem);
dim3 grid(N), block(WARP);
flash_attn_kernel<<<grid, block, smem>>>(dQ, dK, dV, dO, N, d, scale);
check(cudaGetLastError(), "launch");
check(cudaDeviceSynchronize(), "sync");
// 计时 (20 次迭代)
cudaEvent_t e0, e1;
cudaEventCreate(&e0); cudaEventCreate(&e1);
cudaEventRecord(e0);
for (int it = 0; it < 20; it++)
flash_attn_kernel<<<grid, block, smem>>>(dQ, dK, dV, dO, N, d, scale);
cudaEventRecord(e1);
check(cudaDeviceSynchronize(), "sync2");
float ms = 0.f; cudaEventElapsedTime(&ms, e0, e1); ms /= 20.0f;
check(cudaMemcpy(hO, dO, bytes, cudaMemcpyDeviceToHost), "copy O");
naive_ref(hQ, hK, hV, hRef, N, d, scale);
double maxerr = 0.0;
for (size_t i = 0; i < bytes / 4; i++) maxerr = fmax(maxerr, fabs(hO[i] - hRef[i]));
const double flops = 4.0 * (double)N * N * d; // QK^T + PV, 各 2N²d
printf("N=%d d=%d 时间=%.3f ms %.2f GFLOP/s max|err|=%.3e\n",
N, d, ms, flops / (ms * 1e6), maxerr);
printf("naive 版本需要物化 %.2f MB 的 N×N 分数矩阵; 融合版本需要 0 MB\n",
(double)N * N * 4 / 1e6);
return 0;
}
【代码做什么？】
flash_attn_kernel用一个 warp 负责一行 query：blockIdx.x是行号，32 个 lane 各负责 $d/32$ 个 feature 分量。- 外层循环每次搬入 $B_J=64$ 个 key 的 K/V tile 到共享内存（
extern __shared__，__syncwarp()保证可见性）。 - 对 tile 内每个 key，lane 各自算部分点积，再通过
__shfl_xor_sync做 warp 内归约得到分数 $s$。 - 用 online softmax 递推更新 running max
m、归一化和l，以及输出累加器acc[p]：$\alpha = e^{m_{old}-m_{new}}$，l = l·α + e^{s-m_new}，acc = acc·α + e^{s-m_new}·v。 - 全部 key 处理完后一次性除以
l写出结果——整个过程中 $N\times N$ 的分数矩阵从未被创建。 - 主机端用
naive_ref物化 $N\times N$ 矩阵并三遍计算作为参考，比较最大绝对误差，并用 CUDA event 计时 20 次迭代。
【并行机制与性能解说】
- 并行机制：$N$ 个 thread block，每块 32 个线程（一个 warp），映射到 SM 的 SIMD lane 上。warp 内归约用 5 步 shuffle（$\log_2 32$），无需共享内存与 barrier；块内的 K/V tile 复用发生在共享内存（256 KB/SM 的片上存储）；
acc[]与q[]常驻寄存器。关键点：$N\times N$ 的中间结果被”稀释”成 $d$ 宽的累加器——存储复杂度从 $O(N^2)$ 降到 $O(d)$。 - Work（总工作量）：分数矩阵 $2N^2d$ FLOPs + 输出 $2N^2d$ FLOPs = $4N^2d$ FLOPs；$N=1024, d=64$ 时为 $2.68\times10^{8}$ FLOPs $=0.268$ GFLOP。此外每个 tile 要对 $d$ 维累加器做一次重标定：每行约 $(N/B_J)\cdot d$ 次额外乘法，占总量的比例约 $1/(2B_J) \approx 0.8\%$（$B_J=64$）。
- Span（关键路径）：online 形式把 key 循环变成了线性的依赖链——第 $j+1$ 个 key 的
m、l、acc都依赖第 $j$ 个 key 的结果。单个 query 行的关键路径 $L \approx N \times (d/32 + 5 + c)$（每个 key 需要 $d/32$ 次串行乘加 + 5 步 shuffle + 若干次指数），取 $c\approx3$：$L \approx 1024 \times (2+5+3) = 10240$。 - 并行度 = Work / Span $= 2.68\times10^{8} / 10240 \approx 2.6\times10^{4}$。作为对比，如果允许用”两遍”的经典 softmax（先全算 $S$ 再归约），$N$ 维归约可用树形，span 可降到 $O(\log N + \log d)$，并行度提高几个数量级——融合/流式是以”提高关键路径长度”换取”减少片上存储与 DRAM 流量”的典型交易（讲义明确指出 chunked 版本”比原版有额外计算”）。
- 瓶颈分析：
- 并行工作不足：只有 $N=1024$ 个 warp，而 H100 有 144 个 SM、每 SM 最多 64 个 warp（约 9216 个 warp 容量）⇒ 单个 head、单条序列根本填不满机器（呼应讲义 “Need a lot of parallel work to fill the machine”：$N=1$ 的卷积只有 2 MB 输出，而 $N=32$ 时有 1 GB 输出）。工程上必须靠 batch × head 维度并行。
- 共享内存带宽与容量：$d=64$ 时 shared memory $= 2 \times 64 \times 64 \times 4 = 32$ KB/块；若 $d$ 增大要检查 48 KB 的默认上限并显式
cudaFuncSetAttribute。 - 同步：只用
__syncwarp()（warp 内屏障，几乎零成本），没有__syncthreads()，这是”一个 warp 负责一行”这个映射带来的额外好处。 - 每线程寄存器压力：
q[8]+acc[8]+ 标量，$d=256$ 时接近DPT_MAX上限；这也是生产实现改用”多 warp 协作 + 跨 warp 归约”的原因（代价是引入__syncthreads()与更多共享内存）。
4. 性能模型与复杂度分析
本节把 Roofline、算术强度、work-span 与能耗账本放在一起做定量推理。所有参数都在算例中显式给出假设。
4.1 Roofline、ridge point 与”算力提升反而更容易带宽受限”
假设一台机器：16 核 × 3.0 GHz × 8 宽 SIMD × 2（FMA 计 2 FLOP）= 768 GFLOP/s 峰值；DRAM 带宽 20 GB/s。 则 ridge point 为 \(I^* = \frac{\pi_{peak}}{\beta} = \frac{768\times10^{9}}{20\times10^{9}} = 38.4\ \text{FLOP/Byte}.\)
- 一个算术强度 $I=1$ 的 kernel（例如逐元素相加）受带宽限制，吞吐只有 $1 \times 20\ \text{GB/s} = 20$ GFLOP/s，仅用到峰值的 2.6%。
- 融合变换把 $I$ 从 $1/3$（三遍循环）提高到 $3/5$（融成一遍）——讲义的原例：两遍 add/mul 加一遍 add，每次遍历 2 load + 1 store / 1 op（$I=1/3$）；融合后 4 load + 1 store / 3 ops（$I=3/5$）。吞吐因此提高 $1.8\times$，而代码没换、机器没换。
- 若把机器升级为算力翻倍（1536 GFLOP/s，带宽不变），ridge point 变成 76.8 FLOP/Byte，$I=3/5$ 的 kernel 依然是带宽受限，吞吐一点不涨。这就是讲义那句总结的定量含义：“提高算力让程序更容易带宽受限；提高算术强度让程序更容易计算受限。”
4.2 GEMM 分块的算术强度：为什么必须做层次化 blocking
设三级块（A、B、C 的三个 $b\times b$ 块同时驻留在某一级存储中），则每个”块周期”执行 $2b^3$ FLOP、搬运 $3b^2$ 个元素。于是 \(OI(b) = \frac{2b^3}{3b^2} = \frac{2b}{3}\ \text{FLOP/元素} = \frac{b}{6}\ \text{FLOP/Byte}\quad(\text{fp32}).\)
| 块大小 $b$ | 所在层次（假设） | $OI$ (FLOP/Byte) | vs ridge point 38.4 |
|---|---|---|---|
| 8 | 寄存器微内核 | 1.3 | 极度带宽受限 |
| 32 | L1 | 5.3 | 带宽受限 |
| 64 | L1/L2 | 10.7 | 带宽受限 |
| 128 | L2 | 21.3 | 仍带宽受限 |
| 256 | L2 | 42.7 | 计算受限 ✅ |
结论（数值化）：要在 768 GFLOP/s / 20 GB/s 的机器上跑满算力，仅靠单级分块需要 $b \ge 230$，而三个 $256^2$ 的 fp32 块需要 $3 \times 256^2 \times 4\text{B} = 786$ KB——单核私有 L2 通常没有这么大。因此必须层次化分块：寄存器级用 $8\times8$ 微内核保证每次 load 复用 8 次，L1/L2 级用大块摊薄 DRAM 流量，两者叠加才能让有效算术强度跨过 38.4。这正是讲义 “Hierarchical blocked matrix mult” 与 “Not shown: final level of blocking for register locality” 的动机。
4.3 explicit GEMM 的代价：DRAM 流量放大 $R\times S$ 倍（数值算例）
假设参数：$N_{batch}=1$，$H=W=256$，$C=64$，$K=64$，$R=S=3$，fp32，机器带宽 $\beta=900$ GB/s（V100 级 HBM）。
- 计算量：$2 \times (P\,Q) \times K \times (R\,S\,C) = 2 \times 65536 \times 64 \times 576 = 4.83$ GFLOP。在 A100 的 19.5 TFLOPS fp32 上理论耗时 $4.83\times10^{9}/(19.5\times10^{12}) = 0.25$ ms。
- im2col 矩阵大小：$(P\,Q) \times (R\,S\,C) \times 4\text{B} = 65536 \times 576 \times 4 = 151$ MB。写一次 + 读一次 = 302 MB 流量。
- 时间：$302\times10^{6} / 900\times10^{9} = 0.336$ ms。
⇒ 光是把卷积矩阵物化一次，搬运时间就超过了计算时间本身（0.336 ms vs 0.25 ms），并且需要 151 MB 的额外存储（是输入的 $R\times S$ 倍）。这就是 implicit GEMM 的全部理由：只在片上共享内存里物化一个 sub-block，不增加 DRAM 流量也不需要额外 off-chip 存储。
4.4 融合的收益：以 Conv+Scale/Bias+Pool 与 softmax 为例
算例 A（讲义 slide 56 的场景）：设卷积输出为 1 GB。
- 不融合：写 conv 输出 1 GB → 读回做 scale/bias 1 GB → 写 1 GB → 读回做 pool 1 GB → 写 0.25 GB。总流量 ≈ 4.25 GB。
- 融合：写最终 pooled 输出 0.25 GB（conv 中间结果永不落 DRAM），总流量 ≈ 0.25 GB。
- 在 $\beta = 900$ GB/s 下：$4.25/0.9 \approx 4.72$ ms vs $0.25/0.9 \approx 0.28$ ms，节省约 4.4 ms。
算例 B（softmax，讲义 slide 59）：naive 实现读 $5MN+2M$ 个元素、写 $3MN+2M$ 个元素；融合实现读 $MN$、写 $MN$（前提是一行的工作集能装进片上存储）。
- 取 $M=N=4096$，fp32：$MN = 16.78$ M 元素 = 67.1 MB。
- naive：读 $5\times67.1 + 2\times0.0168 \approx 335.5$ MB，写 $3\times67.1 \approx 201.3$ MB，共 536.8 MB。
- 融合：读 67.1 MB + 写 67.1 MB = 134.2 MB。
- 流量降为 1/4；在 900 GB/s 下从 $0.596$ ms 降到 $0.149$ ms。
4.5 硬件加速器的定量账本
(a) 能效阶梯的算术：假设 CPU 方案 100 W 达到 768 GFLOP/s（7.7 GFLOP/s/W）。若按讲义的”10×（GPU）/20×（域专用）/100–1000×（ASIC）”，则同样功耗下：GPU ≈ 77 GFLOP/s/W，域专用加速器 ≈ 154 GFLOP/s/W，ASIC 可达 770–7700 GFLOP/s/W。
(b) TPU 式 systolic array 的峰值与能效（假设参数：256×256 = 65536 个 PE，每 PE 每周期 1 次 MAC，时钟 0.7 GHz，整芯片功耗 40 W）： \(\text{峰值} = 65536 \times 0.7\times10^{9} \times 2\ \text{op} = 9.18\times10^{13}\ \text{op/s} \approx 92\ \text{TOPS (8-bit)}\) \(\text{能效} = 92\times10^{12} / 40 = 2.3\ \text{TOPS/W} = 0.43\ \text{pJ/op}.\) 把 $0.43$ pJ/op 与讲义给出的”整数运算约 1 pJ”对比可知：systolic array 通过把指令流开销摊薄到零，其每 op 能耗已落到”纯逻辑运算”量级甚至更低——这正是”专业化 = 能效”的物理来源。
(c) 数据移动 vs 计算的能耗（讲义 ballpark：整数运算 1 pJ，浮点运算 20 pJ，1 mm 外小 SRAM 读 64 bit 26 pJ，LPDDR 读 64 bit 1200 pJ）：
- LPDDR：$1200/8 = 150$ pJ/Byte。
- 一个 kernel 做 1 GFLOP fp32（$5\times10^{8}$ 次 FMA）并读写 1 GB DRAM：
- 计算能耗 $= 1\times10^{9} \times 20\ \text{pJ} = 0.02$ J。
- 访存能耗 $= 1.07\times10^{9} \times 150\ \text{pJ} = 0.16$ J。
- 访存是计算的 8 倍 ⇒ 即使算力提升 8 倍，总能耗也几乎不变。这就是”降低数据移动”压倒一切的定量理由。
- 反过来看算术强度：要让访存能耗与计算能耗持平（fp32），需要 $OI \approx 20\ \text{pJ} / 150\ \text{pJ/Byte} \approx 0.13$ FLOP/Byte——能耗意义上的 ridge point 远低于性能意义上的 ridge point，所以”能效受限”往往比”性能受限”更容易先撞墙。
(d) Tensor Core 把 ridge point 推向何处：H100 上 fp16 tensor core 峰值约 989 TFLOPS，HBM 带宽约 3 TB/s（H100 SXM 量级）⇒ $I^* \approx 989\times10^{12}/3\times10^{12} \approx 330$ FLOP/Byte。也就是说，要用满张量核，kernel 必须每读 1 字节 DRAM 就完成 330 次浮点运算——只有把数据全部 tiling 进片上存储（shared memory / TMEM）并做算子融合才可能达到。这解释了为什么现代 AI kernel 的形态是”小块 + 异步搬运 + 全融合”，而不是”大矩阵 + 一次算完”。
4.6 work-span 视角下的三种实现对比（汇总）
| 实现 | Work（FLOPs） | Span（关键路径） | 并行度 | 主要瓶颈 |
|---|---|---|---|---|
| 直接卷积（3.1） | $2N H W K R S C = 1.21$ GFLOP | ≈ 1167 次串行运算 | $5.2\times10^{5}$ | 循环顺序错 ⇒ DRAM 流量 ×64；无 SIMD ⇒ 利用率低 |
| 分块 GEMM（3.2） | $2MNK = 2.15$ GFLOP | $K = 1024$ 次 FMA（树形可降到 10） | $1.05\times10^{6}$ | DRAM 带宽（B 被重读 $M/MB$ 次）；寄存器压力 |
| 融合注意力（3.3） | $4N^2d = 0.268$ GFLOP（+0.8% 重标定） | ≈ $N(d/32+8) = 10240$ | $2.6\times10^{4}$ | 并行工作不足（仅 $N$ 个 warp）；共享内存容量 |
三行合起来给出本讲的性能方法论：先看 $W/L$ 判断并行度是否够填机器；再看 $OI$ 与 ridge point 判断是计算受限还是带宽受限；最后用融合把中间结果留在片上，用分块把每一级存储的复用榨干。
5. 关键要点
- DNN 的性能问题，90% 是关于数据复用的 GEMM 问题。 全连接层、卷积层、Transformer 的 attention 块最终都是（隐式）矩阵乘；矩阵乘的算术强度随分块尺寸线性增长（$OI \propto n$），因此”分块 + 融合 + 低精度”这三件软件层面的事，加上”更高带宽 + 更专的硬件”这两件硬件层面的事，几乎就是全部优化空间。
- 一切设计围绕”减少数据移动”。 片外 DRAM 访问的能量成本比一次运算高 1–2 个数量级（1 mm SRAM 读 64 bit ≈ 26 pJ，LPDDR 读 64 bit ≈ 1200 pJ，而一次整数运算 ≈ 1 pJ），所以 tiled tensor（16×16/32×32）、异步计算/异步访存/异步片间通信、以及计算单元之间的直接通信（融合与流水）成为理想加速器的三大特征。
- 能效来自专业化，而不是更多的核。 相对 CPU 上的高质量 C 代码：吞吐型处理器（GPU）约 10× perf/W，域专用加速器（TPU 式）约 20×，固定功能 ASIC 可达 100–1000×；代价是可编程性的丧失与数千万至上亿美元的设计成本。TPU 把 30% 面积给算术单元、只保留 5 条指令、用 weight-stationary 的 systolic array 让通信退化为邻居之间的局部传递，是这一原则的教科书实现。
- 融合（fusion）是软件侧收益率最高的单点变换。 Conv+Scale/Bias+Pool 的融合可把 DRAM 流量从 4.25 GB 降到 0.25 GB；逐行 softmax 的融合把流量降为 1/4；FlashAttention 用 online softmax 的结合律换来”永不物化 $N^2$ 矩阵”，代价仅是不到 1% 的额外重标定计算。融合的本质是把”计算顺序”从内存往返中解耦。
- 抽象层次与自动化是现代性能工程的出路。 从寄存器/共享内存级的手写 CUDA，到 tile 级原语（CUTLASS/Triton/Thunderkittens），再到 Halide 式的”算法 / 调度”分离、自动调度搜索（ML 代价模型 + 树搜索，与人类最佳调度相当），再到 LLM agent 的”生成-剖析-反思-修改”循环——每一层抽象都在把”性能专家的隐性知识”变成可复用、可搜索、可自动生成的东西；而”硬件彩票”提醒我们，最终胜出的算法往往是因为它契合了当下的硬件，而非因为它普遍更优。
6. 常见陷阱与注意事项
- 把卷积物化成 im2col 矩阵再用库函数：explicit GEMM 会把 DRAM 流量放大 $R\times S$ 倍（3×3 卷积即 9 倍）并需要与激活同量级的额外存储。在上述算例中，物化一次 151 MB 的卷积矩阵（读+写 302 MB）的耗时（0.336 ms）已经超过计算本身的 0.25 ms。除非模型很小或库实现经过专门调优，否则应使用 implicit GEMM（在片上共享内存里物化 sub-block）。
- 只做单层 blocking，或盲目加大 block size：单层分块下 $OI = b/6$ FLOP/Byte（fp32），要达到 ridge point 38.4 需要 $b \approx 230$，单核私有缓存根本装不下；必须用”寄存器微内核 + L1/L2 大块”的层次化分块。反过来，块越大越好是错的：$MR\times NR = 4\times16$ 已经需要约 70 个向量寄存器，超过 AVX2 的 16 个 ymm 就会 spill，性能反而下降（讲义的自检问题正是问这个）。
- 忽略算子之间的内存往返：把
Conv → Scale/Bias → MaxPool写成三个独立的 kernel，会产生数倍于必要量的 DRAM 流量（算例 A：4.25 GB vs 0.25 GB）。同理，softmax 的 naive 实现读写量是融合实现的 4 倍。判断标准很简单：中间张量有没有必要离开片上存储？ - 误以为”峰值 TFLOPS 更高就一定更快”：Tensor Core 把 ridge point 推到 300+ FLOP/Byte，此时任何”每字节算不到 300 次运算”的 kernel 都仍然带宽受限。低精度也会改变账本：fp16 让同一份数据的字节数减半（算术强度翻倍），但如果累加器也用 fp16，长归约链的舍入误差会毁掉数值结果——正确做法是 A、B 低精度、累加用 fp32，并且 bf16（指数位与 fp32 相同、尾数只有 7 位）与 fp16（范围小、精度高）不能混用，fp8 还要区分 E4M3（范围 0–448）与 E5M2（范围 0–57344）。
- 忽视小 batch / 短序列导致的并行度不足：$N=1$、$P=Q=64$ 的卷积只有 524K 个输出（2 MB），融合注意力的一个 head 只有 $N$ 个 warp（$N=1024$ 也仅占 H100 约 9216 个 warp 容量的 11%），systolic array 更会因矩阵维度不整除阵列尺寸而大量空转。这类负载应该靠 batch × head × 层间流水来补足并行度，而不是奢望单 kernel 打满硬件。
- 并行维度选错导致伪共享与负载不均：并行化最内层（或在多个线程间共享输出 cache line）会引入伪共享；把
#pragma omp parallel for放在N（不整除线程数）而非collapse(3)的平坦迭代空间上会造成尾部负载不均；在 GEMM 里把线程切在i（列）方向上会让不同线程写同一行的相邻元素——正确做法是把并行切分点放在最外层块（如 3.2 中的jb），并保证每个线程的写区间按 cache line（64 字节 = 16 个 fp32）对齐。
7. 思考题（带答案）
问题 1
某推理服务对单个样本做一次 $y = Wx$，其中 $W$ 是 $4096\times4096$ 的 fp16 权重矩阵。请用 Roofline 定量解释：为什么这块 GPU 的 989 TFLOPS 张量核算力实际只能用到百分之几？如果改成一次处理 32 个样本（batch=32），情况如何变化？
【答案】
- batch=1 时，这是一个矩阵-向量乘（GEMV）。计算量：$2 \times 4096 \times 4096 = 3.36\times10^{7}$ FLOP。必须搬运的数据：权重 $4096^2 \times 2\text{B} = 33.6$ MB（每个权重只用一次），加上输入 8 KB、输出 8 KB。
- 算术强度 $I = 3.36\times10^{7} / 3.36\times10^{7} \approx 1.0$ FLOP/Byte（每个权重字节只支撑约 1 次浮点运算）。
- 在 $\beta \approx 3$ TB/s 的 HBM 上，吞吐上限 $= I\times\beta \approx 1 \times 3\times10^{12} = 3$ TFLOPS，仅为 989 TFLOPS 的 0.3%。ridge point 是 $989/3 \approx 330$ FLOP/Byte，GEMV 差了两个数量级。
- batch=32 时，同一个权重矩阵被 32 个样本复用：计算量变成 $32\times3.36\times10^{7} = 1.07\times10^{9}$ FLOP，而权重流量不变（33.6 MB），$I$ 提高到约 32 FLOP/Byte，吞吐上限升到约 96 TFLOPS（约 10%）。若 batch 提到 256，$I\approx256$，接近 ridge point，算力利用率可达 70% 以上。
- 结论：小 batch 推理在 GPU 上是典型的带宽受限（”权重流式读取”）负载；提高 batch、做权重量化（int8/fp8 直接把字节数减半）、或用批内多流并行，才是正确的优化方向。这也解释了讲义 slide 49 为什么要对比 $N=1$（2 MB 输出）与 $N=32$（1 GB 输出）两种情形。
问题 2
(chunked / online) softmax 为什么能在不物化 $N\times N$ 矩阵的前提下，给出与”全量 softmax”完全相同的数学结果？请写出递推公式，说明数值稳定性来自哪里，并定量说明它多付出的计算量有多大。
【答案】
- 关键在于 softmax 的两步归约（取最大值 $m(x)=\max_i x_i$ 与求和 $l(x)=\sum_i e^{x_i-m(x)}$）都是可结合的。把行向量 $x$ 切成两块 $x^{(1)}, x^{(2)}$： \(m(x)=\max\big(m(x^{(1)}),\,m(x^{(2)})\big),\qquad f(x)=\big[e^{m(x_1)-m(x)}f(x^{(1)}),\; e^{m(x_2)-m(x)}f(x^{(2)})\big],\) \(l(x)=e^{m(x^{(1)})-m(x)}\,l(x^{(1)})+e^{m(x^{(2)})-m(x)}\,l(x^{(2)}),\qquad \mathrm{softmax}(x)=f(x)/l(x).\) 逐块处理时，把新块的统计量 $(m_{tile}, l_{tile}, P_{tile}V_{tile})$ 与已累积的 $(m, l, O)$ 合并即可：$m_{new}=\max(m, m_{tile})$，$\alpha=e^{m-m_{new}}$，$l_{new}=l\alpha + l_{tile}e^{m_{tile}-m_{new}}$，$O_{new}=O\alpha + P_{tile}V_{tile}e^{m_{tile}-m_{new}}$。这与”先算完整行再归一化”在实数域上完全等价（只是浮点舍入次序不同），因此结果一致。
- 数值稳定性来自”始终减去当前的最大值”：$e^{x_i-m}$ 的指数永不为正，避免上溢；一旦某个块的分数更大，就用 $\alpha<1$ 把旧累加器按比例缩小再合并，而不是先算完再统一缩放。这就是讲义中 $f(x)=[e^{x_1-m(x)},\dots]$、$m(x)=\max_i(x_i)$ 的定义在分块情形下的直接推广。
- 代价：每处理一个 key tile，都要对已有的输出累加器（$d$ 维）做一次乘以 $\alpha$ 的重标定。每行的额外乘法约 $(N/B_J)\cdot d$ 次，而主计算是 $2Nd + 2Nd = 4Nd$ FLOPs（这里按乘加算 1 次乘法），因此额外开销比例为 $\frac{(N/B_J)d}{2N^2 d/(\text{每行} N \text{个 key 的 } 2d)} \approx 1/(2B_J)$。取 $B_J=64$ 时约为 0.8%——用不到 1% 的额外计算，换掉 $O(N^2)$ 的存储与 $O(N^2 d)$ 级别的 DRAM 往返。这正是讲义所说 “Save memory footprint: never materialize $N^2$ matrix；Save memory bandwidth: high arithmetic intensity”，以及”存在额外计算（必须重标定 O 的旧值）”的定量版本。
问题 3
假设你要为一个”单用户、单序列、batch=1、要求极低延迟（<10 ms）”的推理服务选择硬件。请结合”效率 vs 可编程性”曲线与 systolic array 的特性，说明为什么”能效最高的加速器”未必是最优选择，并给出你的决策步骤。
【答案】
- 先做 Roofline 定位：如问题 1 所示，batch=1 的负载算术强度极低（$\approx1$ FLOP/Byte），属于带宽受限。此时”峰值算力”（无论是 GPU 的 tensor core 还是 ASIC 的 PE 阵列）几乎不影响性能，唯一重要的是有效带宽与权重能否留在片上。
- 再看硬件形态的匹配度：systolic array 是 weight-stationary 的稠密矩阵乘机器——它擅长的是”输入矩阵足够大、能填满整阵列、每拍都有新数据灌入”的场景；batch=1 的 GEMV 会让阵列的绝大多数 PE 在每一拍只做一半有效工作（甚至因为维度不整除 256×256 而空转），能效优势根本无法体现。相反，通用 GPU/CPU 的动态调度能灵活处理不规则形状、支持任意控制流、能容纳快速迭代的模型结构。
- 再算总拥有成本：ASIC 的设计/验证/流片成本是”数千万到上亿美元”级别，且不可编程——一旦模型结构变化（AI 领域每几个月一次）就要重新设计；FPGA 相对灵活但编程困难（讲义称之为 ~50×？”jury still out”）；域专用加速器（DSL 可编程）约 20×，但仍受 DSL 表达力限制。
- 决策步骤（可操作）：① 测量负载的算术强度与 batch/序列分布，确定是带宽受限还是计算受限；② 若带宽受限，优先在通用硬件上做量化（int8/fp8 减半字节数）、权重常驻片上/近存、kernel 融合与批内并行，而不是换芯片；③ 若计算受限且形状规整、QPS 稳定、模型结构固定，再把”能效更高”的域专用加速器纳入比较，并用”每美元每秒查询数 + 迭代速度 + 迁移成本”而非单纯的 perf/W 决策；④ 始终警惕”硬件彩票”：今天的模型之所以长这样，部分原因是它契合了今天的硬件——把模型结构锁死在一块 ASIC 上，等于同时锁死了未来算法的选择空间。