diff --git a/python/sglang/srt/mem_cache/pool_host/common.py b/python/sglang/srt/mem_cache/pool_host/common.py index 43616e769..60d6bd4bb 100644 --- a/python/sglang/srt/mem_cache/pool_host/common.py +++ b/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: diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index f19d7b6c1..b060958bb 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/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, diff --git a/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py b/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py index 547ebe401..73499e35c 100644 --- a/test/registered/kernels/ops/kvcache/test_hicache_page_first_write_back.py +++ b/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,