FlashSinkhorn: IO-Aware Entropic Optimal Transport

2026-02-03

AI Summary

This work addresses the inefficiency of high-bandwidth memory (HBM) access in large-scale GPU-based entropy-regularized optimal transport (EOT) solvers, where conventional approaches suffer from excessive I/O overhead due to dense matrix operations or generic reduction kernels. We propose the first integration of FlashAttention’s I/O-aware tiling and fusion strategy into EOT, reformulating stabilized log-domain Sinkhorn updates as row-wise LogSumExp reductions of bias-augmented dot-product scores—inspired by Transformer attention mechanisms. By fusing computations and leveraging on-chip SRAM through streaming tiling, our method processes data in a single pass while updating dual potentials. This yields highly efficient computation with linear memory complexity, enabling scalable first- and second-order optimization. On an A100 GPU, our approach achieves up to 32× speedup in forward passes and 161× end-to-end training acceleration, significantly enhancing the scalability of point cloud optimal transport and its downstream applications.

Technology Category

• Machine Learning: Hardware-aware ML
• Search and Optimization: Learning to Search
• Computer Vision: Learning & Optimization for CV

Application Category

• Graph Algorithms and Modeling for the Web: Efficient manipulation of static and dynamic Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Security and Privacy: Privacy-enhancing technologies

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.

Problem

• Entropic Optimal Transport
• Sinkhorn iterations
• GPU efficiency
• IO bottleneck
• large-scale optimization

Innovation

• FlashSinkhorn
• IO-aware
• entropic optimal transport
• fused kernel
• LogSumExp reduction