FlashSinkhorn: IO-Aware Entropic Optimal Transport
This work addresses a performance bottleneck for researchers and practitioners using entropic optimal transport in machine learning, though it is incremental as it optimizes an existing method rather than introducing a new paradigm.
The paper tackled the inefficiency of GPU solvers for entropic optimal transport at scale by introducing FlashSinkhorn, which rewrites Sinkhorn updates to enable fused kernels and tiling, achieving up to 32x forward-pass and 161x end-to-end speedups over state-of-the-art baselines on point-cloud optimal transport.
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/ot_triton.