From 85398c354e2be72dccdb47c5ca0a8261a0166aaa Mon Sep 17 00:00:00 2001 From: GodKickMyAss <1195570512@qq.com> Date: Sat, 11 Jul 2026 07:39:24 +0000 Subject: [PATCH 1/2] perf: optimize inference kernels and collectives --- src/infinicore/tensor/copy.cc | 9 +++ .../nvidia/random_sample_kernel.cuh | 58 +++++++++++++------ 2 files changed, 49 insertions(+), 18 deletions(-) diff --git a/src/infinicore/tensor/copy.cc b/src/infinicore/tensor/copy.cc index 1297d9f8c..cb54a9fb9 100644 --- a/src/infinicore/tensor/copy.cc +++ b/src/infinicore/tensor/copy.cc @@ -50,6 +50,15 @@ void TensorImpl::copy_from(Tensor src) { context::memcpyH2D(local_src->data(), src->data(), copy_size); op::rearrange_(Tensor(const_cast(this)->shared_from_this()), local_src); } + } else { + 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(this)->shared_from_this()), local_src); + } } } } diff --git a/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh b/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh index 3c2696f6b..3b92e7fa1 100644 --- a/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh +++ b/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh @@ -3,6 +3,7 @@ #include #include #include +#include namespace op::random_sample::nvidia { @@ -50,6 +51,23 @@ static cudaError inclusiveSum( } // ↑↑↑ 重新封装 cub api，减少模板参数，方便调用 +// ↓↓↓ random sampling keeps token indices in 32-bit workspace and casts only at the output boundary. + +template +struct InternalSampleIndex { + using Type = Tidx; +}; + +template <> +struct InternalSampleIndex { + using Type = int32_t; +}; + +template <> +struct InternalSampleIndex { + using Type = uint32_t; +}; + // ↓↓↓ 计算 workspace // 地址对齐到 256 @@ -59,6 +77,7 @@ static constexpr size_t align256(size_t size) { template utils::Result calculateWorkspace(size_t n_) { + using TworkIdx = typename InternalSampleIndex::Type; const auto n = static_cast(n_); size_t argmax; @@ -70,14 +89,14 @@ utils::Result 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( + CHECK_CUDA((radixSort( nullptr, size_radix_sort, nullptr, nullptr, nullptr, nullptr, @@ -158,9 +177,9 @@ static __global__ void setSoftmaxMaxKernel( // 直接 for 循环遍历采样 // 这个 kernel 仅用于避免将数据拷贝到 cpu -template +template 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(indices_out[i]); return; } } @@ -218,6 +237,7 @@ struct Algo { void *stream_) const { using Tval = typename CudaTval::Type; + using TworkIdx = typename InternalSampleIndex::Type; auto stream = (cudaStream_t)stream_; auto logits = (Tval *)probs; @@ -226,14 +246,14 @@ struct Algo { auto workspace = reinterpret_cast(workspace_); auto workspace_end = workspace + workspace_size; - auto indices = reinterpret_cast(workspace); - workspace += align256(sizeof(Tidx) * n); + auto indices = reinterpret_cast(workspace); + workspace += align256(sizeof(TworkIdx) * n); auto sorted = reinterpret_cast(workspace); workspace += align256(sizeof(Tval) * n); - auto indices_out = reinterpret_cast(workspace); - workspace += align256(sizeof(Tidx) * n); + auto indices_out = reinterpret_cast(workspace); + workspace += align256(sizeof(TworkIdx) * n); workspace_ = reinterpret_cast(workspace); workspace_size = workspace_end - workspace; @@ -244,23 +264,25 @@ struct Algo { #endif auto grid = (n + block - 1) / block; // sort - fillIndices<<(grid), static_cast(block), 0, stream>>>(indices, static_cast(n)); - CHECK_CUDA(radixSort( + fillIndices<<(grid), static_cast(block), 0, stream>>>( + indices, static_cast(n)); + CHECK_CUDA((radixSort( workspace_, workspace_size, logits, sorted, indices, indices_out, static_cast(n), - stream)); + stream))); // softmax - partialSoftmaxKernel<<(grid), static_cast(block), 0, stream>>>(sorted, static_cast(n), temperature); - setSoftmaxMaxKernel<<<1, 1, 0, stream>>>(sorted); + partialSoftmaxKernel<<(grid), static_cast(block), 0, stream>>>( + sorted, static_cast(n), temperature); + setSoftmaxMaxKernel<<<1, 1, 0, stream>>>(sorted); // sum - CHECK_CUDA(inclusiveSum( - workspace_, workspace, + CHECK_CUDA(inclusiveSum( + workspace_, workspace_size, sorted, static_cast(n), stream)); // sample - randomSampleKernel<<<1, 1, 0, stream>>>( + randomSampleKernel<<<1, 1, 0, stream>>>( result, sorted, indices_out, n, random_val, topp, topk); From 913da748dc0651a8246fe96f3086e896153cbceb Mon Sep 17 00:00:00 2001 From: qinyiqun Date: Fri, 31 Jul 2026 06:42:39 +0000 Subject: [PATCH 2/2] feat(runtime): support asynchronous token handoff --- include/infinicore/context/context.hpp | 2 +- include/infinicore/tensor.hpp | 7 ++ python/infinicore/tensor.py | 4 ++ src/infinicore/context/context_impl.cc | 4 +- src/infinicore/context/runtime/runtime.cc | 10 ++- src/infinicore/context/runtime/runtime.hpp | 2 +- src/infinicore/device_event.cc | 18 +++++ src/infinicore/pybind11/context.hpp | 4 +- src/infinicore/pybind11/device_event.hpp | 3 +- src/infinicore/pybind11/tensor.hpp | 9 ++- src/infinicore/tensor/copy.cc | 65 +++++++++++++++++-- .../nvidia/random_sample_kernel.cuh | 4 +- 12 files changed, 116 insertions(+), 16 deletions(-) diff --git a/include/infinicore/context/context.hpp b/include/infinicore/context/context.hpp index 82b8cb309..dcff27804 100644 --- a/include/infinicore/context/context.hpp +++ b/include/infinicore/context/context.hpp @@ -29,7 +29,7 @@ std::shared_ptr allocateHostMemory(size_t size); std::shared_ptr 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); diff --git a/include/infinicore/tensor.hpp b/include/infinicore/tensor.hpp index b1650417d..68c04766c 100644 --- a/include/infinicore/tensor.hpp +++ b/include/infinicore/tensor.hpp @@ -186,6 +186,13 @@ class TensorImpl : public std::enable_shared_from_this { */ 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. diff --git a/python/infinicore/tensor.py b/python/infinicore/tensor.py index bbe801f93..670a993a1 100644 --- a/python/infinicore/tensor.py +++ b/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) diff --git a/src/infinicore/context/context_impl.cc b/src/infinicore/context/context_impl.cc index 00ed9a986..045065b59 100644 --- a/src/infinicore/context/context_impl.cc +++ b/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) { diff --git a/src/infinicore/context/runtime/runtime.cc b/src/infinicore/context/runtime/runtime.cc index 65d244f41..ea22b1022 100644 --- a/src/infinicore/context/runtime/runtime.cc +++ b/src/infinicore/context/runtime/runtime.cc @@ -74,7 +74,7 @@ std::shared_ptr Runtime::allocatePinnedHostMemory(size_t size) { } std::byte *data_ptr = pinned_host_memory_allocator_->allocate(size); return std::make_shared( - 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) { diff --git a/src/infinicore/context/runtime/runtime.hpp b/src/infinicore/context/runtime/runtime.hpp index 7dbb5185f..90f7112d7 100644 --- a/src/infinicore/context/runtime/runtime.hpp +++ b/src/infinicore/context/runtime/runtime.hpp @@ -45,7 +45,7 @@ class Runtime { std::shared_ptr reinstantiateBlob(std::shared_ptr 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); diff --git a/src/infinicore/device_event.cc b/src/infinicore/device_event.cc index 347e7effd..420d78287 100644 --- a/src/infinicore/device_event.cc +++ b/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); + } } } diff --git a/src/infinicore/pybind11/context.hpp b/src/infinicore/pybind11/context.hpp index 9c24a322e..ec357cdb1 100644 --- a/src/infinicore/pybind11/context.hpp +++ b/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()); m.def("sync_device", &syncDevice, "Synchronize the current device"); // Graph diff --git a/src/infinicore/pybind11/device_event.hpp b/src/infinicore/pybind11/device_event.hpp index f482422bb..1827a6f26 100644 --- a/src/infinicore/pybind11/device_event.hpp +++ b/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()) .def("query", &DeviceEvent::query, "Check if the event has been completed") diff --git a/src/infinicore/pybind11/tensor.hpp b/src/infinicore/pybind11/tensor.hpp index 8f117591b..a251288e7 100644 --- a/src/infinicore/pybind11/tensor.hpp +++ b/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()) + .def( + "to", + [](const Tensor &tensor, const Device &device) { return tensor->to(device); }, + py::call_guard()) .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); }) diff --git a/src/infinicore/tensor/copy.cc b/src/infinicore/tensor/copy.cc index cb54a9fb9..47ac38069 100644 --- a/src/infinicore/tensor/copy.cc +++ b/src/infinicore/tensor/copy.cc @@ -22,6 +22,14 @@ 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(this)->shared_from_this()), src); } else { @@ -29,16 +37,15 @@ void TensorImpl::copy_from(Tensor src) { 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(this)->shared_from_this()), local_src); } } else if (src->device().getType() == Device::Type::CPU) { @@ -51,6 +58,10 @@ void TensorImpl::copy_from(Tensor src) { op::rearrange_(Tensor(const_cast(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); @@ -63,6 +74,52 @@ void TensorImpl::copy_from(Tensor 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(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); + } +} + Tensor TensorImpl::contiguous() const { if (is_contiguous()) { return Tensor(const_cast(this)->shared_from_this()); diff --git a/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh b/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh index 3b92e7fa1..870ebbbff 100644 --- a/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh +++ b/src/infiniop/ops/random_sample/nvidia/random_sample_kernel.cuh @@ -1,9 +1,9 @@ #include "../../../devices/nvidia/nvidia_kernel_common.cuh" #include "infinicore.h" +#include #include #include #include -#include namespace op::random_sample::nvidia { @@ -51,7 +51,7 @@ static cudaError inclusiveSum( } // ↑↑↑ 重新封装 cub api，减少模板参数，方便调用 -// ↓↓↓ random sampling keeps token indices in 32-bit workspace and casts only at the output boundary. +// ↓↓↓ Random sampling keeps token indices in a 32-bit workspace and casts only at the output boundary. template struct InternalSampleIndex {