企业官网建设流程全解析

1. 为什么要在写 Tiled GEMM 之前先算一笔账

AIInfra 工程师的含金量，很多时候不是体现在你会不会调一个 70B 模型，而是体现在你能不能把一个看起来最没技术含量、人人都说“网上有现成库”的算子做到极致。Tiled GEMM 就是最典型的一个。大模型训练、推理的每一步都压在矩阵乘法上，而手写 GEMM 并不是把两层循环套上 shared memory 那么简单；真正决定性能的，是你在动键盘之前有没有把一笔账算清楚：每个 block 占多少寄存器、shared memory 能不能双缓冲、整个 grid 会产生多少 wave、每条 FMA 指令背后拖了多少条辅助指令。这篇笔记是 AIInfra 笔记系列第二篇，不展示最终优化完的完整代码，而是把“动手前怎么用性能模型选参数”的完整过程讲透，全程以 A100 + FP32 + M=N=K=4096 为例。适合已经跑通基础 CUDA、想从“照着教程改”升级到“自己定设计”的读者。

1.1 这篇笔记想解决什么问题

很多人第一次写 Tiled GEMM 时，都会面对一屏幕的玄学问题：tile size 到底选 64 还是 128？线程块用 128 还是 256？每个线程算 4x4、8x8 还是 8x4？为什么同样的代码在 4090 上表现不错，换到 A100 反而变慢？这些问题如果一个个去试，一轮调参就是编译加 profiling，半小时起步，试十组就是半天。性能模型的价值不是替你做决定，而是先帮你把不可能到达的区间划掉，把真正值得实验的候选变小。比如寄存器容量和 shared memory 容量能直接排除掉一批 tile；wave 数能预测 tail effect；每秒 FMA 指令数和辅助指令数的比例能提前判断一个 kernel 的天花板。把这些约束组合起来，你会发现好的设计不是“调出来的”，是“算出来的”。

1.2 一张“性能预算表”需要哪些输入

动手算之前，至少要拿到三组数字。第一组是硬件的执行带宽：这里的 A100 SXM 有几个关键参数，FP32 的 CUDA core 理论峰值是 19.5 TFLOPS，SM 数量是 108，每个 SM 的 shared memory 上限是 164KB，寄存器文件是 64K 个 32-bit 寄存器，DRAM 带宽大约 2TB/s。第二组是你的问题规模：M、N、K 分别多大，数据是 FP32 还是 FP16，这决定了总 FLOP 和最低内存流量。第三组是你打算用哪一种硬件路径：CUDA core 做 FP32 FMA，还是 Tensor Core 走 TF32/FP16，这两者的峰值差出一个量级，模型结论完全不同。这篇笔记先默认你写的是纯 CUDA core 的 FP32 SGEMM，不做 Tensor Core，因为这样能把 Tiled GEMM 的指令级问题看得更清楚。

2. 性能模型拆解：三层判断

2.1 第一层判断：算力密集型还是带宽密集型

先算全局的算术强度。M=N=K=4096 的 FP32 GEMM，总计算量是 2 * 4096^3 = 137.44 GFLOP。假设 A、B、C 三个矩阵都只在 DRAM 里过一遍，最小内存流量是 3 * 4096 * 4096 * 4 字节，约 201MB。算术强度就是 137.44e9 / 201e6，约 684 FLOP/byte。A100 的平衡点大约是 19.5 TFLOPS / 2TB/s，约 9.75 FLOP/byte。两者一对比，684 远大于 9.75，所以这个规模的 GEMM 是绝对的计算密集型，不是带宽密集型。这个结论非常重要：如果你的 kernel 实测只有 6 TFLOPS，不要急着优化 global memory 的读取次数，问题几乎肯定出在计算流水本身，比如指令发射、寄存器阻塞、bank conflict、wave tail。反过来，如果矩阵很小，比如 M=N=64、K=4096，算术强度会急剧下降，结论可能变成带宽瓶颈，那优化的思路就完全反过来了。

2.2 第二层判断：片上带宽和指令发射约束

全局 DRAM 不构成瓶颈，不代表 shared memory 和寄存器不构成瓶颈。每个 K 迭代，一个 TM x TN 的输出 tile 需要从 shared memory 读入 A 的 TMTK 个元素和 B 的 TKTN 个元素，产生的计算量是 2TMTNTK。忽略常数后，每读一个字节能换来多少 FLOP，可以粗略写成 TMTN / (2*(TM+TN))。如果 tile 是 64x64，这个比例是 16 FLOP/byte；如果 tile 是 128x128，比例是 32 FLOP/byte。这个数字说明，tile 越大，每次 shared memory load 能支撑的计算量越大，片上带宽压力的缓解越明显。但这只是分析的一半，另一半是指令发射。A100 每个 SM 每周期可以执行 4 条 warp 指令，但 FP32 FMA 要打到峰值，每周期需要 2 条 warp-FMA 指令，剩下只有 2 个指令槽给 load、store、地址计算和循环控制。如果你的代码里每条 FMA 平均还拖着一条甚至两条辅助指令，那理论峰值就要打对折。这是很多 kernel 性能卡在 50% 以下的核心原因，而性能模型能提前把这个天花板算出来。

2.3 第三层判断：wave quantization 和占用率

第三个容易被忽略的约束是 wave 数。一个 A100 有 108 个 SM，如果你的 block tile 是 128x128，那么 M=N=4096 时一共有 32*32=1024 个 block，平均每个 SM 要处理 9.48 个 block。实际上 block 是按整数轮分配的，某个 SM 可能分到 10 轮，最后一个 wave 只有部分 SM 在工作，这就是 tail effect。9.48 wave 对应的浪费大约是 10/9.48，即 5% 左右。如果 tile 是 64x64，block 数变成 4096，每 SM 平均 37.9 wave，ceil 之后几乎没有损失。但 tile 变小又会降低 2.2 节里的片上效率，所以这里存在一个真实的权衡，不是“越大越好”或者“越小越好”。另外还要看每个 SM 能同时放几个 block，shared memory 和寄存器都会限制 resident block 数。经验做法是至少让每个 SM 能放 2 到 4 个 block，这样 global load 延迟可以被 block 之间的调度掩盖，而不是完全依赖双缓冲。

3. 用模型把 A100 上的 Tiled GEMM 设计空间过一遍

3.1 先定理论时间，再定目标

用前面的数字，A100 在纯 CUDA core FP32 下的理论最短时间是 137.44e9 / 19.5e12，约 7.05ms。这不是一个可以真正达到的数字，因为任何指令级开销都会把它拉长。按照 2.2 节的判断，一个结构良好的 kernel 能把 FMA 有效发射率做到 70% 到 85%，所以合理目标区间是 8.5ms 到 10ms，换算成 TFLOPS 就是 13.7 到 16.2。如果实测连 5ms 都跑不到，说明你根本不是在逼近峰值，而是在某个结构性设计上出了问题。把这个目标写死在设计文档里，后面所有参数选择都围绕它来判断。很多人在这一步就会犯错误：没有先算 7.05ms，于是调参时根本不知道自己的 kernel 离理论极限还有多远，误以为 10ms 已经不错了，其实对于 A100 的 CUDA core 来说，10ms 只是及格线。

3.2 用容量约束筛掉候选 tile

把常见的候选方案放进模型里，一条条过容量和 wave 约束。128x128 的 tile 搭配 256 线程，每个线程负责 8x8=64 个输出，C 矩阵的寄存器就要 64 个，加上 A 片段和 B 片段以及地址索引，每个线程大约 90 个寄存器，一个 block 约 23000 个寄存器，两个 block 约 46000，三个 block 就超标了，所以 resident block 数最多是 2。shared memory 方面，如果 TK=32，单 stage 的 A tile 和 B tile 分别是 16KB，双缓冲是 64KB，两个 block 占 128KB，也刚好能放。这个方案看起来不错，但 wave 数是 9.48，tail 损失约 5%。

再看 64x64 tile，256 线程，每个线程只负责 4x4=16 个输出，寄存器需求很低，shared 双缓冲 32KB，resident block 数可以到 4 甚至更多，wave 数 38，tail 几乎为零。但代价是片上计算密度低，每条 shared load 对应的 FMA 少，辅助指令占比反而升高，实际很难冲到高位。

128x64 的不对称 tile 在这个模型下显得尤其有意思：block 数量是 32*64=2048，wave 数约 19，tail 只有零点几个百分点；寄存器方面每个线程负责 8x4=32 个输出，约 50 到 55 个寄存器，三个 block 也能放下；shared 双缓冲 48KB，三个 block 占 144KB，在 164KB 内。这个方案同时满足了寄存器不超标、shared 够放、wave 尾效应小、片上密度也不错，是我现在上手新 shape 时优先尝试的几个结构之一。

3.3 从模型到具体线程形状

把 128x64 的 tile 展开成线程映射。一个 block 内用 256 个线程，组织成 16x16 的形状，每个线程负责 8 行 4 列的输出 C 小片。沿着 K 方向，每个 K_TILE=32 的迭代里，每个线程需要从 shared memory 里拿到对应的 8 个 A 元素和 4 个 B 元素，然后做 8432 次 FMA 累加。注意这里的 A 片段和 B 片段不一定要全部放在寄存器里，可以在 K 迭代内部再分小段，比如每次取 4 个 A 和 2 个 B，这样寄存器压力更小，编译器调度也更友好。不要小看这个映射方式，它决定了后续指令的排布：C tile 越大，FMA 之间越容易互相隐藏延迟，但不能大到 255 个寄存器上限；C tile 越小，寄存器压力越小但辅助指令比例越高。128x64 + 8x4 的组合在 A100 上是一个经过多次验证的甜点区，这就是性能模型给出的方向，不是凭感觉猜出来的。

4. 让预算兑现的实现细节

4.1 向量化搬运与双缓冲

模型算完，接下来要让每一笔预算真正落实到代码里。第一个细节是搬运必须用 float4。global memory 一次读 16 字节，shared memory 一次写 16 字节，不仅减少 load/store 指令数量，也降低地址计算次数。很多新手在这里用 float 数组一个一个搬，同一份代码指令数多出三倍，性能直接掉一半。shared memory 的 A tile 和 B tile 最好声明成连续数组，并且 K_TILE 方向上保证 16 字节对齐。第二个细节是双缓冲的节奏：先提前用 cp.async 发起下一个 K_TILE 的拷贝，再计算当前 K_TILE，最后等待异步拷贝完成并同步。cp.async 是 Ampere 之后比较推荐的方案，它不会占用寄存器，也不会阻塞计算流水。需要注意的是，double buffer 的 stage buffer 大小会直接挤占 shared memory，模型里的 48KB 双缓冲已经是把这个因素算进去的。

4.2 计算核心与指令瘦身

计算循环内部，C 的结果一定要留在寄存器里，不要每算一个输出就写回 shared memory。常见的问题是有些人会在内层循环里反复读写 shared memory 的 C tile，导致 shared load 指令数量爆炸。正确的姿态是：每个线程的 8x4 C 小片全部用局部数组 c[8][4] 保存，内层只做 FMA 和从 shared 加载 A、B 片段，最后一次性写回 global memory。指令瘦身的另一个重点是地址计算。在 K 内层循环里，不要在每次迭代都重新算 shared memory 地址，应该用指向 As[stage][row][k] 和 Bs[stage][k][col] 的指针，每轮迭代做一次指针偏移。编译器在 FMA 足够多时会把地址计算调度到空闲指令槽里，但前提是地址模式足够简单，别引入乘除法。还有一个小技巧是循环展开，通常把 K_TILE=32 的内部循环完全展开，让编译器看到足够多的独立 FMA 来隐藏流水线延迟，同时让 cp.async 的发射和计算重叠得更自然。

4.3 别忘了 grid swizzle

最后一块容易超预算的是 L2 缓存行为。GEMM 的读写如果按 row-major 顺序访问，同一个 block 处理的两个输出 tile 在 L2 里的复用并不理想。一个简单的改进是给 grid 的 blockIdx 做 swizzle，比如把 32x32 的 block 网格按对角线或者按列交错映射，让相邻 block 访问的 A/B 分片在时间和空间上更接近，提高 L2 hit。这个问题在模型里不体现，但在实测中往往能带来 5% 到 15% 的收益。执行起来也不复杂，通常只需要在 kernel 开头用 blockIdx.x 和 blockIdx.y 做一次位运算或查表，把逻辑 tile 坐标重映射到物理 block 坐标。不要把它当成可有可无的优化，AIInfra 场景下矩阵规模动辄上万维度，L2 流量相当可观，swizzle 就是压住内存流量的最后一道闸。

5. 实测中的常见问题与排查

5.1 性能远低于模型预测

最典型的状况是，模型说你的设计天花板有 15 TFLOPS，实测却只有 5 TFLOPS。这时候先别怀疑硬件，按顺序排查。第一步看 ncu 报告里的 Actual Occupancy，如果远低于模型里的 resident block 预期，检查是不是编译器给了超过预期的寄存器数，导致占用被压低。第二步看 sm__inst_executed 里 FMA 指令的比例，如果 FMA 只占 30%，说明地址计算和 load/store 指令太多，回到 4.2 节做指令瘦身。第三步看 shared memory 的 bank conflict，方法我下一节细说。还有一个日常踩坑：很多人在 ncu 里看到 DRAM Throughput 很低，就以为是访存优化不够，但在计算密集型 GEMM 里 DRAM 低是正常现象，说明你的数据已经被 L2 和 shared 吃掉了，真正该盯的是 compute pipe utilization。

5.2 bank conflict 怎么定位

shared memory 的 bank conflict 在 GEMM 里是个经典问题。A100 的 shared memory 有 32 个 bank，每个 bank 每周期能返回 4 字节，同一个 warp 访问相同地址会广播，访问打散到不同 bank 就是一次生效，如果多个地址落在同一个 bank 就会冲突。典型错误是把 B tile 声明成 [TK][TN]，每行按 K 递增读取，结果 warp 内不同线程访问同一个 bank 的不同行，冲突率很高。解决办法是给 B tile 的每行末尾加 padding，或者把整个 tile 打成 swizzle 布局。判断是否命中，用 ncu 的 shared_load_throughput 和 smsp__warp_issue_stalled_shared 这类指标，如果 stall 占比高，基本就是冲突。不要凭感觉改，用 profiler 确认之后再调整。

5.3 wave tail 与小 shape 修正

大矩阵 4096 的 wave 尾效应只有 5% 左右，但实际业务里经常会遇到 256、512 这种中等规模，比如某些推理场景的 batch size 很小。如果 M、N 只有 512，你用 128x128 tile 的话 block 数只有 16 个，108 个 SM 里大部分空闲，wave quantization 直接吃掉一半性能。这时候性能模型要反过来用：要么把 tile 调小到 64x64 增加 block 数，要么干脆做 split-K，把一个 K 维度拆成多段分配给不同 block，最后再用原子加或者二次 kernel 合并结果。split-K 的代价是 C 的多次写回，会增加一小部分 DRAM 流量，但在小 shape 场景下收益远大于损失。这也是为什么不要背一套固定 tile 走天下，模型的价值就在于每次遇到新 shape 都能重新推导。

5.4 用 profiler 核对四项指标

拿到一张 ncu 报告，我通常会先看四项指标。第一项是 SM busy 或者 compute pipe utilization，它告诉你 SM 的计算单元有没有真正在工作。第二项是 achieved occupancy，和模型里的 resident block 预期对比，如果差距大，说明寄存器或者 shared memory 超预算了。第三项是 shared memory throughput，如果长期接近 100%，需要把 tile 调大来稀释 shared 访问，或者重新审视加载 pattern。第四项是 L2 hit rate，太低的话检查 grid swizzle 和 global load 的合并情况。不要一上来就看 DRAM throughput，计算密集型算子很容易被这个指标带偏。把 ncu 的输出分别记录为 baseline 和优化后，保存成本地笔记，这样下一次调参会有历史数据可以参考。

5.5 一张快速自检清单

我在写完一个 Tiled GEMM 版本后，会拿模型和实测做一轮快速交叉验证：先确认理论总 FLOP 和理论最短时间，再对照 ncu 里实际报告的 FMA 执行数，两者相差超过 10% 说明写错了或者编译器没按预期生成代码；然后确认 shared memory 占用和寄存器占用没有超过不同 resident block 数量的预算；最后用 wave 数估算 tail loss，并决定要不要调整 tile。这套流程走完，大部分性能问题都能定位到具体层级，而不是继续盲调。把这个自检清单保存在项目根目录，每次重写 kernel 都跑一遍，比任何优化技巧都更省时间。