ARTICLE · INTELLIGENCE

战地情报 · 详情页

来自尧图项目组的一线实战观察与深度解析

性能模型驱动Tiled GEMM优化:从算术强度到Roofline实战

性能模型驱动Tiled GEMM优化:从算术强度到Roofline实战

1. 性能模型到底在指导什么

1.1 为什么大家都在死磕 GEMM

在 AIInfra 这个圈子里,绕不开的一个词就是 GEMM,通用矩阵乘法。你去看任何一个推理引擎的 profile,排在耗时前几名的算子,绝大多数都能归约成 GEMM:卷积的 im2col 转化、Attention 里的 QK^T 和 AV、全连接层、MoE 的专家路由分发,全是矩阵乘。夸张点说,把 GEMM 做好了,一个推理引擎的“地基”就稳了一大半。

但为什么是“Tiled GEMM”?因为现实中的矩阵远大于任何一级缓存的容量,比如一次大模型推理的 prefill 阶段,M 维度可能是序列长度乘以 batch,N 和 K 动辄几千甚至上万。整块 A 乘整块 B,数据根本放不进缓存,CPU/GPU 跑起来大部分时间都浪费在等数据上。Tiled 的核心思路就是把它切成一块块能“装进”片上高速存储的小块,让数据在寄存器、共享内存、L1/L2 里反复复用,而不是一次次去 DRAM 大海捞针。

这块工程师的含金量,恰好就体现在这里:不是会调一个 cutlass 模板就能叫懂,而是面对一个新硬件、一个新 shape,你能不能在动手写 kernel 之前,就用一张纸、几个公式把“瓶颈在哪、能不能跑到峰值、该往哪个方向调”推出来。性能模型就是这张纸。

1.2 性能模型不是“跑分”,是设计前的数字推演

很多初学者有个误区,觉得性能模型就是事后拿 ncu 或者 torch.profiler 测一堆指标,然后对照数据找瓶颈。但真正的性能模型要在你写第一行代码之前就用起来。它是一组数学关系,把“算力峰值、带宽、数据复用次数、寄存器数量、共享内存容量”这些硬件约束,与“矩阵维度、tile 大小、数据类型”这些算法参数关联起来。

我特别喜欢拿装修来打比方。你请师傅贴瓷砖,师傅会先量房子尺寸、算瓷砖规格、算水泥用量,然后告诉你 600×600 还是 800×800 更省料、更好看。他不会直接买一卡车瓷砖回来铺开试试。性能模型做的就是这个“量房算料”的活:先算清楚一块 tile 能带来多少算力、消耗多少带宽、需要多少寄存器,再去写 kernel。这样做的最大价值不是预测一个精确到小数点的性能数字,而是帮你迅速地排除掉一大批明显不行的方案。

比如你脑补了一个 128×128 的 tile,感觉挺大,数据复用率应该不错,但一算共享内存发现一个块就要吃掉 192KB,而硬件一个 SM 总共也只有 228KB,还要双缓冲,直接就爆了。这种问题如果等到代码写完再查,至少浪费一个下午。

2. Tile 结构:把大矩阵切成能装进高速缓存的小块

2.1 为什么非要分块

我们从头推导一下。GEMM 要算 C = A × B,其中 A 是 M×K,B 是 K×N。理论上每个 A 的元素要被读 N 次(跟 B 的每一列乘),每个 B 的元素要被读 M 次。也就是说,如果不做任何缓存优化,数据从 DRAM 到计算单元要重复搬运 M×N 次量级,DRAM 带宽瞬间被打满。

怎么减少重复搬运?答案是让数据在离计算单元更近的地方待着,多复用几次。寄存器离计算单元最近,其次是共享内存,再往外是 L1/L2 缓存。Tiled GEMM 的基本动作是:先把 A 的一个纵向块(BM×BK)和 B 的一个横向块(BK×BN)搬到共享内存,然后让这一小块数据在寄存器里完成尽可能多次的乘累加,写回 C 的对应分块。这样每个 DRAM 数据只被读一次,但计算复用了很多次,算术强度就上来了。

你可以把 tile 理解成“批处理”:每次从慢速的大仓库领一批货到高速工作台旁边,接下来所有工人都从工作台拿料干活,不用再跑仓库。至于批多大,取决于工作台有多大,也就是片上存储有多大。

2.2 六个关键参数:BM、BN、BK、TM、TN、线程数

Tiled GEMM 的参数看着很多,其实逻辑上是两层的嵌套结构:

  • 第一层是线程块(block)切分,对应 BM、BN、BK。BM 和 BN 决定了一个 block 负责计算的 C 分块大小,BK 决定了每次从原矩阵中切多厚的 K 维度切片。
  • 第二层是线程内的寄存器分块(register tiling),对应 TM、TN。TM 是一个线程在 M 维度上负责的 C 元素个数,TN 是在 N 维度上负责的个数,所以每个线程的累加器数量就是 TM×TN。

这两个层级的参数不是随便拍的,它们受到一组硬约束:

  • 共享内存容量:至少 (BM×BK + BK×BN)×type_bytes,如果做双缓冲还要翻倍。
  • 寄存器数量:每线程至少 TM×TN 个寄存器用来存累加器,加上地址计算、数据缓存,超过 255 个就会 spill 到本地内存,性能直接崩。
  • 线程数:block 内线程数 = (BM/TM)×(BN/TN),这个值要能被 32(warp 大小)整除,且不能超过硬件最大线程数限制。

把这些约束列出来,你就能理解为什么 tile 不是越大越好:BM、BN 增大能提高数据复用,但会挤压线程数和可用寄存器;BK 增大会摊薄外层循环开销,但会成倍增加共享内存占用和同步延迟。

2.3 一个 128×128×8 的参数示例

纸上谈兵没意思,直接来一个具体的数值推演。假设我打算设计一个 FP32 的 Tiled GEMM,选用 BM=128、BN=128、BK=8,每个线程 TM=8、TN=8:

  • 线程数 = (128/8)×(128/8) = 256,是 warp 大小的整数倍,合规。
  • 每线程累加器 = 8×8 = 64 个寄存器,加上索引和临时缓存大概到 80~100 个,还没爆掉 255 的上限。
  • 共享内存 = 128×8 + 8×128 = 2048 个 float,乘以 4 字节等于 8KB,双缓冲也就 16KB,在多数 GPU 上完全够用。

再看计算量:这个 block 的 K 切片长度为 8,所以要做 2×128×128×8 = 262144 FLOPs。分摊给 256 个线程,每个线程要做 1024 FLOPs,也就是 512 次 FMA。如果硬件 SM 峰值是 384 GFLOPS(比如主频 1.5GHz 的 128 FMA lane),这个 block 的理论耗时约 0.68 微秒。这个数字本身不代表最终性能,但它能帮你反推:要跑满一个 140 SM 的 GPU,需要多少个这样的 block 并发、需要多少批 tile 才能喂饱所有计算单元,从而估算整体延迟。这就是模型的价值:它不是回答“能跑多快”,而是回答“如果一切理想,最快要多久,离这个目标的差距是从哪来的”。

3. 用性能模型算出“瓶颈在哪”

3.1 算术强度:从 DRAM 视角看 GEMM

性能建模最核心的一个概念叫算术强度(Arithmetic Intensity,AI),定义是“每从内存读一个字节,能换来多少次浮点运算”,单位是 FLOPs/Byte。这个概念我每次讲都会强调,因为它直接决定了一个 kernel 是受计算瓶颈还是访存瓶颈。

考虑一个理想情况下的 tile 计算:block 负责 BM×BN 的 C 分块,叠加 BK 的 K 维度切片。整个过程中需要从 DRAM 读取 A 分块(BM×BK 个元素)、B 分块(BK×BN 个元素),最终写回 C 分块(BM×BN 个元素)。所以:

  • 计算量 = 2×BM×BN×BK
  • 访存量 = (BM×BK + BK×BN + BM×BN)×type_bytes
  • 算术强度 ≈ 2×BM×BN×BK / (BM×BK + BK×BN + BM×BN)

这个分子和分母的结构很有意思:当 BM 和 BN 不变,BK 越大,分子线性涨,分母主要是 BM×BN 这一项在拖后腿,AI 会变大;当 M=N=K=4096、一次性读全量数据时,AI ≈ 2×4096/3 ≈ 2730 FLOPs/Byte,非常惊人。所以大 GEMM 天然就是计算受限。而当你处理 M=N=K=128 的小矩阵时,AI 只有约 21 FLOPs/Byte,数据搬移的成本远超计算本身。优化小矩阵 GEMM 本质上就是想办法把 tile 做大、让数据多复用,把 AI 往上推。

3.2 计算受限还是访存受限:Roofline 临界点

光有 AI 还不够,要拿它跟硬件能力比对。每个平台都有一个“临界点”,叫 Ridge Point,公式是 硬件峰值算力 ÷ 硬件峰值带宽。比如一块典型的现代 AI 加速卡,FP16 稠密算力 989 TFLOPS,HBM 带宽 3.35 TB/s,那 Ridge = 989e12 / 3.35e12 ≈ 295 FLOPs/Byte。

  • 如果你的 kernel AI < 295,就跑在访存受限区,再怎么优化指令排布也没用,得想办法减少 DRAM 流量,相当于给房子“减重”。
  • 如果 AI > 295,就跑在计算受限区,优化的重点变成压满 FMA 流水线、减少指令开销,相当于给工人“提速”。

Tiled GEMM 设计里很多看似玄学的选择,在这个框架下一目了然。比如 FP32 换成 TF32/BF16 或 FP16 后 type_bytes 减半,AI 直接翻倍,这就是为什么混合精度在同等硬件上能获得接近两倍的乘加吞吐。再比如同样一个 shape,BK 从 8 提到 16 后 AI 有所提升,但共享内存占用同步翻倍,可能使双缓冲放不下,这就是典型的需要用模型做权衡的场景。

3.3 共享内存、寄存器和延迟隐藏建模

Roofline 模型只能判断“DRAM 这一层”的瓶颈,真正写 kernel 时,还要建模“片上”的资源瓶颈:共享内存带宽、寄存器文件读写、warp 调度延迟。

拿共享内存举例。如果代码写成传统三层循环,内层每次从共享内存取一个 A 元素和一个 B 元素,再对 TM×TN 个累加器做乘加,那么每 TM×TN 次乘加就要读 2 次共享内存。假设 warp 内 32 个线程同时执行,SM 的共享内存 load 带宽是有限的,很快你就能算出共享内存会成为瓶颈。解法是把从共享内存取值提升到寄存器:每处理一个 kk,先一次性把这一行 A 的 TM 个元素和 B 的 TN 个元素搬进寄存器,再做 TM×TN 次乘加,相当于把访存量从 2×(TM×TN) 降到 TM+TN。这就是“寄存器分块”在性能模型中的意义。

延迟隐藏也类似。GPU 靠大量并发 warp 来掩盖访存延迟,一个 warp 等数据的时候,调度器立刻切换到另一个 warp。所以性能和占用率直接相关。如果寄存器占用太高,导致一个 SM 只塞得下一个 block,那共享内存加载阶段所有 warp 都在等,计算单元就空转了。善用性能模型的做法是:先按寄存器数、共享内存容量分别算出理论最大驻留 block 数,取较小值,再看是否足够让每个 SM 至少驻留 4~6 个 block 以掩盖延迟。

4. 实操案例:用模型推导一套 Tiled GEMM

4.1 硬件假设与数字准备

为了把方法论落地,我假设一个“典型加速卡”环境来做演示。以下参数故意简化,实际工程上请用 nvidia-smi、Nsight Compute 或者厂商文档替换成真实数值:

参数数值
SM 数量140
SM 内 FP32 FMA 吞吐128 次/时钟
主频1.5 GHz
SM 峰值算力384 GFLOPS = 1.5G×128×2
DRAM 带宽1 TB/s
L2 带宽5 TB/s
每 SM 共享内存上限228 KB
每 SM 寄存器文件256 KB
每线程寄存器上限255 个

目标是设计一个 FP32 的 Tiled GEMM kernel,用来算一个中等规模的矩阵乘:M=N=K=4096。

4.2 从全局到 tile:推导设计参数

第一步,先算全局算术强度。总计算量是 2×4096³ ≈ 137.4 GFLOPs。如果 A、B、C 都是 FP32,每个矩阵 67MB,三者约 201MB。AI ≈ 137.4e9 / 201e6 ≈ 683 FLOPs/Byte。这个数值大于我们假设的 Ridge(约 295),所以从 DRAM 角度讲,这个 GEMM 是计算受限的,理论上限约为 140×384 = 53.8 TFLOPS。

第二步,选 tile 参数。我以 BM=128、BN=128、BK=8、TM=8、TN=8 为例继续推演:

  • block 内线程数 256,每线程 64 个累加器,寄存器约 100 个,占寄存器文件约 100×256×4 = 102KB。一个 SM 的 256KB 寄存器文件能容纳约 2 个 block。
  • 共享内存双缓冲 16KB,228KB 上限能容纳很多 block,所以寄存器是驻留上限的约束。每个 SM 驻留 2 个 block,即 512 线程,对于掩盖延迟来说偏少,可以考虑减小 BM/BN 到 64×128,让 block 更小、驻留更多。
  • 全局内存载入的 tile 数量:M=N=K=4096,按 BK=8 切分 K 维度共 512 次外循环。主循环中每个 block 处理 128×128 的 C 分块,所以 M、N 各自切 32 份,共 1024 个 block,跑在 140 个 SM 上大概每个 SM 要处理 7~8 个 block,负载相对均匀。

第三步,用模型预估一个 block 的计算耗时。LM 块总 FLOPs = 262144,SM 峰值 384 GFLOPS,理论耗时约 0.68 微秒。配合双缓冲,理想情况下共享内存加载可以和计算完全重叠,那 512 次外循环、每次 0.68 微秒也就是 350 微秒左右。但寄存器占用限制了并发度,实际调度器需要在每个 SM 的 2 个 block 之间切换,延迟隐藏效果需要实测验证。

你会发现,整套推导没有一行代码,却已经把“选多大 tile、驻留多少 block、理论时间多少”全摆出来了。这正是性能模型最重要的使用场景:先在脑子里把方案逼到极限,再动手。

4.3 伪代码与落地时的三个注意点

模型只是指导,最终还得落到实现。下面给一个简化的 CUDA 风格伪代码,展示寄存器分块和双缓冲的基本骨架:

// 线程块负责计算 C[BM][BN],每线程 TM×TN 个累加器 // BM=BN=128, BK=8, TM=TN=8, blockDim = (16, 16) __shared__ float As[2][BK][BM]; // 双缓冲 A 切片 __shared__ float Bs[2][BK][BN]; // 双缓冲 B 切片 float acc[TM][TN] = {0}; int cur = 0; for (int k0 = 0; k0 < K; k0 += BK) { // 预加载第 k0+BK 个切片到共享内存(可与计算重叠,实际用 cp.async) load_A_tile(As[cur ^ 1], A, k0 + BK); load_B_tile(Bs[cur ^ 1], B, k0 + BK); __syncthreads(); // 等当前块的共享内存数据就绪 // 计算当前切片 for (int kk = 0; kk < BK; kk++) { float a_reg[TM], b_reg[TN]; // 从共享内存读入寄存器,注意边界、bank conflict for (int i = 0; i < TM; i++) a_reg[i] = As[cur][kk][thread_y * TM + i]; for (int j = 0; j < TN; j++) b_reg[j] = Bs[cur][kk][thread_x * TN + j]; // 累加 for (int i = 0; i < TM; i++) for (int j = 0; j < TN; j++) acc[i][j] += a_reg[i] * b_reg[j]; } cur ^= 1; __syncthreads(); // 防止覆盖共享内存时和计算冲突 } store_C_tile(C, acc);

三个注意点,都是模型算不出来但实测必踩的坑:

第一,共享内存 bank conflict。上面这段代码中,如果 As 数组按行存储,warp 内相邻线程访问的是同一 kk 下的不同行,但它们的列索引相同,很容易产生几个线程命中同一 bank 的冲突。一个惯用解法是在共享内存声明的列维度加 1 个 padding,打乱地址映射。

第二,全局内存对齐。load_A_tile 和 load_B_tile 的访存模式必须让一个 warp 访问的地址尽量落在同一段连续地址上,否则无法向量化。每行起始地址要按 16 字节甚至 32 字节对齐,分块尺寸最好也是对齐宽度的整数倍。很多新手 kernel 慢不是算不快,而是 load 没合并。

第三,边界处理别用复杂分支。M、N、K 如果是 8 的整数倍,可以用干净的主循环;如果矩阵维度不能整除 tile,就要单独写边界清理 kernel,或者在主循环里加 if 判断。分支本身没多贵,但分支引起的 warp 发散和额外指令会让 FMA 流水线空转,性能模型很难预测这一点,靠经验更靠谱。

5. 常见问题与排查技巧实录

5.1 为什么模型算得挺好,实测上不去

我见过很多同行的第一版 Tiled GEMM 跑出来只有理论峰值的 30%,回来问我“模型是不是算错了”。模型通常没有错,错在它只覆盖了 DRAM 流量、共享内存容量、寄存器数量这几层约束,但真实硬件的坑远不止这些。

第一个最常见的因素是同步开销。每次 __syncthreads() 都有代价,如果 BK 设得太小,外循环次数特别多,同步次数也跟着变多,GPU 在排队等走上消耗的时间会非常可观。用模型算的时候可以顺手把同步次数乘个平均延迟,看看是不是占了整个执行时间的 10% 以上。

第二个是地址对齐和向量化问题。编译器没有把 load 向量化成 128 位指令时,你会发现指令数比理想情况多了几倍,算术管线根本喂不满。打开 SASS 或者 PTX 看一眼,确认每个线程的全局 load 是 LDG.E.128 级别而不是 32 位,这能省掉大量调试时间。

第三个是 L2 命中率。模型里假设每个数据从 DRAM 只读一次,但如果 tile 调度顺序不好,A 片段的某些数据会在 L2 里被多次驱逐和重取。解决办法是用一个简单的“块级循环顺序”思维:让 K 维度的循环尽量连续,避免不同 block 之间交叉访问同一段 B 数据。

5.2 问题排查速查表

现象可能原因检查手段常见解法
性能只有峰值 30% 左右全局 load 未合并或未向量化看 SASS,确认 LDG 宽度对齐每行地址,用 float4 读取
共享内存访问慢bank conflict用 profiler 看 shared load 冲突率数组加 padding,或调整线程与数据映射
寄存器占用过高,占用率低TM/TN 过大编译报告或 occupancy API减小 TM/TN,或改为 6×8 类非对称分块
算力高但 DRAM 流量超预期tile 调度顺序不合理Nsight 看 dram__bytes 总和调整 block 的遍历顺序,让数据尽量留在 L2
不同 shape 表现差异极大边界分支导致 warp 发散观察在小整除 shape 上是否正常主循环 + 单独边界 kernel
双缓冲没起作用预加载没有和计算重叠看 stall 是否为 long_scoreboard使用异步拷贝指令或在加载前做独立计算

这张表是我踩过多次坑之后整理的,不一定覆盖所有平台,但排查思路是通用的:先把一个极端简单的 shape(比如全整除的小矩阵)调好,再往外扩展,问题会好定位得多。

5.3 踩过几次坑之后的习惯

我现在的习惯是,拿到一个新优化任务,先花 20 分钟把硬件参数查清楚,把 M、N、K 代入公式算出 AI 和理论峰值,再往 Roofline 图上标一下,确定这是访存受限还是计算受限。然后才会考虑 tile 结构、寄存器分块和共享内存方案。

如果流程反过来,先写 kernel 再找瓶颈,最常见的结果是:代码改了一整版,发现瓶颈根本不在你改的地方。用性能模型最大的收益不是“预知精确性能”,而是“提前建立正确的预期”,让你在实测数据偏差时能立刻知道往哪个方向找问题。这种能力,恰恰是 AIInfra 工程师含金量的核心体现。模型不是万能的,真实硬件还有指令发射、功耗墙、时钟关系等一大堆复杂因素,但如果没有模型,你连该怀疑谁都不知道。

最后再分享一个小技巧:性能模型其实可以很轻量,不一定要搞什么复杂框架,一个 Python 脚本、一张纸、甚至一堆命令行计算就够。关键是养成“动手前先算清楚”的习惯。我很多次靠这个习惯在评审会上直接指出某个热门方案的不合理之处,省下了整整两周的开发试错成本。对做 AIInfra 的人来说,这会是你最有性价比的一项基本功。

RELATED READING

延伸阅读

更多一线实战笔记与深度复盘,助您持续精进