FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

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 densen×min-teractions, while existing online backends avoid O(nm)storage but still realize each Sinkhorn update via general-purpose tiled map-reduce reductions with limited fusion. We present FlashSinkhorn, an IO-aware EOT solver for squared Euclidean cost that rewrites stabilized log-domain Sinkhorn updates as row-wise Log-Sum-Exp 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, maintain online max/sumexp statistics, and update dual potentials in a single pass, substantially reducing HBM IO per iteration while retaining linear-memory operation. We further provide streaming kernels for transport application and Hessian–vector products, enabling scalable first- and second-order optimization without forming the transport plan. On A100 GPUs, FlashSinkhorn achieves up to 32× forward-pass and 161× end-to-end speedups over state-of-the-art online baselines on point-cloud OT, OTDD gradient flows, and OT-based shuffled regression.