Skip to content

Fix Sinkhorn CUDA graph caching - #11

Open
maokovski wants to merge 3 commits into
mainfrom
fix/graph
Open

maokovski wants to merge 3 commits into
mainfrom
fix/graph

Conversation

@maokovski

@maokovski maokovski commented Oct 21, 2026 •

Copy link
Copy Markdown
Collaborator

Sinkhorn previously captured a new CUDA graph on every long fused solve. This change caches the graph and reuses it across compatible calls.

  • Refresh inputs and reset solver state before replay, then copy results into each call’s output buffers.
  • Rebuild the cache when the device, stream, tensor layout, or kernel configuration changes.
  • Exclude warmup from the requested iteration count.
  • Propagate graph errors instead of retrying with modified state.

Add regression tests for cache reuse with new inputs, exact iteration counts, output lifetime, configuration and stream changes, and error propagation.

Changed files:

  • torchmorph/optimal_transport.py
  • test/test_optimal_transport.py

Validation: Python syntax and diff whitespace checks passed. Runtime tests are delegated to the server CI; CUDA execution and results remain to be confirmed. No performance benchmark results are claimed.

@maokovski maokovski changed the title perf: reuse cached CUDA graphs for Sinkhorn iterations Fix Sinkhorn CUDA graph caching Oct 21, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant