Repository navigation
UCT/CUDA_COPY/MD: Don't widen VMM range crossing a handle boundary - #12186
Conversation
cuMemGetAddressRange() returns only the cuMem handle containing the address. The guard added for multi-handle VMM compared only lengths, so a range shorter than a handle that starts inside one handle and ends in the next was still widened to the first handle. The dmabuf was then exported for that handle only and ibv_reg_dmabuf_mr() failed with EINVAL. Keep the requested extent unless the handle contains the whole range.
|
🤖 Starting review — findings will be posted here when done. |
|
Residual coverage gap: both multi-handle VMM tests need a GPU with VMM support; they silently skip ( Related (out of scope for this PR): |
|
🤖 Starting review — findings will be posted here when done. |
|
Both multi-handle VMM tests require a GPU with VMM support and silently skip via If a driver returns the full reserved VA range from Duplicate detection against the PR conversation could not be performed with the tools available here; it was done only against the previously reported findings supplied in the prompt. Please re-run the inline/issue/review comment fetch before submitting. |
|
🤖 Starting review — findings will be posted here when done. |
|
Residual coverage gap: both multi-handle VMM tests and the host-VMM test need a GPU with VMM / host-NUMA VMM support and skip silently otherwise, so a CI machine without such a GPU gives no signal on this fix (already noted in earlier rounds). Process note: only |
|
🤖 Starting review — findings will be posted here when done. |
|
🤖 Starting review — findings will be posted here when done. |
|
|
|
🤖 Starting review — findings will be posted here when done. |
|
🤖 Review complete — no issues found. |
|
Did you or your organization sign the CLA? |
|
🤖 Starting review — findings will be posted here when done. |
|
🤖 Starting review — findings will be posted here when done. |
Rename check_mem_query_range to check_mem_query_exact_range and also check the query is extended to the whole handle when the range is contained in it.
| ADD_FAILURE() << "locality domain 1 allocation failed even though " | ||
| "domain 0 succeeded"; | ||
| return; | ||
| } |
There was a problem hiding this comment.
The moved copy of vmm_locality_domain_mem_type is an older revision than the one it replaces (see test/gtest/uct/cuda/test_cuda_ipc_md.cc:790-805). The current version probes domain 1 first and skips when it's absent (since domain 1 is present only on devices with several locality domains), then only calls ADD_FAILURE() if domain 0 fails. This PR's version inverts the order: it skips only when domain 0 fails and calls ADD_FAILURE() when domain 1 fails. On a GPU with a single locality domain, domain 0 succeeds and domain 1 fails, so the moved test turns a legitimate skip into a hard test failure (gtest failure on any single-locality-domain GPU built with CU_MEM_LOCATION_TYPE_DEVICE_LOCALITY_DOMAIN available).
Suggested fix:
| } | |
| ucs_sys_device_t sys_dev0, sys_dev1; | |
| /* Domain 1 is present only on a device with several locality domains, so | |
| * its absence is what tells a single-domain device apart */ | |
| if (!check_domain(1, sys_dev1)) { | |
| UCS_TEST_SKIP_R("single locality-domain device"); | |
| } | |
| if (!check_domain(0, sys_dev0)) { | |
| ADD_FAILURE() << "locality domain 0 allocation failed even though " | |
| "domain 1 succeeded"; | |
| return; | |
| } |
Please move the current version of this test (test_cuda_ipc_md.cc:790) instead of the older one.
|
🤖 Starting review — findings will be posted here when done. |
Resolve the conflict with openucx#12176 by keeping the moved test_cuda_copy_md tests in test_mem_alloc_device, and take its vmm_locality_domain_mem_type domain probe order.
| /* Host-located VMM reports a zero base address */ | ||
| if (base_address == 0) { | ||
| goto out_ctx_pop; | ||
| } |
There was a problem hiding this comment.
what about combining this condition and condition below?
| /* Host-located VMM reports a zero base address */ | |
| if (base_address == 0) { | |
| goto out_ctx_pop; | |
| } | |
| /* For multi-handle VMM, the physical allocation containing the address | |
| * may not cover the whole requested range */ | |
| if ((base_address == 0) || /* host-located VMM reports a zero base */ | |
| (is_vmm && ((base_address + alloc_length) < (address + length)))) { | |
| goto out_ctx_pop; | |
| } |
| return true; | ||
| }; | ||
|
|
||
| ucs_sys_device_t sys_dev0, sys_dev1; |
There was a problem hiding this comment.
pls rebase — master now probes domain 1 first, because domain 0 also succeeds on a single-locality-domain device and this ADD_FAILURE() then fires spuriously.
The vmm_locality_domain_mem_type test is moved here in its pre-rebase form: the base branch now contains a different version of this block (in test_cuda_ipc_md.cc) that probes domain 1 first, with the comment "Domain 1 is present only on a device with several locality domains, so its absence is what tells a single-domain device apart". The version being moved probes domain 0 first, so on a device that supports locality domains but has only one, domain 0 succeeds, domain 1 fails, and ADD_FAILURE() fires spuriously. Since this PR deletes that block from test_cuda_ipc_md.cc, a rebase will conflict there and the stale logic can silently land in the new location.
| ucs_sys_device_t sys_dev0, sys_dev1; | |
| ucs_sys_device_t sys_dev0, sys_dev1; | |
| /* Domain 1 is present only on a device with several locality domains, so | |
| * its absence is what tells a single-domain device apart */ | |
| if (!check_domain(1, sys_dev1)) { | |
| UCS_TEST_SKIP_R("single locality-domain device"); | |
| } | |
| if (!check_domain(0, sys_dev0)) { | |
| ADD_FAILURE() << "locality domain 0 allocation failed even though " | |
| "domain 1 succeeded"; | |
| return; | |
| } |
|
🤖 Starting review — findings will be posted here when done. |
| if (is_vmm && (alloc_length < length)) { | ||
| /* For multi-handle VMM, the physical allocation containing the address | ||
| * may not cover the whole requested range */ | ||
| if (is_vmm && ((base_address + alloc_length) < (address + length))) { |
There was a problem hiding this comment.
Why keep is_vmm here? The new check is just "the reported range must cover the requested range", which also holds for legacy allocations.
|
|
|
🤖 Starting review — findings will be posted here when done. |
| /* For multi-handle VMM, the physical allocation may be smaller than the | ||
| * mapped virtual range; preserve the caller's requested extent. */ | ||
| if (is_vmm && (alloc_length < length)) { | ||
| /* For multi-handle VMM, the physical allocation containing the address |
There was a problem hiding this comment.
minor: the zero-base case is unrelated to multi-handle VMM (and not VMM-specific), pls split it out so each condition carries its own comment:
| /* For multi-handle VMM, the physical allocation containing the address | |
| if (base_address == 0) { | |
| /* Host-located VMM reports a zero base address */ | |
| goto out_ctx_pop; | |
| } | |
| /* For multi-handle VMM, the physical allocation containing the address | |
| * may not cover the whole requested range */ | |
| if (is_vmm && ((base_address + alloc_length) < (address + length))) { | |
| goto out_ctx_pop; | |
| } |
|
@d4l3k can you pls backport it to 1.23.x branch? This will enable this fix in NIXL 1.6 |
|
@d4l3k can you pls confirm you have CLA signed? |
|
🤖 Starting review — findings will be posted here when done. |
|
@brminich I've signed the CLA |
|
🤖 Review complete — no issues found. |
What?
Don't widen a CUDA VMM registration to the containing physical allocation unless that allocation covers the whole requested range.
Fixes #12185.
Why?
cuMemGetAddressRange()returns only thecuMemCreatehandle containing the address. The guard from #11461 compares lengths only, so a range shorter than a handle that starts inside one handle and ends in the next is widened to the first handle. The dmabuf is exported for that handle only, andibv_reg_dmabuf_mr()fails with EINVAL. This hits PyTorchexpandable_segments(20 MiB handles in one VA range) via NIXL.This is needed for PyTorch: with the CUDA caching allocator in
expandable_segmentsmode, tensors routinely span handle boundaries, so NIXL/UCX transfers of PyTorch tensors (e.g. torchstore weight sync, meta-pytorch/torchstore#249) fail to register. The only workaround today isUCX_CUDA_COPY_REG_WHOLE_ALLOC=off. NIXL wheels bundle UCX, so a v1.23.x backport would let PyTorch users pick this up.How?
Compare the end of the containing allocation with the end of the requested range instead of comparing lengths. Also skip widening when
cuMemGetAddressRange()reports a zero base address, as it does for host-located VMM, so theCUDA_COPY_REG_WHOLE_ALLOC=offworkaround in thehost_vmm_*tests is dropped.Test plan:
test_cuda_copy_md.vmm_multi_handle_unaligned_rangeandtest_mem_alloc_device.host_vmm_mem_rangefail without the fix and pass with it; all 61*cuda_cpy*gtests pass on GB200, CUDA 13.0.expandable_segmentspass over dmabuf; the unpatched plugin fails withibv_reg_dmabuf_mrEINVAL.cc @tvegas1 @nicolasnoble @tomerg-nvidia