triton之softmax

一 计算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的各行是独立的，所以我们在这些行上进行并行处理