提出FlashSinkhorn,让GPU高效计算大规模最优传输。
FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU
- 将Sinkhorn迭代重写为类似注意力的行级操作,支持高效融合与分块。
- A100上比现有在线方法快32倍(前向)和161倍(端到端)。
- 适合需要大规模最优传输的生成模型与优化任务。
基于Sinkhorn迭代的熵正则最优传输(EOT)广泛应用于现代机器学习,但现有GPU求解器在大规模场景下效率低下。张量化实现因稠密n×m交互导致二次级显存流量,而现有在线后端虽避免存储稠密矩阵,仍依赖通用分块归约核函数且融合能力有限。本文提出FlashSinkhorn,一种面向平方欧氏距离代价的输入输出感知型EOT求解器,将稳定化的对数域Sinkhorn更新重写为带偏置点积得分的行级LogSumExp归约,其归一化形式与Transformer注意力一致。该设计支持FlashAttention式融合与分块:融合的Triton内核通过片上SRAM流式处理数据块,并单次遍历更新对偶变量,显著降低每轮迭代的显存读写开销,同时保持线性内存操作。进一步提供用于运输任务的流式核函数,支持可扩展的一阶与二阶优化。在A100 GPU上,FlashSinkhorn在点云最优传输任务中相较最先进在线基线,前向计算提速达32倍,端到端加速达161倍,显著提升基于最优传输的下游任务可扩展性。为保证可复现性,代码已开源至https://github.com/ot-triton-lab/flash-sinkhorn。
原文摘要 · Abstract (English)
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/flash-sinkhorn .
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。