Faster log-sum-exp computation speeds up AI on GPUs and CPUs
Speeding up Log-Sum-Exp: Kernel Fusion at the Memory Wall, Integer Arithmetic at the Compute Wall
Distributed, Parallel, and Cluster Computing
Summary
Calculating the Log-Sum-Exp function is important for many machine learning tasks but can be slow on common computer processors like CPUs and GPUs. The authors propose two new ways to speed up this calculation: one combines steps on GPUs to reduce memory use and overhead, and the other uses simpler integer math on CPUs. Their approach runs much faster than the usual methods, making machine learning routines more efficient without changing existing software. This helps computers run AI tasks more quickly, especially when handling lots of data.
What this means in practice
- •For machine learning engineers: Accelerate critical log-domain calculations in AI models by replacing default GPU and CPU routines with faster fused or integer-based kernels.
- •For high performance computing teams: Reduce memory bottlenecks and improve integer arithmetic usage in scientific workloads requiring stable log-sum-exp computations on CPUs and GPUs.
Authors
Lingyun Yao, Martin Andraud, Niki Loppi, Andrea Pilzer, Anji Liu, Guy Van den Broeck, Martin Trapp
Abstract
While the Log-Sum-Exp (LSE) function is a foundational numerical primitive for stable log-domain computation in modern machine learning, its fast and efficient computation on generic platforms (\ie, CPUs and GPUs) has been underexplored, potentially creating execution bottlenecks for large workloads. On GPU, the de facto implementation torch.logsumexp dispatches a chain of nine sub-kernels per call, incurring redundant High Bandwidth Memory (HBM) traffic and launch overhead. On CPU, exp and log operations are hardware-hungry, limiting the computing efficiency. To address these bottlenecks, we propose two drop-in kernels to significantly accelerate LSE computation on both GPUs and CPUs: FuseLSE, a single-pass fused kernel that evaluates LSE's exp and log functions on the GPU's Special Function Unit (SFU); and IntLSE, which further replaces the LSE computation with approximated integer arithmetic for CPU computation. Our experiments, validated by benchmarks across workload sizes and hardware platforms, show that both kernels are ${\sim}10-11\times$ faster than PyTorch at small workloads and ${\sim}6\times$ faster on larger ones on GPUs, where the speed-up is attributed to reduced launch overhead and HBM traffic. On CPU, where no SFU is available, LSE computation dominates the cost, allowing IntLSE to gain ${\sim}3-6\times$ execution speed over FuseLSE.