softmax

Implementation of Online normalizer calculation for softmax for learning purpose.

Benchmark on RTX 5090

To reproduce:

cmake -B build && cmake --build build && ./build/softmax_bench | python3 plot_bench.py -o bench.png

Future directions

• Implement Softmax + TopK fusion
• Add robust tests (FP numerical analysis)
• Study Pytorch's Softmax.cu

References

• Online normalizer calculation for softmax
• Milakov's implementation

Note

• Tests and plot_bench.py are Claude generated.