OpenAI Triton 初探

Triton 2021年发布了1.0，我在调研GPU使用方法的时候知道了有这个东西，但是当时还不了解OpenAI，觉得这个项目太新太小众，并没有深究。现在GPT大火之后，再回过头看看他们的这个东西。

现在相关文档还是很少，pip安装后发现版本已经默默升到了2.0.0.post1。

Triton的概念模型相较于cuda的模型，把基于Thread的模型扩展为基于block的。我觉得从线性代数的观点来看，就是把基于矩阵元素的粒度，变成了分块矩阵的粒度。

要想把他用好，可能需要复习一下线性代数课程里关于分块矩阵的部分，分块矩阵对于简化问题讨论有很大帮助，理论价值仅限于此。但是对于实际的计算问题来说，分块矩阵就更有用了。在科学计算，计算方法，高性能计算等课程中，经常会把一个大矩阵分块之后放到集群的不同节点上去计算。我认为Triton的思路是类似的。

对于GPU这样的设备，每个SM是有局部缓存的。分块之后，数据可以保存在缓存中，可以大大提高命中率。

对于两个矩阵的乘法

C （M，N） =A （M，K） * B（K，N）

传统的写法是用三重循环

for(int m = 0; m < M; m++)
for(int n = 0; n < N; n++){
  float acc = 0;
  for(int k = 0; k < K; k++)
    acc += A[m, k] * B[k, n];
  C[m, n] = acc;
}

CUDA的做法通常是对c的每个元素分配一个Thread来计算

Triton的做法是分块，然后每个块对应一个program

#pragma parallel
for(int m = 0; m < M; m += MB)
#pragma parallel
for(int n = 0; n < N; n += NB){
  float acc[MB, NB] = 0;
  for(int k = 0; k < K; k += KB)
    acc +=  A[m:m+MB, k:k+KB]
          @ B[k:k+KB, n:n+NB];
  C[m:m+MB, n:n+NB] = acc;
}

伪代码从概念上看非常清晰

官方提供的案例 03-matrix-multiplication ，示例代码有一些问题。

autotune前两个配置在我的卡上有问题，需要注掉

通过help查看triton.testing.do_bench 的 quantiles 参数已经变成percentiles，但是文档还没更新

torch.allclose比对结果，两个矩阵大体相同，只有一个数差距比较大

最后的效果看起来还不错

另外我还试了cupy，结果和cuBLAS那条线是一样的。

最后是我修改后的文件