GEMM

NT,TN,NN,TT 不同厂商的规则不一样，我们现在是把以下情况称之为 NN，也就是 A 是行优先，B 是列优先的时候。 再具体点就是 A 和 B 都是 K 方向连续的时候，称之为 NN。 当 A 或者 B 的连续方向改变，就变成了 NT, TT, TN。用下表总结一下，我把最后一个维度作为连续的维度。 上图对应的就是 id=0 的记录。 id A B 注释 0 [M, K] [N, K] NN 1 [M, K] [K, N] NT 2 [K, M] [N, K] TN 3 [K, M] [K, N] TT 从数据加载的角度来看，A 和 B 都是 K 方向连续的时候性能是最好的，因为这时可以尽量合并使用 b128 的指令。 而且对于减少 L1 的 set 冲突也是很有效的。 GROUPED-GEMM 用 NN 的数据布局来举例，A 的形状是 [M, K], B 的形状是 [NUM_GROUPS, N, K]，结果 C 的形状是 [M, N]。伴随 A 的有一个输入是 m_indicies，它是一个长度为 M 的一维数组，用来指定 A 的每一行使用 B 的哪一个 group 的权重来进行矩阵乘。 ...

THREADBLOCK-SWIZZLE (GROUP-M, GROUP-N) 下面这张图可以直观地解释 threadblock-swizzle 想要优化的方向。假设一共有 6 个 Compute Unit；左图是 raster along M 同时 swizzle=1，右图是 raster along M 同时 swizzle=2。 看左图，出一波 6 个输出块需要加载 7 个输入块（6A + 1B）；而右图，同样出一波 6 个输出块只需要加载 5 个输入快（3A + 2B）。如果不考虑 cache hit 的话，显然右侧对带宽的需求更低，性能更好。 这个也好理解，小学就学过，当面积相等的时候，正方形的周长最短，所以越接近方形的分组，需要加载的输入就越少。 不过在实际情况下，需要考虑 CU 的个数，L2 缓存的大小等因素，很难从第一性原理出发计算出最好的 raster 方向和 swizzle_size，一般是作为参数去 autoune。 比方说假设共有 6 个 CU，那么基本可以按照右图的顺序计算 tiles：第一波先计算 0-5，第二波计算 6-11，以此类推；但是如果有 12 个 CU，那么第一波此就可以计算 0-11，左图和右图就没有区别了。在这种情况下，我们会希望 CU/swizzle_size < M tiles，这样不同组的 tiles 确实会在不同的波次中计算。 PERSISTENT-KERNEL Persistent kernel 指的是每一个 cu 只起一个 threadblock, 这个 tg 一直执行同一个 kernel。比方说共有 4 个 cu，那么共起 4 个 tg，每个 cu 一个单独的 tg。 ...

实验数据 性能的总体影响 由于之前的 streamk 的实现，导致了约一半的 GEMM MNK 组合出现了性能的下降；因此，我把 data-parallel，splitk 和 streamk (dp + 1 streamk 的版本) 都加入了 triton 的 autoune 的范围。在四千多个 MNK GEMM 组合上进行了测试，结果如下： 可以看到 99.2% 的组合都是性能提升的，也出现了一些性能下降的。这里需要强调一下，这次的测试除了引入 streamk 之外，也做了一些数据精度上的调整。之前的 GEMM 在 K 方向的累加是用 float16 实现的，这个跟甲方确认过没有问题。但是在实现 streamk 的过程中，因为 k 方向的累加次数可能会很多，另外也是为了使得 GEMM 的精度更高，因此统一都更换到了 float32 进行累加，这个改动会使得性能下降很多（大概比 float16 精度下降约一半）。因此这里的性能的改变并不是单单由于 streamk 导致的，实际上如果单独考量的话，streamk 的提升肯定会更多。 下面这个分布图可以直观地看出来性能提升的分布，可以看到平均数和中位数都大约是 120%-130%，另外在 150%-200% 之间也有分布。最多的提升到原来的 266%，而最差的是原来性能的 46% （因为累加精度的提升）。 性能下降的奇怪 pattern 在第一张图中，注意到一个奇怪的现象，性能下降的 MNK 似乎遵循某种 pattern，排列得非常有规律。下图很明显地体现出来，当 M * N 等于三个数的时候，性能是有较多下降的。这张图中横坐标是 M * N， 纵坐标是性能百分比，虚线表示 100% 的位置。至于原因，我研究了半天也没有发现为什么，怀疑可能跟用 atomic_add 写出 float32 的结果矩阵有关系，但是没有完全解释清楚。 ...

原理 StreamK 这个算法的目的是将 GEMM 中的 K 方向的 MAC 均分到每一个 Compute Unit，使得所有 CU 可以同时结束，从而减少 idle 的时间。 GEMM 最常见的调度算法是每一个输出的块由一个 TG 负责计算，这个 TG 把 K 方向的累加全部算完了。这个算法以前是没有什么问题的，性能很好。但是现在 CU 性能越来越强大，就使得切块约切越大，对于同一个 MKN，需要用到的 TG 越来越少；而同时 CU 数量越来越多，这就使得很多情况下，CU 没有被打满，浪费了算力。 当然可以减少切块的大小，但是这样会使得 cache 的利用率降低。也有一种 split-k 的算法，将一个输出的块平均分给 N 个 TG 来计算，最后再累加到一起。Split-k 的切分是通过在 K 方向均匀地切分，每个 TG 计算相等的一部分的 K 方向的 MAC。这个算法提高了 CU 的利用率，但是不能完全消灭 CU 的 idle 的时间。Stream-k 的算法就是进阶版的 split-k，K 方向不进行均匀切分，而是按需分配，相切多少切多少，满足的目标就是所有 CU 分配到的 K 方向的 MAC 是完全相等的，这样就完全没有 idle 的浪费。 Xprof 显示 80 个 CU 的耗时基本是平均的。 ...