# FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

arXiv: https://arxiv.org/abs/2602.03067 · HTML: https://mlsanity.com/paper/2602.03067/

## Key facts

- Cited by 11 papers in the map (data through 14 Oct 2026).
- Builds on 37 of its 56 references that are in the map.
- First posted on arXiv 3 Feb 2026.

## Builds on (most cited first, 30 of 37)

- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [FlashAttention-2: Faster Attention with Better Parallelism and Work Partitioning](https://mlsanity.com/paper/2307.08691/) (2023)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Computational Optimal Transport](https://mlsanity.com/paper/1803.00567/) (2018)
- [How to Escape Saddle Points Efficiently](https://mlsanity.com/paper/1703.00887/) (2017)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [Escaping From Saddle Points --- Online Stochastic Gradient for Tensor Decomposition](https://mlsanity.com/paper/1503.02101/) (2015)
- [Near-linear time approximation algorithms for optimal transport via Sinkhorn iteration](https://mlsanity.com/paper/1705.09634/) (2017)
- [Fast Computation of Wasserstein Barycenters](https://mlsanity.com/paper/1310.4375/) (2013)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Geometric Dataset Distances via Optimal Transport](https://mlsanity.com/paper/2002.02923/) (2020)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [Optimal Transport Tools (OTT): A JAX Toolbox for all things Wasserstein](https://mlsanity.com/paper/2201.12324/) (2022)
- [Entropic estimation of optimal transport maps](https://mlsanity.com/paper/2109.12004/) (2021)
- [Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems](https://mlsanity.com/paper/1610.06519/) (2016)
- [Escaping Saddles with Stochastic Gradients](https://mlsanity.com/paper/1803.05999/) (2018)
- [A Unified Framework for Implicit Sinkhorn Differentiation](https://mlsanity.com/paper/2205.06688/) (2022)
- [Dataset Dynamics via Gradient Flows in Probability Space](https://mlsanity.com/paper/2010.12760/) (2020)
- [TritonForge: Profiling-Guided Framework for Automated Triton Kernel Optimization](https://mlsanity.com/paper/2512.09196/) (2025)
- [Sinkhorn Natural Gradient for Generative Models](https://mlsanity.com/paper/2011.04162/) (2020)
- [ESPFormer: Doubly-Stochastic Attention with Expected Sliced Transport Plans](https://mlsanity.com/paper/2502.07962/) (2025)
- [Scaled-Dot-Product Attention as One-Sided Entropic Optimal Transport](https://mlsanity.com/paper/2508.08369/) (2025)
- [A Hypergradient Approach to Robust Regression without Correspondence](https://mlsanity.com/paper/2012.00123/) (2020)
- [A Truncated Newton Method for Optimal Transport](https://mlsanity.com/paper/2504.02067/) (2025)
- [Coreset selection for the Sinkhorn divergence and generic smooth divergences](https://mlsanity.com/paper/2504.20194/) (2025)
- [LOTFormer: Doubly-Stochastic Linear Attention via Low-Rank Optimal Transport](https://mlsanity.com/paper/2509.23436/) (2025)
- [You Need Better Attention Priors](https://mlsanity.com/paper/2601.15380/) (2026)

## Built on by (most cited first, 11 of 11)

- [TileMaxSim: IO-Aware GPU MaxSim Scoring with Dimension Tiling and Fused Product Quantization](https://mlsanity.com/paper/2606.26439/) (2026)
- [Block-Wise Differentiable Sinkhorn Attention: Tail-Refinement Gradients with a Gap-Aware Dustbin Bridge](https://mlsanity.com/paper/2605.08123/) (2026)
- [ASAP: Amortized Doubly-Stochastic Attention via Sliced Dual Projection](https://mlsanity.com/paper/2605.12879/) (2026)
- [Spherical Harmonic Optimal Transport: Application to Climate Models Comparisons](https://mlsanity.com/paper/2605.18389/) (2026)
- [From Cone Geometry to Monge Structure: Local Defects and Global Regret in High-Dimensional Optimal Transport](https://mlsanity.com/paper/2606.04695/) (2026)
- [DrainSinkhorn: Safe Elimination for Batched Entropic Optimal Transport](https://mlsanity.com/paper/2607.24741/) (2026)
- [SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport](https://mlsanity.com/paper/2608.28262/) (2026)
- [AI as a Compiler: Compiling Triton kernels without the Triton compiler](https://mlsanity.com/paper/2609.36800/) (2026)
- [A Unified Dual Method for Matching Problems](https://mlsanity.com/paper/2609.39339/) (2026)
- [FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport](https://mlsanity.com/paper/2610.02395/) (2026)
- [Scaling Online Optimal Transport Minibatches for Flow Matching](https://mlsanity.com/paper/2610.18838/) (2026)

## Abstract (arXiv)

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× 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× forward-pass and 161× 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 .