Skip to content

UCT/CUDA_IPC: Skip zero-length elements in SGL zcopy #12219

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
dmitrygx merged 12 commits into from
Nov 8, 2026
Merged

UCT/CUDA_IPC: Skip zero-length elements in SGL zcopy #12219

Show file tree
Hide file tree
Changes from 2 commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
68ba8c2
GTEST/UCT: Check SGL zcopy with a zero-length element
dmitrygx Nov 3, 2026
b7e2134
UCT/CUDA_IPC: Skip zero-length elements in SGL zcopy
dmitrygx Nov 3, 2026
5f70e97
GTEST/UCT: Check SGL zcopy with all elements of zero length
dmitrygx Nov 3, 2026
eefee38
GTEST/UCT: Check SGL zcopy with empty first and last elements
dmitrygx Nov 3, 2026
b661636
UCT/CUDA_IPC: Compact SGL arrays only if some element is empty
dmitrygx Nov 3, 2026
cf3ae7e
UCT/API: Document zero-length elements of SGL zcopy
dmitrygx Nov 3, 2026
b022817
UCT/API: Drop the zero-length rule from SGL zcopy doc
dmitrygx Nov 3, 2026
bf15f67
UCT/CUDA_IPC: Size SGL mapping by the elements with data
dmitrygx Nov 3, 2026
b118b78
GTEST/UCT: Add min_zcopy() helper to SGL test
dmitrygx Nov 3, 2026
c3082be
UCT/CUDA_IPC: Fill SGL copy arrays in the mapping loop
dmitrygx Nov 3, 2026
3ff81b8
Merge branch 'master' into topic/ucx/cuda-ipc-sgl-zero-length
dmitrygx Nov 4, 2026
f0aabce
UCT/CUDA_IPC: Find first non-empty SGL element in the sum loop
dmitrygx Nov 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
43 changes: 30 additions & 13 deletions src/uct/cuda/cuda_ipc/cuda_ipc_ep.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -288,14 +288,16 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
uct_cuda_ipc_unpacked_rkey_t *key;
uct_cuda_ipc_event_desc_t *cuda_ipc_event;
uct_cuda_ipc_sgl_mapping_t *mapping;
uct_cuda_ipc_sgl_entry_t *entry;
uct_cuda_queue_desc_t *q_desc;
CUmemcpyAttributes attr;
void *mapped_addr;
const void *mapped_base_addr;
CUdeviceptr *mapped_addrs;
CUdeviceptr *mapped_addrs, *local_addrs;
CUdeviceptr *dsts, *srcs;
size_t *copy_lengths;
size_t attrs_idx;
size_t i, total_length;
size_t i, first, total_length;
CUdevice cuda_device;
CUstream *stream;
ucs_status_t status;
Expand All @@ -308,7 +310,12 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
}

total_length = 0;
first = count;
for (i = 0; i < count; i++) {
if ((first == count) && (lengths[i] != 0)) {

@svc-nvidia-pr-review svc-nvidia-pr-review Nov 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: maybe find the first non-empty element right after the total_length == 0 check (for (first = 0; lengths[first] == 0; ++first) {})? The sum loop stays branchless and the first = count sentinel is not needed.

@dmitrygx dmitrygx Nov 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in b661636: the first element with data is found right after the total_length == 0 check, the sum loop has no branch, and the first = count sentinel is gone. The search is a while loop rather than a for with an empty body.

first = i;
}

total_length += lengths[i];
}

Expand All @@ -318,15 +325,15 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
return UCS_OK;
}

status = uct_cuda_ipc_check_and_push_ctx((CUdeviceptr)buffers[0],
status = uct_cuda_ipc_check_and_push_ctx((CUdeviceptr)buffers[first],
&cuda_device, &is_ctx_pushed);
if (ucs_unlikely(status != UCS_OK)) {
return status;
}

mapping = ucs_malloc(sizeof(*mapping) +
count * (sizeof(*mapping->entries) +
sizeof(CUdeviceptr)),
(2 * sizeof(CUdeviceptr)) + sizeof(size_t)),

@svc-nvidia-pr-review svc-nvidia-pr-review Nov 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

maybe keep the common case cheap: when no length is zero we could pass buffers/lengths directly as before, and skip allocating and filling the two extra arrays per element.

@dmitrygx dmitrygx Nov 3, 2026 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in b661636: the sum loop also counts the elements with data. If none is empty, buffers and lengths are passed to cuMemcpyBatchAsync as on master, and the allocation has the same size as on master. Only if some element is empty, the allocation also has room for the compacted local addresses and lengths, which a separate loop fills.

"cuda_ipc_sgl_mapping");
if (ucs_unlikely(mapping == NULL)) {
ucs_error("Failed to allocate cuda_ipc sgl mapping");
Expand All @@ -337,12 +344,19 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
mapping->count = 0;
mapping->entries = (uct_cuda_ipc_sgl_entry_t *)(mapping + 1);
mapped_addrs = (CUdeviceptr *)(mapping->entries + count);
local_addrs = mapped_addrs + count;
copy_lengths = (size_t *)(local_addrs + count);

memset(&attr, 0, sizeof(attr));
attr.srcAccessOrder = CU_MEMCPY_SRC_ACCESS_ORDER_STREAM;
attrs_idx = 0;

for (i = 0; i < count; i++) {
/* cuMemcpyBatchAsync does not accept copies of zero length */
Comment thread
dmitrygx marked this conversation as resolved.
Resolved
if (lengths[i] == 0) {
continue;
}

key = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[i];

status = uct_cuda_ipc_get_remote_address(&key->super, remote_addrs[i],
Expand All @@ -352,23 +366,26 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
goto out_unmap;
}

mapping->entries[i].pid = key->super.super.pid;
mapping->entries[i].pid_ns = key->super.pid_ns;
mapping->entries[i].d_bptr = (uintptr_t)key->super.super.d_bptr;
mapping->entries[i].mapped_addr = mapped_base_addr;
mapped_addrs[i] = (CUdeviceptr)mapped_addr;
entry = &mapping->entries[mapping->count];
entry->pid = key->super.super.pid;
entry->pid_ns = key->super.pid_ns;
entry->d_bptr = (uintptr_t)key->super.super.d_bptr;
entry->mapped_addr = mapped_base_addr;
mapped_addrs[mapping->count] = (CUdeviceptr)mapped_addr;
local_addrs[mapping->count] = (CUdeviceptr)buffers[i];
copy_lengths[mapping->count] = lengths[i];
mapping->count++;
}

if (direction == UCT_CUDA_IPC_PUT) {
dsts = mapped_addrs;
srcs = (CUdeviceptr*)buffers;
srcs = local_addrs;
} else {
dsts = (CUdeviceptr*)buffers;
dsts = local_addrs;
srcs = mapped_addrs;
}

key = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[0];
key = (uct_cuda_ipc_unpacked_rkey_t *)rkeys[first];
status = uct_cuda_ipc_get_stream_and_event(iface, key->stream_id,
&q_desc, &stream,
&cuda_ipc_event);
Expand All @@ -377,7 +394,7 @@ uct_cuda_ipc_post_cuda_sgl_async_copy(uct_ep_h tl_ep, void * const *buffers,
}

status = UCT_CUDADRV_FUNC_LOG_ERR(
cuMemcpyBatchAsync(dsts, srcs, (size_t *)lengths, count, &attr,
cuMemcpyBatchAsync(dsts, srcs, copy_lengths, mapping->count, &attr,
Comment thread
dmitrygx marked this conversation as resolved.
Resolved
&attrs_idx, 1, *stream));
if (ucs_unlikely(status != UCS_OK)) {
ucs_mpool_put(cuda_ipc_event);
Expand Down
35 changes: 33 additions & 2 deletions test/gtest/uct/v2/test_uct_sgl.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -130,6 +130,30 @@ class test_uct_sgl : public uct_test {
test_sgl(std::vector<size_t>(count, elem_size(UCS_KBYTE)), &comp);
}

void test_sgl_zero_length_element()
{
static const size_t count = 3;
static const size_t empty_index = 1;

@svc-nvidia-pr-review svc-nvidia-pr-review Nov 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pls cover a zero length in the first element too (e.g. loop the index over 0..count-1) - the new first selection in cuda_ipc only matters when lengths[0] == 0, and with empty_index = 1 it is never exercised.

@dmitrygx dmitrygx Nov 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Added zero_length_first_last_elements with {0, 2} in eefee38: the first element with data is element 1, and the last element is empty too. I did not loop over the index, to keep one named case per scenario; with {1} and {0, 1, 2}, each position is empty in some test. All elements in the test have valid buffers and rkeys on the same GPU, so the test runs the first > 0 path and checks the data, but it would also pass if the context and the stream were taken from element 0.

🤖 Generated with Claude Code


connect_or_skip();

if (max_count() < count) {
UCS_TEST_SKIP_R("not enough sgl elements");
}

if (elem_size(0) > 0) {

@svc-nvidia-pr-review svc-nvidia-pr-review Nov 3, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

minor: elem_size(0) > 0 as a min_zcopy check is hard to read; can we add a small min_zcopy() helper and use it here and in elem_size()?

@dmitrygx dmitrygx Nov 3, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Done in b118b78: added a min_zcopy() helper, used here and in elem_size().

UCS_TEST_SKIP_R("zero length is not supported");
}

sgl_arrays sgl(*m_sender, *m_receiver, first_supported_mem_type(), op(),
std::vector<size_t>(count, elem_size(UCS_KBYTE)));
sgl.lengths[empty_index] = 0;

ASSERT_UCS_OK_OR_INPROGRESS(sgl_op(sgl));
m_sender->flush();
check_sgl(sgl);
}

private:
enum sgl_op_t {
SGL_OP_PUT,
Expand Down Expand Up @@ -261,10 +285,13 @@ class test_uct_sgl : public uct_test {
void check_sgl(const sgl_arrays &sgl)
{
for (size_t i = 0; i < sgl.lengths.size(); ++i) {
/* The target of a zero-length element keeps its initial data */
uint64_t seed = (sgl.lengths[i] == 0) ? SEED2 : (SEED1 + i);

if (op() == SGL_OP_PUT) {
sgl.remote[i]->pattern_check(SEED1 + i);
sgl.remote[i]->pattern_check(seed);
} else {
sgl.local[i]->pattern_check(SEED1 + i);
sgl.local[i]->pattern_check(seed);
}
}
}
Expand Down Expand Up @@ -327,4 +354,8 @@ UCS_TEST_P(test_uct_sgl, with_callback) {
test_sgl_with_callback();
}

UCS_TEST_P(test_uct_sgl, zero_length_element) {
test_sgl_zero_length_element();
}

UCT_INSTANTIATE_TEST_CASE(test_uct_sgl)
Loading