feat(runtime): add explicit asynchronous tensor copy support

include/infinicore/context/context.hpp

@@ -29,7 +29,7 @@ std::shared_ptr<Memory> allocateHostMemory(size_t size);
 std::shared_ptr<Memory> allocatePinnedHostMemory(size_t size);
 
 void memcpyH2D(void *dst, const void *src, size_t size, bool async = true);
-void memcpyD2H(void *dst, const void *src, size_t size);
+void memcpyD2H(void *dst, const void *src, size_t size, bool async = false);
 void memcpyD2D(void *dst, const void *src, size_t size, bool async = true);
 void memcpyH2H(void *dst, const void *src, size_t size);
 

include/infinicore/tensor.hpp

@@ -186,6 +186,13 @@ class TensorImpl : public std::enable_shared_from_this<TensorImpl> {
      */
     void copy_from(Tensor src);
 
+    /**
+     * Queue a non-blocking copy from another tensor on the active device stream.
+     * Cross-device copies require contiguous tensors. D2H destinations and H2D
+     * sources must use page-locked host memory.
+     */
+    void copy_from_async(Tensor src);
+
     /**
      * Return a tensor with the same data in contiguous arrangement as current tensor.
      * If this tensor is already contiguous, the original tensor is returned.

python/infinicore/tensor.py

@@ -81,6 +81,10 @@ def is_pinned(self):
     def copy_(self, src):
         self._underlying.copy_(src._underlying)
 
+    def copy_async_(self, src):
+        """Queue an explicit non-blocking copy on the active device stream."""
+        self._underlying.copy_async_(src._underlying)
+
     def to(self, *args, **kwargs):
         return Tensor(
             self._underlying.to(*tuple(arg._underlying for arg in args), **kwargs)

src/infinicore/context/context_impl.cc

@@ -138,8 +138,8 @@ void memcpyH2D(void *dst, const void *src, size_t size, bool async) {
     return ContextImpl::singleton().getCurrentRuntime()->memcpyH2D(dst, src, size, async);
 }
 
-void memcpyD2H(void *dst, const void *src, size_t size) {
-    return ContextImpl::singleton().getCurrentRuntime()->memcpyD2H(dst, src, size);
+void memcpyD2H(void *dst, const void *src, size_t size, bool async) {
+    return ContextImpl::singleton().getCurrentRuntime()->memcpyD2H(dst, src, size, async);
 }
 
 void memcpyD2D(void *dst, const void *src, size_t size, bool async) {

src/infinicore/context/runtime/runtime.cc

@@ -74,7 +74,7 @@ std::shared_ptr<Memory> Runtime::allocatePinnedHostMemory(size_t size) {
     }
     std::byte *data_ptr = pinned_host_memory_allocator_->allocate(size);
     return std::make_shared<Memory>(
-        data_ptr, size, device_,
+        data_ptr, size, Device::cpu(),
         [alloc = pinned_host_memory_allocator_.get()](std::byte *p) {
             alloc->deallocate(p);
         },
@@ -110,8 +110,12 @@ void Runtime::memcpyH2D(void *dst, const void *src, size_t size, bool async) {
     }
 }
 
-void Runtime::memcpyD2H(void *dst, const void *src, size_t size) {
-    INFINICORE_CHECK_ERROR(infinirtMemcpy(dst, src, size, INFINIRT_MEMCPY_D2H));
+void Runtime::memcpyD2H(void *dst, const void *src, size_t size, bool async) {
+    if (async) {
+        INFINICORE_CHECK_ERROR(infinirtMemcpyAsync(dst, src, size, INFINIRT_MEMCPY_D2H, stream_));
+    } else {
+        INFINICORE_CHECK_ERROR(infinirtMemcpy(dst, src, size, INFINIRT_MEMCPY_D2H));
+    }
 }
 
 void Runtime::memcpyD2D(void *dst, const void *src, size_t size, bool async) {

src/infinicore/context/runtime/runtime.hpp

@@ -45,7 +45,7 @@ class Runtime {
     std::shared_ptr<Memory> reinstantiateBlob(std::shared_ptr<Memory> blob);
 
     void memcpyH2D(void *dst, const void *src, size_t size, bool async = true);
-    void memcpyD2H(void *dst, const void *src, size_t size);
+    void memcpyD2H(void *dst, const void *src, size_t size, bool async = false);
     void memcpyD2D(void *dst, const void *src, size_t size, bool async = true);
 
     void setDeviceMemory(void *ptr, int value, size_t count);

src/infinicore/device_event.cc

@@ -42,7 +42,16 @@ DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept {
     if (this != &other) {
         // Clean up current resources
         if (event_ != nullptr) {
+            Device current_device = context::getDevice();
+            if (current_device != device_) {
+                context::setDevice(device_);
+            }
+
             context::destroyEvent(event_);
+
+            if (current_device != device_) {
+                context::setDevice(current_device);
+            }
         }
 
         // Transfer ownership
@@ -59,7 +68,16 @@ DeviceEvent &DeviceEvent::operator=(DeviceEvent &&other) noexcept {
 
 DeviceEvent::~DeviceEvent() {
     if (event_ != nullptr) {
+        Device current_device = context::getDevice();
+        if (current_device != device_) {
+            context::setDevice(device_);
+        }
+
         context::destroyEvent(event_);
+
+        if (current_device != device_) {
+            context::setDevice(current_device);
+        }
     }
 }
 

src/infinicore/pybind11/context.hpp

@@ -22,7 +22,9 @@ inline void bind(py::module &m) {
     m.def("get_stream", &getStream, "Get the current stream");
 
     // Synchronization
-    m.def("sync_stream", &syncStream, "Synchronize the current stream");
+    m.def("sync_stream", &syncStream,
+          "Synchronize the current stream",
+          py::call_guard<py::gil_scoped_release>());
     m.def("sync_device", &syncDevice, "Synchronize the current device");
 
     // Graph

src/infinicore/pybind11/device_event.hpp

@@ -21,7 +21,8 @@ inline void bind(py::module &m) {
              "Record the event on a specific stream", py::arg("stream"))
 
         .def("synchronize", &DeviceEvent::synchronize,
-             "Wait for the event to complete (blocking)")
+             "Wait for the event to complete (blocking)",
+             py::call_guard<py::gil_scoped_release>())
         .def("query", &DeviceEvent::query,
              "Check if the event has been completed")
 

src/infinicore/pybind11/tensor.hpp

@@ -29,7 +29,14 @@ inline void bind(py::module &m) {
         .def("debug", [](const Tensor &tensor, const std::string &filename) { return tensor->debug(filename); })
 
         .def("copy_", [](Tensor &tensor, const Tensor &other) { tensor->copy_from(other); })
-        .def("to", [](const Tensor &tensor, const Device &device) { return tensor->to(device); })
+        .def(
+            "copy_async_",
+            [](Tensor &tensor, const Tensor &other) { tensor->copy_from_async(other); },
+            py::call_guard<py::gil_scoped_release>())
+        .def(
+            "to",
+            [](const Tensor &tensor, const Device &device) { return tensor->to(device); },
+            py::call_guard<py::gil_scoped_release>())
         .def("contiguous", [](const Tensor &tensor) { return tensor->contiguous(); })
 
         .def("as_strided", [](const Tensor &tensor, const Shape &shape, const Strides &strides) { return tensor->as_strided(shape, strides); })

src/infinicore/tensor/copy.cc

@@ -22,23 +22,30 @@ void TensorImpl::copy_from(Tensor src) {
         throw std::runtime_error(
             "Cannot copy from tensor with different shape. Src: " + src->info() + " Dst: " + this->info());
     }
+    if (src->dtype() != this->dtype()) {
+        throw std::runtime_error(
+            "Cannot copy from tensor with different dtype. Src: " + src->info() + " Dst: " + this->info());
+    }
+    if (src->nbytes() != this->nbytes()) {
+        throw std::runtime_error(
+            "Cannot copy from tensor with different byte size. Src: " + src->info() + " Dst: " + this->info());
+    }
     if (this->device() == src->device()) {
         op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), src);
     } else {
         if (!src->is_contiguous()) {
             src = src->contiguous();
         }
 
-        // Use nbytes() to get the actual tensor size, not the full memory size
-        size_t copy_size = std::min(this->nbytes(), src->nbytes());
+        const size_t copy_size = this->nbytes();
         if (this->device().getType() == Device::Type::CPU) {
             if (this->is_contiguous()) {
                 context::setDevice(src->device());
-                context::memcpyD2H(this->data(), src->data(), copy_size);
+                context::memcpyD2H(this->data(), src->data(), copy_size, false);
             } else {
                 auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device());
                 context::setDevice(src->device());
-                context::memcpyD2H(local_src->data(), src->data(), copy_size);
+                context::memcpyD2H(local_src->data(), src->data(), copy_size, false);
                 op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), local_src);
             }
         } else if (src->device().getType() == Device::Type::CPU) {
@@ -50,7 +57,66 @@ void TensorImpl::copy_from(Tensor src) {
                 context::memcpyH2D(local_src->data(), src->data(), copy_size);
                 op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), local_src);
             }
+        } else {
+            if (this->device().getType() != src->device().getType()) {
+                throw std::runtime_error(
+                    "Cannot copy directly between different accelerator backends. Src: " + src->info() + " Dst: " + this->info());
+            }
+            context::setDevice(this->device());
+            if (this->is_contiguous()) {
+                context::memcpyD2D(this->data(), src->data(), copy_size);
+            } else {
+                auto local_src = Tensor::empty(this->shape(), this->dtype(), this->device());
+                context::memcpyD2D(local_src->data(), src->data(), copy_size);
+                op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), local_src);
+            }
+        }
+    }
+}
+
+void TensorImpl::copy_from_async(Tensor src) {
+    if (src->shape() != this->shape()) {
+        throw std::runtime_error(
+            "Cannot copy asynchronously from tensor with different shape. Src: " + src->info() + " Dst: " + this->info());
+    }
+    if (src->dtype() != this->dtype()) {
+        throw std::runtime_error(
+            "Cannot copy asynchronously from tensor with different dtype. Src: " + src->info() + " Dst: " + this->info());
+    }
+    if (src->nbytes() != this->nbytes()) {
+        throw std::runtime_error(
+            "Cannot copy asynchronously from tensor with different byte size. Src: " + src->info() + " Dst: " + this->info());
+    }
+
+    if (this->device() == src->device()) {
+        op::rearrange_(Tensor(const_cast<TensorImpl *>(this)->shared_from_this()), src);
+        return;
+    }
+    if (!this->is_contiguous() || !src->is_contiguous()) {
+        throw std::runtime_error(
+            "Asynchronous cross-device copy requires contiguous tensors. Src: " + src->info() + " Dst: " + this->info());
+    }
+
+    const size_t copy_size = this->nbytes();
+    if (this->device().getType() == Device::Type::CPU) {
+        if (!this->is_pinned()) {
+            throw std::runtime_error("Asynchronous D2H copy requires pinned destination memory");
+        }
+        context::setDevice(src->device());
+        context::memcpyD2H(this->data(), src->data(), copy_size, true);
+    } else if (src->device().getType() == Device::Type::CPU) {
+        if (!src->is_pinned()) {
+            throw std::runtime_error("Asynchronous H2D copy requires pinned source memory");
+        }
+        context::setDevice(this->device());
+        context::memcpyH2D(this->data(), src->data(), copy_size, true);
+    } else {
+        if (this->device().getType() != src->device().getType()) {
+            throw std::runtime_error(
+                "Cannot copy asynchronously between different accelerator backends. Src: " + src->info() + " Dst: " + this->info());
         }
+        context::setDevice(this->device());
+        context::memcpyD2D(this->data(), src->data(), copy_size, true);
     }
 }
 

src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh

@@ -1,5 +1,6 @@
 #include "../../../devices/nvidia/nvidia_kernel_common.cuh"
 #include "infinicore.h"
+#include <cstdint>
 #include <cub/device/device_radix_sort.cuh>
 #include <cub/device/device_reduce.cuh>
 #include <cub/device/device_scan.cuh>
@@ -50,6 +51,23 @@ static cudaError inclusiveSum(
 }
 
 // ↑↑↑ 重新封装 cub api，减少模板参数，方便调用
+// ↓↓↓ Random sampling keeps token indices in a 32-bit workspace and casts only at the output boundary.
+
+template <class Tidx>
+struct InternalSampleIndex {
+    using Type = Tidx;
+};
+
+template <>
+struct InternalSampleIndex<int64_t> {
+    using Type = int32_t;
+};
+
+template <>
+struct InternalSampleIndex<uint64_t> {
+    using Type = uint32_t;
+};
+
 // ↓↓↓ 计算 workspace
 
 // 地址对齐到 256
@@ -59,6 +77,7 @@ static constexpr size_t align256(size_t size) {
 
 template <class Tidx, class Tval>
 utils::Result<size_t> calculateWorkspace(size_t n_) {
+    using TworkIdx = typename InternalSampleIndex<Tidx>::Type;
     const auto n = static_cast<int>(n_);
 
     size_t argmax;
@@ -70,14 +89,14 @@ utils::Result<size_t> calculateWorkspace(size_t n_) {
     argmax += 256;
 
     // indices
-    size_t size_random = align256(sizeof(Tidx) * n);
+    size_t size_random = align256(sizeof(TworkIdx) * n);
     // sorted
     size_random += align256(sizeof(Tval) * n);
     // indices_out
-    size_random += align256(sizeof(Tidx) * n);
+    size_random += align256(sizeof(TworkIdx) * n);
     // cub device api
     size_t size_radix_sort;
-    CHECK_CUDA((radixSort<Tval, Tidx>(
+    CHECK_CUDA((radixSort<Tval, TworkIdx>(
         nullptr, size_radix_sort,
         nullptr, nullptr,
         nullptr, nullptr,
@@ -158,9 +177,9 @@ static __global__ void setSoftmaxMaxKernel(
 
 // 直接 for 循环遍历采样
 // 这个 kernel 仅用于避免将数据拷贝到 cpu
-template <class Tval, class Tidx>
+template <class Tval, class Tout, class Tidx>
 static __global__ void randomSampleKernel(
-    Tidx *__restrict__ result,
+    Tout *__restrict__ result,
     const Tval *__restrict__ sorted,
     const Tidx *__restrict__ indices_out,
     size_t n,
@@ -174,7 +193,7 @@ static __global__ void randomSampleKernel(
 #endif
     for (size_t i = 0;; ++i) {
         if ((sorted[i]) >= p) {
-            *result = indices_out[i];
+            *result = static_cast<Tout>(indices_out[i]);
             return;
         }
     }
@@ -218,6 +237,7 @@ struct Algo {
         void *stream_) const {
 
         using Tval = typename CudaTval<Tval_>::Type;
+        using TworkIdx = typename InternalSampleIndex<Tidx>::Type;
 
         auto stream = (cudaStream_t)stream_;
         auto logits = (Tval *)probs;
@@ -226,14 +246,14 @@ struct Algo {
         auto workspace = reinterpret_cast<size_t>(workspace_);
         auto workspace_end = workspace + workspace_size;
 
-        auto indices = reinterpret_cast<Tidx *>(workspace);
-        workspace += align256(sizeof(Tidx) * n);
+        auto indices = reinterpret_cast<TworkIdx *>(workspace);
+        workspace += align256(sizeof(TworkIdx) * n);
 
         auto sorted = reinterpret_cast<Tval *>(workspace);
         workspace += align256(sizeof(Tval) * n);
 
-        auto indices_out = reinterpret_cast<Tidx *>(workspace);
-        workspace += align256(sizeof(Tidx) * n);
+        auto indices_out = reinterpret_cast<TworkIdx *>(workspace);
+        workspace += align256(sizeof(TworkIdx) * n);
 
         workspace_ = reinterpret_cast<void *>(workspace);
         workspace_size = workspace_end - workspace;
@@ -244,23 +264,25 @@ struct Algo {
 #endif
         auto grid = (n + block - 1) / block;
         // sort
-        fillIndices<<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(indices, static_cast<int>(n));
-        CHECK_CUDA(radixSort(
+        fillIndices<TworkIdx><<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(
+            indices, static_cast<int>(n));
+        CHECK_CUDA((radixSort<Tval, TworkIdx>(
             workspace_, workspace_size,
             logits, sorted,
             indices, indices_out,
             static_cast<int>(n),
-            stream));
+            stream)));
         // softmax
-        partialSoftmaxKernel<<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(sorted, static_cast<int>(n), temperature);
-        setSoftmaxMaxKernel<<<1, 1, 0, stream>>>(sorted);
+        partialSoftmaxKernel<Tval><<<static_cast<unsigned int>(grid), static_cast<unsigned int>(block), 0, stream>>>(
+            sorted, static_cast<int>(n), temperature);
+        setSoftmaxMaxKernel<Tval><<<1, 1, 0, stream>>>(sorted);
         // sum
-        CHECK_CUDA(inclusiveSum(
-            workspace_, workspace,
+        CHECK_CUDA(inclusiveSum<Tval>(
+            workspace_, workspace_size,
             sorted, static_cast<int>(n),
             stream));
         // sample
-        randomSampleKernel<<<1, 1, 0, stream>>>(
+        randomSampleKernel<Tval, Tidx, TworkIdx><<<1, 1, 0, stream>>>(
             result,
             sorted, indices_out, n,
             random_val, topp, topk);