FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

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

Many machine-learning tasks compare two collections of data points, such as cells from two single-cell experiments, image datasets, or the numerical representations produced by different AI models. Optimal transport is a principled way to make these comparisons, and Sinkhorn iterations are the standard scalable algorithm. But on GPUs (the specialized chips that train and run most modern AI), it becomes slow at large scale: the algorithm builds a huge table of all pairwise distances between the two collections and repeatedly shuffles it between the small, fast memory inside the chip and the slower memory outside. Much of the cost comes from moving this table through GPU memory, not from the arithmetic itself. We observed that the core step has the same mathematical structure as the "attention" mechanism inside modern transformer AI models. That let us adapt FlashAttention, a widely used technique for speeding up transformers, to optimal transport: stream small tiles of the table through fast on-chip memory and never write the full table out. Our implementation, FlashSinkhorn, is up to 32× faster per step and 161× faster end-to-end than state-of-the-art online GPU baselines, while using far less memory. We release the code as open source so researchers can reproduce the results and build on the method.