🤖 AI

FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

FlashSinkhorn is an IO-aware GPU solver for entropic optimal transport that leverages FlashAttention-style fusion and tiling to drastically reduce HBM memory traffic, achieving up to 161× speedups over state-of-the-art baselines while enabling scalable optimization for large-scale point-cloud tasks.

Original authors: Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

Published 2026-05-22

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Felix X. -F. Ye, Xingjie Li, An Yu, Ming-Ching Chang, Linsong Chu, Davis Wertheimer

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to match two huge crowds of people. One crowd is standing on one side of a field (the "source"), and the other is on the opposite side (the "target"). Your goal is to figure out the most efficient way to pair everyone up so that the total distance everyone has to walk is minimized. This is a classic math problem called Optimal Transport.

In modern machine learning, we often add a little bit of "fuzziness" to this matching process to make the math easier to handle. This is called Entropic Optimal Transport. To solve it, computers use a method called Sinkhorn iterations, which is like a game of "hot potato" where the computer keeps passing notes back and forth between the two crowds, refining the matches over and over until it finds the best solution.

The Problem: The Traffic Jam

The paper explains that while this method works well for small crowds, it hits a massive wall when the crowds get huge (like tens of thousands of people).

Think of the computer's memory like a city:

HBM (High Bandwidth Memory): This is the city's main highway. It's huge and can hold a lot of data, but it's slow to get to.
SRAM (On-chip Memory): This is a tiny, super-fast private office right inside the computer's processor. It's incredibly fast but very small.

Older methods for solving this matching problem were like a delivery truck that had to drive from the highway (HBM) to the office (SRAM) and back every single time it needed to check a single pair of people. Because there are millions of possible pairs, the truck was stuck in traffic jams on the highway, constantly moving data back and forth. The computer spent more time waiting for data than actually doing the math.

The Solution: FlashSinkhorn

The authors created a new tool called FlashSinkhorn. They realized that the math behind this matching problem looks exactly like the math used in Transformers (the technology behind AI chatbots like the one you are talking to).

In Transformers, there is a clever trick called FlashAttention that solves a similar traffic jam. Instead of driving the truck back and forth, FlashAttention loads a whole "tile" (a small batch) of data into the fast office, does all the necessary calculations there, and only writes the final result back to the highway.

FlashSinkhorn takes this same "tile-based" strategy and applies it to the matching problem:

No More Full Maps: Instead of writing down the entire map of every possible connection (which would be too big to fit in memory), it calculates connections on the fly, one small tile at a time.
The "Office" Strategy: It keeps the current batch of calculations in the fast, small office (SRAM). It updates the "match scores" right there without ever needing to write the massive intermediate list back to the slow highway.
Streaming: It streams through the data like a conveyor belt, processing and discarding the heavy lifting as it goes, keeping the highway clear.

The Results: Speed and Scale

The paper tested this on powerful GPUs (specifically the A100). The results were dramatic:

Speed: It was up to 32 times faster for the initial calculation and up to 161 times faster for the full process (including learning from mistakes) compared to the best existing online methods.
Memory: While older methods would crash (run out of memory) when trying to match crowds of 30,000 people, FlashSinkhorn could handle 50,000 people easily because it never tried to store the whole map at once.
Real-World Use: They showed it works on real tasks like comparing huge datasets (like thousands of images) and solving complex regression problems where the order of data is mixed up.

The Bottom Line

FlashSinkhorn is like upgrading from a delivery truck stuck in traffic to a high-speed drone. It doesn't change the destination (the math answer is still exact), but it changes how the data is moved. By keeping the heavy lifting inside the computer's fast "office" and only using the slow "highway" for the final results, it makes solving massive matching problems practical and fast, turning a task that used to take hours or crash the computer into something that takes seconds.

Technical Summary: FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

Problem Statement

Entropic Optimal Transport (EOT), solved via Sinkhorn iterations, is a fundamental tool in modern machine learning for comparing probability measures, ranging from generative modeling to distributional alignment. However, scaling EOT to large point clouds (e.g., $n \approx 70,000$ ) on GPUs remains a persistent bottleneck.

Existing GPU solvers face two primary inefficiencies:

Tensorized Implementations: These materialize the dense $n \times m$ cost matrix, leading to quadratic High-Bandwidth Memory (HBM) traffic ($O(nm)$ reads/writes) and memory capacity issues that cause Out-of-Memory (OOM) errors as $n$ grows.
Online Backends: While these avoid storing the dense matrix by evaluating interactions on-the-fly (e.g., via KeOps or OTT-JAX), they rely on generic tiled map-reduce reduction kernels. These kernels suffer from limited fusion, preventing the exploitation of tensor-core GEMM structures and resulting in suboptimal compute utilization and high runtime despite reduced memory usage.

The core systems bottleneck is that current methods are either memory-bound (due to excessive HBM traffic) or compute-bound with poor instruction efficiency, failing to scale to the repeated large-scale solves required in downstream pipelines.

Methodology

The authors introduce FlashSinkhorn, an IO-aware GPU solver for EOT with squared Euclidean cost. The methodology draws a direct parallel between stabilized log-domain Sinkhorn updates and the attention mechanism in Transformers, leveraging the principles of FlashAttention to optimize memory access patterns.

1. Attention-Form Reformulation

The key theoretical insight (Proposition 1) is that for squared Euclidean cost, the stabilized log-domain Sinkhorn updates can be rewritten as row-wise LogSumExp (LSE) reductions of a biased dot-product score matrix.

Specifically, the update $f_i \leftarrow -\varepsilon \text{LSE}_j [(g_j - C_{ij})/\varepsilon + \log b_j]$ is equivalent to computing the LSE of a matrix $S_X = (QK^\top + \mathbf{1}(\hat{g} + \delta))/\varepsilon$ , where $Q=\sqrt{2}X$ , $K=\sqrt{2}Y$ , and $\delta, \gamma$ are precomputed log-weight biases.
This reformulation transforms the Sinkhorn iteration into the exact same normalization statistic used in scaled dot-product attention, enabling the application of IO-aware streaming techniques.

2. Streaming Kernels and Tiling

FlashSinkhorn implements this reformulation using fused Triton kernels that employ a two-level memory model (slow HBM and fast on-chip SRAM):

No Materialization: The $n \times m$ cost matrix is never materialized in HBM. Instead, tiles of $Q$ and $K$ are streamed.
Online LSE: The kernels maintain running row-wise max and sum-exp statistics on-chip (SRAM) as they process tiles of the score matrix. This allows the computation of exact LSE reductions with significantly fewer HBM accesses.
Fusion: The entire update (score computation, bias addition, LSE reduction, and potential update) is fused into a single kernel pass. This eliminates intermediate writes to HBM and reduces kernel launch overhead.
Row-Stationary Traversal: Unlike FlashAttention, which often loops over keys/values in the outer loop, FlashSinkhorn adopts a row-stationary traversal (fixing a query row block in SRAM and streaming all key blocks) to better match the row-wise reduction nature of Sinkhorn updates.

3. Differentiation and Second-Order Optimization

The paper extends this IO-aware strategy to differentiation and second-order methods:

Transport Matrix Application: The barycentric projection and gradient computation are expressed as attention-like operations (weighted averages), allowing them to be computed via the same streaming kernels without materializing the transport matrix $P^*$ .
Hessian-Vector Products (HVP): The authors derive a streaming HVP oracle that decomposes the Hessian-vector product into a sequence of transport-vector and transport-matrix products. This avoids materializing the $O(n^2 d^2)$ Hessian tensor or the $O(nm)$ transport matrix, reducing working memory to $O((n+m)d)$ .

Key Contributions

Attention-Form Sinkhorn Update: The authors express stabilized Sinkhorn updates as biased dot-product LSE reductions, enabling exact FlashAttention-style streaming for EOT.
IO-Aware EOT Solver: They provide fused Triton kernels for streamed Sinkhorn iterations and transport-matrix-free streaming operators for differentiation, substantially reducing HBM IO volume.
Scalable Second-Order Optimization: They introduce streaming HVPs that enable scalable Newton-CG and trust-region methods for EOT-based losses, previously infeasible at large scales due to memory constraints.
Large-Scale Evaluation: The paper demonstrates substantial speedups and memory efficiency on A100 GPUs across synthetic benchmarks and downstream tasks like Optimal Transport Dataset Distance (OTDD) and shuffled regression.

Experimental Results

Evaluated on NVIDIA A100-80GB GPUs, FlashSinkhorn demonstrates significant performance gains over state-of-the-art baselines (GeomLoss with KeOps/Tensorized backends and OTT-JAX):

Speed:
- Forward Pass: Achieves up to 32× speedup over KeOps and 5.1× over OTT-JAX.
- End-to-End (Forward + Backward): Achieves up to 161× speedup over KeOps at high dimensions ( $d=512$ ).
- HVP: Achieves 3–25× speedup depending on dimensionality, making Newton-CG practical at scales where other methods time out.
Memory Efficiency:
- FlashSinkhorn maintains linear memory complexity $O((n+m)d)$ , whereas tensorized baselines scale as $O(nm)$ or $O(n^2)$ .
- It successfully processes point clouds with $n=50,000$ and dimensions up to $d=1024$ , where tensorized methods run out of memory and online baselines (KeOps) exceed time limits.
Downstream Tasks:
- OTDD: FlashSinkhorn scales OTDD distance computation and gradient flow to $n=60,000$ (MNIST $\leftrightarrow$ Fashion-MNIST) with peak memory under 1 GB, while tensorized baselines OOM at $n > 20,000$ .
- Saddle Escape: In shuffled linear regression, the streaming HVP enables monitoring the minimum Hessian eigenvalue to detect saddle points, allowing a hybrid Adam-Newton optimizer to escape saddles and converge rapidly.

Significance and Claims

The paper claims that FlashSinkhorn makes large-scale EOT practical in the memory-bound regime. By rewriting Sinkhorn updates to match the IO-aware structure of FlashAttention, the authors bridge the gap between optimal transport and modern transformer acceleration techniques.

The significance lies in:

System-Level Optimization: The gains are attributed to kernel-level specialization (fusion and tiling) rather than algorithmic changes to the Sinkhorn iteration itself.
Enabling New Workloads: The reduction in memory footprint and runtime enables the use of EOT in downstream tasks (like dataset distance and implicit differentiation) at scales previously considered computationally prohibitive.
Reproducibility: The authors release an open-source implementation to facilitate further research in scalable optimal transport.

The authors modestly note that this approach currently applies to costs with a specific structure (squared Euclidean and cosine distance after normalization). Costs without this structure (e.g., raw Euclidean with square roots or learned neural costs) are identified as future work.