Use cudaMemcpyBatchAsync for RMM copies

cpp/CMakeLists.txt

@@ -1,6 +1,6 @@
 # =============================================================================
 # cmake-format: off
-# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION.
+# SPDX-FileCopyrightText: Copyright (c) 2018-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 # cmake-format: on
 # =============================================================================
@@ -80,6 +80,7 @@ add_library(
   rmm
   src/aligned.cpp
   src/cuda_device.cpp
+  src/cuda_memcpy.cpp
   src/cuda_stream.cpp
   src/cuda_stream_pool.cpp
   src/cuda_stream_view.cpp

cpp/include/rmm/detail/cuda_memcpy.hpp

@@ -0,0 +1,24 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <rmm/detail/cuda_stream.hpp>
+#include <rmm/detail/export.hpp>
+
+#include <cuda_runtime_api.h>
+
+#include <cstddef>
+
+RMM_NAMESPACE_BEGIN
+namespace detail {
+
+[[nodiscard]] RMM_EXPORT cudaError_t memcpy_async(void* dst,
+                                                  void const* src,
+                                                  std::size_t count,
+                                                  cuda::stream_ref stream);
+
+}  // namespace detail
+RMM_NAMESPACE_END

cpp/include/rmm/detail/cuda_stream.hpp

@@ -0,0 +1,18 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#pragma once
+
+#include <rmm/detail/export.hpp>
+
+#include <cuda/stream_ref>
+
+RMM_NAMESPACE_BEGIN
+namespace detail {
+
+[[nodiscard]] bool is_default_stream(cuda::stream_ref stream) noexcept;
+
+}  // namespace detail
+RMM_NAMESPACE_END

cpp/include/rmm/device_uvector.hpp

@@ -6,6 +6,7 @@
 #pragma once
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/detail/cuda_memcpy.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/detail/exec_check_disable.hpp>
 #include <rmm/detail/export.hpp>
@@ -216,8 +217,8 @@ class device_uvector {
   {
     RMM_EXPECTS(
       element_index < size(), "Attempt to access out of bounds element.", rmm::out_of_range);
-    RMM_CUDA_TRY(cudaMemcpyAsync(
-      element_ptr(element_index), &value, sizeof(value), cudaMemcpyDefault, stream.value()));
+    RMM_CUDA_TRY(
+      rmm::detail::memcpy_async(element_ptr(element_index), &value, sizeof(value), stream));
   }
 
   // We delete the r-value reference overload to prevent asynchronously copying from a literal or
@@ -306,8 +307,8 @@ class device_uvector {
     RMM_EXPECTS(
       element_index < size(), "Attempt to access out of bounds element.", rmm::out_of_range);
     value_type value;
-    RMM_CUDA_TRY(cudaMemcpyAsync(
-      &value, element_ptr(element_index), sizeof(value), cudaMemcpyDefault, stream.value()));
+    RMM_CUDA_TRY(
+      rmm::detail::memcpy_async(&value, element_ptr(element_index), sizeof(value), stream));
     stream.synchronize();
     return value;
   }

cpp/src/cuda_memcpy.cpp

@@ -0,0 +1,29 @@
+/*
+ * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+#include <rmm/detail/cuda_memcpy.hpp>
+
+RMM_NAMESPACE_BEGIN
+namespace detail {
+
+cudaError_t memcpy_async(void* dst, void const* src, std::size_t count, cuda::stream_ref stream)
+{
+  if (count == 0) { return cudaSuccess; }
+
+#if defined(CUDART_VERSION) && CUDART_VERSION >= 13000
+  if (!is_default_stream(stream)) {
+    cudaMemcpyAttributes attrs{};
+    attrs.srcAccessOrder = cudaMemcpySrcAccessOrderStream;
+    attrs.flags          = cudaMemcpyFlagPreferOverlapWithCompute;
+    std::size_t attr_idx = 0;
+    return cudaMemcpyBatchAsync(&dst, &src, &count, 1, &attrs, &attr_idx, 1, stream.get());
+  }
+#endif
+
+  return cudaMemcpyAsync(dst, src, count, cudaMemcpyDefault, stream.get());
+}
+
+}  // namespace detail
+RMM_NAMESPACE_END

cpp/src/cuda_stream_view.cpp

@@ -4,6 +4,7 @@
  */
 
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/detail/cuda_stream.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/detail/export.hpp>
 
@@ -34,14 +35,23 @@ bool cuda_stream_view::is_per_thread_default() const noexcept
 }
 
 bool cuda_stream_view::is_default() const noexcept
+{
+  return detail::is_default_stream(static_cast<cuda::stream_ref>(*this));
+}
+
+namespace detail {
+
+bool is_default_stream(cuda::stream_ref stream) noexcept
 {
 #ifdef CUDA_API_PER_THREAD_DEFAULT_STREAM
-  return *this == cuda_stream_legacy;
+  return stream == cudaStreamLegacy;
 #else
-  return *this == cuda_stream_legacy || value() == nullptr;
+  return stream == cudaStreamLegacy || stream == cudaStream_t{};
 #endif
 }
 
+}  // namespace detail
+
 void cuda_stream_view::synchronize() const { RMM_CUDA_TRY(cudaStreamSynchronize(stream_)); }
 
 void cuda_stream_view::synchronize_no_throw() const noexcept

cpp/src/device_buffer.cpp

@@ -4,6 +4,7 @@
  */
 
 #include <rmm/aligned.hpp>
+#include <rmm/detail/cuda_memcpy.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/device_buffer.hpp>
 #include <rmm/error.hpp>
@@ -138,7 +139,7 @@ void device_buffer::copy_async(void const* source, std::size_t bytes)
     RMM_EXPECTS(nullptr != source, "Invalid copy from nullptr.");
     RMM_EXPECTS(nullptr != _data, "Invalid copy to nullptr.");
 
-    RMM_CUDA_TRY(cudaMemcpyAsync(_data, source, bytes, cudaMemcpyDefault, stream().value()));
+    RMM_CUDA_TRY(rmm::detail::memcpy_async(_data, source, bytes, stream()));
   }
 }
 
@@ -149,7 +150,7 @@ void device_buffer::reserve(std::size_t new_capacity, cuda_stream_view stream)
     cuda_set_device_raii dev{_device};
     auto tmp            = device_buffer{new_capacity, alignment(), stream, _mr};
     auto const old_size = size();
-    RMM_CUDA_TRY(cudaMemcpyAsync(tmp.data(), data(), size(), cudaMemcpyDefault, stream.value()));
+    RMM_CUDA_TRY(rmm::detail::memcpy_async(tmp.data(), data(), size(), stream));
     *this = std::move(tmp);
     _size = old_size;
   }
@@ -165,7 +166,7 @@ void device_buffer::resize(std::size_t new_size, cuda_stream_view stream)
   } else {
     cuda_set_device_raii dev{_device};
     auto tmp = device_buffer{new_size, alignment(), stream, _mr};
-    RMM_CUDA_TRY(cudaMemcpyAsync(tmp.data(), data(), size(), cudaMemcpyDefault, stream.value()));
+    RMM_CUDA_TRY(rmm::detail::memcpy_async(tmp.data(), data(), size(), stream));
     *this = std::move(tmp);
   }
 }

cpp/tests/cuda_stream_tests.cpp

@@ -1,10 +1,11 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <rmm/cuda_stream.hpp>
 #include <rmm/cuda_stream_view.hpp>
+#include <rmm/detail/cuda_stream.hpp>
 #include <rmm/device_buffer.hpp>
 
 #include <cuda/stream_ref>
@@ -53,6 +54,21 @@ TEST_F(CudaStreamTest, StreamRefConsistentWithView)
   EXPECT_EQ(ref_from_stream, ref_from_view);
 }
 
+TEST_F(CudaStreamTest, IsDefaultStream)
+{
+  rmm::cuda_stream stream;
+
+  EXPECT_FALSE(rmm::detail::is_default_stream(stream));
+
+  EXPECT_EQ(rmm::cuda_stream_default.is_default(),
+            rmm::detail::is_default_stream(rmm::cuda_stream_default));
+  EXPECT_EQ(rmm::cuda_stream_legacy.is_default(),
+            rmm::detail::is_default_stream(rmm::cuda_stream_legacy));
+  EXPECT_EQ(rmm::cuda_stream_per_thread.is_default(),
+            rmm::detail::is_default_stream(rmm::cuda_stream_per_thread));
+  EXPECT_EQ(stream.view().is_default(), rmm::detail::is_default_stream(stream));
+}
+
 TEST_F(CudaStreamTest, MoveConstructor)
 {
   rmm::cuda_stream stream_a;

cpp/tests/device_buffer_tests.cu

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2019-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -650,6 +650,14 @@ TEST(DeviceBufferAlignmentTest, DefaultConstructedResizeLarger)
   EXPECT_EQ(buff.size(), 100);
 }
 
+TEST(DeviceBufferAlignmentTest, DefaultConstructedResizeLargerOnNonDefaultStream)
+{
+  rmm::cuda_stream stream;
+  rmm::device_buffer buff;
+  EXPECT_NO_THROW(buff.resize(100, stream.view()));
+  EXPECT_EQ(buff.size(), 100);
+}
+
 TEST(DeviceBufferAlignmentTest, DefaultConstructedReserveLarger)
 {
   rmm::device_buffer buff;

cpp/tests/device_uvector_tests.cpp

@@ -1,10 +1,11 @@
 
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2020-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
 #include <rmm/aligned.hpp>
+#include <rmm/cuda_stream.hpp>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/detail/error.hpp>
 #include <rmm/device_uvector.hpp>
@@ -258,6 +259,15 @@ TYPED_TEST(TypedUVectorTest, GetSetElementAsync)
   }
 }
 
+TEST(DeviceUVectorMemcpyTest, GetSetElementOnNonDefaultStream)
+{
+  rmm::cuda_stream stream;
+  rmm::device_uvector<int> vec(1, stream.view());
+  int const value = 42;
+  vec.set_element_async(0, value, stream.view());
+  EXPECT_EQ(vec.element(0, stream.view()), value);
+}
+
 TYPED_TEST(TypedUVectorTest, SetElementZeroAsync)
 {
   auto const size{100};