chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

Add clear error message when OOM with symmetric memory (#17038)

Browse Source
This commit is contained in:
Nicolas Castet
2026-01-16 06:44:37 -06:00
committed by GitHub
parent daea51385d
commit 9d8bbd4223
1 changed files with 12 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
python/sglang/srt/distributed/device_communicators/pynccl_allocator.py
View File

@@ -34,10 +34,20 @@ ncclResult_t ncclCommWindowRegister(ncclComm_t comm, void* buff, size_t size, n
ncclResult_t ncclMemAlloc(void** ptr, size_t size);
ncclResult_t ncclMemFree(void *ptr);
const char* ncclGetErrorString(ncclResult_t result);
#define NCCLCHECK(cmd) do { \
ncclResult_t res = cmd; \
if (res != ncclSuccess) { \
fprintf(stderr, "ERROR: NCCL symmetric memory allocation failed. Most likely out of device memory. '%s'\\n", \
ncclGetErrorString(res)); \
return NULL; \
} \
} while(0)
void* nccl_alloc_plug(size_t size, int device, void* stream) {
void* ptr;
ncclResult_t err = ncclMemAlloc(&ptr, size);
NCCLCHECK(ncclMemAlloc(&ptr, size));
const char *str_val = getenv("SGLANG_TMP_NCCL_COMM_VALUE");
char *endptr;
@@ -45,7 +55,7 @@ void* nccl_alloc_plug(size_t size, int device, void* stream) {
ncclComm_t comm = (ncclComm_t)(int_val);
ncclWindow_t win;
ncclResult_t err2 = ncclCommWindowRegister(comm, ptr, size, &win, NCCL_WIN_COLL_SYMMETRIC);
NCCLCHECK(ncclCommWindowRegister(comm, ptr, size, &win, NCCL_WIN_COLL_SYMMETRIC));
return ptr;
}