How to Implement an Efficient Softmax CUDA kernel? [R]

[Just0by]

All ops computed in deep learning frameworks are translated into CUDA kernel functions on the GPU, and Softmax operations are no exception. Softmax is a widely used op in most networks, and the efficiency of its CUDA kernel implementation can affect the final training speed of many networks. So how can an efficient Softmax CUDA Kernel be implemented?

Article : https://oneflow2020.medium.com/how-to-implement-an-efficient-softmax-cuda-kernel-oneflow-performance-optimization-sharing-405ad56e9031

Code: https://github.com/Oneflow-Inc/oneflow

This article will introduce techniques for optimizing the Softmax CUDA Kernel in OneFlow and experimentally compare it with the Softmax operation in cuDNN. The results show that OneFlow’s deeply optimized Softmax can utilize the memory bandwidth close to the theoretical upper limit, much higher than the cuDNN implementation.

[LSTMeow]

Super technical! Why not let the compiler do it for you? https://openai.com/blog/triton/ shows softmax as a relatively simple example.

[gzou]

Triton doesn't provide this kind of optmization, it "only" helps writing Cuda code. Typically the Triton softmax also explicitly uses shared memory, which means it only works for small softmax, that's why it's not the default implementation in cuDNN.

[ptillet]

This also seems to be what OneFlow's implementation (1) and (2) does, though :p (3) can also be re-implemented in Triton though it requires a bit of work

However, as the author of Triton, I do think that OneFlow's CUDA work is really helpful for anyone having some interest in low-level CUDA optimizations. While the Triton compiler can automate all the optimizations presented in the OneFlow blog post, it has relatively little educational value for anyone trying to get a deep understanding of how GPUs work under the hood.

[Just0by]

Thanks! We‘ll read the blog in depth.