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
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
_CUDART = None
def _get_cudart():
global _CUDART
if _CUDART is 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
from cuda.bindings import runtime as _cudart
except ImportError: # pragma: no cover - cuda-python is always present in images
_cudart = None
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
to the safe host path.
"""
cudart = _get_cudart()
if cudart is None:
if _cudart is None:
# Cannot tell; assume device so behavior stays unchanged.
return True
attr = _CudaPointerAttributes()
ret = cudart.cudaPointerGetAttributes(
_ctypes.byref(attr), _ctypes.c_void_p(int(ptr))
)
if ret != 0:
err, attr = _cudart.cudaPointerGetAttributes(int(ptr))
if int(err) != 0:
# Clear the error so subsequent CUDA calls are not poisoned.
cudart.cudaGetLastError()
_cudart.cudaGetLastError()
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]:
if length <= 0:
return b""
if _is_device_pointer(addr):
cudart = _get_cudart()
buf = (_ctypes.c_char * length)()
buf = bytearray(length)
# cudaMemcpyDeviceToHost = 2; synchronous default-stream copy.
ret = cudart.cudaMemcpy(
buf, _ctypes.c_void_p(int(addr)), _ctypes.c_size_t(length), 2
err, = _cudart.cudaMemcpy(
_ctypes.addressof((_ctypes.c_char * length).from_buffer(buf)),
int(addr),
length,
2,
)
if ret != 0:
if int(err) != 0:
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 bytes(buf)
@@ -174,15 +152,14 @@ def _write_bytes_to_address(addr: int, data: bytes) -> bool:
if not data:
return True
if _is_device_pointer(addr):
cudart = _get_cudart()
buf = _ctypes.create_string_buffer(data, len(data))
# cudaMemcpyHostToDevice = 1; synchronous default-stream copy.
ret = cudart.cudaMemcpy(
_ctypes.c_void_p(int(addr)), buf, _ctypes.c_size_t(len(data)), 1
err, = _cudart.cudaMemcpy(
int(addr), _ctypes.addressof(buf), len(data), 1
)
if ret != 0:
if int(err) != 0:
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)}"
)
return False