MadSpace stream handling for external users

[spinjo]

Several modifications to make the code more flexible for external users of the madspace package, focusing on the application of https://arxiv.org/abs/2608.23022. Also tried to generally make the stream handling in madspace cleaner/safer.

On a NVIDIA RTX PRO 500 Blackwell GPU this made a batchsize=1k evaluation of 3-body TPropagatorMapping.map_forward 2x faster (384mus -> 177mus, because cudaFreeAsync) and significantly reduced host occupation (3.21ms -> 0.37ms at batchsize=100k, because no stream sync). I didn't find any scenario that gets slower.

I used claude for the code edits.

Bug fixes

• MadSpace released memory while external kernels still read it; use cudaFreeAsync instead of cudaFree to fix this
• Imported inputs were not properly ordered; fixed by always passing the DLPack stream argument
• DLPack protocol ignored arguments stream, copy, max_version; now they are used
• GPU calls after CUDA runtime shutdown cause issues; now fixed
• Several smaller bugs in gradient handling with external tensors

New features

• External users can set the madspace stream through a context manager, for instance with madspace.stream(torch.cuda.current_stream().cuda_stream): ...; avoids overhead from stream syncs
• fork_streams and join_streams keep each stream clean until passing it back to the caller
• Stream ordering across the DLPack boundary
• release_inputs() to give manual control about when tensors held in madspace are cleaned

Comments

• Tested the changes on CPU and CUDA, not on HIP and ROCm
• Code only supports single-GPU; mismatched dl_device raises an error
• The _prev_caches, _prev_caches_backward, update_cached_tensors, MemPool::reset and _backward_wait_events functions are called only from commented-out code

[theoheimel]

Hi Jonas, thanks a lot! There is another PR on the way which touches memory management, so I'll review this one in detail once the other one is merged. On the first glance, it looks good to me.

I have one suggestion what could be improved: The DLPACK API defines a function current_work_stream that madspace could use to retrieve the CUDA stream. If we use that, we don't need our own context manager and it should work with the native with torch.cuda.stream(stream) (and the torch default stream). That should make the integration into torch even more seamless.