mooncake: probe pointers via cuda.bindings instead of ctypes libcudart

The ctypes-based probing loaded a second libcudart copy into the process,
which corrupted CUDA/runtime state and segfaulted the scheduler processes
shortly after registration (reproducible on prefill ranks, intermittent on
decode). cuda.bindings.runtime.cudaPointerGetAttributes is a properly
typed binding and boots cleanly on both roles.
This commit is contained in:
2026-09-21 12:46:25 +08:00
parent b963295489
commit c74a4037fb
@@ -101,32 +101,12 @@ FAILED_SESSION_RECOVERIES = Counter(
import ctypes as _ctypes import ctypes as _ctypes
class _CudaPointerAttributes(_ctypes.Structure):
_fields_ = [
("type", _ctypes.c_int),
("device", _ctypes.c_int),
("devicePointer", _ctypes.c_void_p),
("hostPointer", _ctypes.c_void_p),
]
_CUDA_MEMORY_TYPE_DEVICE = 2 _CUDA_MEMORY_TYPE_DEVICE = 2
_CUDART = None
try:
def _get_cudart(): from cuda.bindings import runtime as _cudart
global _CUDART except ImportError: # pragma: no cover - cuda-python is always present in images
if _CUDART is None: _cudart = None
for name in ("libcudart.so", "libcudart.so.13", "libcudart.so.12"):
try:
_CUDART = _ctypes.CDLL(name)
break
except OSError:
continue
else:
_CUDART = False
return _CUDART or None
def _is_device_pointer(ptr: int) -> bool: def _is_device_pointer(ptr: int) -> bool:
@@ -136,34 +116,32 @@ def _is_device_pointer(ptr: int) -> bool:
segment addresses). Returns False on any error so the caller falls back segment addresses). Returns False on any error so the caller falls back
to the safe host path. to the safe host path.
""" """
cudart = _get_cudart() if _cudart is None:
if cudart is None:
# Cannot tell; assume device so behavior stays unchanged. # Cannot tell; assume device so behavior stays unchanged.
return True return True
attr = _CudaPointerAttributes() err, attr = _cudart.cudaPointerGetAttributes(int(ptr))
ret = cudart.cudaPointerGetAttributes( if int(err) != 0:
_ctypes.byref(attr), _ctypes.c_void_p(int(ptr))
)
if ret != 0:
# Clear the error so subsequent CUDA calls are not poisoned. # Clear the error so subsequent CUDA calls are not poisoned.
cudart.cudaGetLastError() _cudart.cudaGetLastError()
return False return False
return attr.type == _CUDA_MEMORY_TYPE_DEVICE return int(attr.type) == _CUDA_MEMORY_TYPE_DEVICE
def _read_bytes_from_address(addr: int, length: int) -> Optional[bytes]: def _read_bytes_from_address(addr: int, length: int) -> Optional[bytes]:
if length <= 0: if length <= 0:
return b"" return b""
if _is_device_pointer(addr): if _is_device_pointer(addr):
cudart = _get_cudart() buf = bytearray(length)
buf = (_ctypes.c_char * length)()
# cudaMemcpyDeviceToHost = 2; synchronous default-stream copy. # cudaMemcpyDeviceToHost = 2; synchronous default-stream copy.
ret = cudart.cudaMemcpy( err, = _cudart.cudaMemcpy(
buf, _ctypes.c_void_p(int(addr)), _ctypes.c_size_t(length), 2 _ctypes.addressof((_ctypes.c_char * length).from_buffer(buf)),
int(addr),
length,
2,
) )
if ret != 0: if int(err) != 0:
logger.error( logger.error(
f"cudaMemcpy D2H failed (ret={ret}) for addr {hex(addr)} len {length}" f"cudaMemcpy D2H failed (err={err}) for addr {hex(addr)} len {length}"
) )
return None return None
return bytes(buf) return bytes(buf)
@@ -174,15 +152,14 @@ def _write_bytes_to_address(addr: int, data: bytes) -> bool:
if not data: if not data:
return True return True
if _is_device_pointer(addr): if _is_device_pointer(addr):
cudart = _get_cudart()
buf = _ctypes.create_string_buffer(data, len(data)) buf = _ctypes.create_string_buffer(data, len(data))
# cudaMemcpyHostToDevice = 1; synchronous default-stream copy. # cudaMemcpyHostToDevice = 1; synchronous default-stream copy.
ret = cudart.cudaMemcpy( err, = _cudart.cudaMemcpy(
_ctypes.c_void_p(int(addr)), buf, _ctypes.c_size_t(len(data)), 1 int(addr), _ctypes.addressof(buf), len(data), 1
) )
if ret != 0: if int(err) != 0:
logger.error( logger.error(
f"cudaMemcpy H2D failed (ret={ret}) for addr {hex(addr)} " f"cudaMemcpy H2D failed (err={err}) for addr {hex(addr)} "
f"len {len(data)}" f"len {len(data)}"
) )
return False return False