DrainSinkhorn: Instance-Resident Batched Entropic Optimal Transport on GPUs

Many learning pipelines call entropic optimal transport (OT) inside a loop, and each call is a batch of small, independent problems: one per image in detection label assignment, per patch pair in point-cloud registration, per frame pair in tracking, and per tile in domain decomposition. General GPU solvers launch every Sinkhorn update over the whole batch and run the batch until its slowest problem stops, so practitioners write a kernel for each shape or drop OT. We present DrainSinkhorn, a batched Sinkhorn solver that makes each problem of the batch, an instance, the unit of GPU execution. Wherever an instance fits on chip, one GPU thread block solves it from start to finish and exits as soon as the instance converges, freeing its slot for a pending instance; with a fixed iteration count, several instances share a thread block, and the forward map and its exact gradient take one launch each. Larger instances leave later launches once they converge. On tracking and domain-decomposition workloads, DrainSinkhorn is 2.73–10.75× faster than OTT-JAX, the fastest general batched GPU solver in our comparison, and 5.3–34× faster than POT; the gain is largest on calls with thousands of instances. In training layers it is 1.05–2.8× faster than OTT-JAX and 9.9–24× faster than the Sinkhorn modules of OTA and GeoTransformer, cutting label assignment from 16.3% to 4.7% of each GPU's training step and making GeoTransformer's patch matching 10.4× faster in a training step.