[Bugfix] Fix ncclCommQueryProperties heap overflow with NCCL >= 2.31 (#53008)
Signed-off-by: Xuan Jiang <xuanj@amazon.com> Co-authored-by: Vladimir Aerov <vaerov@amazon.com>
Xuan-1998 committed
Aug 28, 2026 at 17:25 UTC
026d5af7f410e59e3910a4d66af1f691db29ca1d
2 files changed
+19
-3
vllm/distributed/device_communicators/pynccl_wrapper.py
+15
-2
@@ -51,8 +51,15 @@ class ncclUniqueId(ctypes.Structure):
51
_fields_ = [("internal", ctypes.c_byte * 128)]
52
53
54
-# NCCL 2.30+ ncclCommProperties_t. Only fields through ginType are read;
55
-# trailing fields keep the layout aligned with NCCL's versioned structure.
54
+# Mirror of NCCL's versioned ncclCommProperties_t (nccl_device/core.h,
55
+# v2.31.2-1, 144 bytes). ncclCommQueryProperties fills fields gated by the
56
+# version the caller declares in props.version, not by props.size, so never
57
+# declare a version newer than this layout: clamp to
58
+# NCCL_COMM_PROPERTIES_LAYOUT_VERSION. To use fields added in a newer NCCL,
59
+# extend the layout and bump the constant.
60
+NCCL_COMM_PROPERTIES_LAYOUT_VERSION = 23102 # NCCL_VERSION(2, 31, 2)
61
+
62
+
63
class ncclCommProperties(ctypes.Structure):
64
_fields_ = [
65
("size", ctypes.c_size_t),
@@ -68,6 +75,12 @@ class ncclCommProperties(ctypes.Structure):
75
("nLsaTeams", ctypes.c_int),
76
("hostRmaSupport", ctypes.c_bool),
77
("railedGinType", ctypes.c_int),
78
+ # Filled only when the declared version is >= NCCL_VERSION(2, 31, 0).
79
+ ("commHash", ctypes.c_uint64),
80
+ ("ginMinStride", ctypes.c_int),
81
+ ("ginConnectionType", ctypes.c_int),
82
+ ("ginSupport", ctypes.c_bool * 64),
83
+ ("devCommRuntimeVersionSize", ctypes.c_size_t),
84
]
85
86
vllm/utils/nccl.py
+4
-1
@@ -86,6 +86,7 @@ def find_nccl_library_paths() -> list[str] | None:
86
def query_nccl_gin_type(group: torch.distributed.ProcessGroup) -> int | None:
87
"""Return the GIN type for an initialized group, or ``None`` on failure."""
88
from vllm.distributed.device_communicators.pynccl_wrapper import (
89
+ NCCL_COMM_PROPERTIES_LAYOUT_VERSION,
90
NCCLLibrary,
91
ncclCommProperties,
92
)
@@ -114,7 +115,9 @@ def query_nccl_gin_type(group: torch.distributed.ProcessGroup) -> int | None:
115
ctypes.memset(ctypes.addressof(props), 0, ctypes.sizeof(props))
116
props.size = ctypes.sizeof(props)
117
props.magic = 0xCAFEBEEF
117
- props.version = nccl.ncclGetRawVersion()
118
+ props.version = min(
119
+ nccl.ncclGetRawVersion(), NCCL_COMM_PROPERTIES_LAYOUT_VERSION
120
+ )
121
result = query_fn(ctypes.c_void_p(comm_ptr), ctypes.byref(props))
122
except Exception:
123
logger.warning("Failed to query NCCL communicator properties", exc_info=True)