从 tile、warp、CTA 到性能提升：一次 GPU kernel 优化复盘

先说明一下：这篇内容是我自己的技术整理，AI 只帮我做了少量语句润色。正文重点放在 kernel 调优的排查顺序和验证方法，不写未经测量的绝对加速数字。
最近整理了一个包含 K 维循环和二维输出 tile 的 GPU kernel。调优过程中，我把问题拆成三个层次：tile 负责数据复用，warp 负责线程内协作，CTA 负责整体资源平衡。
1. 先确认 baseline 的瓶颈
最初的实现比较直接：一个 CTA 负责一块输出区域，每个线程计算一个或几个输出元素，输入数据直接从 global memory 读取。
这种写法容易验证正确性，但同一个 tile 内的输入值会被多个线程重复加载。调优前我会先记录这些指标：

• kernel 平均耗时、P50 和 P95
• global memory throughput
• registers per thread
• shared memory per CTA
• achieved occupancy
• 访存是否连续、是否发生明显的 bank conflict

如果 profile 显示 global memory 访问占主要时间，应该先处理数据复用；如果计算单元已经接近饱和，继续扩大 tile 可能没有收益。
2. tile：先减少重复的 global load
第二步把输出划分成 block tile，常见参数是 BM、BN 和 BK：

1. BM = 输出 tile 的行数
2. BN = 输出 tile 的列数
3. BK = 每次处理的 K 维大小

复制代码

每个 CTA 负责一个 BM x BN 的输出块。输入的 A、B 子块由线程协作加载到 shared memory，然后在 K 维循环中重复使用。
典型流程是：

1. for k0 in range(0, K, BK):
2.     cooperative load A tile
3.     cooperative load B tile
4.     synchronize
6.     for k in range(BK):
7.         accumulate from shared memory
9.     synchronize

复制代码

这里最容易出错的是边界和访存布局：

• 相邻线程尽量访问相邻地址，让 global memory load 合并；
• 边界 tile 使用 mask，不能为了消除分支而读取越界地址；
• shared memory 的排列要检查 bank conflict；
• 两次同步要包住 tile 的写入和读取，不能只凭感觉删除 __syncthreads()。

Tile 变大以后，global memory 读取可能减少，但 shared memory 占用也会增加。能同时驻留的 CTA 数量下降后，整体性能反而可能变差，所以 BM、BN、BK 需要和 occupancy 一起看。
3. warp：让每个 warp 负责一个 micro-tile
完成 block tiling 后，再设计 warp 内部的工作分配。一个 CTA 可以拆成多个 warp，每个 warp 负责一个更小的 micro-tile，中间结果放在寄存器里。
可以先把映射写清楚：

1. CTA tile:    BM x BN
2. warp tile:   WM x WN
3. thread tile: TM x TN

复制代码

这样做的目标是让同一个 warp 内的线程以规律的方式读取 shared memory，同时减少中间结果反复写回。
这里有一个明显的取舍：thread tile 太小，索引和协作开销偏高；thread tile 太大，寄存器压力会上升，occupancy 下降。实际调参时，我会同时记录 registers/thread 和 kernel time，而不是只看理论上的每线程计算量。
4. CTA：线程数更大不一定更快
最后比较不同的 CTA 配置，例如 128、256 和 512 threads/CTA。每组配置都要一起记录：

• BM、BN、BK
• registers/thread
• shared memory/CTA
• active warps
• achieved occupancy
• kernel time

大 CTA 有时能减少调度和边界开销，但也可能带来更高的寄存器和 shared memory 使用量。对于某个输入 shape 最好的配置，换到另一个 shape 后可能就不再合适。
所以 CTA 调优更像资源平衡问题，而不是简单地把线程数调大。
5. 几个常见的失败尝试
Tile 继续放大
重复加载减少了，但 shared memory 占用过高，同时驻留的 CTA 变少，最终速度下降。
增大 thread tile
每个线程计算更多输出元素，看起来可以摊薄索引开销，但寄存器压力上升后，occupancy 下降，收益被抵消。
直接使用 vectorized load
对齐且整除的 shape 上可能有效，但边界 shape 需要额外处理。没有把 remainder 路径验证清楚前，这种改动很容易只对单个 benchmark 有效。
只看平均耗时
平均值容易掩盖抖动。更稳妥的做法是固定输入和 warmup 次数，重复多轮，同时记录 mean、P50 和 P95。
6. 我会怎样验证最终版本
性能方面，至少保留一张这样的表：

1. variant              mean       P50        P95
2. baseline              ...        ...        ...
3. block tiling          ...        ...        ...
4. warp micro-tile       ...        ...        ...
5. CTA tuning            ...        ...        ...
6. final                 ...        ...        ...

复制代码

正确性方面，要覆盖随机输入、边界尺寸、非整除的 K 维，以及不同 dtype。至少记录 max absolute error 和 max relative error，不能只看一次输出相等。
7. 总结
这类 kernel 优化最好按顺序推进：先用 profile 确认瓶颈，再用 tile 提高数据复用，然后调整 warp 的 micro-tile，最后平衡 CTA 的寄存器、shared memory 和 occupancy。
Tile 解决的是数据重复读取，warp 映射决定线程如何协作，CTA 配置决定资源能否被有效利用。三者需要放在同一个 benchmark 里比较，单独看某一个参数很容易得出错误结论。
大家在调 BM、BN、BK 时，通常是先依赖硬件经验，还是会用自动搜索？对于边界很多的 shape，大家更倾向于使用 masked load，还是单独维护 remainder kernel？