From 375cda0282e273b2234dd6b24c01a0d5a610c287 Mon Sep 17 00:00:00 2001
From: Bradley Dice <bdice@bradleydice.com>
Date: Sun, 16 Aug 2026 02:04:26 -0500
Subject: [PATCH] Select batched memcpy flag by copy size

Signed-off-by: Bradley Dice <bdice@bradleydice.com>
---
 cpp/include/cudf/detail/utilities/cuda_memcpy.hpp |  8 +++-----
 cpp/src/utilities/cuda_memcpy.cu                  | 14 ++++++++++----
 2 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp b/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp
index 3764a1f345b..6c4997f5504 100644
--- a/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp
+++ b/cpp/include/cudf/detail/utilities/cuda_memcpy.hpp
@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -28,10 +28,8 @@ void cuda_memcpy_async_impl(
  * satisfied, but for host memory the caller must ensure the source is not freed before the stream
  * is synchronized.
  *
- * All copies share a single attribute entry (`cudaMemcpySrcAccessOrderStream` +
- * `cudaMemcpyFlagPreferOverlapWithCompute`). Per-copy attributes are not supported by this
- * wrapper; callers requiring different attributes per copy should call `cudaMemcpyBatchAsync`
- * directly.
+ * A batch uses `cudaMemcpyFlagPreferOverlapWithCompute` when every copy is 128 KiB or less. If
+ * any copy is larger, the batch uses `cudaMemcpyFlagDefault`.
  *
  * @param dsts Host pointer to a list of destination pointers.
  * @param srcs Host pointer to a list of source pointers.
diff --git a/cpp/src/utilities/cuda_memcpy.cu b/cpp/src/utilities/cuda_memcpy.cu
index 8e961a9cf63..d4f103fd9d4 100644
--- a/cpp/src/utilities/cuda_memcpy.cu
+++ b/cpp/src/utilities/cuda_memcpy.cu
@@ -21,6 +21,8 @@ namespace cudf::detail {
 
 namespace {
 
+constexpr std::size_t prefer_overlap_threshold = 128 * 1024;
+
 // Simple kernel to copy between device buffers
 CUDF_KERNEL void copy_kernel(char const* __restrict__ src, char* __restrict__ dst, size_t n)
 {
@@ -90,10 +92,14 @@ cudaError_t memcpy_batch_async(void* const* dsts,
       count = valid_dsts.size();
     }
 
-    cudaMemcpyAttributes attrs = {.srcAccessOrder = cudaMemcpySrcAccessOrderStream,
-                                  .flags          = cudaMemcpyFlagPreferOverlapWithCompute};
-    std::size_t attrs_idxs     = 0;
-    return cudaMemcpyBatchAsync(dsts, srcs, sizes, count, &attrs, &attrs_idxs, 1, stream.value());
+    auto const flags =
+      std::ranges::any_of(std::ranges::views::iota(std::size_t{0}, count),
+                          [&](auto i) { return sizes[i] > prefer_overlap_threshold; })
+        ? cudaMemcpyFlagDefault
+        : cudaMemcpyFlagPreferOverlapWithCompute;
+    cudaMemcpyAttributes attrs = {.srcAccessOrder = cudaMemcpySrcAccessOrderStream, .flags = flags};
+    std::size_t attrs_idx      = 0;
+    return cudaMemcpyBatchAsync(dsts, srcs, sizes, count, &attrs, &attrs_idx, 1, stream.value());
   }
 #endif  // CUDART_VERSION >= 13000
   for (std::size_t i = 0; i < count; ++i) {