[HiCache][AMD] Add UMBP tiered DRAM + SSD L3 storage backend with hugepage host allocator (#25377)

Co-authored-by: TianDi101 ditian12@amd.com
Co-authored-by: Niko Ma nima@amd.com
Co-authored-by: Wu, Yutong yutong.wu@amd.com
Co-authored-by: figo fizhang@amd.com
Co-authored-by: AMD-yanfeiwang <yanfei.wang@amd.com>
Co-authored-by: Zhangheng <hzh0425@apache.org>

python/sglang/srt/managers/cache_controller.py

@@ -483,7 +483,7 @@ class HiCacheController:
 
             if (
                 self.storage_backend_type
-                in ["hf3fs", "mooncake", "eic", "nixl", "simm"]
+                in ["hf3fs", "mooncake", "eic", "nixl", "simm", "mori"]
             ) or (
                 self.storage_backend_type == "dynamic"
                 and bool(self.storage_config.extra_config.get("interface_v1", 0))

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

@@ -40,6 +40,20 @@ def get_allocator_from_storage(allocator_type):
                 "Fallback to use default allocator."
             )
             return HostTensorAllocator()
+    elif allocator_type == "mori":
+        try:
+            from sglang.srt.mem_cache.storage.umbp.umbp_host_allocator import (
+                UMBPHostTensorAllocator,
+            )
+
+            return UMBPHostTensorAllocator()
+        except (ImportError, RuntimeError) as exc:
+            logger.warning(
+                "UMBPHostTensorAllocator unavailable (%s). "
+                "Falling back to torch.empty-based allocator.",
+                exc,
+            )
+            return HostTensorAllocator()
     else:
         return HostTensorAllocator()
 

python/sglang/srt/mem_cache/storage/backend_factory.py

@@ -185,6 +185,8 @@ class StorageBackendFactory:
             return backend_class(storage_config, mem_pool_host)
         elif backend_name == "simm":
             return backend_class(storage_config, mem_pool_host)
+        elif backend_name == "mori":
+            return backend_class(storage_config, mem_pool_host)
         else:
             raise ValueError(f"Unknown built-in backend: {backend_name}")
 
@@ -229,3 +231,9 @@ StorageBackendFactory.register_backend(
     "sglang.srt.mem_cache.storage.simm.hicache_simm",
     "HiCacheSiMM",
 )
+
+StorageBackendFactory.register_backend(
+    "mori",
+    "sglang.srt.mem_cache.storage.umbp.umbp_store",
+    "UMBPStore",
+)

python/sglang/srt/mem_cache/storage/umbp/umbp_host_allocator.py

@@ -0,0 +1,142 @@
+import ctypes
+import logging
+import math
+import os
+from typing import Any, Dict
+
+import torch
+
+from sglang.srt.mem_cache.pool_host.common import HostTensorAllocator
+
+logger = logging.getLogger(__name__)
+
+
+def _bool_env(name: str, default: bool) -> bool:
+    raw = os.getenv(name)
+    if raw is None:
+        return default
+    return raw.strip().lower() in ("1", "true", "yes", "on")
+
+
+def _int_env(name: str, default: int) -> int:
+    raw = os.getenv(name)
+    return int(raw) if raw is not None and raw != "" else default
+
+
+class UMBPHostTensorAllocator(HostTensorAllocator):
+    """Allocate the HiCache L2 host tensor from mori's UMBPHostMemAllocator."""
+
+    def __init__(self) -> None:
+        super().__init__()
+        try:
+            import mori.umbp as umbp_mod
+        except ImportError as exc:
+            raise RuntimeError(
+                "mori.umbp is not available. Build mori with BUILD_UMBP=ON "
+                "or fall back to the default torch host allocator."
+            ) from exc
+
+        self._mod = umbp_mod
+        self._allocator = umbp_mod.UMBPHostMemAllocator()
+
+        self._use_hugepage = _bool_env("SGLANG_HICACHE_HOST_HUGEPAGE", True)
+        self._hugepage_size = _int_env(
+            "SGLANG_HICACHE_HOST_HUGEPAGE_SIZE", 2 * 1024 * 1024
+        )
+        self._numa_node = _int_env("SGLANG_HICACHE_HOST_NUMA_NODE", -1)
+        self._prefault = _bool_env("SGLANG_HICACHE_HOST_PREFAULT", True)
+        self._handles: Dict[int, Any] = {}
+
+    def allocate(
+        self, dims: tuple, dtype: torch.dtype, device: str = "cpu"
+    ) -> torch.Tensor:
+        if device != "cpu":
+            raise ValueError(
+                "UMBPHostTensorAllocator only supports CPU host memory, "
+                f"got device={device}"
+            )
+
+        self.dims = dims
+        self.dtype = dtype
+
+        element_size = torch.empty((), dtype=dtype).element_size()
+        nbytes = math.prod(int(dim) for dim in dims) * element_size
+
+        requested_backing = (
+            self._mod.UMBPHostBufferBacking.AnonymousHugetlb
+            if self._use_hugepage
+            else self._mod.UMBPHostBufferBacking.Anonymous
+        )
+
+        handle = self._allocator.alloc(
+            nbytes,
+            requested_backing,
+            self._hugepage_size,
+            self._numa_node,
+            self._prefault,
+        )
+        if not handle:
+            raise RuntimeError(
+                f"UMBPHostMemAllocator.alloc({nbytes} bytes) failed "
+                f"(requested_backing={requested_backing}, "
+                f"numa_node={self._numa_node})."
+            )
+        self._handles[int(handle.ptr)] = handle
+
+        c_array = (ctypes.c_byte * nbytes).from_address(handle.ptr)
+        tensor = torch.frombuffer(c_array, dtype=torch.uint8, count=nbytes)
+
+        if dtype != torch.uint8:
+            tensor = tensor.view(dtype)
+
+        logger.info(
+            "UMBPHostTensorAllocator: allocated %.2f GB at 0x%x "
+            "requested_backing=%s actual_backing=%s actual_alignment=%d "
+            "mapped_size=%d numa_node=%d",
+            nbytes / 1e9,
+            handle.ptr,
+            requested_backing,
+            handle.actual_backing,
+            handle.actual_alignment,
+            handle.mapped_size,
+            self._numa_node,
+        )
+        if (
+            self._use_hugepage
+            and handle.actual_backing == self._mod.UMBPHostBufferBacking.Anonymous
+        ):
+            logger.warning(
+                "UMBPHostTensorAllocator: requested AnonymousHugetlb backing "
+                "but kernel demoted to Anonymous (4 KiB pages). Check "
+                "vm.nr_hugepages and HugePages_Free in /proc/meminfo. "
+                "Performance and AINIC MR-size benefits will not apply."
+            )
+
+        return tensor.view(dims)
+
+    def mapped_size_for(self, ptr: int) -> int:
+        """Actual mmap size for the allocation whose base address is *ptr*."""
+        handles = getattr(self, "_handles", None)
+        if handles is None:
+            return 0
+        h = handles.get(ptr)
+        return int(h.mapped_size) if h is not None else 0
+
+    @property
+    def mapped_size(self) -> int:
+        """Largest mapped_size across all live allocations, or 0."""
+        handles = getattr(self, "_handles", None)
+        if not handles:
+            return 0
+        return max(int(h.mapped_size) for h in handles.values())
+
+    def __del__(self) -> None:
+        try:
+            handles = getattr(self, "_handles", None)
+            allocator = getattr(self, "_allocator", None)
+            if handles and allocator is not None:
+                for h in handles.values():
+                    allocator.free(h)
+                self._handles.clear()
+        except Exception:
+            pass

python/sglang/srt/server_args.py

@@ -1923,6 +1923,7 @@ class ServerArgs:
                 "dynamic",
                 "eic",
                 "simm",
+                "mori",
             ],
         ),
     ] = None