超级会员免费看
本文介绍了如何使用CUDA的计算内核实现softmax操作，特别是针对OpenAI Triton的入门应用。内核处理过程中，每个线程块加载输入矩阵X的一行进行归一化，并将结果写入输出Y。由于Triton要求块大小为2的幂次方，因此需要对输入进行padding和边界检查。此外，强调了CUDA中load操作的特性，即加载tensor数据并涉及数据从全局RAM到SRAM的传输。
一 计算kernel
【BBuf的CUDA笔记】十三，OpenAI Triton 入门笔记一 - 知乎 (zhihu.com)
https://zhuanlan.zhihu.com/p/679232270我们的softmax kernel的工作方式如下：每个程序加载输入矩阵X的一行，对其进行归一化处理，然后将结果写回到输出Y中。需要注意的是，Triton的一个重要限制是每个块必须包含2的幂次方个元素，因此如果我们想处理任何可能的输入形状，我们需要在内部对每行进行“pad”以及对内存访问操作进行保护（也就是防止越界）：
订阅专栏 解锁全文
105
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
