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
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
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 219 to 223 (a,c) and ε from 0.001 to 0.02 at n=222 (b,d). A cross marks a GL return with residual R>τ=0.005 . (e) Peak GPU memory allocated by PyTorch. (f) On clustered n=222 , 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 τ and the arrow marks the first fine update. (g,h) Change in the grouped objective U over twenty optimization steps at n=221 , 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=221 , ε=0.001 problem. Bold marks solves that meet τ=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 ε until its full residual first meets τ .
Table 4
Figure 3: FS2 on Quijote. (a) Time of a complete warm call of FS2 at τ=0.05 . The cross marks GL’s instrumented call at 224 . (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 τ , and RMSE is the displacement error of the 227 solution at τ=0.01 and of a zero-displacement predictor.
Appendix figures & tables11 assets
Supplementary material from the paper’s appendix.
Appendix
δrel
νi over
∥TεA−Tε∥
Gradient
Cloud
factor
Admitted
Max νi
Thm. 1 bound
Max error
over νidiam(Y)
rel. ℓ2
Clustered
1
6.33%
6.7×10−6
0.005
2.7×10−6
0.48
1.6×10−7
Clustered
64
5.32%
3.7×10−4
0.005
1.4×10−4
0.48
1.3×10−5
Heterog.
1
60.54%
2.9×10−5
0.023
2.4×10−6
0.23
9.1×10−6
Heterog.
64
49.24%
3.0×10−3
0.038
1.7×10−4
0.23
1.8×10−3
Appendix
Table 4: Proposition 2 at FS2’s returned potentials, n=m=220 , ε=0.001 : largest values over all rows. The νi column divides by the Theorem 1 bound MBδrelai/ri , the error column by νidiam(Y) , and the gradient compares 2ai(xi−TεA(xi)) with 2ai(xi−Tε(xi)) .
n=m
ε
FS2 time (s)
FS2 R
GL time (s)
GL R
Clustered
219
0.001
5.638
0.004879
1.521
0.054122 †
219
0.01
0.290
0.002918
0.026
0.072730 †
220
0.001
4.176
0.004791
5.500
0.053802 †
220
0.01
0.275
0.003245
0.027
0.071918 †
221
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 221 , ε=0.001 as medians of three calls.
Clustered, ε=0.001
Heterogeneous, ε=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=221 problems, truncation 5, rerun with FS2 in the same session. Bold marks solves that meet τ=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 ε
Problem
GL setting
R
K∗
Time (s)
R
ε=0.001
.90, default
.0542
–
–
.0203 at K=512
.99, default
.00687
118
421.10
.004997
.90, matched
.0597
–
–
.0160 at K=2048
ε=0.01
.90, default
.0731
–
–
.0372 at K=512
.99, default
.0373
–
–
.0372 at K=512
Appendix
Table 7: GeomLoss continued at the target ε on clustered n=m=221 , truncation 5, on the GPU of Table 6 . K∗ is the smallest number K of appended target- ε entries whose full residual meets τ=0.005 . Time is the median of three clean calls at K∗ , which at ε=0.001 are the GL + calls of Table 1 a. Bold marks solves that meet τ . Without a K∗ , the last column gives the lowest scored residual. FS2 takes 15.44 s with R=0.00483 at ε=0.001 ( Table 1 a) and, in the run of this table, 0.817 s with R=0.00267 at ε=0.01 .
Problem
n
ε
τ
Mask
R
RA
Ω
RA+Ω
Clustered
219
0.001
0.005
sparse
0.004879
0.004879
4.7×10−6
0.004884
Clustered
219
0.01
0.005
sparse
0.002918
0.002918
2.7×10−5
0.002945
Clustered
220
0.001
0.005
sparse
0.004791
0.004791
4.7×10−6
0.004796
Clustered
220
0.01
0.005
sparse
0.003245
0.003246
2.5×10−5
0.003271
Clustered
221
0.001
0.005
sparse
0.004828
0.004828
4.8×10−6
0.004833
Clustered
221
0.002
0.005
sparse
0.004094
0.004094
7.5×10−6
0.004101
Appendix
Table 8: The certificate RA+Ω of Section 3.2 at returned potentials, with RA the larger of the row and column residuals of PA in IEEE FP32 and Ω 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, Ω=0 and RA is the full residual. R is the recorded full residual, and bold certifies R≤τ .
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 ε . 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=220 , ε=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×1012
2.58×1011
2.58×1011
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=221 , ε=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
ε
U0
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=221 . A decrease is the percentage drop of U from the common initial score U0 over the twenty steps, and final scores within 3×10−6 tie. Times cover the twenty updates, including compilation and autotuning.
method
n=m
τ
configuration
time (s)
peak (GiB)
light validation [U3SE]
full max
FS2
224
.05
B64/t47/c128/h80
146.905 W
12.053
.046836 [.048531]
.046631 ✓
FS2
225
.05
B64/t45/c128/h48
374.839 W
19.472
.049304 [.049758]
—
FS2
226
.05
B128/t45/c128/h72
694.630 W
33.131
.048340 [.048783]
—
FS2
227
.05
B256/t45/c128/h80
1351.589 W
64.129
.049009 [.050797]
.048965 ✓
FS2
224
.01
B64/t47/c800/h432
736.393 W
—
.009538 [.009924]
.009433 ✓
FS2
227
.01
B256/t45/c800/h560
8832.994 F
—
.009585 [.009954]
.009154 ✓
Appendix
Table 13: Quijote forward solves, with the τ=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 225 and 226 , 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 227 -particle FS2 solution at τ=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.
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(N2) operations, which restricts its use on large-scale datasets (e.g. N≥104). 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.
Ian Hsieh, Soumya Snigdha Kundu, Tom Vercauteren +1
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
Entropic regularized optimal transport (OT) via the Sinkhorn algorithm has become a fundamental tool in machine learning, yet existing implementations either suffer from numerical instability for small regularization parameters or incur significant overhead from deep learning frameworks. We present FastSinkhorn, a lightweight, native CUDA implementation of the log-domain Sinkhorn algorithm that combines warp-level shuffle reductions with shared-memory tiling to achieve high GPU utilization without sacrificing numerical stability. Our solver operates entirely in the log-domain, enabling robust computation for regularization parameters as small as epsilon = 10^{-4} where standard-domain methods fail. On dense OT problems with n = m = 8192, our implementation achieves 12x speedup over the widely-used POT library and 5.9x speedup over GPU-accelerated PyTorch baselines, while consuming only 256 MB of GPU memory. We validate our solver on image color transfer, 3D point cloud matching, and convergence analysis, demonstrating that native CUDA kernels with careful numerical treatment provide a practical and efficient foundation for large-scale optimal transport computation.
We present GPU-accelerated Bregman Douglas--Rachford splitting algorithm (BDRS) for discrete optimal transport problem in three input formats: an explicit cost matrix, a point cloud with a ground cost between them, and a separable cost on a regular grid. For each input format, we propose hardware-aware designs of mathematically equivalent representations for the BDRS iterations to enhance numerical stability and empirical runtime. We benchmark the three proposed implementations against eight GPU baseline solvers from the literature on the same device. We demonstrate that our implementations of BDRS achieve state-of-the-art performance on their respective input formats. To the best of our knowledge, this is the first cross-solver study of GPU DOT solvers with a unified measure of optimality.
Yifan Xu, Shiqian Ma
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