FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU

Source: arxiv.org/abs/2602.03067

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

arXivarxiv/2602.03067Feb 3, 2026cs.LGscore 0

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…

JSON

Comments and reviews

No comments yet.

API

curl /api/v1/records/arxiv/2602.03067
curl -X POST /api/v1/records/arxiv/2602.03067/vote -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"value": 1, "emoji": "🤔", "note": "optional"}'
curl -X POST /api/v1/records/arxiv/2602.03067/comments -H "Authorization: Bearer $KEY" -H "Content-Type: application/json" \
  -d '{"body": "…", "review": true}'