PaperlayerView exact revision on arXiv

Source-linked AI summary

Speeding up Log-Sum-Exp: Kernel Fusion at the Memory Wall, Integer Arithmetic at the Compute Wall

Lingyun Yao, Martin Andraud, Niki Loppi, Andrea Pilzer, Anji Liu, Guy Van den Broeck, Martin Trapp

arXiv:2610.33647v3cs.DC

TL;DR

LSE is a widely used primitive for stable log-domain computation, but its efficient computation on generic CPUs and GPUs has been underexplored. The paper presents FuseLSE and IntLSE to address launch, memory-traffic, and arithmetic costs; both speed up GPU workloads over PyTorch, while IntLSE outperforms FuseLSE on CPU benchmarks.

  • Problem

    LSE is widely used for stable log-domain computation, but optimization on generic CPUs and GPUs has received limited attention despite PyTorch’s multi-kernel implementation and redundant HBM traffic.

    Source: §1 · p. 2

  • Method

    FuseLSE fuses the LSE sub-kernel chain, while IntLSE additionally replaces exp and log with integer approximations.

    Source: §1 · p. 2, §3 · p. 3

  • Results

    Both kernels achieve ∼10−11× speedup over torch.logsumexp at small GPU workloads and ∼6× at large workloads; on CPU, IntLSE is 3–6× faster than FuseLSE on the microbenchmark.

    Source: §1 · p. 2

  • Takeaways & Limitations

    For memory-bound GPU reductions, kernel fusion is the critical lever, while integer arithmetic wins on platforms without a fast SFU.

    Source: §6 · p. 10

    IntLSE’s integer approximation incurs cumulative error that may be unsuitable for the most accuracy-critical workloads, and reduced-precision variants are not integrated into application benchmarks.

    Source: §6 · p. 10

Abstract

from arXiv · show

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 PyTorch 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.

1 Introduction

LSE is widely used for numerically stable log-domain computation, but has received limited optimization attention on generic CPUs and GPUs. FuseLSE and IntLSE target GPU launch and memory overheads, while experiments report speedups over PyTorch across workload sizes and platforms.

Motivation

  • LSE supports stable log-domain computation in applications including LLM token log-probabilities, particle filtering, probabilistic inference, and entropy-regularized optimal transport.The cited examples include per-token LLM log-probabilities and particle filters.Source: §1 · p. 1
  • Despite broad use in machine learning, LSE is rarely an optimization target on generic CPUs and GPUs, and PyTorch can be ∼6× slower than torch.sum on the same GPU tensor.The cited comparison is 1.95 ms for torch.logsumexp versus 0.31 ms for torch.sum.Source: §1 · p. 2

Approach

  • FuseLSE fuses the sub-kernel chain to reduce launch overhead and HBM traffic, while IntLSE additionally replaces exp and log with integer approximations.The baseline makes roughly six input-sized HBM transfers per call across four full-tensor operations.Source: §1 · p. 2

Results

  • ∼10−11× speedups at small workloads and ∼6× at large workloads over torch.logsumexp were reported for both kernels on GPUs; IntLSE was 3–6× faster than FuseLSE on CPUs.The GPU results were confirmed by single-kernel benchmarks and three ML applications; CPU comparisons also report IntLSE as 2.8–5.2× faster than torch.logsumexp for particle filters.Source: §1 · p. 2

2 Preliminaries

LSE is used to aggregate probabilities or scores in the log domain, including in RLHF, probabilistic circuits, and particle filtering. The section describes PyTorch’s multi-kernel baseline and existing approaches that fuse operations or compute LSE in a streaming pass.

2 Preliminaries

  • LSE is used in RLHF log-probability calculations, probabilistic-circuit inference, and particle filtering, where it can be called across large vocabularies, deep computation graphs, or repeated filter steps.The cited settings include vocabulary sizes from 32k to 256k in RLHF, and millions of LSE evaluations in a deep probabilistic circuit.Source: §2 · p. 2

2.1 Baseline: PyTorch’s torch.logsumexp

  • At (B, N) = (512, 128256), torch.logsumexp expands its six-step max-subtract recipe into nine CUDA sub-kernels, with four operations accounting for six input-sized HBM transfers.The caption profiles the call on a V100-PCIe-32 GB and identifies the four heavy operations as amax, sub, exp, and sum.Source: §2.1 · p. 3, Figure 1 · p. 3

2.2 Existing accelerating methods

  • torch.compile reduces HBM traffic from roughly 6× to 2× the input size on the profiled workload, while online softmax computes LSE in one streaming pass by carrying a maximum and exp-sum.Other systems fuse LSE with upstream matrix operations, while standalone settings may still need an optimized LSE primitive.Source: §2.2 · p. 3

3 Implementation of the proposed LSE acceleration kernels

The proposed kernels target the launch and HBM costs of the baseline with a fused reduction, then offer an integer-approximation variant. FuseLSE uses pairwise LSE recursively; IntLSE replaces that pair operation with fixed-point integer arithmetic and a correction lookup table.

3 Implementation of the proposed LSE acceleration kernels

  • Fusing nine launches into one is intended to reduce HBM traffic from roughly 6× to 1× the input size, while replacing exp and log with integer arithmetic motivates IntLSE.These two opportunities motivate FuseLSE and IntLSE as replacement kernels.Source: §3 · p. 3

3.1 FuseLSE: Single-Pass Fused Kernel

  • FuseLSE recursively applies the associative, commutative pairwise LSE operation, processing each row in one HBM pass with SFU exp and log operations.Each thread maintains a register accumulator; the kernel reads each input once and writes one output scalar per row.Source: §3.1 · pp. 3–4

3.2 IntLSE: Approximate integer LSE

  • IntLSE approximates the pairwise correction with integer adds and shifts; its correction lookup table reduces maximum pointwise error from 0.085 to 0.0068.Inputs are converted to base-2 log space in Q16.16 fixed point, and the final integer result is converted back to float32 base-e.Source: §3.2 · pp. 4–5

4 Individual kernel experiments and benchmarking

The section benchmarks FuseLSE and IntLSE across GPU and CPU workloads, profiles GPU time, and evaluates reduced-bit GPU variants. It reports workload-dependent speedups and examines their performance and accuracy trade-offs.

4 Individual kernel experiments and benchmarking

  • The standalone experiments compare CUDA and CPU implementations across workload sizes and input ranges, with accuracy measured against an FP64 baseline.The GPU benchmarks use a V100; CPU implementations use a Xeon Gold 6134.Source: §4 · p. 5, §4.1 · p. 5

4.1 Kernel benchmark

  • On GPU, FuseLSE and IntLSE are ∼10−11× faster than torch.logsumexp for N ≤1024 and ∼6× faster for N ≥8192; on CPU, IntLSE is 3.1−6.1× faster than FuseLSE.At N = 128,256, GPU speedups are 6.24× for FuseLSE and 6.18× for IntLSE.Source: Figure 3 · p. 5, §4.1 · p. 5
  • FuseLSE matches torch at float precision (median relative error <10−4%), while LUT-corrected IntLSE error ranges from 0.04% to 2.38% at N = 128,256, depending on input range.Without the LUT, IntLSE’s large-N absolute error is ∼20−80× above the LUT-corrected variant on wider ranges.Source: §4.1 · p. 6
  • At small N, launch overhead explains the GPU speedup; at large N, reduced HBM traffic explains it, while CPU gains reflect the cost of exp and log versus integer arithmetic.IntLSE is faster than torch on all but the narrowest input range, even when using one thread against torch’s four.Source: §4.1 · p. 6

4.2 Cross-GPU time decomposition

  • On V100 and H100, total GPU time approaches the HBM lower bound at large N, and torch.logsumexp moves ∼6× more bytes than FuseLSE and IntLSE.IntLSE’s SM-compute exceeds FuseLSE’s by ∼50% on V100 and ∼26% on H100, but the discussion attributes GPU performance to HBM traffic rather than integer arithmetic.Source: Figure 4 · p. 6, §4.2 · pp. 6–7
  • FuseLSE is 6.3−6.9× faster than eager PyTorch and 1.8−2.1× faster than torch.compile and XLA across four GPUs.The compilers move 1.9−2.0× the input volume per call, compared with ∼6× for eager PyTorch and 1.01× for FuseLSE.Source: §4.2 · p. 7

4.3 Reduced-bit LSE

  • At N = 128,256, FuseLSE16, IntLSE-Q10.6, and FuseLSE reach 12.06×, 10.73×, and 6.25× speedups over torch.logsumexp, respectively; FuseLSE16 maintains mean absolute error below 3×10−3.IntLSE-Q10.6 error reaches ∼1% at that size; the authors favor FuseLSE16 on this V100 workload but do not pursue reduced-bit variants in downstream applications because surrounding operations run in FP32.Source: Figure 5 · p. 7, §4.3 · p. 7

5 Applications

The applications evaluate LSE kernels in GPU workloads for RLHF and probabilistic-circuit inference, and in a CPU particle filter. Results report faster execution against application baselines, with accuracy comparisons to torch.logsumexp or an FP64 reference.

5 Applications

  • The GPU applications replace LSE in vocabulary reductions and probabilistic-circuit sum layers, while the CPU application replaces NumPy LSE in particle filtering.FuseLSE is recommended for the two GPU applications; IntLSE is used for the CPU particle-filter benchmark.Source: §5 · p. 8, §5.1 · p. 8, §5.2 · p. 9

5.1 RLHF log-probability computation on GPU

  • 5.96–6.30× speedups over torch.logsumexp across four vocabulary sizes were achieved by both GPU kernels in RLHF SLS.The kernels were within ∼3.3% of each other; FuseLSE error was 1.1×10−6–1.9×10−5, while IntLSE error was 5.4×10−3–1.5×10−2. The benchmark uses a GPT-2 backbone and vocabularies from 32K to 256K.Source: §5.1 · pp. 8–9, Figure 6 · p. 8

5.2 Probabilistic circuit inference on GPU

  • ∼1.7× forward-pass speedups over PyJuice default mode were achieved by FuseLSE and IntLSE for all four latent counts K.FuseLSE error was 3.0–4.2×10−4; IntLSE error was 0.45–1.26 nats. The benchmark evaluates HCLT inference on MNIST at batch size 512.Source: §5.2 · p. 9, Figure 7 · p. 9

5.3 Particle filter on CPU

  • IntLSE was 9.1–15.4× faster than SciPy, 1.5–2.1× faster than NumPy, and 2.8–5.2× faster than PyTorch for CPU LSE calls.Its absolute error stayed ∼10−2 across particle counts; timings are primitive-level, not complete-filter runtime.Source: §5.3 · p. 10, Figure 8 · p. 10

6 Conclusion and discussion

FuseLSE and IntLSE replace torch.logsumexp, with reported speedups across microbenchmarks and three ML applications. The discussion attributes GPU gains to fusion reducing HBM traffic and identifies integer arithmetic as advantageous when fast SFUs are absent.

6 Conclusion and discussion

  • ∼6−12× speedups over torch.logsumexp were delivered by FuseLSE and IntLSE across microbenchmarks and three ML applications.The paper identifies kernel fusion as the critical lever for memory-bound GPU reductions, cutting redundant HBM traffic.Source: §6 · p. 10
  • On platforms without a fast SFU, computation becomes the bottleneck and integer arithmetic wins, while GPU LSE and softmax reductions are memory-bound.The broader discussion says redundant HBM traffic dominates memory-bound reductions and that fast integer pipelines should also be exposed on GPUs.Source: §6 · p. 10
  • IntLSE’s integer approximation incurs cumulative error that may be unsuitable for the most accuracy-critical workloads.Reduced-precision FuseLSE16 and IntLSE-Q10.6 variants are not yet integrated into application benchmarks.Source: §6 · p. 10

A.1 IntLSE Approximation-error analysis

The approximation-error analysis splits the expected accumulated error into an exact component and an approximation component. The latter is summed over periods.

A.1 IntLSE Approximation-error analysis

  • The expected accumulated error is split into an exact part, f(z) = log2(1+2^z), and an approximation part, f̂(z) = (1+Fz)2^Iz.The exact component is labeled (I), and the approximation component (II).Source: §A.1 · p. 14
  • The approximation component is summed over periods using t = z + n in [0, 1).Source: §A.1 · p. 14
  • Subtracting approximation component (II) from exact component (I) yields the next expression in the analysis.Source: §A.1 · p. 14

A.2 Additional experimental results

The appendix captions describe tables covering kernel variants, robustness, accuracy, CPU timing, GPU latency, and GPU profiling. They specify evaluation settings but do not provide the table cell results.

A.2 Additional experimental results

  • Table 2 organizes kernels by variant and arithmetic and identifies the sections where they are used.Source: Table 2 · p. 14
  • Table 3 covers edge-case handling, backward support, fallback rules, and robust-kernel overhead at specified GPU and CPU settings.Source: Table 3 · p. 14
  • Table 4 reports p95 absolute error versus FP64 across input distributions for FuseLSE on V100, IntLSE on CPU, and torch FP32 as reference.Source: Table 4 · p. 15
  • Table 5 reports CPU time across input distributions on Xeon Gold 6148 with one thread and defines underflow by values below -87.3 relative to the maximum.Source: Table 5 · p. 15
  • Table 6 reports LSE latency and FuseLSE speedups at (B, N) = (512, 128,256), excluding compilation, with unsupported B300 toolchains marked.Source: Table 6 · p. 15
  • Table 7 profiles GPU traffic, achieved DRAM and SFU throughput, and pipe stalls at (B, N) = (512, 128,256).Traffic is expressed relative to input size, and multi-kernel values are time-weighted.Source: Table 7 · p. 15

NeurIPS Paper Checklist

The checklist is presented as a required component addressing reproducibility, transparency, research ethics, and societal impact. The authors provide affirmative or not-applicable responses with justifications and point to sections for supporting details.

NeurIPS Paper Checklist

Loading 2610.33647v3…