From 48f3264807eac6bdef161aeaf5f272de6e4bda6f Mon Sep 17 00:00:00 2001
From: Chao Shi <chao.shi@alibaba-inc.com>
Date: Tue, 26 May 2026 17:44:15 +0800
Subject: [PATCH] [HiCache]: Check return code of cudaHostRegister (#26301)

---
 python/sglang/srt/mem_cache/memory_pool_host.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/mem_cache/memory_pool_host.py b/python/sglang/srt/mem_cache/memory_pool_host.py
index a872e24ea..d5d9d5655 100644
--- a/python/sglang/srt/mem_cache/memory_pool_host.py
+++ b/python/sglang/srt/mem_cache/memory_pool_host.py
@@ -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