diff --git a/madspace/include/madspace/driver/backend.hpp b/madspace/include/madspace/driver/backend.hpp
index 792d5828b5..7d859eb4d1 100644
--- a/madspace/include/madspace/driver/backend.hpp
+++ b/madspace/include/madspace/driver/backend.hpp
@@ -19,6 +19,7 @@ class Runtime {
         const std::vector<bool>& eval_grad,
         bool return_contiguous_grads = false
     ) = 0;
+    virtual void release_inputs() {}
     friend std::unique_ptr<Runtime>
     build_runtime(const Function& function, ContextPtr context, bool concurrent);
 
diff --git a/madspace/include/madspace/driver/context.hpp b/madspace/include/madspace/driver/context.hpp
index c58985541b..f6898f7a64 100644
--- a/madspace/include/madspace/driver/context.hpp
+++ b/madspace/include/madspace/driver/context.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
-#include <stdint.h>
+#include <cstdint>
+#include <optional>
 #include <unordered_map>
 
 #include "madspace/compgraphs.hpp"
@@ -200,6 +201,10 @@ ContextPtr default_cuda_context(std::size_t index = 0);
 ContextPtr default_hip_context(std::size_t index = 0);
 ContextPtr default_device_context(DevicePtr device);
 
+// thread-local stream to run on, empty to use our own and synchronize. 0 keeps ours
+std::optional<std::uintptr_t> caller_stream();
+void set_caller_stream(std::optional<std::uintptr_t> stream);
+
 inline std::string prefixed_name(const std::string& prefix, const std::string& name) {
     return prefix == "" ? name : std::format("{}.{}", prefix, name);
 }
diff --git a/madspace/include/madspace/driver/tensor.hpp b/madspace/include/madspace/driver/tensor.hpp
index 2ef40b2605..33e890e966 100644
--- a/madspace/include/madspace/driver/tensor.hpp
+++ b/madspace/include/madspace/driver/tensor.hpp
@@ -9,6 +9,7 @@
 #include <cstdint>
 #include <functional>
 #include <initializer_list>
+#include <optional>
 #include <vector>
 
 namespace madspace {
@@ -199,10 +200,14 @@ inline bool needs_zero_init(AllocHint hint) {
 
 class Device {
 public:
+    static constexpr bool stream_ordered_alloc = false;
+
     virtual ~Device() = default;
     virtual std::pair<void*, Tensor>
     allocate(std::size_t size, AllocHint hint) const = 0;
     virtual void free(void* ptr) const = 0;
+    virtual void free_on_stream(void* ptr, void* stream) const { free(ptr); }
+    virtual void order_streams(void* from, void* to) const {}
     virtual void memcpy(void* to, void* from, std::size_t size) const = 0;
     virtual void tensor_copy(const Tensor& source, Tensor& target) const = 0;
     virtual void tensor_zero(Tensor& tensor) const = 0;
@@ -386,7 +391,12 @@ class Tensor {
         );
     }
 
-    ~Tensor() { reset(); }
+    ~Tensor() {
+        try {
+            reset();
+        } catch (...) {
+        }
+    }
 
     Tensor& operator=(const Tensor& other) {
         reset();
@@ -398,7 +408,10 @@ class Tensor {
     }
 
     Tensor& operator=(Tensor&& other) noexcept {
-        reset();
+        try {
+            reset();
+        } catch (...) {
+        }
         impl = other.impl;
         other.impl = nullptr;
         return *this;
@@ -474,6 +487,18 @@ class Tensor {
         check_impl();
         return impl->device;
     }
+    std::optional<std::uintptr_t> stream() const {
+        return impl == nullptr ? std::nullopt : storage()->stream;
+    }
+    void set_stream(std::optional<std::uintptr_t> stream) {
+        if (impl == nullptr) {
+            return;
+        }
+        TensorImpl* owner = storage();
+        if (owner->stream_ordered) {
+            owner->stream = stream;
+        }
+    }
     std::size_t index_value() const {
         check_impl();
         if (impl->batch_sizes.size() > 0) {
@@ -499,21 +524,26 @@ class Tensor {
 
     std::size_t byte_size() const { return dtype_size() * shape().product(); }
 
-    void reset() {
+    void reset() { reset_on_stream(stream().value_or(0)); }
+
+    template <typename D>
+    void reset(const D& device) {
         if (impl == nullptr) {
             return;
         }
-        impl->reset(*impl->device);
+        // released before the free, which can throw
+        TensorImpl* owner = impl;
         impl = nullptr;
+        owner->reset(device);
     }
 
-    template <typename D>
-    void reset(const D& device) {
+    void reset_on_stream(std::uintptr_t stream) {
         if (impl == nullptr) {
             return;
         }
-        impl->reset(device);
+        TensorImpl* owner = impl;
         impl = nullptr;
+        owner->reset_on_stream(stream);
     }
 
     Tensor select(std::size_t axis, std::size_t index) const;
@@ -634,6 +664,8 @@ class Tensor {
         Sizes stride;
         std::size_t contiguous_dims;
         SizeVec batch_sizes;
+        bool stream_ordered = false;
+        std::optional<std::uintptr_t> stream;
 
         template <typename D>
         void reset(const D& device) {
@@ -651,9 +683,36 @@ class Tensor {
             delete this;
         }
 
+        void reset_on_stream(std::uintptr_t stream) {
+            if (ref_count.fetch_sub(1, std::memory_order_acq_rel) != 1) {
+                return;
+            }
+            if (owns_data && data != nullptr) {
+                if (stream_ordered) {
+                    device->free_on_stream(data, reinterpret_cast<void*>(stream));
+                } else {
+                    device->free(data);
+                }
+                --Tensor::tensor_count;
+            } else if (data_owner != nullptr) {
+                data_owner->reset_on_stream(stream);
+            } else if (external_reset) {
+                (*external_reset)();
+            }
+            delete this;
+        }
+
         void incref() { ref_count.fetch_add(1, std::memory_order_relaxed); }
     };
 
+    TensorImpl* storage() const {
+        TensorImpl* item = impl;
+        while (item->data_owner != nullptr) {
+            item = item->data_owner;
+        }
+        return item;
+    }
+
     Tensor(TensorImpl* _impl) : impl(_impl) {
         if (impl->data_owner != nullptr) {
             impl->data_owner->incref();
@@ -677,6 +736,10 @@ class Tensor {
             impl->data_owner = parent.impl;
         } else if (data != nullptr) {
             ++tensor_count;
+            if constexpr (D::stream_ordered_alloc) {
+                impl->stream_ordered = true;
+                impl->stream = reinterpret_cast<std::uintptr_t>(device.stream());
+            }
         }
     }
 
diff --git a/madspace/include/madspace/driver/thread_pool.hpp b/madspace/include/madspace/driver/thread_pool.hpp
index a8c34ca356..242859b977 100644
--- a/madspace/include/madspace/driver/thread_pool.hpp
+++ b/madspace/include/madspace/driver/thread_pool.hpp
@@ -72,15 +72,17 @@ class ThreadResource {
         std::optional<std::function<void(T&)>> destructor = std::nullopt
     ) :
         _pool(&pool),
+        _constructor(std::move(constructor)),
         _destructor(destructor),
-        _listener_id(pool.add_listener([this, constructor](std::size_t thread_count) {
+        _listener_id(pool.add_listener([this](std::size_t thread_count) {
             while (_resources.size() < thread_count) {
-                _resources.push_back(constructor());
+                _resources.emplace_back();
             }
         })) {
         for (std::size_t i = 0; i == 0 || i < pool.thread_count(); ++i) {
-            _resources.push_back(constructor());
+            _resources.emplace_back();
         }
+        get();
     }
     ~ThreadResource() {
         reset();
@@ -88,6 +90,7 @@ class ThreadResource {
     ThreadResource(ThreadResource&& other) noexcept :
         _pool(std::move(other._pool)),
         _resources(std::move(other._resources)),
+        _constructor(std::move(other._constructor)),
         _listener_id(std::move(other._listener_id)),
         _destructor(std::move(other._destructor)) {
         other._pool = nullptr;
@@ -97,6 +100,7 @@ class ThreadResource {
         reset();
         _pool = std::move(other._pool);
         _resources = std::move(other._resources);
+        _constructor = std::move(other._constructor);
         _listener_id = std::move(other._listener_id);
         _destructor = std::move(other._destructor);
         other._pool = nullptr;
@@ -104,13 +108,21 @@ class ThreadResource {
     }
     ThreadResource(const ThreadResource&) = delete;
     ThreadResource& operator=(const ThreadResource&) = delete;
-    T& get() { return _resources.at(ThreadPool::thread_index()); }
-    const T& get() const { return _resources.at(ThreadPool::thread_index()); }
+    T& get() const {
+        auto& [flag, item] = _resources.at(ThreadPool::thread_index());
+        std::call_once(flag, [&] {
+            std::unique_lock<std::mutex> lock(construction_mutex());
+            item.emplace(_constructor());
+        });
+        return *item;
+    }
     void reset() {
         if (_pool) {
             if (_destructor) {
-                for (auto& item : _resources) {
-                    _destructor.value()(item);
+                for (auto& [flag, item] : _resources) {
+                    if (item) {
+                        _destructor.value()(*item);
+                    }
                 }
             }
             _pool->remove_listener(_listener_id);
@@ -118,8 +130,14 @@ class ThreadResource {
     }
 
 private:
+    static std::mutex& construction_mutex() {
+        static std::mutex mutex;
+        return mutex;
+    }
+
     ThreadPool* _pool = nullptr;
-    std::vector<T> _resources;
+    mutable std::deque<std::pair<std::once_flag, std::optional<T>>> _resources;
+    std::function<T()> _constructor;
     std::size_t _listener_id;
     std::optional<std::function<void(T&)>> _destructor;
 };
diff --git a/madspace/madspace/_madspace_py_loader.py b/madspace/madspace/_madspace_py_loader.py
index 66c69baf21..3e5278b44f 100644
--- a/madspace/madspace/_madspace_py_loader.py
+++ b/madspace/madspace/_madspace_py_loader.py
@@ -1,3 +1,4 @@
+import contextlib
 import ctypes
 import logging
 import os
@@ -17,6 +18,20 @@
 from ._madspace_py import *
 
 
+@contextlib.contextmanager
+def stream(handle):
+    """
+    Run the calls on the given cuda stream instead of madspace's own, without
+    synchronizing. Keep it alive until everything madspace handed out has been dropped.
+    """
+    previous = get_stream()
+    set_stream(handle)
+    try:
+        yield
+    finally:
+        set_stream(previous)
+
+
 def _init():
     """
     Monkey-patch classes for a more pythonic experience.
@@ -27,7 +42,7 @@ def call_and_convert(runtime, args):
         if len(args) == 0:
             tensorlib = "numpy"
         else:
-            tensorlib = type(args[0]).__module__
+            tensorlib = type(args[0]).__module__.partition(".")[0]
         outputs = runtime.call(args)
         # Convert outputs, lazy-loading torch or numpy
         if tensorlib == "torch":
@@ -122,11 +137,19 @@ def log_handler(level, message):
             case Logger.level_error:
                 py_logger.error(message)
 
+    def release_inputs(self):
+        for name in ("runtime", "forward_runtime", "inverse_runtime"):
+            if hasattr(self, name):
+                getattr(self, name).release_inputs()
+
     FunctionRuntime.__call__ = runtime_call
     Function.__call__ = function_call
     FunctionGenerator.__call__ = function_generator_call
     Mapping.map_forward = map_forward
     Mapping.map_inverse = map_inverse
+    Function.release_inputs = release_inputs
+    FunctionGenerator.release_inputs = release_inputs
+    Mapping.release_inputs = release_inputs
     Tensor.numpy = tensor_numpy
     Tensor.torch = tensor_torch
     # Logger.set_log_handler(log_handler)
diff --git a/madspace/madspace/torch.py b/madspace/madspace/torch.py
index 9e0b4d35e0..43aec8619b 100644
--- a/madspace/madspace/torch.py
+++ b/madspace/madspace/torch.py
@@ -29,11 +29,14 @@ def __init__(
             ),
         )
 
+    def release_inputs(self) -> None:
+        self.runtime.release_inputs()
+
     def forward(self, *args: torch.Tensor) -> list[torch.Tensor]:
         if torch.is_grad_enabled():
             return AutogradWrapper.apply(self, self.dummy, *args)
         else:
-            outputs = self.runtime.call(args)
+            outputs = self.runtime.call([arg.detach() for arg in args])
             if len(outputs) == 1:
                 return torch.from_dlpack(outputs[0])
             else:
@@ -53,9 +56,8 @@ def forward(
         )
         ctx.module = module
         ctx.eval_grad = eval_grad
-        ctx.save_for_backward(
-            *(None if grad is None else torch.from_dlpack(grad) for grad in local_grads)
-        )
+        ctx.stream = me.get_stream()
+        ctx.stored_locals = local_grads
         if len(outputs) == 1:
             return torch.from_dlpack(outputs[0])
         else:
@@ -64,9 +66,10 @@ def forward(
     @staticmethod
     @once_differentiable
     def backward(ctx: FunctionCtx, *output_grads: torch.Tensor):
-        input_grads, global_grads = ctx.module.runtime.call_backward(
-            output_grads, ctx.saved_tensors, ctx.eval_grad
-        )
+        with me.stream(ctx.stream):
+            input_grads, global_grads = ctx.module.runtime.call_backward(
+                output_grads, ctx.stored_locals, ctx.eval_grad
+            )
         for name, grad in global_grads:
             if grad is None:
                 continue
diff --git a/madspace/src/cpu/runtime.cpp b/madspace/src/cpu/runtime.cpp
index 873897a4a8..5f068dbaa2 100644
--- a/madspace/src/cpu/runtime.cpp
+++ b/madspace/src/cpu/runtime.cpp
@@ -1172,8 +1172,22 @@ CpuRuntime::CpuRuntime(const Function& function, ContextPtr context, bool concur
         );
     }
 
+    std::vector<bool> grad_accumulated(function.locals().size());
+    SizeVec output_counts(function.locals().size());
+    for (auto& instr : function.instructions()) {
+        for (auto& in : instr.inputs) {
+            grad_accumulated.at(in.local_index) = true;
+        }
+    }
+    for (auto& out : function.outputs()) {
+        ++output_counts.at(out.local_index);
+    }
     for (auto& out : function.outputs()) {
         _output_indices.push_back(out.local_index);
+        _copy_output_grads.push_back(
+            grad_accumulated.at(out.local_index) ||
+            output_counts.at(out.local_index) > 1
+        );
     }
 }
 
@@ -1294,8 +1308,19 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_single(
     auto& device = CpuDevice::instance();
     TensorVec local_grads(stored_locals.size());
     TensorVec locals(stored_locals);
-    for (auto [index, grad] : zip(_output_indices, output_grads)) {
-        local_grads[index] = grad;
+    for (auto [index, grad, copy] :
+         zip(_output_indices, output_grads, _copy_output_grads)) {
+        auto& local_grad = local_grads[index];
+        if (!grad) {
+            continue;
+        }
+        if (local_grad) {
+            local_grad.add(grad);
+        } else if (copy) {
+            local_grad = grad.copy(AllocHint::local_grad);
+        } else {
+            local_grad = grad;
+        }
     }
 
     Tensor all_global_grads(
@@ -1304,6 +1329,9 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_single(
     // all_global_grads.zero();
     TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes);
     for (auto [index, grad] : zip(_grad_global_indices, global_grads)) {
+        if (local_grads[index]) {
+            grad.add(local_grads[index]);
+        }
         local_grads[index] = grad;
     }
 
@@ -1474,8 +1502,19 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_concurrent(
     auto& thread_pool = _context->thread_pool();
     TensorVec local_grads(stored_locals.size());
     TensorVec locals(stored_locals);
-    for (auto [index, grad] : zip(_output_indices, output_grads)) {
-        local_grads[index] = grad;
+    for (auto [index, grad, copy] :
+         zip(_output_indices, output_grads, _copy_output_grads)) {
+        auto& local_grad = local_grads[index];
+        if (!grad) {
+            continue;
+        }
+        if (local_grad) {
+            local_grad.add(grad);
+        } else if (copy) {
+            local_grad = grad.copy(AllocHint::local_grad);
+        } else {
+            local_grad = grad;
+        }
     }
 
     Tensor all_global_grads(
@@ -1484,6 +1523,9 @@ std::pair<TensorVec, TensorVec> CpuRuntime::run_backward_concurrent(
     // all_global_grads.zero();
     TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes);
     for (auto [index, grad] : zip(_grad_global_indices, global_grads)) {
+        if (local_grads[index]) {
+            grad.add(local_grads[index]);
+        }
         local_grads[index] = grad;
     }
 
diff --git a/madspace/src/cpu/runtime.hpp b/madspace/src/cpu/runtime.hpp
index 9b70fdbfc0..eb1a21ecbe 100644
--- a/madspace/src/cpu/runtime.hpp
+++ b/madspace/src/cpu/runtime.hpp
@@ -76,6 +76,7 @@ class CpuRuntime : public Runtime {
 
     std::vector<Instruction> _instructions;
     SizeVec _output_indices;
+    std::vector<bool> _copy_output_grads;
     std::size_t _input_count;
     TensorVec _locals_init;
     std::vector<bool> _requires_grad_init;
diff --git a/madspace/src/driver/context.cpp b/madspace/src/driver/context.cpp
index c279c09798..38c44ba2fd 100644
--- a/madspace/src/driver/context.cpp
+++ b/madspace/src/driver/context.cpp
@@ -3,6 +3,7 @@
 #include <dlfcn.h>
 #include <filesystem>
 #include <nlohmann/json.hpp>
+#include <stdexcept>
 #include <unordered_map>
 
 #include "madspace/driver/io.hpp"
@@ -12,6 +13,7 @@ using json = nlohmann::json;
 
 namespace {
 UmamiStatus umami_key_query_not_implemented(bool const**, int*) { return UMAMI_ERROR_NOT_IMPLEMENTED; }
+thread_local std::optional<std::uintptr_t> current_caller_stream;
 } // namespace
 
 MatrixElementApi::MatrixElementApi(
@@ -94,13 +96,16 @@ MatrixElementApi::MatrixElementApi(
         );
     }
 
-    _instances = ThreadResource<InstanceType>(thread_pool, [&, device] {
+    _instances = ThreadResource<InstanceType>(thread_pool, [this, device, param_card] {
         device->activate();
         void* instance;
         check_umami_status(_initialize(&instance, param_card.c_str()));
         return InstanceType(instance, [this, device](void* proc) {
-            device->activate();
-            _free(proc);
+            try {
+                device->activate();
+                _free(proc);
+            } catch (...) {
+            }
         });
     });
 }
@@ -349,7 +354,7 @@ ContextPtr madspace::default_hip_context(std::size_t index) {
 }
 
 ContextPtr madspace::default_device_context(DevicePtr device) {
-    static std::unordered_map<DevicePtr, ContextPtr> default_contexts;
+    static auto& default_contexts = *new std::unordered_map<DevicePtr, ContextPtr>;
     if (auto search = default_contexts.find(device); search != default_contexts.end()) {
         return search->second;
     } else {
@@ -358,3 +363,16 @@ ContextPtr madspace::default_device_context(DevicePtr device) {
         return context;
     }
 }
+
+std::optional<std::uintptr_t> madspace::caller_stream() {
+    return current_caller_stream;
+}
+
+void madspace::set_caller_stream(std::optional<std::uintptr_t> stream) {
+    if (stream && *stream == 2) {
+        throw std::invalid_argument("the per-thread default stream is not supported");
+    }
+    // cuda spells the legacy stream 1 and the per-thread one 2, we spell ours 0
+    current_caller_stream =
+        stream && *stream == 1 ? std::optional<std::uintptr_t>(0) : stream;
+}
diff --git a/madspace/src/gpu/device.cu b/madspace/src/gpu/device.cu
index 46e4bd85b0..0c8feff203 100644
--- a/madspace/src/gpu/device.cu
+++ b/madspace/src/gpu/device.cu
@@ -1,11 +1,27 @@
 #include "../kernels/kernels.hpp"
 #include "device.hpp"
+#include "madspace/driver/context.hpp"
 #include "tensor.cuh"
 
 using namespace madspace;
 using namespace madspace::gpu;
 using namespace madspace::kernels;
 
+namespace {
+template <typename... T>
+void wait_for(const T&... tensors) {
+    auto caller = caller_stream();
+    if (caller && *caller != 0) {
+        check_error(gpuStreamSynchronize(reinterpret_cast<gpuStream_t>(*caller)));
+    }
+    for (const Tensor* tensor : {&tensors...}) {
+        if (auto stream = tensor->stream()) {
+            check_error(gpuStreamSynchronize(reinterpret_cast<gpuStream_t>(*stream)));
+        }
+    }
+}
+} // namespace
+
 std::pair<void*, Tensor> GpuDevice::allocate(std::size_t size, AllocHint hint) const {
     activate();
     void* ptr;
@@ -21,6 +37,25 @@ void GpuDevice::free(void* ptr) const {
     check_error(gpuFree(ptr));
 }
 
+void GpuDevice::free_on_stream(void* ptr, void* stream) const {
+    activate();
+    check_error(gpuFreeAsync(ptr, static_cast<gpuStream_t>(stream)));
+}
+
+void GpuDevice::order_streams(void* from, void* to) const {
+    activate();
+    static thread_local std::vector<gpuEvent_t> events;
+    if (events.size() <= static_cast<std::size_t>(_index)) {
+        events.resize(_index + 1);
+    }
+    gpuEvent_t& event = events.at(_index);
+    if (!event) {
+        check_error(gpuEventCreate(&event));
+    }
+    check_error(gpuEventRecord(event, static_cast<gpuStream_t>(from)));
+    check_error(gpuStreamWaitEvent(static_cast<gpuStream_t>(to), event));
+}
+
 void GpuDevice::memcpy(void* to, void* from, std::size_t size) const {
     activate();
     check_error(gpuMemcpy(to, from, size, gpuMemcpyDefault));
@@ -28,24 +63,28 @@ void GpuDevice::memcpy(void* to, void* from, std::size_t size) const {
 
 void GpuDevice::tensor_copy(const Tensor& source, Tensor& target) const {
     activate();
+    wait_for(source, target);
     AsyncGpuDevice(*this, gpuStreamPerThread, 0).tensor_copy(source, target);
     check_error(gpuStreamSynchronize(gpuStreamPerThread));
 }
 
 void GpuDevice::tensor_zero(Tensor& tensor) const {
     activate();
+    wait_for(tensor);
     AsyncGpuDevice(*this, gpuStreamPerThread, 0).tensor_zero(tensor);
     check_error(gpuStreamSynchronize(gpuStreamPerThread));
 }
 
 void GpuDevice::tensor_add(const Tensor& source, Tensor& target) const {
     activate();
+    wait_for(source, target);
     AsyncGpuDevice(*this, gpuStreamPerThread, 0).tensor_add(source, target);
     check_error(gpuStreamSynchronize(gpuStreamPerThread));
 }
 
 void GpuDevice::tensor_cpu(const Tensor& source, Tensor& target) const {
     activate();
+    wait_for(source);
     check_error(
         gpuMemcpy(target.data(), source.data(), source.byte_size(), gpuMemcpyDefault)
     );
@@ -64,6 +103,7 @@ void GpuDevice::adam_step(
     double weight_decay
 ) const {
     activate();
+    wait_for(gradient, parameter, exp_avg, exp_avg_sq);
     AsyncGpuDevice device(*this, gpuStreamPerThread, 0);
     tensor_foreach_dynamic<kernel_adam_step<GpuTypes>, 1, 3>(
         {&gradient},
@@ -86,7 +126,7 @@ MemPool::MemPool(
         cached_sizes_and_tensors,
     gpuStream_t stream
 ) :
-    _device(device) {
+    _device(device), _stream(stream) {
     std::size_t pool_count = 0;
     for (auto& [pool_index, size, parent_tensor, zero_init] :
          cached_sizes_and_tensors) {
@@ -123,7 +163,10 @@ MemPool::~MemPool() {
             for (auto& [size, item] : stream_free_pointers) {
                 auto& [ptr, parent] = item;
                 if (!parent) {
-                    check_error(gpuFree(ptr));
+                    try {
+                        _device.free_on_stream(ptr, _stream);
+                    } catch (...) {
+                    }
                 }
             }
         }
diff --git a/madspace/src/gpu/device.hpp b/madspace/src/gpu/device.hpp
index 15231af77e..de46ee5f3f 100644
--- a/madspace/src/gpu/device.hpp
+++ b/madspace/src/gpu/device.hpp
@@ -32,6 +32,10 @@ inline void check_error(gpuError_t error) {
 
 inline void check_error() { check_error(gpuGetLastError()); }
 
+inline void ignore_error(gpuError_t) {}
+inline void ignore_error(gpublasStatus_t) {}
+inline void ignore_error(gpurandStatus_t) {}
+
 class GpuDevice : public Device {
 public:
 #ifdef __CUDACC__
@@ -42,6 +46,8 @@ class GpuDevice : public Device {
     virtual std::pair<void*, Tensor>
     allocate(std::size_t size, AllocHint hint) const override;
     void free(void* ptr) const override;
+    void free_on_stream(void* ptr, void* stream) const override;
+    void order_streams(void* from, void* to) const override;
     void memcpy(void* to, void* from, std::size_t size) const override;
 
     void tensor_copy(const Tensor& source, Tensor& target) const override;
@@ -124,10 +130,13 @@ class MemPool {
     std::vector<PoolItem> _pools;
     std::unordered_map<void*, AllocItem> _allocs;
     const GpuDevice& _device;
+    gpuStream_t _stream;
 };
 
 class AsyncGpuDevice {
 public:
+    static constexpr bool stream_ordered_alloc = true;
+
     AsyncGpuDevice(
         const GpuDevice& device,
         gpuStream_t stream,
diff --git a/madspace/src/gpu/gpu_abstraction.cuh b/madspace/src/gpu/gpu_abstraction.cuh
index 4677439b99..ad09c84d79 100644
--- a/madspace/src/gpu/gpu_abstraction.cuh
+++ b/madspace/src/gpu/gpu_abstraction.cuh
@@ -1,5 +1,11 @@
 #pragma once
 
+// our lanes are blocking streams, ordered against the legacy stream but not this one
+#if defined(CUDA_API_PER_THREAD_DEFAULT_STREAM) ||                                     \
+    defined(__HIP_API_PER_THREAD_DEFAULT_STREAM__)
+#error "madspace cannot be built with the per-thread default stream"
+#endif
+
 #ifdef __CUDACC__
 
 #include <cub/cub.cuh>
@@ -32,6 +38,9 @@
 #define gpuEventDestroy cudaEventDestroy
 #define gpuStreamWaitEvent cudaStreamWaitEvent
 #define gpuEventRecord cudaEventRecord
+#define gpuEventQuery cudaEventQuery
+#define gpuEventSynchronize cudaEventSynchronize
+#define gpuErrorNotReady cudaErrorNotReady
 #define gpuDeviceSynchronize cudaDeviceSynchronize
 
 #define gpublasStatus_t cublasStatus_t
@@ -90,6 +99,9 @@
 #define gpuEventDestroy hipEventDestroy
 #define gpuStreamWaitEvent(stream, event) hipStreamWaitEvent(stream, event, 0)
 #define gpuEventRecord hipEventRecord
+#define gpuEventQuery hipEventQuery
+#define gpuEventSynchronize hipEventSynchronize
+#define gpuErrorNotReady hipErrorNotReady
 #define gpuDeviceSynchronize hipDeviceSynchronize
 
 #define gpublasStatus_t rocblas_status
diff --git a/madspace/src/gpu/runtime.cu b/madspace/src/gpu/runtime.cu
index 83d6ca4c67..f90d6e7ed1 100644
--- a/madspace/src/gpu/runtime.cu
+++ b/madspace/src/gpu/runtime.cu
@@ -1504,11 +1504,27 @@ private:
     std::vector<bool> _sync_matrix;
 };
 
+struct StreamGuard {
+    gpuStream_t main_stream;
+    const std::vector<gpuStream_t>& streams;
+    bool dismissed = false;
+
+    ~StreamGuard() {
+        if (!dismissed) {
+            ignore_error(gpuStreamSynchronize(main_stream));
+            for (auto stream : streams) {
+                ignore_error(gpuStreamSynchronize(stream));
+            }
+        }
+    }
+};
+
 } // namespace
 
 GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
     _context(context),
     _input_count(function_arg.inputs().size()),
+    _last_stream(context->thread_pool(), []() { return std::optional<gpuStream_t>{}; }),
     _gpublas_handle(
         context->thread_pool(),
         []() {
@@ -1516,7 +1532,7 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
             check_error(gpublasCreate(&handle));
             return handle;
         },
-        [](gpublasHandle_t handle) { check_error(gpublasDestroy(handle)); }
+        [](gpublasHandle_t handle) { ignore_error(gpublasDestroy(handle)); }
     ),
     _gpurand_generator(
         context->thread_pool(),
@@ -1527,10 +1543,23 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
             check_error(gpurandSetPseudoRandomGeneratorSeed(handle, rand_dev()));
             return handle;
         },
-        [](gpurandGenerator_t handle) { check_error(gpurandDestroyGenerator(handle)); }
+        [](gpurandGenerator_t handle) { ignore_error(gpurandDestroyGenerator(handle)); }
     ),
     _prev_caches(context->thread_pool(), []() { return TensorVec{}; }),
-    _prev_caches_backward(context->thread_pool(), []() { return TensorVec{}; }) {
+    _prev_caches_backward(context->thread_pool(), []() { return TensorVec{}; }),
+    _held_inputs(
+        context->thread_pool(),
+        []() { return HeldInputs{}; },
+        [](HeldInputs& held) {
+            for (auto& [event, tensors] : held.items) {
+                ignore_error(gpuEventSynchronize(event));
+                ignore_error(gpuEventDestroy(event));
+            }
+            for (auto event : held.free_events) {
+                ignore_error(gpuEventDestroy(event));
+            }
+        }
+    ) {
     if (context->device()->device_type() != GpuDevice::gpu_device_type) {
         throw std::runtime_error("Context has incompatible device");
     }
@@ -1774,9 +1803,22 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
         );
     }
 
+    std::vector<bool> grad_accumulated(function.locals().size());
+    SizeVec output_counts(function.locals().size());
+    for (auto& instr : function.instructions()) {
+        for (auto& in : instr.inputs) {
+            grad_accumulated.at(in.local_index) = true;
+        }
+    }
+    for (auto& out : function.outputs()) {
+        ++output_counts.at(out.local_index);
+    }
     for (auto& out : function.outputs()) {
         _output_indices.push_back(out.local_index);
-        update_sync(out.local_index, 0, _wait_events);
+        _copy_output_grads.push_back(
+            grad_accumulated.at(out.local_index) ||
+            output_counts.at(out.local_index) > 1
+        );
     }
 
     _streams = ThreadResource<std::vector<gpuStream_t>>(
@@ -1784,17 +1826,25 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
         [stream_count]() {
             std::vector<gpuStream_t> streams(stream_count);
             for (auto& item : streams) {
+                // blocking, a caller on the legacy default stream relies on it
                 check_error(gpuStreamCreate(&item));
             }
             return streams;
         },
         [](auto& streams) {
             for (auto item : streams) {
-                check_error(gpuStreamDestroy(item));
+                ignore_error(gpuStreamDestroy(item));
             }
         }
     );
     std::size_t max_event_count = std::max(event_count, backward_event_count);
+    _stream_switch_event = max_event_count++;
+    if (stream_count > 1) {
+        _fork_event = max_event_count++;
+        for (std::size_t stream = 1; stream < stream_count; ++stream) {
+            _join_events.push_back(max_event_count++);
+        }
+    }
     _events = ThreadResource<std::vector<gpuEvent_t>>(
         context->thread_pool(),
         [max_event_count]() {
@@ -1806,24 +1856,119 @@ GpuRuntime::GpuRuntime(const Function& function_arg, ContextPtr context) :
         },
         [](auto& events) {
             for (auto item : events) {
-                check_error(gpuEventDestroy(item));
+                ignore_error(gpuEventDestroy(item));
             }
         }
     );
 }
 
+void GpuRuntime::fork_streams(
+    gpuStream_t main_stream,
+    const std::vector<gpuStream_t>& streams,
+    const std::vector<gpuEvent_t>& events
+) const {
+    if (!_fork_event) {
+        return;
+    }
+    check_error(gpuEventRecord(events.at(*_fork_event), main_stream));
+    for (std::size_t stream = 1; stream < streams.size(); ++stream) {
+        check_error(gpuStreamWaitEvent(streams.at(stream), events.at(*_fork_event)));
+    }
+}
+
+void GpuRuntime::join_streams(
+    gpuStream_t main_stream,
+    const std::vector<gpuStream_t>& streams,
+    const std::vector<gpuEvent_t>& events
+) const {
+    for (std::size_t stream = 1; stream < streams.size(); ++stream) {
+        gpuEvent_t event = events.at(_join_events.at(stream - 1));
+        check_error(gpuEventRecord(event, streams.at(stream)));
+        check_error(gpuStreamWaitEvent(main_stream, event));
+    }
+}
+
+void GpuRuntime::release_inputs() {
+    auto& held = _held_inputs.get();
+    gpuError_t error = gpuSuccess;
+    held.items.erase(
+        std::remove_if(
+            held.items.begin(),
+            held.items.end(),
+            [&](auto& item) {
+                auto status = gpuEventQuery(item.first);
+                if (status != gpuSuccess) {
+                    if (status != gpuErrorNotReady) {
+                        error = status;
+                    }
+                    return false;
+                }
+                held.free_events.push_back(item.first);
+                return true;
+            }
+        ),
+        held.items.end()
+    );
+    check_error(error);
+}
+
+void GpuRuntime::hold_inputs(
+    const TensorVec& inputs, gpuStream_t stream, bool legacy_caller
+) {
+    release_inputs();
+    auto handle = reinterpret_cast<std::uintptr_t>(stream);
+    TensorVec kept;
+    for (auto& input : inputs) {
+        auto input_stream = input.stream();
+        if (input && input_stream != handle && !(legacy_caller && input_stream == 0)) {
+            kept.push_back(input);
+        }
+    }
+    if (kept.empty()) {
+        return;
+    }
+    auto& held = _held_inputs.get();
+    gpuEvent_t event;
+    if (held.free_events.empty()) {
+        check_error(gpuEventCreate(&event));
+    } else {
+        event = held.free_events.back();
+        held.free_events.pop_back();
+    }
+    check_error(gpuEventRecord(event, stream));
+    held.items.emplace_back(event, std::move(kept));
+}
+
+// the cublas and curand handles are shared, so a new stream waits for the last one
+void GpuRuntime::switch_stream(
+    gpuStream_t main_stream, const std::vector<gpuEvent_t>& events
+) {
+    auto& last_stream = _last_stream.get();
+    if (last_stream && *last_stream != main_stream) {
+        check_error(gpuStreamWaitEvent(main_stream, events.at(_stream_switch_event)));
+    }
+    last_stream = main_stream;
+}
+
 TensorVec GpuRuntime::run(const TensorVec& inputs) {
     auto& gpu_device = *static_cast<const GpuDevice*>(_context->device());
+    gpu_device.activate();
     auto& streams = _streams.get();
     auto& events = _events.get();
-    gpu_device.activate();
     auto locals = _locals_init;
     std::copy(inputs.begin(), inputs.end(), locals.begin());
-    gpuStream_t main_stream = streams.at(0);
-    MemPool mem_pool(gpu_device, load_pool_size_cache(false), main_stream);
+    auto caller = caller_stream();
+    gpuStream_t main_stream =
+        caller && *caller != 0 ? reinterpret_cast<gpuStream_t>(*caller) : streams.at(0);
+    auto stream_handle = reinterpret_cast<std::uintptr_t>(main_stream);
+    TensorVec outputs;
+    switch_stream(main_stream, events);
+    MemPool mem_pool(gpu_device, load_pool_size_cache(false, !caller), main_stream);
+    StreamGuard stream_guard{main_stream, streams};
+    fork_streams(main_stream, streams, events);
 
     for (auto& instr : _instructions) {
-        gpuStream_t stream = streams.at(instr.stream);
+        gpuStream_t stream = instr.stream == 0 ? main_stream : streams.at(instr.stream);
         AsyncGpuDevice device(gpu_device, stream, instr.stream, &mem_pool);
         for (auto event : instr.wait_events) {
             check_error(gpuStreamWaitEvent(stream, events.at(event)));
@@ -1838,16 +1983,25 @@ TensorVec GpuRuntime::run(const TensorVec& inputs) {
             check_error(gpuEventRecord(events.at(instr.record_event), stream));
         }
     }
-    for (auto event : _wait_events) {
-        check_error(gpuStreamWaitEvent(main_stream, events.at(event)));
-    }
+    join_streams(main_stream, streams, events);
     update_pool_size_cache(mem_pool.total_sizes(), false);
     //update_cached_tensors(mem_pool.reset(main_stream), false);
-    TensorVec outputs;
     for (auto index : _output_indices) {
         outputs.push_back(locals[index]);
+        outputs.back().set_stream(caller);
     }
-    check_error(gpuStreamSynchronize(main_stream));
+    for (auto& local : locals) {
+        local.reset_on_stream(stream_handle);
+    }
+    if (caller) {
+        check_error(gpuEventRecord(events.at(_stream_switch_event), main_stream));
+        hold_inputs(inputs, main_stream, *caller == 0);
+    } else {
+        check_error(gpuStreamSynchronize(main_stream));
+        _last_stream.get().reset();
+        release_inputs();
+    }
+    stream_guard.dismissed = true;
     return outputs;
 }
 
@@ -1855,9 +2009,9 @@ std::tuple<TensorVec, TensorVec, std::vector<bool>> GpuRuntime::run_with_grad(
     const TensorVec& inputs, const std::vector<bool>& input_requires_grad
 ) {
     auto& gpu_device = *static_cast<const GpuDevice*>(_context->device());
+    gpu_device.activate();
     auto& streams = _streams.get();
     auto& events = _events.get();
-    gpu_device.activate();
     auto locals = _locals_init;
     auto requires_grad = _requires_grad_init;
     std::vector<bool> store_local(locals.size());
@@ -1866,11 +2020,17 @@ std::tuple<TensorVec, TensorVec, std::vector<bool>> GpuRuntime::run_with_grad(
     std::copy(
         input_requires_grad.begin(), input_requires_grad.end(), requires_grad.begin()
     );
-    gpuStream_t main_stream = streams.at(0);
-    MemPool mem_pool(gpu_device, load_pool_size_cache(false), main_stream);
+    auto caller = caller_stream();
+    gpuStream_t main_stream =
+        caller && *caller != 0 ? reinterpret_cast<gpuStream_t>(*caller) : streams.at(0);
+    TensorVec outputs;
+    switch_stream(main_stream, events);
+    MemPool mem_pool(gpu_device, load_pool_size_cache(false, !caller), main_stream);
+    StreamGuard stream_guard{main_stream, streams};
+    fork_streams(main_stream, streams, events);
 
     for (auto [instr, instr_eval_grad] : zip(_instructions, eval_grad)) {
-        gpuStream_t stream = streams.at(instr.stream);
+        gpuStream_t stream = instr.stream == 0 ? main_stream : streams.at(instr.stream);
         AsyncGpuDevice device(gpu_device, stream, instr.stream, &mem_pool);
         if (instr.differentiable) {
             for (auto input_index : instr.input_indices) {
@@ -1907,16 +2067,24 @@ std::tuple<TensorVec, TensorVec, std::vector<bool>> GpuRuntime::run_with_grad(
             check_error(gpuEventRecord(events.at(instr.record_event), stream));
         }
     }
-    for (auto event : _wait_events) {
-        check_error(gpuStreamWaitEvent(main_stream, events.at(event)));
-    }
+    join_streams(main_stream, streams, events);
     update_pool_size_cache(mem_pool.total_sizes(), false);
     //update_cached_tensors(mem_pool.reset(main_stream), false);
-    TensorVec outputs;
     for (auto index : _output_indices) {
         outputs.push_back(locals[index]);
     }
-    check_error(gpuStreamSynchronize(main_stream));
+    for (std::size_t i = _input_count; i < locals.size(); ++i) {
+        locals[i].set_stream(caller);
+    }
+    if (caller) {
+        check_error(gpuEventRecord(events.at(_stream_switch_event), main_stream));
+        hold_inputs(inputs, main_stream, *caller == 0);
+    } else {
+        check_error(gpuStreamSynchronize(main_stream));
+        _last_stream.get().reset();
+        release_inputs();
+    }
+    stream_guard.dismissed = true;
     return {outputs, locals, eval_grad};
 }
 
@@ -1927,18 +2095,32 @@ std::pair<TensorVec, TensorVec> GpuRuntime::run_backward(
     bool return_contiguous_grads
 ) {
     auto& gpu_device = *static_cast<const GpuDevice*>(_context->device());
+    gpu_device.activate();
     auto& streams = _streams.get();
     auto& events = _events.get();
-    gpu_device.activate();
     TensorVec local_grads(stored_locals.size());
     TensorVec locals(stored_locals);
-    for (auto [index, grad] : zip(_output_indices, output_grads)) {
-        local_grads[index] = grad;
-    }
-    gpuStream_t main_stream = streams.at(0);
-    MemPool mem_pool(gpu_device, load_pool_size_cache(true), main_stream);
-
+    auto caller = caller_stream();
+    gpuStream_t main_stream =
+        caller && *caller != 0 ? reinterpret_cast<gpuStream_t>(*caller) : streams.at(0);
+    switch_stream(main_stream, events);
+    MemPool mem_pool(gpu_device, load_pool_size_cache(true, !caller), main_stream);
+    StreamGuard stream_guard{main_stream, streams};
     AsyncGpuDevice init_device(gpu_device, main_stream, 0, &mem_pool);
+    for (auto [index, grad, copy] :
+         zip(_output_indices, output_grads, _copy_output_grads)) {
+        auto& local_grad = local_grads[index];
+        if (!grad) {
+            continue;
+        }
+        if (local_grad) {
+            local_grad.add(grad, init_device);
+        } else if (copy) {
+            local_grad = grad.copy(init_device, AllocHint::local_grad);
+        } else {
+            local_grad = grad;
+        }
+    }
     Tensor all_global_grads(
         DataType::dt_float,
         {_grad_global_total_size},
@@ -1948,6 +2130,9 @@ std::pair<TensorVec, TensorVec> GpuRuntime::run_backward(
     // all_global_grads.zero(init_device);
     TensorVec global_grads = all_global_grads.split_and_reshape(_grad_global_shapes);
     for (auto [index, grad] : zip(_grad_global_indices, global_grads)) {
+        if (local_grads[index]) {
+            grad.add(local_grads[index], init_device);
+        }
         local_grads[index] = grad;
     }
 
@@ -1987,7 +2172,21 @@ std::pair<TensorVec, TensorVec> GpuRuntime::run_backward(
     }*/
     update_pool_size_cache(mem_pool.total_sizes(), true);
     //update_cached_tensors(mem_pool.reset(main_stream), true);
-    check_error(gpuStreamSynchronize(main_stream));
+    for (auto& grad : local_grads) {
+        grad.set_stream(caller);
+    }
+    all_global_grads.set_stream(caller);
+    if (caller) {
+        check_error(gpuEventRecord(events.at(_stream_switch_event), main_stream));
+        TensorVec held(output_grads);
+        held.insert(held.end(), stored_locals.begin(), stored_locals.end());
+        hold_inputs(held, main_stream, *caller == 0);
+    } else {
+        check_error(gpuStreamSynchronize(main_stream));
+        _last_stream.get().reset();
+        release_inputs();
+    }
+    stream_guard.dismissed = true;
     return {
         {local_grads.begin(), local_grads.begin() + _input_count},
         return_contiguous_grads ? TensorVec{all_global_grads} : global_grads
@@ -1995,14 +2194,14 @@ std::pair<TensorVec, TensorVec> GpuRuntime::run_backward(
 }
 
 std::vector<std::tuple<std::size_t, std::size_t, Tensor, bool>>
-GpuRuntime::load_pool_size_cache(bool backward) {
+GpuRuntime::load_pool_size_cache(bool backward, bool synchronizes) {
     auto cache = backward ? _pool_size_cache_backward.load() : _pool_size_cache.load();
     std::vector<std::tuple<std::size_t, std::size_t, Tensor, bool>> ret;
     if (cache) {
         //auto& thread_prev_caches =
             //backward ? _prev_caches_backward.get() : _prev_caches.get();
         for (auto [pool_index, size] : *cache) {
-            Tensor new_cache = _context->cached_tensor(size);
+            Tensor new_cache = synchronizes ? _context->cached_tensor(size) : Tensor();
             /*if (pool_index < thread_prev_caches.size()) {
                 Tensor& prev_cache = thread_prev_caches.at(pool_index);
                 if (prev_cache && prev_cache.is_only_reference()) {
diff --git a/madspace/src/gpu/runtime.hpp b/madspace/src/gpu/runtime.hpp
index 2383e6e254..a571330710 100644
--- a/madspace/src/gpu/runtime.hpp
+++ b/madspace/src/gpu/runtime.hpp
@@ -7,6 +7,7 @@
 
 #include <memory>
 #include <map>
+#include <optional>
 
 namespace madspace {
 namespace gpu {
@@ -42,13 +43,20 @@ class GpuRuntime : public Runtime {
         const std::vector<bool>& eval_grad,
         bool return_contiguous_grads
     ) override;
+    void release_inputs() override;
     Context& context() { return *_context; }
     gpublasHandle_t gpublas_handle() { return _gpublas_handle.get(); }
     gpurandGenerator_t gpurand_generator() { return _gpurand_generator.get(); }
 
 private:
+    struct HeldInputs {
+        std::vector<std::pair<gpuEvent_t, TensorVec>> items;
+        std::vector<gpuEvent_t> free_events;
+    };
+
+    // a cached block can only be handed out again if the call that used it synchronized
     std::vector<std::tuple<std::size_t, std::size_t, Tensor, bool>>
-    load_pool_size_cache(bool backward);
+    load_pool_size_cache(bool backward, bool synchronizes);
     void update_pool_size_cache(
         const std::vector<std::pair<std::size_t, std::size_t>>& total_sizes,
         bool backward
@@ -56,8 +64,21 @@ class GpuRuntime : public Runtime {
     void update_cached_tensors(
         const std::vector<std::pair<std::size_t, Tensor>>& tensors, bool backward
     );
+    void fork_streams(
+        gpuStream_t main_stream,
+        const std::vector<gpuStream_t>& streams,
+        const std::vector<gpuEvent_t>& events
+    ) const;
+    void join_streams(
+        gpuStream_t main_stream,
+        const std::vector<gpuStream_t>& streams,
+        const std::vector<gpuEvent_t>& events
+    ) const;
+    void switch_stream(gpuStream_t main_stream, const std::vector<gpuEvent_t>& events);
+    void hold_inputs(const TensorVec& inputs, gpuStream_t stream, bool legacy_caller);
     std::vector<Instruction> _instructions;
     SizeVec _output_indices;
+    std::vector<bool> _copy_output_grads;
     std::size_t _input_count;
     TensorVec _locals_init;
     std::vector<bool> _requires_grad_init;
@@ -67,7 +88,10 @@ class GpuRuntime : public Runtime {
     ContextPtr _context;
     ThreadResource<std::vector<gpuStream_t>> _streams;
     ThreadResource<std::vector<gpuEvent_t>> _events;
-    std::vector<std::size_t> _wait_events;
+    ThreadResource<std::optional<gpuStream_t>> _last_stream;
+    std::size_t _stream_switch_event;
+    std::optional<std::size_t> _fork_event;
+    std::vector<std::size_t> _join_events;
     std::vector<std::size_t> _backward_wait_events;
     ThreadResource<gpublasHandle_t> _gpublas_handle;
     ThreadResource<gpurandGenerator_t> _gpurand_generator;
@@ -77,6 +101,8 @@ class GpuRuntime : public Runtime {
         _pool_size_cache_backward;
     ThreadResource<TensorVec> _prev_caches;
     ThreadResource<TensorVec> _prev_caches_backward;
+    // last, so it is destroyed first: it waits for the work using the members above
+    ThreadResource<HeldInputs> _held_inputs;
 };
 
 extern "C" Runtime*
diff --git a/madspace/src/python/function_runtime.cpp b/madspace/src/python/function_runtime.cpp
index e4a311e5e5..7d9ff4f693 100644
--- a/madspace/src/python/function_runtime.cpp
+++ b/madspace/src/python/function_runtime.cpp
@@ -1,10 +1,13 @@
 #include "function_runtime.hpp"
 
+#include <algorithm>
+#include <array>
 #include <format>
 #include <ranges>
 #include <stdexcept>
 
 #include "dlpack.h"
+#include "madspace/driver/context.hpp"
 
 using namespace madspace_py;
 using namespace pybind11::literals;
@@ -16,13 +19,52 @@ struct ManagerContext {
     std::vector<int64_t> stride;
     std::vector<int64_t> batch_sizes;
     Tensor tensor;
+    std::uintptr_t stream = 0;
 };
 
-void deleter(struct DLManagedTensor* self) {
-    delete static_cast<ManagerContext*>(self->manager_ctx);
+void deleter(struct DLManagedTensor* self) noexcept {
+    ManagerContext* context = static_cast<ManagerContext*>(self->manager_ctx);
+    try {
+        std::uintptr_t free_stream = context->tensor.stream().value_or(0);
+        if (free_stream != context->stream) {
+            context->tensor.device()->order_streams(
+                reinterpret_cast<void*>(context->stream),
+                reinterpret_cast<void*>(free_stream)
+            );
+        }
+    } catch (...) {
+    }
+    delete context;
     delete self;
 };
 
+std::uintptr_t consumer_stream(std::optional<std::int64_t> stream, int device_type) {
+    if (!stream) {
+        return 0;
+    }
+    if (*stream == 2 || (device_type == kDLROCM && *stream == 1)) {
+        throw py::buffer_error(
+            std::format("dlpack stream {} is not supported", *stream)
+        );
+    }
+    return *stream <= 1 ? 0 : static_cast<std::uintptr_t>(*stream);
+}
+
+bool orders_whole_stream(PyTypeObject* type) {
+    static std::array<py::handle, 2> torch_types;
+    if (!torch_types[0]) {
+        auto modules = py::module_::import("sys").attr("modules");
+        if (!modules.contains("torch")) {
+            return false;
+        }
+        py::object torch = modules["torch"];
+        torch_types[0] = py::object(torch.attr("Tensor")).release();
+        torch_types[1] = py::object(torch.attr("nn").attr("Parameter")).release();
+    }
+    PyObject* obj = reinterpret_cast<PyObject*>(type);
+    return obj == torch_types[0].ptr() || obj == torch_types[1].ptr();
+}
+
 Runtime* get_runtime(FunctionRuntime& func_runtime, DevicePtr expected_device) {
     Runtime* runtime;
     if (!expected_device) {
@@ -69,11 +111,18 @@ std::tuple<std::vector<Tensor>, Runtime*> check_and_convert_args(
     }
     std::vector<Tensor> inputs;
     DevicePtr expected_device = nullptr;
+    std::vector<PyTypeObject*> ordered_producers;
     for (int i = 0; i < n_args; ++i) {
         auto& arg = args.at(i);
         auto& input_type = func_runtime._function.inputs().at(i).type;
-        auto tensor =
-            dlpack_to_tensor(arg, input_type, i, expected_device, dlpack_version_cache);
+        auto tensor = dlpack_to_tensor(
+            arg,
+            input_type,
+            i,
+            expected_device,
+            dlpack_version_cache,
+            &ordered_producers
+        );
         if (expected_device == nullptr && tensor &&
             tensor.dtype() != DataType::batch_sizes) {
             expected_device = tensor.device();
@@ -100,16 +149,22 @@ std::tuple<int, int> madspace_py::dlpack_device(Tensor tensor) {
 
 py::object madspace_py::tensor_to_dlpack(
     Tensor tensor,
-    std::optional<int> stream,
+    std::optional<std::int64_t> stream,
     std::optional<std::tuple<int, int>> max_version,
-    std::optional<int> dl_device,
+    std::optional<std::tuple<int, int>> dl_device,
     std::optional<bool> copy
 ) {
-    // TODO: do something with the arguments
     if (!tensor) {
         return py::none();
     }
 
+    if (copy && *copy) {
+        throw py::buffer_error("dlpack export cannot copy the tensor");
+    }
+    if (dl_device && *dl_device != dlpack_device(tensor)) {
+        throw py::buffer_error("dlpack export cannot change the device");
+    }
+
     DLManagedTensor* dl_tensor;
     if (tensor.dtype() == DataType::batch_sizes) {
         ManagerContext* context = new ManagerContext{
@@ -141,13 +196,22 @@ py::object madspace_py::tensor_to_dlpack(
         default:
             break;
         }
+        auto [device_type, device_id] = dlpack_device(tensor);
+        std::uintptr_t consumer = consumer_stream(stream, device_type);
+        if (stream && *stream == -1) {
+            consumer = tensor.stream().value_or(consumer);
+        } else if (auto producer = tensor.stream(); producer && *producer != consumer) {
+            tensor.device()->order_streams(
+                reinterpret_cast<void*>(*producer), reinterpret_cast<void*>(consumer)
+            );
+        }
         ManagerContext* context = new ManagerContext{
             {tensor.shape().begin(), tensor.shape().end()},
             {tensor.stride().begin(), tensor.stride().end()},
             {},
-            tensor
+            tensor,
+            consumer
         };
-        auto [device_type, device_id] = dlpack_device(tensor);
         dl_tensor = new DLManagedTensor{
             {context->tensor.data(),
              DLDevice{static_cast<DLDeviceType>(device_type), device_id},
@@ -184,7 +248,8 @@ Tensor madspace_py::dlpack_to_tensor(
     std::optional<Type> expected_type,
     std::size_t arg_index,
     DevicePtr expected_device,
-    bool* dlpack_version_cache
+    bool* dlpack_version_cache,
+    std::vector<PyTypeObject*>* ordered_producers
 ) {
     if (tensor.is_none()) {
         return {};
@@ -193,17 +258,48 @@ Tensor madspace_py::dlpack_to_tensor(
     py::object dlpack_func = tensor.attr("__dlpack__");
     py::object capsule_obj;
 
+    py::dict stream_arg;
+    PyTypeObject* producer = Py_TYPE(tensor.ptr());
+    if (!(expected_type && expected_type->dtype == DataType::batch_sizes) &&
+        !(expected_device && expected_device->device_type() == DeviceType::cpu) &&
+        !(ordered_producers &&
+          std::find(ordered_producers->begin(), ordered_producers->end(), producer) !=
+              ordered_producers->end())) {
+        std::uintptr_t stream = madspace::caller_stream().value_or(0);
+        std::tuple<int, int> dl_device;
+        try {
+            dl_device = tensor.attr("__dlpack_device__")().cast<std::tuple<int, int>>();
+        } catch (const py::cast_error&) {
+            throw std::invalid_argument(
+                std::format(
+                    "Argument {}: __dlpack_device__ must return a pair of ints",
+                    arg_index + 1
+                )
+            );
+        }
+        auto [device_type, device_id] = dl_device;
+        if (device_id == 0 && (device_type == kDLCUDA || device_type == kDLROCM)) {
+            stream_arg["stream"] = device_type == kDLCUDA && stream == 0
+                ? 1
+                : static_cast<std::int64_t>(stream);
+            if (ordered_producers && orders_whole_stream(producer)) {
+                ordered_producers->push_back(producer);
+            }
+        }
+    }
+
     // catching exceptions is extremely expensive so we cache whether to use the new or
     // old version of the dlpack protocol
     if (dlpack_version_cache == nullptr || !*dlpack_version_cache) {
         try {
             capsule_obj = dlpack_func(
                 "max_version"_a =
-                    std::make_tuple(DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION)
+                    std::make_tuple(DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION),
+                **stream_arg
             );
         } catch (py::error_already_set& e) {
             if (e.matches(PyExc_TypeError)) {
-                capsule_obj = dlpack_func();
+                capsule_obj = dlpack_func(**stream_arg);
                 if (dlpack_version_cache != nullptr) {
                     *dlpack_version_cache = true;
                 }
@@ -213,11 +309,15 @@ Tensor madspace_py::dlpack_to_tensor(
         }
     } else {
         try {
-            capsule_obj = dlpack_func();
+            capsule_obj = dlpack_func(**stream_arg);
         } catch (py::error_already_set& e) {
+            if (!e.matches(PyExc_TypeError)) {
+                throw;
+            }
             capsule_obj = dlpack_func(
                 "max_version"_a =
-                    std::make_tuple(DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION)
+                    std::make_tuple(DLPACK_MAJOR_VERSION, DLPACK_MINOR_VERSION),
+                **stream_arg
             );
             *dlpack_version_cache = false;
         }
@@ -444,6 +544,12 @@ Tensor madspace_py::dlpack_to_tensor(
     return ret_tensor;
 }
 
+void FunctionRuntime::release_inputs() {
+    for (auto& entry : _runtimes) {
+        entry.second->release_inputs();
+    }
+}
+
 std::vector<Tensor> FunctionRuntime::call(std::vector<py::object> args) {
     auto [inputs, runtime] =
         check_and_convert_args(args, *this, &_dlpack_version_cache);
@@ -477,31 +583,34 @@ FunctionRuntime::call_backward(
     const std::vector<py::object>& stored_locals,
     const std::vector<bool>& eval_grad
 ) {
-    std::vector<Tensor> arg_out;
     DevicePtr expected_device = nullptr;
+    std::vector<PyTypeObject*> ordered_producers;
     std::size_t arg_index = 0;
-    for (auto& grad : output_grads) {
-        auto tensor = dlpack_to_tensor(
-            grad, std::nullopt, arg_index, expected_device, &_dlpack_version_cache
-        );
+    auto convert = [&](const py::object& arg) {
+        Tensor tensor = py::isinstance<Tensor>(arg)
+            ? arg.cast<Tensor>()
+            : dlpack_to_tensor(
+                  arg,
+                  std::nullopt,
+                  arg_index,
+                  expected_device,
+                  &_dlpack_version_cache,
+                  &ordered_producers
+              );
         if (expected_device == nullptr && tensor &&
             tensor.dtype() != DataType::batch_sizes) {
             expected_device = tensor.device();
         }
-        arg_out.push_back(tensor);
         ++arg_index;
+        return tensor;
+    };
+    std::vector<Tensor> arg_out;
+    for (auto& grad : output_grads) {
+        arg_out.push_back(convert(grad));
     }
     std::vector<Tensor> arg_locals;
     for (auto& local : stored_locals) {
-        auto tensor = dlpack_to_tensor(
-            local, std::nullopt, arg_index, expected_device, &_dlpack_version_cache
-        );
-        if (expected_device == nullptr && tensor &&
-            tensor.dtype() != DataType::batch_sizes) {
-            expected_device = tensor.device();
-        }
-        arg_locals.push_back(tensor);
-        ++arg_index;
+        arg_locals.push_back(convert(local));
     }
     // TODO: checks here
     Runtime* runtime = get_runtime(*this, expected_device);
diff --git a/madspace/src/python/function_runtime.hpp b/madspace/src/python/function_runtime.hpp
index f3bf205334..6d189be7d4 100644
--- a/madspace/src/python/function_runtime.hpp
+++ b/madspace/src/python/function_runtime.hpp
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <cstdint>
 #include <pybind11/pybind11.h>
 #include <unordered_map>
 #include <vector>
@@ -15,9 +16,9 @@ namespace madspace_py {
 std::tuple<int, int> dlpack_device(Tensor tensor);
 py::object tensor_to_dlpack(
     Tensor tensor,
-    std::optional<int> stream = std::nullopt,
+    std::optional<std::int64_t> stream = std::nullopt,
     std::optional<std::tuple<int, int>> max_version = std::nullopt,
-    std::optional<int> dl_device = std::nullopt,
+    std::optional<std::tuple<int, int>> dl_device = std::nullopt,
     std::optional<bool> copy = std::nullopt
 );
 Tensor dlpack_to_tensor(
@@ -25,7 +26,8 @@ Tensor dlpack_to_tensor(
     std::optional<Type> expected_type = std::nullopt,
     std::size_t arg_index = 0,
     DevicePtr expected_device = nullptr,
-    bool* dlpack_version_cache = nullptr
+    bool* dlpack_version_cache = nullptr,
+    std::vector<PyTypeObject*>* ordered_producers = nullptr
 );
 
 struct FunctionRuntime {
@@ -49,6 +51,7 @@ struct FunctionRuntime {
         const std::vector<py::object>& stored_locals,
         const std::vector<bool>& eval_grad
     );
+    void release_inputs();
 
     Function _function;
     ContextPtr _context;
diff --git a/madspace/src/python/madspace.cpp b/madspace/src/python/madspace.cpp
index 4d5d31c889..204bd08967 100644
--- a/madspace/src/python/madspace.cpp
+++ b/madspace/src/python/madspace.cpp
@@ -236,6 +236,7 @@ PYBIND11_MODULE(_madspace_py, m) {
         .def(
             "__dlpack__",
             &tensor_to_dlpack,
+            py::kw_only(),
             py::arg("stream") = std::nullopt,
             py::arg("max_version") = std::nullopt,
             py::arg("dl_device") = std::nullopt,
@@ -281,13 +282,16 @@ PYBIND11_MODULE(_madspace_py, m) {
     m.def("default_context", &default_context);
     m.def("default_cuda_context", &default_cuda_context, py::arg("index") = 0);
     m.def("default_hip_context", &default_hip_context, py::arg("index") = 0);
+    m.def("get_stream", &caller_stream);
+    m.def("set_stream", &set_caller_stream, py::arg("stream"));
 
     py::classh<FunctionRuntime>(m, "FunctionRuntime", py::dynamic_attr())
         .def(py::init<Function>(), py::arg("function"))
         .def(py::init<Function, ContextPtr>(), py::arg("function"), py::arg("context"))
         .def("call", &FunctionRuntime::call)
         .def("call_with_grad", &FunctionRuntime::call_with_grad)
-        .def("call_backward", &FunctionRuntime::call_backward);
+        .def("call_backward", &FunctionRuntime::call_backward)
+        .def("release_inputs", &FunctionRuntime::release_inputs);
 
     auto& fb =
         py::classh<FunctionBuilder>(m, "FunctionBuilder")
@@ -1000,10 +1004,16 @@ PYBIND11_MODULE(_madspace_py, m) {
             "add_data",
             [](DiscreteOptimizer& opt, std::vector<py::object> values_and_counts) {
                 TensorVec input_tensors;
+                std::vector<PyTypeObject*> ordered_producers;
                 for (std::size_t i = 1; auto& input : values_and_counts) {
-                    input_tensors.push_back(
-                        dlpack_to_tensor(input, i % 2 == 0 ? batch_int : batch_float, i)
-                    );
+                    input_tensors.push_back(dlpack_to_tensor(
+                        input,
+                        i % 2 == 0 ? batch_int : batch_float,
+                        i,
+                        nullptr,
+                        nullptr,
+                        &ordered_producers
+                    ));
                     ++i;
                 }
                 opt.add_data(input_tensors);
@@ -1056,11 +1066,17 @@ PYBIND11_MODULE(_madspace_py, m) {
                 TensorVec tensors;
                 tensors.reserve(inputs.size());
                 bool dlpack_version_cache = false;
+                std::vector<PyTypeObject*> ordered_producers;
                 for (std::size_t i = 0;
                      auto [input, type] : zip(inputs, opt.input_types())) {
-                    tensors.push_back(
-                        dlpack_to_tensor(input, type, i, device, &dlpack_version_cache)
-                    );
+                    tensors.push_back(dlpack_to_tensor(
+                        input,
+                        type,
+                        i,
+                        device,
+                        &dlpack_version_cache,
+                        &ordered_producers
+                    ));
                     ++i;
                 }
                 return opt.step(tensors);