FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

TL;DR

FlashSinkhorn introduces an IO-aware GPU solver for entropic optimal transport that significantly accelerates computations by fusing operations and reducing memory traffic, enabling scalable applications.

Contribution

It presents a novel IO-aware Sinkhorn algorithm with fused kernels and on-chip streaming, improving efficiency and scalability for optimal transport on GPUs.

Findings

Achieves up to 32x speedup in forward-pass

Achieves up to 161x end-to-end speedup

Enhances scalability for OT-based downstream tasks

Abstract

Entropic optimal transport (EOT) via Sinkhorn iterations is widely used in modern machine learning, yet GPU solvers remain inefficient at scale. Tensorized implementations suffer quadratic HBM traffic from dense $n\times m$ interactions, while existing online backends avoid storing dense matrices but still rely on generic tiled map-reduce reduction kernels with limited fusion. We present \textbf{FlashSinkhorn}, an IO-aware EOT solver for squared Euclidean cost that rewrites stabilized log-domain Sinkhorn updates as row-wise LogSumExp reductions of biased dot-product scores, the same normalization as transformer attention. This enables FlashAttention-style fusion and tiling: fused Triton kernels stream tiles through on-chip SRAM and update dual potentials in a single pass, substantially reducing HBM IO per iteration while retaining linear-memory operations. We further provide streaming kernels for transport application, enabling scalable first- and second-order optimization. On A100 GPUs, FlashSinkhorn achieves up to $32\times$ forward-pass and $161\times$ end-to-end speedups over state-of-the-art online baselines on point-cloud OT, improves scalability on OT-based downstream tasks. For reproducibility, we release an open-source implementation at https://github.com/ot-triton-lab/flash-sinkhorn .