[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,

test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py

@@ -15,9 +15,15 @@ import torch
 
 from sglang.kernels.ops.kvcache.hicache import can_use_write_back_jit_kernel
 from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
+from sglang.srt.mem_cache.pool_host import common as pool_host_common
 from sglang.srt.mem_cache.pool_host.common import (
     ALLOC_MEMORY_FUNCS,
+    HostTensorAllocator,
+    _cuda_host_unregister,
+    _resolve_device_accessible_ptr_fn,
+    alloc_with_host_register,
     alloc_with_pin_memory,
+    make_kernel_ptr_table,
 )
 from sglang.srt.mem_cache.pool_host.mha import MHATokenToKVPoolHost
 from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
@@ -267,8 +273,46 @@ def test_page_first_staged_write_back_mla(element_dim: int, page_count: int) ->
     _run_mla(element_dim, page_count)
 
 
+@pytest.mark.skipif(
+    is_hip(),
+    reason="ROCm maps registered host memory at a distinct device address.",
+)
+def test_registered_mmap_kernel_ptr_table_fallback_matches_device_alias(
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    """CUDA maps registered host memory at the host address itself;
+    ``make_kernel_ptr_table``'s raw-host-address fallback depends on it."""
+    if _resolve_device_accessible_ptr_fn() is None:
+        pytest.skip(
+            "installed sglang-kernel has no get_device_accessible_ptr; "
+            "build it from python/sglang/kernels/aot to run this test"
+        )
+
+    buffer = alloc_with_host_register(
+        (PAGE_SIZE * 4, 128),
+        torch.bfloat16,
+        "cpu",
+        True,
+        HostTensorAllocator(),
+    )
+    try:
+        aliased = make_kernel_ptr_table([buffer], DEVICE, host_memory_registered=True)
+        monkeypatch.setattr(
+            pool_host_common, "_resolve_device_accessible_ptr_fn", lambda: None
+        )
+        raw = make_kernel_ptr_table([buffer], DEVICE, host_memory_registered=True)
+        assert torch.equal(aliased, raw)
+    finally:
+        _cuda_host_unregister(buffer)
+
+
 def test_registered_mmap_pointer_domains_and_all_layer_transfer() -> None:
-    from sgl_kernel.kvcacheio import get_device_accessible_ptr
+    get_device_accessible_ptr = _resolve_device_accessible_ptr_fn()
+    if get_device_accessible_ptr is None:
+        pytest.skip(
+            "installed sglang-kernel has no get_device_accessible_ptr; "
+            "build it from python/sglang/kernels/aot to run this test"
+        )
 
     device_pool = MLATokenToKVPool(
         size=PAGE_SIZE * 4,