FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport

Felix X. -F. Ye; Yu Chin Fabian Lim; Naigang Wang; Davis Wertheimer

cs.AIOct 1, 2026

FlashSinkhorn 2: Block-Sparse Entropic Optimal Transport

Authors: Felix X. -F. Ye, Yu Chin Fabian Lim, Naigang Wang, Davis Wertheimer

Organizations: Department of Mathematics & Statistics, University at Albany, Albany, NY, USA · IBM T. J. Watson Research Center, Yorktown Heights, NY, USA

Abstract

Streaming GPU solvers for entropic optimal transport (EOT), such as FlashSinkhorn, avoid storing the dense kernel but still evaluate all $n\times 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\times10^8$ -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

Figures & tables

Figure 1: FlashSinkhorn 2 overview ( Algorithm 1 ), drawn with

B{=}4

. Panel (a) is the Morton order both stages share. Panels (b)–(d) are the coarse stage ( Section 3.1 ) and panels (e)–(g) the fine stage ( Section 3.2 ). Red outlines trace source block I4. Yellow and green mark candidate and admitted tiles. Purple denotes on-chip work, and hatched grey denotes HBM storage.

Figure 2: FS2 vs GeomLoss (GL). (a–d) Time of a complete solve on 26 problems:

n

from

2^{19}

2^{23}

(a,c) and

\varepsilon

from 0.001 to 0.02 at

n{=}2^{22}

(b,d). A cross marks a GL return with residual

R>\tau{=}0.005

. (e) Peak GPU memory allocated by PyTorch. (f) On clustered

n{=}2^{22}

, the time at which each new best full residual

R

is reached, for FS2’s potentials along the solve and for GL’s native returns. The residual decreases to the right. The dotted line marks

\tau

and the arrow marks the first fine update. (g,h) Change in the grouped objective

U

over twenty optimization steps at

n{=}2^{21}

, against cumulative end-to-end (E2E) time, including compilation and autotuning.

	FS2	GL	GL +	FS
Time (s)	15.44	21.45	421.10	311.98
coarse stage	2.94	0.005	0.022	–
fine stage	12.30	21.40	420.67	311.98
Fine updates	18	12	225	44
per update (s)	0.68	1.78	1.87	7.09
Memory (GiB)	2.09	0.26	0.26	0.20

Table 1: FS2 against GeomLoss (GL) and FlashSinkhorn (FS) on the clustered

n{=}m{=}2^{21}

\varepsilon{=}0.001

problem. Bold marks solves that meet

\tau{=}0.005

. In (a), stage times are GPU kernel time from a separate profiled call, with sampled checks counted in the fine stage. GL + is GL at scaling 0.99, continued at the target

\varepsilon

until its full residual first meets

\tau

Table 4

Figure 3: FS2 on Quijote. (a) Time of a complete warm call of FS2 at

\tau{=}0.05

. The cross marks GL’s instrumented call at

2^{24}

. (b) Peak GPU memory allocated by PyTorch, and the selected block size

B

. (c) Best sampled residual versus time. Diamonds mark the first fine update and dotted lines the targets. (d) Key results. Orange marks a return with residual above

\tau

, and RMSE is the displacement error of the

2^{27}

solution at

\tau{=}0.01

and of a zero-displacement predictor.

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

	$\delta_{\mathrm{rel}}$			$\nu_{i}$ over		$\\|T^{\mathcal{A}}_{\varepsilon}-T_{\varepsilon}\\|$	Gradient
Cloud	factor	Admitted	Max $\nu_{i}$	Thm. 1 bound	Max error	over $\nu_{i}\operatorname{diam}(\mathbf{Y})$	rel. $\ell_{2}$
Clustered	$1$	6.33%	$6.7\times 10^{-6}$	0.005	$2.7\times 10^{-6}$	0.48	$1.6\times 10^{-7}$
Clustered	$64$	5.32%	$3.7\times 10^{-4}$	0.005	$1.4\times 10^{-4}$	0.48	$1.3\times 10^{-5}$
Heterog.	$1$	60.54%	$2.9\times 10^{-5}$	0.023	$2.4\times 10^{-6}$	0.23	$9.1\times 10^{-6}$
Heterog.	$64$	49.24%	$3.0\times 10^{-3}$	0.038	$1.7\times 10^{-4}$	0.23	$1.8\times 10^{-3}$

Appendix

Table 4: Proposition 2 at FS2’s returned potentials,

n=m=2^{20}

\varepsilon=0.001

: largest values over all rows. The

\nu_{i}

column divides by the Theorem 1 bound

M_{B}\delta_{\mathrm{rel}}a_{i}/r_{i}

, the error column by

\nu_{i}\operatorname{diam}(\mathbf{Y})

, and the gradient compares

2a_{i}(x_{i}-T^{\mathcal{A}}_{\varepsilon}(x_{i}))

with

2a_{i}(x_{i}-T_{\varepsilon}(x_{i}))

Clustered
$n=m$	$\varepsilon$	FS2 time (s)	FS2 $R$	GL time (s)	GL $R$
$2^{19}$	0.001	5.638	0.004879	1.521	0.054122 †
$2^{19}$	0.01	0.290	0.002918	0.026	0.072730 †
$2^{20}$	0.001	4.176	0.004791	5.500	0.053802 †
$2^{20}$	0.01	0.275	0.003245	0.027	0.071918 †
$2^{21}$	0.001	15.593	0.004828	21.448	0.054161 †

Appendix

Table 5: The 32-problem forward grid. A dagger marks

R>0.005

. FS2 meets the target on all 32 problems and GeomLoss on 10, all heterogeneous. Each entry is one call. Table 1 a and Table 6 remeasure clustered

2^{21}

\varepsilon=0.001

as medians of three calls.

		Clustered, $\varepsilon{=}0.001$		Heterogeneous, $\varepsilon{=}0.01$
Solver	Cells	Time (s)	$R$	Time (s)	$R$
FS2	auto	15.37	.00483	0.273	.00095
GL, scaling .90	default	21.46	.0542	57.80	.00084
GL, scaling .95	default	39.48	.0310	86.59	.00094
GL, scaling .99	default	200.38	.00687	404.03	.00020
GL, scaling .90	matched	0.361	.0597	0.391	.00076

Appendix

Table 6: GeomLoss settings on two

n{=}m{=}2^{21}

problems, truncation 5, rerun with FS2 in the same session. Bold marks solves that meet

\tau{=}0.005

. Default is GeomLoss’s cluster scale (430 and 2033 source cells), and matched gives as many source cells as FS2’s coarse stage, to within one (23,799 against 23,800, and 28,332). FS2 and matched times are medians of three calls, the others single calls.

R

is the full two-sided residual.

		$K{=}0$	Continued at the target $\varepsilon$
Problem	GL setting	$R$	$K^{*}$	Time (s)	$R$
$\varepsilon{=}0.001$	.90, default	.0542	–	–	.0203 at $K{=}512$
	.99, default	.00687	118	421.10	.004997
	.90, matched	.0597	–	–	.0160 at $K{=}2048$
$\varepsilon{=}0.01$	.90, default	.0731	–	–	.0372 at $K{=}512$
	.99, default	.0373	–	–	.0372 at $K{=}512$

Appendix

Table 7: GeomLoss continued at the target

\varepsilon

on clustered

n{=}m{=}2^{21}

, truncation 5, on the GPU of Table 6 .

K^{*}

is the smallest number

K

of appended target-

\varepsilon

entries whose full residual meets

\tau{=}0.005

. Time is the median of three clean calls at

K^{*}

, which at

\varepsilon{=}0.001

are the GL + calls of Table 1 a. Bold marks solves that meet

\tau

. Without a

K^{*}

, the last column gives the lowest scored residual. FS2 takes 15.44 s with

R{=}0.00483

\varepsilon{=}0.001

( Table 1 a) and, in the run of this table, 0.817 s with

R{=}0.00267

\varepsilon{=}0.01

Problem	$n$	$\varepsilon$	$\tau$	Mask	$R$	$R_{\mathcal{A}}$	$\Omega$	$R_{\mathcal{A}}+\Omega$
Clustered	$2^{19}$	0.001	0.005	sparse	0.004879	0.004879	$4.7{\times}10^{-6}$	0.004884
Clustered	$2^{19}$	0.01	0.005	sparse	0.002918	0.002918	$2.7{\times}10^{-5}$	0.002945
Clustered	$2^{20}$	0.001	0.005	sparse	0.004791	0.004791	$4.7{\times}10^{-6}$	0.004796
Clustered	$2^{20}$	0.01	0.005	sparse	0.003245	0.003246	$2.5{\times}10^{-5}$	0.003271
Clustered	$2^{21}$	0.001	0.005	sparse	0.004828	0.004828	$4.8{\times}10^{-6}$	0.004833
Clustered	$2^{21}$	0.002	0.005	sparse	0.004094	0.004094	$7.5{\times}10^{-6}$	0.004101

Appendix

Table 8: The certificate

R_{\mathcal{A}}+\Omega

of Section 3.2 at returned potentials, with

R_{\mathcal{A}}

the larger of the row and column residuals of

\mathbf{P}^{\mathcal{A}}

in IEEE FP32 and

\Omega

the tile-sum bound on the omitted mass in FP64. Where the mask would admit more than half of the tiles (dense), the update is dense,

\Omega=0

and

R_{\mathcal{A}}

is the full residual.

R

is the recorded full residual, and bold certifies

R\leq\tau

Case	Control	FS2 (s)	Control (s)	Full $R$	Ratio	State
C, 0.001	No coarse stage	4.165	153.330	0.02135	–	misses target
	Dense fine	4.165	33.327	0.00479	8.001	qualified
	Box only	4.167	4.335	0.00479	1.040	qualified
C, 0.01	No coarse stage	0.275	183.132	0.00477	664.818	qualified
	Dense fine	0.276	0.275	0.00325	–	inactive
	Box only	0.276	0.275	0.00325	–	inactive

Appendix

Table 9: Component ablations. C and H denote clustered and heterogeneous inputs, followed by

\varepsilon

. The ratio (control/FS2) is given for active controls that meet the target.

(a) Extra coarse updates, clustered
Extra coarse	Fine cap	Fine updates	Time (s)	$R$	Ratio
2048 (native)	128	13 + 2	4.169	0.00479	1
64	128	128 + 3	24.355	0.01896	–
256	128	128 + 3	24.817	0.01421	–
64	8192	1267 + 38	250.235	0.00478	60.0
256	8192	1077 + 38	215.674	0.00477	51.6

Appendix

Table 10: Coarse-continuation and threshold controls at

n=m=2^{20}

\varepsilon=0.001

. Fine updates count ordinary updates plus terminal previews. (a) Clustered, with the number of extra coarse updates fixed. Ratios divide by the FS2 reference of the same run. (b) Every threshold control meets the target.

Metric	FS (dense)	GL (KeOps, FP32)	FS2 (TF32)
Point pairs per direction	$4.40\times 10^{12}$	$2.58\times 10^{11}$	$2.58\times 10^{11}$
NCU time per direction (ms)	3,600	393.957	251.435
Warm update, both directions (ms)	–	791.073	504.392
Warp instructions / pair	0.34	0.628315	0.278530
Shared-load wavefronts / pair	–	0.0625001	0.0246587
Shared-store wavefronts / pair	–	0.0009806	0.0053709

Appendix

Table 11: Kernel counters for one fine update per direction on clustered

n=m=2^{21}

\varepsilon=0.001

, on the profiling A100. FS replays one dense reduction per direction from its native incoming state, with its counters averaged over the two directions. GL and FS2 run FS2’s first ordinary fine update on the same

B=128

mask, and

f/g

gives the two directions. Per-pair counts divide totals by the point pairs they cover, instructions are executed warp instructions, and dashes mark counters not collected.

Cloud	$\varepsilon$	$U_{0}$	FS2 s	GL s	FS2 decrease	GL decrease
Heterogeneous	.001	.004297	52.80	3215.67	.5829%	.5826%
Heterogeneous	.01	.01501	51.66	1290.01	.2551%	.2550%
Clustered	.001	.05230	79.81	453.22	.3064%	.2906%
Clustered	.01	.08166	42.58	0.80	.1756%	.1682%

Appendix

Table 12: Twenty-step optimization at

n=m=2^{21}

. A decrease is the percentage drop of

U

from the common initial score

U_{0}

over the twenty steps, and final scores within

3\times 10^{-6}

tie. Times cover the twenty updates, including compilation and autotuning.

method	$n=m$	$\tau$	configuration	time (s)	peak (GiB)	light validation $[U_{3\mathrm{SE}}]$	full max
FS2	$2^{24}$	.05	$B64/t47/c128/h80$	146.905 W	12.053	.046836 [.048531]	.046631 ✓
FS2	$2^{25}$	.05	$B64/t45/c128/h48$	374.839 W	19.472	.049304 [.049758]	—
FS2	$2^{26}$	.05	$B128/t45/c128/h72$	694.630 W	33.131	.048340 [.048783]	—
FS2	$2^{27}$	.05	$B256/t45/c128/h80$	1351.589 W	64.129	.049009 [.050797]	.048965 ✓
FS2	$2^{24}$	.01	$B64/t47/c800/h432$	736.393 W	—	.009538 [.009924]	.009433 ✓
FS2	$2^{27}$	.01	$B256/t45/c800/h560$	8832.994 F	—	.009585 [.009954]	.009154 ✓

Appendix

Table 13: Quijote forward solves, with the

\tau=0.05

size sweep in the first four rows.

B/t/c/h

: block size, cell level, fine cap and returned ordinary updates. Peak: PyTorch allocator peak of the size sweep. Light validation: mean absolute relative marginal error on 4096 sampled indices per direction, 65,536 at

2^{25}

and

2^{26}

, with the mean plus three standard errors in brackets. Full max:

R

, with dashes where it is not enumerated. W: warm call, F: first call, I: instrumented call. MS/.95/default/5: multiscale, scaling 0.95, default cluster scale, truncation 5.

Predictor	Particles	RMSE (Mpc/ $h$ )	Normalized RMSE	95% interval
Zero displacement	4096	10.064	1.000	–
FS2 barycentric displacement	4096	6.473	0.643	[0.625, 0.662]
no periodic-face crossing	4034	6.333	–	–
periodic-face crossing	62	12.603	–	–

Appendix

Table 14: Displacement reconstruction from the

2^{27}

-particle FS2 solution at

\tau=0.01

on 4096 sampled particles, with minimum-image scoring. The interval comes from 500 bootstrap resamples over particles at 95% confidence. The subsets split the sample by whether the reference trajectory crosses a periodic face.

Optimal Transport

Optimal transport compares and transforms probability distributions by finding low-cost ways to move mass between them.

See all 284 papers

SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport
Aug 28, 2026Ian Hsieh, Soumya Snigdha Kundu, Tom Vercauteren +1Optimal Transport Entropy Regularization
Sorbonne Université, Institut du Cerveau - Paris Brain Institute - ICM, CNRS, Inria, Inserm, AP-HP, Hôpital de la Pitié Salpêtrière, 75013 Paris, France · King’s College London, UK
Fast Log-Domain Sinkhorn Optimal Transport with Warp-Level GPU Reductions
Apr 4, 2026Hao XiaoOptimal Transport GPU Acceleration
ATLAS AI Lab
GPU-Accelerated Bregman Douglas-Rachford Splitting for Discrete Optimal Transport
Oct 3, 2026Yifan Xu, Shiqian MaOptimal Transport GPU Acceleration
Department of Computational Applied Mathematics and Operations Research, Rice University, Houston, TX, USA · Data Science and AI Institute, and Department of Applied Mathematics and Statistics, Johns Hopkins University, Baltimore, MD, USA

Aug 28, 2026J/K move · Enter open · S save

SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport

Ian Hsieh, Soumya Snigdha Kundu, Tom Vercauteren, Reuben Dorent

Sorbonne Université, Institut du Cerveau - Paris Brain Institute - ICM, CNRS, Inria, Inserm, AP-HP, Hôpital de la Pitié Salpêtrière, 75013 Paris, France · King’s College London, UK

Optimal Transport Entropy Regularization

Open paper AskPDF arXiv

Entropic optimal transport (EOT) has been shown to offer a computationally tractable approximation to exact optimal transport. However, the standard Sinkhorn-Knopp algorithm has two main limitations. First, given discrete measures with

N

points, each iteration requires

O(N^2)

operations, which restricts its use on large-scale datasets (e.g.

N\geq10^4

). Second, it uses the independent coupling as a reference measure for regularisation. This assigns mass to high-cost transport edges at moderate regularisation strengths. We propose SinkSLOT, which addresses both limitations by putting forth the expected sliced lifted transport plan as a natural way to sparsify the Gibbs kernel with a non-independent prior coupling. We prove that: 1) SinkSLOT converges; 2) with

L

slices, each resulting sparse Sinkhorn iteration costs

O(LN)

; and 3) the resulting objective is a divergence requiring no debiasing. Experiments on synthetic benchmarks show that SinkSLOT delivers substantial speedups over state-of-the-art dense and sparse EOT methods. We also demonstrate the applicability of the proposed divergence in a gradient flow experiment. The code is publicly available at https://github.com/cai4cai/SinkSLOT.

Abstract

Optimal Transport

Explore similar work

SinkSLOT: Sinkhorn via Sparse Lifted Optimal Transport

Fast Log-Domain Sinkhorn Optimal Transport with Warp-Level GPU Reductions

GPU-Accelerated Bregman Douglas-Rachford Splitting for Discrete Optimal Transport