STREAMK

实验数据 性能的总体影响 由于之前的 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 的耗时基本是平均的。 ...