[Fix] HiCache startup ImportError on the pinned kernel wheel (#39516)

Co-authored-by: Mohammad Angkad <mohammad.angkad@radixark.ai>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: hnyls2002 <lsyincs@gmail.com>

python/sglang/srt/mem_cache/pool_host/common.py

@@ -4,15 +4,19 @@ import json
 import logging
 import os
 from collections import defaultdict
+from functools import lru_cache
 
 import torch
 
 from sglang.srt.environ import envs
 from sglang.srt.mem_cache.storage.mmap import alloc_mmap
 from sglang.srt.runtime_context import get_memory
+from sglang.srt.utils import is_hip
 
 logger = logging.getLogger(__name__)
 
+_is_hip = is_hip()
+
 _CUDA_HOST_REGISTERED_RANGES_ATTR = "_sglang_cuda_host_registered_ranges"
 
 
@@ -250,6 +254,33 @@ def alloc_with_pin_memory(
     return buffer
 
 
+@lru_cache(maxsize=1)
+def _resolve_device_accessible_ptr_fn():
+    try:
+        from sgl_kernel.kvcacheio import get_device_accessible_ptr
+    except ImportError:
+        get_device_accessible_ptr = None
+    else:
+        if not hasattr(torch.ops.sgl_kernel, "get_device_accessible_ptr"):
+            get_device_accessible_ptr = None
+
+    if get_device_accessible_ptr is None:
+        # CUDA's UVA makes host and device addresses equal; on HIP they differ.
+        if _is_hip:
+            raise ImportError(
+                "sgl_kernel.kvcacheio.get_device_accessible_ptr is missing from the "
+                "installed sglang-kernel. It is required on ROCm, where registered "
+                "host memory carries a distinct device address. Rebuild sglang-kernel "
+                "from python/sglang/kernels/aot (setup_rocm.py)."
+            )
+        logger.warning(
+            "sgl_kernel.kvcacheio.get_device_accessible_ptr is missing from the "
+            "installed sglang-kernel; using raw host addresses for kernel pointer "
+            "tables. Build sglang-kernel from python/sglang/kernels/aot to enable it."
+        )
+    return get_device_accessible_ptr
+
+
 def make_kernel_ptr_table(
     tensors: list[torch.Tensor],
     target_device: torch.device | str,
@@ -257,9 +288,12 @@ def make_kernel_ptr_table(
     host_memory_registered: bool,
 ) -> torch.Tensor:
     device = torch.device(target_device)
-    if host_memory_registered and device.type == "cuda":
-        from sgl_kernel.kvcacheio import get_device_accessible_ptr
-
+    get_device_accessible_ptr = (
+        _resolve_device_accessible_ptr_fn()
+        if host_memory_registered and device.type == "cuda"
+        else None
+    )
+    if get_device_accessible_ptr is not None:
         if device.index is None:
             device_index = torch.cuda.current_device()
         else:

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1902,7 +1902,7 @@ class UnifiedRadixCache(BasePrefixCache):
         storage_hit_count -= storage_hit_count % self.page_size
         return storage_hit_count
 
-    @rank_consensus(same_params=["req_id", "len(new_input_tokens)"])
+    @rank_consensus(same_params=["request.rid", "len(new_input_tokens)"])
     def prefetch_from_storage(
         self,
         request: CacheRequestHandle,