[HiCache]: Check return code of cudaHostRegister (#26301)

2026-05-26 17:44:15 +08:00
parent d25a220fdb
commit 48f3264807
1 changed files with 3 additions and 1 deletions
@@ -187,9 +187,11 @@ def alloc_with_host_register(
    """
    buffer = allocator.allocate(dims, dtype=dtype, device=device)
    if pin_memory:
-        torch.cuda.cudart().cudaHostRegister(
+        ret = torch.cuda.cudart().cudaHostRegister(
            buffer.data_ptr(), buffer.numel() * buffer.element_size(), 0
        )
+        if ret != 0:
+            raise RuntimeError(f"cudaHostRegister failed with error code {ret}")
    return buffer