FlashSinkhorn: GPU上的IO感知熵最优传输
FlashSinkhorn: IO-Aware Entropic Optimal Transport on GPU
AI总结：
本文提出FlashSinkhorn，一种基于GPU的熵最优传输求解器，通过将稳定化的对数域Sinkhorn更新转换为行-wise的LogSumExp归一化，实现了与Transformer注意力相同的归一化方式，从而实现了FlashAttention风格的融合和分块处理，显著降低了HBMIO并保持线性内存操作。
摘要：
熵最优传输（EOT）通过Sinkhorn迭代在现代机器学习中广泛应用，但GPU求解器在大规模情况下仍效率低下。张量化实现因密集的n×m交互导致二次HBM流量，而现有在线后端避免存储密集矩阵但仍然依赖于通用的 tiled map-reduce 减少内核，融合有限。我们提出FlashSinkhorn，一种针对平方欧几里得成本的IO感知EOT求解器，将稳定化的对数域Sinkhorn更新重写为行-wise的LogSumExp归一化，与Transformer注意力相同的归一化方式。这使得FlashAttention风格的融合和分块处理成为可能：融合的Triton内核通过芯片上的SRAM流式传输分块，并在单次通过中更新双潜力，显著减少每个迭代的HBM IO同时保持线性内存操作。我们进一步提供了用于传输应用的流式内核，实现了可扩展的一阶和二阶优化。在A100 GPU上，FlashSinkhorn在点云OT上的前向传递速度比最先进的在线基线快32倍，在端到端速度上快161倍，提高了OT基于下游任务的可扩展性。为了可重复性，我们发布了开源实现，网址为https://github.com/ot-triton-lab/flash-sinkhorn。
英文摘要：
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 .
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
