FlashSinkhorn: IO-Aware Entropic Optimal Transport

TL;DR

FlashSinkhorn introduces an IO-aware GPU solver for Entropic OT with squared Euclidean cost by recasting stabilized Sinkhorn updates as row-wise LogSumExp reductions of biased dot-product scores, enabling online, fused streaming kernels that avoid materializing the dense $n \times m$ score matrix. By leveraging FlashAttention-style tiling and SRAM-based online normalization, it reduces high-bandwidth memory traffic while preserving linear memory in $O((n+m)d)$ and achieving substantial speedups (up to $32\times$ forward and $161\times$ end-to-end) on A100 GPUs. The approach includes streaming Hessian–vector products for second-order optimization and a streaming transport-application operator that yields $P^*V$ without forming $P^*$. Empirical results across synthetic benchmarks and OT-driven downstream tasks demonstrate improved scalability for point-cloud OT, OTDD, and regression problems, with competitive accuracy and reduced memory footprints. This work significantly advances practical, large-scale EOT by combining IO-aware numerical reductions with end-to-end differentiable streaming kernels, though it currently focuses on the squared Euclidean cost and leaves general cost extensions for future work.

Abstract

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 \textbf{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/ot_triton.

Figures (7)

• Figure 1: FlashSinkhorn kernels via IO-aware tiling. Left: streaming $\hat{\mathbf{f}}$ update (\ref{['alg:flash_sinkhorn_hatf']}). In the outer loop (blue arrow), each thread block stages a tile $Q_I$ in on-chip SRAM and maintains per-row online log-sum-exp statistics. In the inner loop (orange arrow), it streams tiles $K_J$ from HBM, adds a broadcast bias row vector, computes the score tile on-the-fly, and updates the running statistics to write $\hat{\mathbf{f}}_I$ back to HBM. Middle: streaming $\hat{\mathbf{g}}$ update with the roles of $Q$ and $K$ swapped. Right: streaming transport application $P^*V$ (\ref{['alg:plan_mat_stream']}). At convergence, it reuses the same streaming row-softmax while additionally streaming tiles of $V$ alongside $K$ to accumulate and write $P^*V$ without materializing $P^*$. Hatched blocks denote tiles staged in on-chip SRAM; solid-color blocks denote tensors resident in HBM; black arrows indicate HBM$\leftrightarrow$SRAM transfers.
• Figure 2: FlashSinkhorn performance benchmarks. (Top) Forward and forward+backward timing versus $n$ at $d{=}64$ and versus $d$ at $n{=}20k$. (Bottom left) Memory scaling at $d{=}1024$: FlashSinkhorn maintains $O(n)$ while tensorized methods scale as $O(n^{1.7}\text{--}n^{1.9})$. (Bottom right) HVP timing; FlashSinkhorn alone scales to $n{=}50k$ and $d{\ge}256$. A100-80GB, $\varepsilon{=}0.1$.
• Figure 3: Scaling OTDD to large datasets (MNIST$\leftrightarrow$Fashion-MNIST, d=512). FlashSinkhorn vs tensorized baseline for (a,b) time and (c,d) memory on OTDD distance computation and gradient flow.
• Figure 4: Saddle escape trajectory comparison: Adam vs Newton $n=40000, \varepsilon=0.1$ (a) Loss: Adam decay slowly in saddle region; post-escape Newton converges in 7 steps vs 90 for Adam continuation. (b) Gradient norm: Adam stalls near 0.1; Newton breaks through rapidly. (c) Minimum Hessian eigenvalue via Lanczos; negative$\leftarrow$positive transition at step 200 triggers Newton switch. (d) Wall time: Adam+Newton (2423s) vs Adam-only (2785s), 2.8× speedup after escape.
• Figure 6: FlashSinkhorn HVP peak memory vs. problem size ($d=64$). Linear scaling confirms O$(nd)$ memory complexity.

Theorems & Definitions (17)

• Proposition 1: Sinkhorn iteration as biased dot-product LSE
• Theorem 2
• Proposition 3: Transport matrix application as an attention output
• Corollary 4
• Theorem 5: Streaming HVP oracle
• proof : Proof sketch
• Lemma 6
• proof
• Theorem 7
• proof