【GPU 算子工程】Softmax、LayerNorm 与逐元素融合
内容提要
本文讨论了softmax和LayerNorm等内存绑定算子的优化方法。softmax通过减去最大值来避免数值溢出，在线softmax则通过增量维护最大值和总和来减少遍数。LayerNorm采用Welford算法进行单遍均值和方差计算。逐元素融合技术将多个操作合并为一个kernel，显著减少访存次数，提高性能。整体优化策略强调数值稳定性和减少访存。
关键要点
-
softmax 和 LayerNorm 是典型的内存绑定算子，优化重点在于减少访存遍数。
-
softmax 通过减去最大值来避免数值溢出，标准做法是三遍遍历数据。
-
在线 softmax 通过增量维护最大值和总和，将遍数减少到两遍，保持数值稳定。
-
LayerNorm 使用 Welford 算法单遍计算均值和方差，避免了数值不稳定的问题。
-
逐元素融合技术将多个操作合并为一个 kernel，显著减少访存次数，提高性能。
-
优化策略强调数值稳定性和减少访存，在线 softmax 和 Welford 算法是关键技术。
延伸解读
内存绑定算子的优化策略
softmax和LayerNorm等算子属于内存绑定类型，优化的关键在于减少访存遍数。通过在线softmax和Welford算法，可以将多次遍历数据减少到两次或一次，从而提高计算效率。这种优化策略在深度学习模型中尤为重要，能够显著提升整体性能。
逐元素融合技术的优势
逐元素融合技术通过将多个操作合并为一个kernel，减少了访存次数，提升了性能。例如，将scale、bias和GELU操作融合后，访存流量从约6次降至2次，实测提速接近3倍。这种技术在深度学习框架中被广泛应用，尤其是在PyTorch和JAX中，成为重要的图优化手段。
数值稳定性的重要性
在优化softmax和LayerNorm时，数值稳定性是首要考虑因素。通过减去最大值和使用Welford算法，可以有效避免数值溢出和不稳定问题。这些技术不仅提高了计算的准确性，也为后续的深度学习模型提供了更可靠的基础。
延伸问答
softmax的数值稳定性如何实现？
softmax通过减去最大值来避免数值溢出，标准做法是三遍遍历数据，但在线softmax通过增量维护最大值和总和，将遍数减少到两遍。
LayerNorm的Welford算法有什么优势？
Welford算法能单遍同时计算均值和方差，数值上比传统方法更稳定，避免了灾难性抵消的问题。
逐元素融合技术如何提高性能？
逐元素融合技术将多个操作合并为一个kernel，显著减少访存次数，从而提高性能。
在线softmax的增量重标定技巧是什么？
在线softmax通过在遍历新元素时同时更新running max和running sum，减少了遍数并保持数值稳定。
如何减少softmax和LayerNorm的访存遍数？
通过在线softmax和Welford算法，分别将softmax的遍数减少到两遍和LayerNorm的遍数减少到一遍。
逐元素融合对内存绑定算子的影响是什么？
逐元素融合显著减少了内存访问次数，从而提高了内存绑定算子的执行效率，实测提速接近3倍。
