diff --git a/benchmark/optimal_transport.py b/benchmark/optimal_transport.py
new file mode 100644
index 0000000..b800933
--- /dev/null
+++ b/benchmark/optimal_transport.py
@@ -0,0 +1,134 @@
+import ot
+import torch
+from torch.utils import benchmark
+
+from torchmorph import SinkhornSolver, build_cost_matrix
+
+
+def _grid_problem(n, H, W, epsilon, max_iter, log_space=False, seed=42):
+    """Random (n, H*W) marginals with a 2-D grid cost matrix."""
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    torch.manual_seed(seed)
+    source = torch.rand(n, H * W, device=device)
+    target = torch.rand(n, H * W, device=device)
+    cost_matrix = build_cost_matrix((H, W), device=device)
+    solver = SinkhornSolver(epsilon=epsilon, max_iter=max_iter, log_space=log_space)
+    return solver, source, target, cost_matrix
+
+
+def run_forward_benchmark(n=1, H=32, W=32, epsilon=1.0, max_iter=100, log_space=False):
+    solver, source, target, cost_matrix = _grid_problem(n, H, W, epsilon, max_iter, log_space)
+    if torch.cuda.is_available():
+        torch.cuda.reset_peak_memory_stats()
+    timer = benchmark.Timer(
+        stmt="solver(source, target, cost_matrix)",
+        globals={
+            "solver": solver,
+            "source": source,
+            "target": target,
+            "cost_matrix": cost_matrix,
+        },
+    )
+    result = timer.blocked_autorange(min_run_time=3)
+    if torch.cuda.is_available():
+        print(f"peak allocated: {torch.cuda.max_memory_allocated() / 1024**2:.2f} MB")
+        print(f"peak reserved:  {torch.cuda.max_memory_reserved() / 1024**2:.2f} MB")
+    return result
+
+
+def run_pot_sinkhorn_benchmark(H=32, W=32, epsilon=1.0, max_iter=100):
+    solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter)
+    source, target, cost_matrix = solver.data_preprocess(source, target, cost_matrix)
+    timer = benchmark.Timer(
+        stmt="ot.sinkhorn(source, target, M, reg=epsilon, numItermax=max_iter, stopThr=1e-5)",
+        globals={
+            "ot": ot,
+            "source": source[0],
+            "target": target[0],
+            "M": cost_matrix,
+            "epsilon": epsilon,
+            "max_iter": max_iter,
+        },
+    )
+    return timer.blocked_autorange(min_run_time=2)
+
+
+@torch.no_grad()
+def run_sinkhorn_relative_error(H=32, W=32, epsilon=1.0, max_iter=100):
+    solver, source, target, cost_matrix = _grid_problem(1, H, W, epsilon, max_iter)
+    plan = solver.plan(source, target, cost_matrix)
+    a, b, cost_matrix = solver.data_preprocess(source, target, cost_matrix)
+    ot_plan = ot.sinkhorn(a[0], b[0], cost_matrix, reg=epsilon, numItermax=max_iter, stopThr=1e-5)
+    return torch.linalg.norm(plan[0] - ot_plan) / torch.clamp(torch.linalg.norm(ot_plan), min=1e-12)
+
+
+def reference_error_check(size=32, epsilon=10.0, max_iter=200, verbose=False):
+    """Compare log-space fused-kernel potentials against POT's sinkhorn_log."""
+    if not torch.cuda.is_available():
+        raise RuntimeError("CUDA device not available.")
+    solver, source, target, cost_matrix = _grid_problem(
+        1, size, size, epsilon, max_iter, log_space=True, seed=2
+    )
+
+    grad_f, grad_g = solver.potentials(source, target, cost_matrix)
+    grad_f, grad_g = grad_f[0], grad_g[0]
+
+    a, b, cost_matrix = solver.data_preprocess(source, target, cost_matrix)
+    _, log = ot.bregman.sinkhorn_log(
+        a[0],
+        b[0],
+        cost_matrix,
+        reg=epsilon,
+        numItermax=max_iter,
+        stopThr=0,
+        log=True,
+    )
+    f = epsilon * log["log_u"]
+    g = epsilon * log["log_v"]
+    f = f - f.mean()
+    g = g - g.mean()
+
+    if verbose:
+        print(f, g, grad_f, grad_g, sep="\n")
+
+    f_rel_l2 = torch.linalg.norm(f - grad_f) / torch.clamp(torch.linalg.norm(f), min=1e-12)
+    g_rel_l2 = torch.linalg.norm(g - grad_g) / torch.clamp(torch.linalg.norm(g), min=1e-12)
+    return {
+        "size": size,
+        "points": size * size,
+        "epsilon": epsilon,
+        "max_iter": max_iter,
+        "f_grad_max_abs": (f - grad_f).abs().max().item(),
+        "f_grad_rel_l2": f_rel_l2.item(),
+        "g_grad_max_abs": (g - grad_g).abs().max().item(),
+        "g_grad_rel_l2": g_rel_l2.item(),
+        "f_allclose": torch.allclose(f, grad_f, rtol=1e-4, atol=1e-4),
+        "g_allclose": torch.allclose(g, grad_g, rtol=1e-4, atol=1e-4),
+    }
+
+
+def print_reference_error_table(
+    sizes=(2, 4, 8, 16, 32, 64), max_iters=(100, 200, 500), epsilon=10.0
+):
+    """Print a size and iteration sweep for POT-vs-CUDA gradient error."""
+    headers = ("size", "N", "eps", "itr", "f_max", "f_rel_l2", "g_max", "g_rel_l2", "close")
+    print("| " + " | ".join(headers) + " |")
+    print("| " + " | ".join(["---"] * len(headers)) + " |")
+    for size in sizes:
+        for max_iter in max_iters:
+            row = reference_error_check(size=size, epsilon=epsilon, max_iter=max_iter)
+            print(
+                f"| {row['size']}x{row['size']} "
+                f"| {row['points']} "
+                f"| {row['epsilon']:.3g} "
+                f"| {row['max_iter']} "
+                f"| {row['f_grad_max_abs']:.3e} "
+                f"| {row['f_grad_rel_l2']:.3e} "
+                f"| {row['g_grad_max_abs']:.3e} "
+                f"| {row['g_grad_rel_l2']:.3e} "
+                f"| {row['f_allclose']}/{row['g_allclose']} |"
+            )
+
+
+if __name__ == "__main__":
+    print_reference_error_table()
diff --git a/test/test_optimal_transport.py b/test/test_optimal_transport.py
new file mode 100644
index 0000000..785e069
--- /dev/null
+++ b/test/test_optimal_transport.py
@@ -0,0 +1,327 @@
+import numpy as np
+import pytest
+import torch
+
+from torchmorph import SinkhornSolver, build_cost_matrix
+
+try:
+    import ot
+except ImportError:  # pragma: no cover - exercised only when POT is absent
+    ot = None
+
+# Every test parametrized over DEVICES runs on CPU and, when available, CUDA,
+# so both the pure-torch and the fused-kernel implementations are covered.
+DEVICES = [
+    "cpu",
+    pytest.param(
+        "cuda",
+        marks=pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available"),
+    ),
+]
+LOG_SPACE = [False, True]
+
+
+def _require_cuda():
+    if not torch.cuda.is_available():
+        pytest.skip("CUDA is required for the fused Sinkhorn kernels")
+
+
+def _require_pot():
+    if ot is None:
+        pytest.skip("POT is required for baseline Sinkhorn tests")
+
+
+def _make_positive(shape, device, seed):
+    generator = torch.Generator(device=device).manual_seed(seed)
+    return torch.rand(shape, generator=generator, device=device) + 0.1
+
+
+def _pot_plan(source, target, cost_matrix, solver):
+    return ot.sinkhorn(
+        source.detach().cpu().numpy().astype(np.float64),
+        target.detach().cpu().numpy().astype(np.float64),
+        cost_matrix.detach().cpu().numpy().astype(np.float64),
+        reg=solver.epsilon,
+        numItermax=solver.max_iter,
+        stopThr=1e-9,
+    )
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_build_cost_matrix(device):
+    # 1-D grid: zero diagonal, symmetry, unit spacing
+    cost_1d = build_cost_matrix((5,), device=device)
+    assert cost_1d.shape == (5, 5)
+    assert cost_1d.device.type == device
+    assert torch.allclose(cost_1d.diag(), torch.zeros(5, device=device))
+    assert torch.allclose(cost_1d, cost_1d.T)
+    assert cost_1d[0, 1].item() == pytest.approx(1.0)
+
+    # 2-D grid: adjacent columns are distance 1 apart
+    cost_2d = build_cost_matrix((3, 4), device=device)
+    assert cost_2d.shape == (12, 12)
+    assert torch.allclose(cost_2d, cost_2d.T)
+    assert cost_2d[0, 1].item() == pytest.approx(1.0)
+
+    # p-norms differ on diagonal moves: grid points (0, 0) and (1, 1)
+    assert build_cost_matrix((2, 2), p=1, device=device)[0, 3].item() == pytest.approx(2.0)
+    assert build_cost_matrix((2, 2), p=2, device=device)[0, 3].item() == pytest.approx(2.0**0.5)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+@pytest.mark.parametrize("n,d", [(1, 4), (3, 6), (2, 16)])
+def test_forward_shapes_and_marginals(n, d, log_space, device):
+    solver = SinkhornSolver(epsilon=1.5, max_iter=300, log_space=log_space)
+
+    source = _make_positive((n, d), device=device, seed=11)
+    target = _make_positive((n, d), device=device, seed=23)
+
+    distance = solver(source, target)
+    plan = solver.plan(source, target)
+    f, g = solver.potentials(source, target)
+    assert distance.shape == (n,)
+    assert distance.device.type == device
+    assert plan.shape == (n, d, d)
+    assert f.shape == (n, d) and g.shape == (n, d)
+
+    a, b, cost_matrix = solver.data_preprocess(source, target)
+    assert torch.allclose(distance, (plan * cost_matrix).sum(dim=(-2, -1)), rtol=1e-5, atol=1e-6)
+    assert torch.allclose(plan.sum(dim=-1), a, atol=1e-4)
+    assert torch.allclose(plan.sum(dim=-2), b, atol=1e-4)
+
+
+@pytest.mark.parametrize("shape", [(4,), (2, 2, 2)])
+def test_forward_rejects_non_2d_inputs(shape):
+    solver = SinkhornSolver(epsilon=1.5, max_iter=10)
+    source = _make_positive(shape, device="cpu", seed=911)
+    target = _make_positive(shape, device="cpu", seed=919)
+
+    with pytest.raises(ValueError, match=r"\(n, d\)"):
+        solver(source, target)
+
+
+def test_forward_rejects_mismatched_cost_matrix():
+    solver = SinkhornSolver(epsilon=1.5, max_iter=10)
+    source = _make_positive((2, 4), device="cpu", seed=5)
+    target = _make_positive((2, 4), device="cpu", seed=7)
+
+    with pytest.raises(ValueError, match="cost_matrix"):
+        solver(source, target, cost_matrix=torch.zeros(3, 3))
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_distance_matches_pot_sinkhorn(log_space, device):
+    _require_pot()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
+
+    source = _make_positive((3, 8), device=device, seed=101)
+    target = _make_positive((3, 8), device=device, seed=211)
+    distance = solver(source, target)
+
+    a, b, cost_matrix = solver.data_preprocess(source, target)
+    cost_np = cost_matrix.cpu().numpy()
+    expected = torch.tensor(
+        [
+            float(np.sum(_pot_plan(a[i], b[i], cost_matrix, solver) * cost_np))
+            for i in range(a.shape[0])
+        ],
+        device=device,
+        dtype=distance.dtype,
+    )
+    assert torch.allclose(distance, expected, rtol=5e-3, atol=1e-5)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_plan_matches_pot_sinkhorn_grid_cost(log_space, device):
+    """A flattened 2-D image with an explicit grid cost matrix matches POT."""
+    _require_pot()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
+    cost_matrix = build_cost_matrix((2, 3), device=device)
+
+    source = _make_positive((1, 6), device=device, seed=307)
+    target = _make_positive((1, 6), device=device, seed=401)
+    plan = solver.plan(source, target, cost_matrix)
+
+    a, b, _ = solver.data_preprocess(source, target, cost_matrix)
+    expected = torch.tensor(
+        _pot_plan(a[0], b[0], cost_matrix, solver), device=device, dtype=plan.dtype
+    )
+    assert torch.allclose(plan[0], expected, rtol=5e-3, atol=1e-5)
+
+
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+@pytest.mark.parametrize("n", [1, 3, 20])  # single-item, partial-tile, and multi-tile batches
+def test_fused_cuda_kernels_match_cpu(log_space, n):
+    """CUDA float32 inputs dispatch to the fused kernels; results must match CPU."""
+    _require_cuda()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
+
+    source = _make_positive((n, 6), device="cpu", seed=503)
+    target = _make_positive((n, 6), device="cpu", seed=607)
+
+    cpu_distance = solver(source, target)
+    cpu_plan = solver.plan(source, target)
+    cpu_f, cpu_g = solver.potentials(source, target)
+    cuda_distance = solver(source.cuda(), target.cuda())
+    cuda_plan = solver.plan(source.cuda(), target.cuda())
+    cuda_f, cuda_g = solver.potentials(source.cuda(), target.cuda())
+
+    assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5)
+    assert torch.allclose(cuda_plan.cpu(), cpu_plan, rtol=1e-4, atol=1e-5)
+    assert torch.allclose(cuda_f.cpu(), cpu_f, rtol=1e-4, atol=1e-4)
+    assert torch.allclose(cuda_g.cpu(), cpu_g, rtol=1e-4, atol=1e-4)
+
+
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_fused_short_run_without_cuda_graph_matches_cpu(log_space):
+    """max_iter below the CUDA-graph threshold takes the plain-launch fused path."""
+    _require_cuda()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=60, log_space=log_space)
+
+    source = _make_positive((2, 6), device="cpu", seed=1409)
+    target = _make_positive((2, 6), device="cpu", seed=1423)
+
+    cpu_distance = solver(source, target)
+    cuda_distance = solver(source.cuda(), target.cuda())
+    assert torch.allclose(cuda_distance.cpu(), cpu_distance, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_dtype_device_combinations_agree(log_space):
+    """float32/float64 on CPU/CUDA must all reach the same solution.
+
+    float64 CUDA inputs exercise the torch fallback on the GPU, float32 CUDA
+    inputs the fused kernels, and CPU inputs the pure-torch paths.
+    """
+    solver = SinkhornSolver(epsilon=1.5, max_iter=2000, log_space=log_space)
+    source = _make_positive((2, 8), device="cpu", seed=41)
+    target = _make_positive((2, 8), device="cpu", seed=43)
+    reference = solver(source.double(), target.double())
+
+    combos = [(source, target)]  # cpu float32
+    if torch.cuda.is_available():
+        combos += [
+            (source.cuda(), target.cuda()),  # cuda float32: fused kernels
+            (source.cuda().double(), target.cuda().double()),  # cuda float64: torch ops
+        ]
+    for src, tgt in combos:
+        distance = solver(src, tgt)
+        tol = 1e-4 if distance.dtype == torch.float32 else 1e-8
+        assert torch.allclose(distance.cpu().double(), reference, rtol=tol, atol=tol)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+def test_log_space_survives_small_epsilon(device):
+    """At this epsilon exp(-C/eps) underflows float32, so only log_space=True is usable.
+
+    On CPU this exercises the torch logsumexp path, on CUDA the fused log kernel.
+    """
+    _require_pot()
+    solver = SinkhornSolver(epsilon=0.1, max_iter=5000, log_space=True)
+    cost_matrix = build_cost_matrix((8, 8), device=device)
+
+    source = _make_positive((1, 64), device=device, seed=31)
+    target = _make_positive((1, 64), device=device, seed=37)
+
+    f, g = solver.potentials(source, target, cost_matrix)
+    assert torch.isfinite(f).all()
+    assert torch.isfinite(g).all()
+
+    plan = solver.plan(source, target, cost_matrix)[0]
+    a, b, _ = solver.data_preprocess(source, target, cost_matrix)
+    assert torch.isfinite(plan).all()
+    assert (plan >= 0).all()
+    assert torch.allclose(plan.sum(dim=0), b[0], atol=1e-5)
+    assert torch.allclose(plan.sum(dim=1), a[0], atol=5e-3)
+
+    cost_np = cost_matrix.cpu().numpy().astype(np.float64)
+    pot_plan = ot.bregman.sinkhorn_log(
+        a[0].cpu().numpy().astype(np.float64),
+        b[0].cpu().numpy().astype(np.float64),
+        cost_np,
+        reg=solver.epsilon,
+        numItermax=solver.max_iter,
+        stopThr=1e-12,
+    )
+    pot_distance = float(np.sum(pot_plan * cost_np))
+    distance = solver(source, target, cost_matrix)
+    assert distance[0].item() == pytest.approx(pot_distance, rel=2e-2)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_backward_matches_numeric_directional_derivative(log_space, device):
+    """Autograd gradients w.r.t. BOTH source and target must match numeric
+    derivatives of the entropic cost.
+
+    Runs in float64, i.e. through the torch paths on both CPU and CUDA.
+    """
+    n, d = 3, 8
+    solver = SinkhornSolver(epsilon=1.5, max_iter=3000, log_space=log_space)
+    cost_matrix = build_cost_matrix((d,), device=device).double()
+
+    source = _make_positive((n, d), device=device, seed=1009).double().requires_grad_(True)
+    target = _make_positive((n, d), device=device, seed=2003).double().requires_grad_(True)
+    direction = _make_positive((n, d), device=device, seed=3001).double() - 0.6
+
+    solver(source, target, cost_matrix).sum().backward()
+    assert source.grad is not None and target.grad is not None
+
+    def entropic_cost(src, tgt):
+        plan = solver.plan(src, tgt, cost_matrix)
+        entropy = (plan * plan.clamp(min=1e-300).log()).sum()
+        return (plan * cost_matrix).sum() + solver.epsilon * entropy
+
+    eps_fd = 1e-4
+    src, tgt = source.detach(), target.detach()
+    for grad, plus, minus in [
+        (source.grad, (src + eps_fd * direction, tgt), (src - eps_fd * direction, tgt)),
+        (target.grad, (src, tgt + eps_fd * direction), (src, tgt - eps_fd * direction)),
+    ]:
+        analytic = (grad * direction).sum()
+        numeric = (entropic_cost(*plus) - entropic_cost(*minus)) / (2 * eps_fd)
+        assert torch.allclose(analytic, numeric, rtol=1e-5, atol=1e-7)
+
+
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_backward_through_fused_kernels_matches_cpu(log_space):
+    """float32 CUDA backward runs through the fused kernels; grads must match CPU."""
+    _require_cuda()
+    solver = SinkhornSolver(epsilon=1.5, max_iter=1000, log_space=log_space)
+
+    source = _make_positive((3, 6), device="cpu", seed=1213)
+    target = _make_positive((3, 6), device="cpu", seed=1217)
+
+    cpu_source = source.clone().requires_grad_(True)
+    cpu_target = target.clone().requires_grad_(True)
+    solver(cpu_source, cpu_target).sum().backward()
+    cuda_source = source.cuda().requires_grad_(True)
+    cuda_target = target.cuda().requires_grad_(True)
+    solver(cuda_source, cuda_target).sum().backward()
+
+    for cuda_grad, cpu_grad in [
+        (cuda_source.grad, cpu_source.grad),
+        (cuda_target.grad, cpu_target.grad),
+    ]:
+        assert torch.isfinite(cuda_grad).all()
+        assert torch.allclose(cuda_grad.cpu(), cpu_grad, rtol=1e-4, atol=1e-5)
+
+
+@pytest.mark.parametrize("device", DEVICES)
+@pytest.mark.parametrize("log_space", LOG_SPACE)
+def test_threshold_early_stopping_reaches_same_solution(log_space, device):
+    """threshold > 0 stops the torch paths early without changing the solution."""
+    source = _make_positive((2, 8), device=device, seed=1301)
+    target = _make_positive((2, 8), device=device, seed=1303)
+    source, target = source.double(), target.double()  # force the torch paths
+
+    full = SinkhornSolver(epsilon=1.5, max_iter=5000, log_space=log_space)
+    early = SinkhornSolver(epsilon=1.5, max_iter=5000, threshold=1e-12, log_space=log_space)
+
+    assert torch.allclose(early(source, target), full(source, target), rtol=1e-8, atol=1e-10)
+    assert torch.allclose(
+        early.plan(source, target), full.plan(source, target), rtol=1e-6, atol=1e-10
+    )
diff --git a/torchmorph/__init__.py b/torchmorph/__init__.py
index 096562b..de23f1c 100644
--- a/torchmorph/__init__.py
+++ b/torchmorph/__init__.py
@@ -19,6 +19,7 @@
     morphological_laplace,
     white_tophat,
 )
+from .optimal_transport import SinkhornSolver, build_cost_matrix
 
 __all__ = [
     "euclidean_distance_transform",
@@ -38,4 +39,6 @@
     "morphological_laplace",
     "white_tophat",
     "black_tophat",
+    "SinkhornSolver",
+    "build_cost_matrix",
 ]
diff --git a/torchmorph/csrc/ot_kernel.cu b/torchmorph/csrc/ot_kernel.cu
new file mode 100644
index 0000000..193d1ac
--- /dev/null
+++ b/torchmorph/csrc/ot_kernel.cu
@@ -0,0 +1,283 @@
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDAGuard.h>
+#include <tuple>
+
+namespace {
+
+constexpr int kThreads = 256;  // multiple of the warp size, assumed below
+constexpr int kBatchTile = 8;  // batch items processed per block
+constexpr int kWarps = kThreads / 32;
+
+__device__ inline float warp_reduce_sum(float v){
+    for (int offset = 16; offset > 0; offset >>= 1){
+        v += __shfl_down_sync(0xffffffff, v, offset);
+    }
+    return v;
+}
+
+// Block-level sum reduction; the result is valid in thread 0 only.
+__device__ inline float block_reduce_sum(float v){
+    __shared__ float partial[kWarps];
+    const int lane = threadIdx.x & 31;
+    const int warp = threadIdx.x >> 5;
+    v = warp_reduce_sum(v);
+    if (lane == 0) partial[warp] = v;
+    __syncthreads();
+    if (warp == 0){
+        v = lane < kWarps ? partial[lane] : 0.0f;
+        v = warp_reduce_sum(v);
+    }
+    __syncthreads();  // let the next call reuse `partial`
+    return v;
+}
+
+// Merge two online-logsumexp states; (m, s) represents m + log(s).
+// s is 0 exactly when m is -inf (the empty state), which the guards below
+// keep NaN-free.
+__device__ inline void lse_merge(float& m, float& s, float m_other, float s_other){
+    const float hi = fmaxf(m, m_other);
+    if (hi == -INFINITY){
+        m = -INFINITY;
+        s = 0.0f;
+        return;
+    }
+    s = (m == -INFINITY ? 0.0f : s * expf(m - hi)) +
+        (m_other == -INFINITY ? 0.0f : s_other * expf(m_other - hi));
+    m = hi;
+}
+
+__device__ inline void warp_reduce_lse(float& m, float& s){
+    for (int offset = 16; offset > 0; offset >>= 1){
+        const float m_other = __shfl_down_sync(0xffffffff, m, offset);
+        const float s_other = __shfl_down_sync(0xffffffff, s, offset);
+        lse_merge(m, s, m_other, s_other);
+    }
+}
+
+// Block-level logsumexp reduction; the result is valid in thread 0 only.
+__device__ inline void block_reduce_lse(float& m, float& s){
+    __shared__ float partial_m[kWarps];
+    __shared__ float partial_s[kWarps];
+    const int lane = threadIdx.x & 31;
+    const int warp = threadIdx.x >> 5;
+    warp_reduce_lse(m, s);
+    if (lane == 0){
+        partial_m[warp] = m;
+        partial_s[warp] = s;
+    }
+    __syncthreads();
+    if (warp == 0){
+        m = lane < kWarps ? partial_m[lane] : -INFINITY;
+        s = lane < kWarps ? partial_s[lane] : 0.0f;
+        warp_reduce_lse(m, s);
+    }
+    __syncthreads();  // let the next call reuse the partial arrays
+}
+
+// One Sinkhorn scaling update. One block per (row, batch tile): the matrix
+// row is streamed once and applied to TILE batch items, so the d^2 matrix is
+// read ceil(n / TILE) times per update instead of n times.
+//   out[b, i] = num[b, i] / (sum_j mat[i, j] * scale[b, j] + 1e-12)
+// The u-update passes (a, K, v); the v-update passes (b, K^T, u).
+template <int TILE>
+__global__ void scaling_update(
+    const float* __restrict__ num,
+    const float* __restrict__ mat,
+    const float* __restrict__ scale,
+    float* __restrict__ out,
+    int d,
+    int n_batch
+){
+    const int row = blockIdx.x;
+    const int batch0 = blockIdx.y * TILE;
+    const int tile = min(TILE, n_batch - batch0);
+    const float* mat_row = mat + (size_t)row * d;
+
+    float acc[TILE];
+#pragma unroll
+    for (int t = 0; t < TILE; t++) acc[t] = 0.0f;
+
+    for (int j = threadIdx.x; j < d; j += blockDim.x){
+        const float k = mat_row[j];
+#pragma unroll
+        for (int t = 0; t < TILE; t++){
+            if (t < tile) acc[t] += k * scale[(size_t)(batch0 + t) * d + j];
+        }
+    }
+
+    for (int t = 0; t < tile; t++){
+        const float total = block_reduce_sum(acc[t]);
+        if (threadIdx.x == 0){
+            const size_t idx = (size_t)(batch0 + t) * d + row;
+            out[idx] = num[idx] / (total + 1e-12f);
+        }
+    }
+}
+
+// One log-domain Sinkhorn update with the same batch tiling:
+//   log_out[b, i] = log_num[b, i] - logsumexp_j(-cost[i, j] / eps + log_scale[b, j])
+// The logsumexp is evaluated online in a single pass over the row (running
+// max with a rescaled running sum), halving the matrix traffic of the
+// classic two-pass max-then-sum evaluation.
+// The u-update passes (log_a, M, log_v); the v-update passes (log_b, M^T, log_u).
+template <int TILE>
+__global__ void log_scaling_update(
+    const float* __restrict__ log_num,
+    const float* __restrict__ cost,
+    const float* __restrict__ log_scale,
+    float* __restrict__ log_out,
+    int d,
+    int n_batch,
+    float inv_eps
+){
+    const int row = blockIdx.x;
+    const int batch0 = blockIdx.y * TILE;
+    const int tile = min(TILE, n_batch - batch0);
+    const float* cost_row = cost + (size_t)row * d;
+
+    float m[TILE], s[TILE];
+#pragma unroll
+    for (int t = 0; t < TILE; t++){
+        m[t] = -INFINITY;
+        s[t] = 0.0f;
+    }
+
+    for (int j = threadIdx.x; j < d; j += blockDim.x){
+        const float c = -cost_row[j] * inv_eps;
+#pragma unroll
+        for (int t = 0; t < TILE; t++){
+            if (t >= tile) continue;
+            const float z = c + log_scale[(size_t)(batch0 + t) * d + j];
+            if (z > m[t]){
+                s[t] = s[t] * expf(m[t] - z) + 1.0f;  // expf(-inf) == 0 covers the first element
+                m[t] = z;
+            } else if (m[t] != -INFINITY){  // skip z == m == -inf, whose contribution is 0
+                s[t] += expf(z - m[t]);
+            }
+        }
+    }
+
+    for (int t = 0; t < tile; t++){
+        block_reduce_lse(m[t], s[t]);
+        if (threadIdx.x == 0){
+            const size_t idx = (size_t)(batch0 + t) * d + row;
+            // m == -inf means the opposite marginal is all zero; pin the
+            // potential instead of producing NaN.
+            log_out[idx] = m[t] == -INFINITY ? -INFINITY : log_num[idx] - m[t] - logf(s[t]);
+        }
+    }
+}
+
+void check_input(const torch::Tensor& t, int64_t numel, const char* name){
+    TORCH_CHECK(t.is_cuda(), name, " must be a CUDA tensor");
+    TORCH_CHECK(t.scalar_type() == torch::kFloat32, name, " must be float32");
+    TORCH_CHECK(t.is_contiguous(), name, " must be contiguous");
+    TORCH_CHECK(t.numel() == numel, name, " must have ", numel, " elements, got ", t.numel());
+}
+
+// Validate batched (n, d) vectors against a shared (d, d) matrix.
+void check_shapes(const torch::Tensor& u, const torch::Tensor& mat){
+    TORCH_CHECK(mat.dim() == 2 && mat.size(0) == mat.size(1), "cost/kernel matrix must be square");
+    TORCH_CHECK(u.dim() == 2 && u.size(1) == mat.size(0),
+                "scaling vectors must be (n, d) with d matching the matrix");
+}
+
+}  // namespace
+
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& K,
+    const torch::Tensor& K_T,
+    torch::Tensor u,
+    torch::Tensor v,
+    int64_t n_iter
+){
+    check_shapes(u, K);
+    const int64_t n = u.size(0);
+    const int64_t d = u.size(1);
+    check_input(a, n * d, "a");
+    check_input(b, n * d, "b");
+    check_input(K, d * d, "K");
+    check_input(K_T, d * d, "K_T");
+    check_input(u, n * d, "u");
+    check_input(v, n * d, "v");
+
+    const at::cuda::CUDAGuard guard(u.device());
+    const auto stream = at::cuda::getCurrentCUDAStream();
+    const bool single = n == 1;
+    const dim3 grid(static_cast<unsigned>(d),
+                    single ? 1u : static_cast<unsigned>((n + kBatchTile - 1) / kBatchTile));
+
+    for (int64_t i = 0; i < n_iter; i++){
+        if (single){
+            scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                a.data_ptr<float>(), K.data_ptr<float>(), v.data_ptr<float>(),
+                u.data_ptr<float>(), (int)d, (int)n);
+            scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                b.data_ptr<float>(), K_T.data_ptr<float>(), u.data_ptr<float>(),
+                v.data_ptr<float>(), (int)d, (int)n);
+        } else {
+            scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                a.data_ptr<float>(), K.data_ptr<float>(), v.data_ptr<float>(),
+                u.data_ptr<float>(), (int)d, (int)n);
+            scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                b.data_ptr<float>(), K_T.data_ptr<float>(), u.data_ptr<float>(),
+                v.data_ptr<float>(), (int)d, (int)n);
+        }
+    }
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    return std::make_tuple(u, v);
+}
+
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
+    const torch::Tensor& log_a,
+    const torch::Tensor& log_b,
+    const torch::Tensor& M,
+    const torch::Tensor& M_T,
+    torch::Tensor log_u,
+    torch::Tensor log_v,
+    int64_t n_iter,
+    double epsilon
+){
+    TORCH_CHECK(epsilon > 0, "epsilon must be positive, got ", epsilon);
+    check_shapes(log_u, M);
+    const int64_t n = log_u.size(0);
+    const int64_t d = log_u.size(1);
+    check_input(log_a, n * d, "log_a");
+    check_input(log_b, n * d, "log_b");
+    check_input(M, d * d, "M");
+    check_input(M_T, d * d, "M_T");
+    check_input(log_u, n * d, "log_u");
+    check_input(log_v, n * d, "log_v");
+
+    const at::cuda::CUDAGuard guard(log_u.device());
+    const auto stream = at::cuda::getCurrentCUDAStream();
+    const float inv_eps = static_cast<float>(1.0 / epsilon);
+    const bool single = n == 1;
+    const dim3 grid(static_cast<unsigned>(d),
+                    single ? 1u : static_cast<unsigned>((n + kBatchTile - 1) / kBatchTile));
+
+    for (int64_t i = 0; i < n_iter; i++){
+        if (single){
+            log_scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                log_a.data_ptr<float>(), M.data_ptr<float>(), log_v.data_ptr<float>(),
+                log_u.data_ptr<float>(), (int)d, (int)n, inv_eps);
+            log_scaling_update<1><<<grid, kThreads, 0, stream>>>(
+                log_b.data_ptr<float>(), M_T.data_ptr<float>(), log_u.data_ptr<float>(),
+                log_v.data_ptr<float>(), (int)d, (int)n, inv_eps);
+        } else {
+            log_scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                log_a.data_ptr<float>(), M.data_ptr<float>(), log_v.data_ptr<float>(),
+                log_u.data_ptr<float>(), (int)d, (int)n, inv_eps);
+            log_scaling_update<kBatchTile><<<grid, kThreads, 0, stream>>>(
+                log_b.data_ptr<float>(), M_T.data_ptr<float>(), log_u.data_ptr<float>(),
+                log_v.data_ptr<float>(), (int)d, (int)n, inv_eps);
+        }
+    }
+    C10_CUDA_KERNEL_LAUNCH_CHECK();
+
+    return std::make_tuple(log_u, log_v);
+}
diff --git a/torchmorph/csrc/torchmorph.cpp b/torchmorph/csrc/torchmorph.cpp
index e553463..669c56d 100644
--- a/torchmorph/csrc/torchmorph.cpp
+++ b/torchmorph/csrc/torchmorph.cpp
@@ -41,6 +41,27 @@ std::tuple<torch::Tensor, torch::Tensor> bfdt_cuda(
     bool return_indices
 );
 
+// Optimal Transport functions
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_fastiter(
+    const torch::Tensor& a,
+    const torch::Tensor& b,
+    const torch::Tensor& K,
+    const torch::Tensor& K_T,
+    torch::Tensor u,
+    torch::Tensor v,
+    int64_t n_iter
+);
+
+std::tuple<torch::Tensor, torch::Tensor> sinkhorn_logiter(
+    const torch::Tensor& log_a,
+    const torch::Tensor& log_b,
+    const torch::Tensor& M,
+    const torch::Tensor& M_T,
+    torch::Tensor log_u,
+    torch::Tensor log_v,
+    int64_t n_iter,
+    double epsilon
+);
 
 PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
     m.def("grey_erosion_cuda", &grey_erosion_cuda,
@@ -82,4 +103,27 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
           py::arg("sampling") = std::vector<float>(),
           py::arg("return_distances") = true,
           py::arg("return_indices") = false);
+
+    // Optimal Transport
+    m.def("sinkhorn_fastiter", &sinkhorn_fastiter,
+          "Sinkhorn scaling-form iterations (CUDA)",
+          py::arg("a"),
+          py::arg("b"),
+          py::arg("K"),
+          py::arg("K_T"),
+          py::arg("u"),
+          py::arg("v"),
+          py::arg("n_iter"));
+
+    m.def("sinkhorn_logiter", &sinkhorn_logiter,
+          "Sinkhorn log-domain iterations (CUDA)",
+          py::arg("log_a"),
+          py::arg("log_b"),
+          py::arg("M"),
+          py::arg("M_T"),
+          py::arg("log_u"),
+          py::arg("log_v"),
+          py::arg("n_iter"),
+          py::arg("epsilon"));
 }
+
diff --git a/torchmorph/optimal_transport.py b/torchmorph/optimal_transport.py
new file mode 100644
index 0000000..3652845
--- /dev/null
+++ b/torchmorph/optimal_transport.py
@@ -0,0 +1,249 @@
+from typing import Optional
+
+import torch
+from torch import Tensor, nn
+
+
+def build_cost_matrix(shape, p=2, device=None) -> Tensor:
+    """Pairwise L^p distance matrix between the points of a spatial grid.
+
+    Returns a (d, d) cost matrix for the d flattened grid points of `shape`.
+    Flatten grid-shaped distributions to rows of an (n, d) tensor and pass the
+    result as `cost_matrix` to :class:`SinkhornSolver`.
+    """
+    coords = torch.stack(
+        torch.meshgrid([torch.arange(s, device=device) for s in shape], indexing="ij"), dim=-1
+    )
+    coords = coords.reshape(-1, len(shape)).float()
+    return torch.cdist(coords, coords, p=p)
+
+
+def _use_fused_kernels(t: Tensor) -> bool:
+    return t.is_cuda and t.dtype == torch.float32
+
+
+def _transpose_or_self(matrix: Tensor) -> Tensor:
+    """Return the contiguous transpose, reusing the input when it is symmetric.
+
+    Cost matrices from build_cost_matrix are always symmetric, so this
+    usually saves a (d, d) copy.
+    """
+    return matrix if torch.equal(matrix, matrix.mT) else matrix.mT.contiguous()
+
+
+_GRAPH_MIN_ITER = 100  # below this, plain launches beat the graph-capture setup cost
+_GRAPH_CHUNK = 25
+
+
+def _run_fused(launch, max_iter, device):
+    """Run `launch(k)`, which enqueues k in-place Sinkhorn iterations, max_iter times in total.
+
+    Long runs are launch-latency-bound (two tiny kernels per iteration), so a
+    chunk of iterations is captured into a CUDA graph once and replayed. The
+    iterations are fixed-point steps on ping-pong buffers, so a fallback that
+    runs extra iterations is always safe.
+    """
+    if max_iter < _GRAPH_MIN_ITER or torch.cuda.is_current_stream_capturing():
+        launch(max_iter)
+        return
+    try:
+        with torch.cuda.device(device):
+            # Warm up on a side stream (required before capture), doing the
+            # first chunk of real iterations in the process.
+            side_stream = torch.cuda.Stream()
+            side_stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(side_stream):
+                launch(_GRAPH_CHUNK)
+            torch.cuda.current_stream().wait_stream(side_stream)
+
+            graph = torch.cuda.CUDAGraph()
+            with torch.cuda.graph(graph):
+                launch(_GRAPH_CHUNK)
+            for _ in range(max_iter // _GRAPH_CHUNK - 1):
+                graph.replay()
+            remainder = max_iter % _GRAPH_CHUNK
+            if remainder:
+                launch(remainder)
+    except RuntimeError:
+        launch(max_iter)  # graphs unavailable; extra iterations are harmless
+
+
+def _scaling_iterations(a, b, cost_matrix, epsilon, max_iter, threshold):
+    """Plain scaling-form Sinkhorn on (n, d) marginals; returns (log_u, log_v)."""
+    K = torch.exp(-cost_matrix / epsilon)
+    if _use_fused_kernels(a):
+        from torchmorph import _C
+
+        a, b = a.contiguous(), b.contiguous()
+        K_T = _transpose_or_self(K)
+        u = torch.ones_like(a)
+        v = torch.ones_like(b)
+        _run_fused(lambda k: _C.sinkhorn_fastiter(a, b, K, K_T, u, v, k), max_iter, a.device)
+        return u.log(), v.log()
+
+    u = torch.ones_like(a)
+    v = torch.ones_like(b)
+    for i in range(max_iter):
+        u_prev = u
+        u = a / (v @ K.mT + 1e-12)
+        v = b / (u @ K + 1e-12)
+        if threshold > 0 and i % 10 == 0:
+            if (u - u_prev).abs().sum(dim=-1).max() <= threshold:
+                break
+    return u.log(), v.log()
+
+
+def _log_iterations(a, b, cost_matrix, epsilon, max_iter, threshold):
+    """Log-domain Sinkhorn on (n, d) marginals; returns (log_u, log_v)."""
+    log_a, log_b = a.log(), b.log()
+    if _use_fused_kernels(a):
+        from torchmorph import _C
+
+        log_a, log_b = log_a.contiguous(), log_b.contiguous()
+        cost_T = _transpose_or_self(cost_matrix)
+        log_u = torch.zeros_like(a)
+        log_v = torch.zeros_like(b)
+        _run_fused(
+            lambda k: _C.sinkhorn_logiter(
+                log_a, log_b, cost_matrix, cost_T, log_u, log_v, k, epsilon
+            ),
+            max_iter,
+            a.device,
+        )
+        return log_u, log_v
+
+    log_K = -cost_matrix / epsilon
+    log_u = torch.zeros_like(a)
+    log_v = torch.zeros_like(b)
+    for i in range(max_iter):
+        log_u_prev = log_u
+        log_u = log_a - torch.logsumexp(log_K + log_v.unsqueeze(-2), dim=-1)
+        log_v = log_b - torch.logsumexp(log_K.mT + log_u.unsqueeze(-2), dim=-1)
+        if threshold > 0 and i % 10 == 0:
+            if (log_u - log_u_prev).abs().sum(dim=-1).max() <= threshold:
+                break
+    return log_u, log_v
+
+
+class _SinkhornDistance(torch.autograd.Function):
+    """Transport distance <P, C> with envelope-theorem gradients.
+
+    The backward pass returns the centered dual potentials, i.e. the exact
+    gradients of the entropic OT cost w.r.t. the (normalized) marginals.
+    """
+
+    @staticmethod
+    def forward(ctx, a, b, cost_matrix, epsilon, max_iter, threshold, log_space):
+        iterate = _log_iterations if log_space else _scaling_iterations
+        log_u, log_v = iterate(a, b, cost_matrix, epsilon, max_iter, threshold)
+        plan = torch.exp(log_u.unsqueeze(-1) - cost_matrix / epsilon + log_v.unsqueeze(-2))
+        distance = (plan * cost_matrix).sum(dim=(-2, -1))
+        f = epsilon * log_u
+        g = epsilon * log_v
+        ctx.save_for_backward(f - f.mean(dim=-1, keepdim=True), g - g.mean(dim=-1, keepdim=True))
+        return distance
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        f, g = ctx.saved_tensors
+        grad = grad_output.unsqueeze(-1)
+        return grad * f, grad * g, None, None, None, None, None
+
+
+class SinkhornSolver(nn.Module):
+    """Entropy-regularized balanced optimal transport as a differentiable module.
+
+    forward(source, target, cost_matrix=None) takes two (n, d) batches of
+    histograms sharing one (d, d) cost matrix and returns the (n,) transport
+    distances <P, C>, with K = exp(-C / epsilon). Gradients w.r.t. source and
+    target are the centered dual potentials (envelope theorem), so the module
+    can be used directly as a loss.
+
+    log_space=True runs the iterations entirely in the log domain, which is
+    numerically stable for small epsilon. The implementation is picked from
+    the input automatically: CUDA float32 tensors use the fused kernels,
+    everything else uses pure torch ops (where `threshold` > 0 enables early
+    stopping; the fused kernels always run `max_iter` iterations).
+    """
+
+    def __init__(self, epsilon=1.0, max_iter=100, threshold=0.0, p=2, log_space=False):
+        super().__init__()
+        if epsilon <= 0:
+            raise ValueError("epsilon must be positive.")
+        if max_iter <= 0:
+            raise ValueError("max_iter must be positive.")
+        if threshold < 0:
+            raise ValueError("threshold must be non-negative.")
+        self.epsilon = epsilon
+        self.max_iter = max_iter
+        self.threshold = threshold
+        self.p = p
+        self.log_space = log_space
+
+    def extra_repr(self):
+        return (
+            f"epsilon={self.epsilon}, max_iter={self.max_iter}, "
+            f"threshold={self.threshold}, log_space={self.log_space}"
+        )
+
+    def data_preprocess(self, source, target, cost_matrix: Optional[Tensor] = None):
+        """Validate (n, d) inputs, clamp negatives, and normalize each row to unit mass.
+
+        All ops are differentiable, so forward gradients flow through the
+        normalization. When `cost_matrix` is None the d bins are treated as
+        points on a line.
+        """
+        if source.ndim != 2 or source.shape != target.shape:
+            raise ValueError("source and target must be (n, d) tensors of the same shape.")
+        if not source.is_floating_point():
+            raise ValueError("source and target must be floating-point tensors.")
+        d = source.shape[1]
+        if cost_matrix is None:
+            cost_matrix = build_cost_matrix((d,), self.p, source.device)
+        elif cost_matrix.shape != (d, d):
+            raise ValueError(
+                f"cost_matrix must have shape ({d}, {d}), got {tuple(cost_matrix.shape)}."
+            )
+        cost_matrix = cost_matrix.to(device=source.device, dtype=source.dtype).contiguous()
+
+        source = source.clamp(min=0)
+        target = target.clamp(min=0)
+        source = source / source.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+        target = target / target.sum(dim=-1, keepdim=True).clamp(min=1e-12)
+        return source, target, cost_matrix
+
+    def forward(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None):
+        """(n,) transport distances between the rows of source and target."""
+        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix)
+        return _SinkhornDistance.apply(
+            source,
+            target,
+            cost_matrix,
+            self.epsilon,
+            self.max_iter,
+            self.threshold,
+            self.log_space,
+        )
+
+    @torch.no_grad()
+    def plan(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None):
+        """(n, d, d) transport plans, reconstructed in log space."""
+        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix)
+        log_u, log_v = self._iterate(source, target, cost_matrix)
+        return torch.exp(log_u.unsqueeze(-1) - cost_matrix / self.epsilon + log_v.unsqueeze(-2))
+
+    @torch.no_grad()
+    def potentials(self, source: Tensor, target: Tensor, cost_matrix: Optional[Tensor] = None):
+        """Centered dual potentials (f, g), each (n, d).
+
+        These are the gradients of the entropic OT cost w.r.t. the marginals.
+        """
+        source, target, cost_matrix = self.data_preprocess(source, target, cost_matrix)
+        log_u, log_v = self._iterate(source, target, cost_matrix)
+        f = self.epsilon * log_u
+        g = self.epsilon * log_v
+        return f - f.mean(dim=-1, keepdim=True), g - g.mean(dim=-1, keepdim=True)
+
+    def _iterate(self, a, b, cost_matrix):
+        iterate = _log_iterations if self.log_space else _scaling_iterations
+        return iterate(a, b, cost_matrix, self.epsilon, self.max_iter, self.threshold)