# FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 25 of its 41 references that are in the map.
- First posted on arXiv 1 Oct 2026.

## Builds on (most cited first, 25 of 25)

- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [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)
- [Flex Attention: A Programming Model for Generating Optimized Attention Kernels](https://mlsanity.com/paper/2412.05496/) (2024)
- [Online normalizer calculation for softmax](https://mlsanity.com/paper/1805.02867/) (2018)
- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](https://mlsanity.com/paper/2004.11127/) (2020)
- [Stabilized Sparse Scaling Algorithms for Entropy Regularized Transport Problems](https://mlsanity.com/paper/1610.06519/) (2016)
- [Low-Rank Sinkhorn Factorization](https://mlsanity.com/paper/2103.04737/) (2021)
- [Massively scalable Sinkhorn distances via the Nystr"om method](https://mlsanity.com/paper/1812.05189/) (2018)
- [Screening Sinkhorn Algorithm for Regularized Optimal Transport](https://mlsanity.com/paper/1906.08540/) (2019)
- [Multiscale Strategies for Computing Optimal Transport](https://mlsanity.com/paper/1708.02469/) (2017)
- [Quantized Gromov-Wasserstein](https://mlsanity.com/paper/2104.02013/) (2021)
- [A Sparse Multi-Scale Algorithm for Dense Optimal Transport](https://mlsanity.com/paper/1510.05466/) (2015)
- [Overrelaxed Sinkhorn-Knopp Algorithm for Regularized Optimal Transport](https://mlsanity.com/paper/1711.01851/) (2017)
- [Fast and Scalable Optimal Transport for Brain Tractograms](https://mlsanity.com/paper/2107.02010/) (2021)
- [Scalable Optimal Transport in High Dimensions for Graph Distances, Embedding Alignment, and More](https://mlsanity.com/paper/2107.06876/) (2021)
- [Accelerating Sinkhorn Algorithm with Sparse Newton Iterations](https://mlsanity.com/paper/2401.12253/) (2024)
- [FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU](https://mlsanity.com/paper/2602.03067/) (2026)
- [Hierarchical Refinement: Optimal Transport to Infinity and Beyond](https://mlsanity.com/paper/2503.03025/) (2025)
- [A note on overrelaxation in the Sinkhorn algorithm](https://mlsanity.com/paper/2012.12562/) (2020)
- [Efficient and Accurate Optimal Transport with Mirror Descent and Conjugate Gradients](https://mlsanity.com/paper/2307.08507/) (2023)
- [A Truncated Newton Method for Optimal Transport](https://mlsanity.com/paper/2504.02067/) (2025)
- [Domain decomposition for entropy regularized optimal transport](https://mlsanity.com/paper/2001.10986/) (2020)
- [Nonequispaced Fast Fourier Transform Boost for the Sinkhorn Algorithm](https://mlsanity.com/paper/2201.07524/) (2022)
- [cuRegOT: A GPU-Accelerated Solver for Entropic-Regularized Optimal Transport](https://mlsanity.com/paper/2605.08793/) (2026)

## Abstract (arXiv)

Streaming GPU solvers for entropic optimal transport (EOT), such as FlashSinkhorn, avoid storing the dense kernel but still evaluate all n× m point pairs in every Sinkhorn iteration. We present FlashSinkhorn~2 (FS2), a solver for squared-Euclidean cost on low-dimensional point clouds that solves large discrete EOT problems to a prescribed marginal residual on a single GPU by coupling two stages. A coarse stage solves on cell centroids, lifts the potentials to every point and, when a sampled marginal check rejects the lift, continues on the centroids, replacing most point-level updates. A block-sparse fine stage then removes the centroid error that coarse updates cannot. Its Morton-ordered blocks support screening and fused tensor-core execution, and a threshold set by the block masses bounds each omitted tile's contribution to every row and column. On synthetic benchmarks, FS2 reaches the target residual on all 32 problems and GeomLoss multiscale on 10. On one A100, FS2 solves discrete EOT between two 1.34×108-particle measures from a cosmological N-body simulation, at an entropic blur equal to the mean interparticle distance, to an all-particle marginal residual below 0.01 in under 2.5 hours. To our knowledge, it is the largest discrete EOT problem solved to this accuracy within hours. For reproducibility, we release an open-source implementation at https://github.com/ot-triton-lab/flash-sinkhorn