Skip to content

Improve NCCL UBR support. Extend cuMem/VMM allocation feature to platforms without fabric allocation support. #152

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
romerojosh merged 2 commits into from
Jul 15, 2026
+128 −32
Merged

Improve NCCL UBR support. Extend cuMem/VMM allocation feature to platforms without fabric allocation support. #152

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
59a44da
Improve NCCL UBR support. Extend cuMem/VMM allocation feature to plat…
romerojosh Jul 7, 2026
3e7ae78
Harden cuMem capability and autotune fallback checks
romerojosh Jul 15, 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
11 changes: 7 additions & 4 deletions docs/env_vars.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,15 +12,18 @@ CUDECOMP_ENABLE_NCCL_UBR
:code:`CUDECOMP_ENABLE_NCCL_UBR` controls whether cuDecomp registers its communication buffers with the NCCL library using :code:`ncclCommRegister`/:code:`ncclCommDeregister` (i.e., user buffer registration).
Registration can improve NCCL send/receive performance in some scenarios. See the `User Buffer Registration <https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/usage/bufferreg.html>`_
section of the NCCL documentation for more details.
This option requires CUDA VMM workspace allocations. Setting :code:`CUDECOMP_ENABLE_NCCL_UBR` also requests the cuMem allocation path used by :code:`cudecompMalloc`; if cuMem support is unavailable, NCCL user buffer registration is disabled. See :code:`CUDECOMP_ENABLE_CUMEM` for more details on cuMem workspace allocations.

Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable this feature.
Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable this feature when cuMem support is available.

CUDECOMP_ENABLE_CUMEM
------------------------
(since v0.5.0, requires CUDA 12.3 driver/toolkit or newer)
(since v0.5.0, requires CUDA 11.3 driver/toolkit or newer)

:code:`CUDECOMP_ENABLE_CUMEM` controls whether cuDecomp uses :code:`cuMem*` APIs to allocate fabric-registered workspace buffers via :code:`cudecompMalloc`. This option can improve the performance of
some MPI distributions on multi-node NVLink (MNNVL) capable systems.
:code:`CUDECOMP_ENABLE_CUMEM` controls whether cuDecomp uses :code:`cuMem*` APIs to allocate VMM workspace buffers via :code:`cudecompMalloc`.
The allocations always request POSIX file-descriptor export support.
When built with CUDA 12.3 or newer, running with a CUDA 12.3 or newer driver, and running on a device that supports fabric handles, cuDecomp also requests fabric export support.
This option can improve the performance of some MPI distributions on multi-node NVLink (MNNVL) capable systems.

Default setting is off (:code:`0`). Setting this variable to :code:`1` will enable this feature.

Expand Down
2 changes: 2 additions & 0 deletions include/internal/halo.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,13 +148,15 @@ void cudecompUpdateHalos_(int ax, const cudecompHandle_t handle, const cudecompG
bool input_has_padding = anyNonzeros(padding);

if (c == 2 && (input_has_padding || haloBackendRequiresNvshmem(grid_desc->config.halo_comm_backend) ||
(haloBackendRequiresNccl(grid_desc->config.halo_comm_backend) && handle->nccl_enable_ubr) ||
(managed && haloBackendRequiresMpi(grid_desc->config.halo_comm_backend)) ||
(handle->cuda_cumem_enable && comm_info.mnnvl_active &&
haloBackendRequiresMpi(grid_desc->config.halo_comm_backend)))) {
// For padded input, always stage to work space.
// For managed memory, always stage to work space if using MPI.
// If using MPI and communicator has MNNVL connections, stage to work space if fabric-allocated.
// For any memory, always stage to workspace if using NVSHMEM.
// For NCCL user buffer registration, stage to the registered workspace instead of application input.
// Can revisit for NVSHMEM if input is NVSHMEM allocated.
c = 1;
}
Expand Down
5 changes: 5 additions & 0 deletions include/internal/transpose.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -374,6 +374,11 @@ static void cudecompTranspose_(int ax, int dir, const cudecompHandle_t handle, c
// in to workspace (which should be nvshmem allocated). Can revisit support for input/output
// arrays allocated with nvshmem.
enable = false;
} else if (transposeBackendRequiresNccl(grid_desc->config.transpose_comm_backend) && handle->nccl_enable_ubr) {
// Note: NCCL user buffer registration requires both the source and destination buffers to be registered.
// cuDecomp only registers the workspace, so keep NCCL communication staged through workspace instead of
// directly using input/output buffers that may not be registered.
enable = false;
} else if (transposeBackendRequiresMpi(grid_desc->config.transpose_comm_backend)) {
// Note: For MPI, disable special cases if input or output pointers are to managed memory
// since MPI performance directly from managed memory is not great
Expand Down
40 changes: 36 additions & 4 deletions src/autotune.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -292,16 +292,32 @@ void autotuneTransposeBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_d
transposeWorkspaceGuard(work_nvshmem, {handle, grid_desc, CUDECOMP_TRANSPOSE_COMM_NVSHMEM});

// Check if there is enough memory for separate non-NVSHMEM allocated work buffer
const bool allow_nvshmem_workspace_fallback = !handle->cuda_cumem_enable && !handle->nccl_enable_ubr;
auto ret = cudaMalloc(&work, work_sz);
if (ret == cudaErrorMemoryAllocation) {
int any_oom = (ret == cudaErrorMemoryAllocation);
CHECK_MPI(MPI_Allreduce(MPI_IN_PLACE, &any_oom, 1, MPI_INT, MPI_LOR, handle->mpi_comm));
if (any_oom) {
if (ret == cudaSuccess) {
CHECK_CUDA(cudaFree(work));
work = nullptr;
} else if (ret == cudaErrorMemoryAllocation) {
cudaGetLastError(); // Reset CUDA error state
} else {
CHECK_CUDA(ret);
}
if (!allow_nvshmem_workspace_fallback) {
THROW_CUDA_ERROR(
"Cannot allocate separate non-NVSHMEM workspace during autotuning while cuMem or NCCL user buffer "
"registration is enabled.");
}
if (handle->rank == 0) {
printf("CUDECOMP:WARN: Cannot allocate separate workspace for non-NVSHMEM backends during "
"autotuning. Using NVSHMEM allocated workspace for all backends, which may cause issues "
"for some MPI implementations. See documentation for more details and suggested workarounds.\n");
}
work = work_nvshmem;
cudaGetLastError(); // Reset CUDA error state
} else {
CHECK_CUDA(ret);
CHECK_CUDA(cudaFree(work));
auto backend = (need_nccl) ? CUDECOMP_TRANSPOSE_COMM_NCCL : CUDECOMP_TRANSPOSE_COMM_MPI_P2P;
tmp = grid_desc->config.transpose_comm_backend;
Expand Down Expand Up @@ -746,16 +762,32 @@ void autotuneHaloBackend(cudecompHandle_t handle, cudecompGridDesc_t grid_desc,
work_nvshmem_guard = haloWorkspaceGuard(work_nvshmem, {handle, grid_desc, CUDECOMP_HALO_COMM_NVSHMEM});

// Check if there is enough memory for separate non-NVSHMEM allocated work buffer
const bool allow_nvshmem_workspace_fallback = !handle->cuda_cumem_enable && !handle->nccl_enable_ubr;
auto ret = cudaMalloc(&work, work_sz);
if (ret == cudaErrorMemoryAllocation) {
int any_oom = (ret == cudaErrorMemoryAllocation);
CHECK_MPI(MPI_Allreduce(MPI_IN_PLACE, &any_oom, 1, MPI_INT, MPI_LOR, handle->mpi_comm));
if (any_oom) {
if (ret == cudaSuccess) {
CHECK_CUDA(cudaFree(work));
work = nullptr;
} else if (ret == cudaErrorMemoryAllocation) {
cudaGetLastError(); // Reset CUDA error state
} else {
CHECK_CUDA(ret);
}
if (!allow_nvshmem_workspace_fallback) {
THROW_CUDA_ERROR(
"Cannot allocate separate non-NVSHMEM workspace during autotuning while cuMem or NCCL user buffer "
"registration is enabled.");
}
if (handle->rank == 0) {
printf("CUDECOMP:WARN: Cannot allocate separate workspace for non-NVSHMEM backends during "
"autotuning. Using NVSHMEM allocated workspace for all backends, which may cause issues "
"for some MPI implementations. See documentation for more details and suggested workarounds.\n");
}
work = work_nvshmem;
cudaGetLastError(); // Reset CUDA error state
} else {
CHECK_CUDA(ret);
CHECK_CUDA(cudaFree(work));
auto backend = (need_nccl) ? CUDECOMP_HALO_COMM_NCCL : CUDECOMP_HALO_COMM_MPI;
tmp = grid_desc->config.halo_comm_backend;
Expand Down
102 changes: 78 additions & 24 deletions src/cudecomp.cc
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -389,44 +389,66 @@ static void gatherGlobalMPIInfo(cudecompHandle_t& handle) {

static void getCudecompEnvVars(cudecompHandle_t& handle) {
// Check CUDECOMP_ENABLE_NCCL_UBR (NCCL user buffer registration)
handle->nccl_enable_ubr = checkEnvVar("CUDECOMP_ENABLE_NCCL_UBR");
bool nccl_ubr_requested = checkEnvVar("CUDECOMP_ENABLE_NCCL_UBR");

// Check CUDECOMP_ENABLE_CUMEM (CUDA VMM allocations for work buffers)
handle->cuda_cumem_enable = checkEnvVar("CUDECOMP_ENABLE_CUMEM");
bool cuda_cumem_requested = checkEnvVar("CUDECOMP_ENABLE_CUMEM");
handle->cuda_cumem_enable = cuda_cumem_requested || nccl_ubr_requested;
if (handle->cuda_cumem_enable) {
#if CUDART_VERSION < 12030
#if CUDART_VERSION < 11030
if (handle->rank == 0) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but CUDA version used for compilation does not "
"support fabric allocations. Disabling this feature.\n");
if (cuda_cumem_requested) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but CUDA version used for compilation does not "
"support CUDA VMM allocations. Disabling this feature.\n");
}
if (nccl_ubr_requested) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_NCCL_UBR is set but CUDA version used for compilation does not "
"support CUDA VMM allocations. Disabling this feature.\n");
}
}
handle->cuda_cumem_enable = false;
#else
int driverVersion;
CHECK_CUDA(cudaDriverGetVersion(&driverVersion));
if (driverVersion < 12030) {
if (driverVersion < 11030) {
if (handle->rank == 0) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but installed driver does not "
"support fabric allocations. Disabling this feature.\n");
if (cuda_cumem_requested) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but installed driver does not "
"support CUDA VMM allocations. Disabling this feature.\n");
}
if (nccl_ubr_requested) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_NCCL_UBR is set but installed driver does not "
"support CUDA VMM allocations. Disabling this feature.\n");
}
}
handle->cuda_cumem_enable = false;
} else {
// Check if fabric allocation type is supported
int dev;
CUdevice cu_dev;
CHECK_CUDA(cudaGetDevice(&dev));
CHECK_CUDA_DRV(cuDeviceGet(&cu_dev, dev));
int flag = 0;
CHECK_CUDA_DRV(cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_dev));
if (!flag) {
if (handle->rank == 0) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but device does not "
"support fabric allocations. Disabling this feature.\n");

int vmm_supported = 0;
int posix_fd_supported = 0;
CHECK_CUDA_DRV(
cuDeviceGetAttribute(&vmm_supported, CU_DEVICE_ATTRIBUTE_VIRTUAL_ADDRESS_MANAGEMENT_SUPPORTED, cu_dev));
CHECK_CUDA_DRV(cuDeviceGetAttribute(&posix_fd_supported,
CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR_SUPPORTED, cu_dev));
handle->cuda_cumem_enable = vmm_supported && posix_fd_supported;
if (!handle->cuda_cumem_enable && handle->rank == 0) {
if (cuda_cumem_requested) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_CUMEM is set but the current device does not support CUDA VMM "
"allocations with POSIX file-descriptor handles. Disabling this feature.\n");
}
if (nccl_ubr_requested) {
printf("CUDECOMP:WARN: CUDECOMP_ENABLE_NCCL_UBR is set but the current device does not support CUDA VMM "
"allocations with POSIX file-descriptor handles. Disabling this feature.\n");
}
handle->cuda_cumem_enable = false;
}
}
#endif
}
handle->nccl_enable_ubr = nccl_ubr_requested && handle->cuda_cumem_enable;

// Check CUDECOMP_ENABLE_CUDA_GRAPHS (CUDA Graphs usage in pipelined backends)
handle->cuda_graphs_enable = checkEnvVar("CUDECOMP_ENABLE_CUDA_GRAPHS");
Expand Down Expand Up @@ -601,7 +623,7 @@ static void cleanupFailedGridDescCreate(cudecompHandle_t handle, cudecompGridDes
releaseUnusedHandleResources(handle, release_streams);
}

#if CUDART_VERSION >= 12030
#if CUDART_VERSION >= 11030
struct cuMemAllocationGuard {
~cuMemAllocationGuard() noexcept {
if (!active) return;
Expand Down Expand Up @@ -1268,16 +1290,28 @@ cudecompResult_t cudecompMalloc(cudecompHandle_t handle, cudecompGridDesc_t grid
#endif
} else {
if (handle->cuda_cumem_enable) {
#if CUDART_VERSION >= 12030
#if CUDART_VERSION >= 11030
int dev;
CUdevice cu_dev;
CHECK_CUDA(cudaGetDevice(&dev));
CHECK_CUDA_DRV(cuDeviceGet(&cu_dev, dev));

int requestedHandleTypes = CU_MEM_HANDLE_TYPE_POSIX_FILE_DESCRIPTOR;
#if CUDART_VERSION >= 12030
int driverVersion;
CHECK_CUDA(cudaDriverGetVersion(&driverVersion));
if (driverVersion >= 12030) {
int fabric_supported = 0;
CHECK_CUDA_DRV(
cuDeviceGetAttribute(&fabric_supported, CU_DEVICE_ATTRIBUTE_HANDLE_TYPE_FABRIC_SUPPORTED, cu_dev));
if (fabric_supported) requestedHandleTypes |= CU_MEM_HANDLE_TYPE_FABRIC;
}
#endif

CUmemAllocationProp prop = {};
prop.type = CU_MEM_ALLOCATION_TYPE_PINNED;
prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE;
prop.requestedHandleTypes = CU_MEM_HANDLE_TYPE_FABRIC;
prop.requestedHandleTypes = static_cast<CUmemAllocationHandleType>(requestedHandleTypes);
prop.location.id = cu_dev;

// Check for RDMA support
Expand All @@ -1286,15 +1320,35 @@ cudecompResult_t cudecompMalloc(cudecompHandle_t handle, cudecompGridDesc_t grid
cuDeviceGetAttribute(&flag, CU_DEVICE_ATTRIBUTE_GPU_DIRECT_RDMA_WITH_CUDA_VMM_SUPPORTED, cu_dev));
if (flag) prop.allocFlags.gpuDirectRDMACapable = 1;

// Align allocation size to required granularity
// Keep the caller-requested size so any retry can realign from the original value.
size_t original_buffer_size_bytes = buffer_size_bytes;
size_t granularity;
CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_MINIMUM));
buffer_size_bytes = (buffer_size_bytes + granularity - 1) / granularity * granularity;
CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
buffer_size_bytes = (original_buffer_size_bytes + granularity - 1) / granularity * granularity;

// Allocate memory
cuMemAllocationGuard cumem_guard;
cumem_guard.size = buffer_size_bytes;
CHECK_CUDA_DRV(cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0));
CUresult err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0);
#if CUDART_VERSION >= 12030
if ((requestedHandleTypes & CU_MEM_HANDLE_TYPE_FABRIC) &&
(err == CUDA_ERROR_NOT_PERMITTED || err == CUDA_ERROR_NOT_SUPPORTED)) {
// Fabric handles are useful when the platform supports them, but regular NCCL user buffer registration only
// requires POSIX FD export support. If Fabric creation is unavailable at runtime, keep VMM enabled and fall
// back to a POSIX-FD-only allocation.
requestedHandleTypes &= ~CU_MEM_HANDLE_TYPE_FABRIC;
prop.requestedHandleTypes = static_cast<CUmemAllocationHandleType>(requestedHandleTypes);
CHECK_CUDA_DRV(cuMemGetAllocationGranularity(&granularity, &prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED));
buffer_size_bytes = (original_buffer_size_bytes + granularity - 1) / granularity * granularity;
cumem_guard.size = buffer_size_bytes;
err = cuFnTable.pfn_cuMemCreate(&cumem_guard.handle, buffer_size_bytes, &prop, 0);
}
#endif
if (CUDA_SUCCESS != err) {
const char* error_str;
cuFnTable.pfn_cuGetErrorString(err, &error_str);
throw cudecomp::CudaError(__FILE__, __LINE__, error_str);
}
cumem_guard.handle_created = true;
CHECK_CUDA_DRV(cuMemAddressReserve(&cumem_guard.ptr, buffer_size_bytes, granularity, 0, 0));
cumem_guard.address_reserved = true;
Expand Down Expand Up @@ -1382,7 +1436,7 @@ cudecompResult_t cudecompFree(cudecompHandle_t handle, cudecompGridDesc_t grid_d

} else {
if (handle->cuda_cumem_enable) {
#if CUDART_VERSION >= 12030
#if CUDART_VERSION >= 11030
if (buffer) {
CUmemGenericAllocationHandle cumem_handle;
CHECK_CUDA_DRV(cuMemRetainAllocationHandle(&cumem_handle, buffer));
Expand Down
Loading