Triton

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

最近在 AMDGPU 上优化用 Triton 实现的 Flash-Attention 算子，有一些优化手段值得记录下来。 通过调整 Block 发射顺序减少 SIMD 的 IDLE 时间 FA 的 Triton 实现中，将 Q 在 M 方向切分为了不同的 block。在前向过程中，如果 causal = True，那么 Q 只有左下三角的元素参与计载。即参与计算的元素在 M 方向从上到下逐渐增加。在默认的实现中，block 是从上到下按序发射的，即先发射负载小的块，再发射负载大的块。由于负载较大的块难以被分配到 SIMD 上，因此导致了较大的 SIMD IDLE。通过从倒序从下到上发射块，即先发射负载大的块，再发射负载小的块，由于负载小的块可以被更均衡地分配到各 SIMD 上，因此可以有效减少 SIMD IDLE。 先发射负载小的块，再发射负载大的块，导致较大的 SIMD IDLE 先发射负载大的块，再发射负载小的块，可以减少 SIMD IDLE 通过实现 chain-dot 减少对 LDS 的访存 在我们的硬件规范下，Q 和 K 矩阵乘的结果 QK 的 Layout 跟 Q 是不同的，因此需要先将 QK 的 Layout 转到跟 Q 一样才可以继续与 V 进行矩阵乘（和 Q 一样作为第一个操作数）。可以通过插入一些寄存器指令对线程之间的数据进行交换，以避免通过写入写出 LDS 来进行 Layout 的转换，这些指令（例如 bpermute, swizzle 等）的开销远小于 LDS 访存。 ...