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那条线是一样的。
最后是我修改后的文件文章来源：https://www.toymoban.com/news/detail-474839.html
文章来源地址https://www.toymoban.com/news/detail-474839.html
到了这里，关于OpenAI Triton 初探的文章就介绍完了。如果您还想了解更多内容，请在右上角搜索TOY模板网以前的文章或继续浏览下面的相关文章，希望大家以后多多支持TOY模板网！
