一 计算kernel
【BBuf的CUDA笔记】十三,OpenAI Triton 入门笔记一 - 知乎 (zhihu.com)https://zhuanlan.zhihu.com/p/679232270我们的softmax kernel的工作方式如下:每个程序加载输入矩阵X的一行,对其进行归一化处理,然后将结果写回到输出Y中。需要注意的是,Triton的一个重要限制是每个块必须包含2的幂次方个元素,因此如果我们想处理任何可能的输入形状,我们需要在内部对每行进行“pad”以及对内存访问操作进行保护(也就是防止越界):
@triton.jit
def softmax_kernel(output_ptr, input_ptr, input_row_stride, output_row_stride, n_cols, BLOCK_SIZE: tl.constexpr):
# softmax的各行是独立的,所以我们在这些行上进行并行处理