feat(hicache): support numa detect to reduce long tail latency (#11028)

Co-authored-by: Zhiqiang Xie <xiezhq@stanford.edu>

python/sglang/srt/mem_cache/hiradix_cache.py

@@ -24,6 +24,7 @@ from sglang.srt.mem_cache.radix_cache import (
     split_node_hash_value,
 )
 from sglang.srt.metrics.collector import StorageMetricsCollector
+from sglang.srt.utils import bind_to_closest_numa_node, is_numa_available
 
 if TYPE_CHECKING:
     from sglang.srt.mem_cache.cache_init_params import CacheInitParams
@@ -43,8 +44,16 @@ class HiRadixCache(RadixCache):
                     "Page first layout is not supported with direct IO backend, switching to page first direct layout"
                 )
 
+        if (
+            not server_args.disable_hicache_numa_detect
+            and is_numa_available()
+            and torch.cuda.is_available()
+        ):
+            bind_to_closest_numa_node()
+
         self.page_size = params.page_size
         self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
+
         if isinstance(self.kv_cache, MHATokenToKVPool):
             self.token_to_kv_pool_host = MHATokenToKVPoolHost(
                 self.kv_cache,

python/sglang/srt/server_args.py

@@ -487,6 +487,7 @@ class ServerArgs:
     hicache_write_policy: str = "write_through"
     hicache_io_backend: str = "kernel"
     hicache_mem_layout: str = "layer_first"
+    disable_hicache_numa_detect: bool = False
     hicache_storage_backend: Optional[str] = None
     hicache_storage_prefetch_policy: str = "best_effort"
     hicache_storage_backend_extra_config: Optional[str] = None
@@ -3891,6 +3892,11 @@ class ServerArgs:
             default=ServerArgs.hicache_mem_layout,
             help="The layout of host memory pool for hierarchical cache.",
         )
+        parser.add_argument(
+            "--disable-hicache-numa-detect",
+            action="store_true",
+            help="Disable binding the process to the NUMA node closest to the active CUDA device when hierarchical cache is enabled.",
+        )
         parser.add_argument(
             "--hicache-storage-backend",
             type=str,

python/sglang/srt/utils/common.py

@@ -3464,19 +3464,27 @@ def check_cuda_result(raw_output):
     return results
 
 
-def get_physical_device_id() -> int:
+def get_physical_device_id(pytorch_device_id: int) -> int:
     """
     Convert PyTorch logical device ID to physical device ID.
+
+    When CUDA_VISIBLE_DEVICES is set, maps the logical device ID (as seen by PyTorch)
+    to the actual physical device ID. If CUDA_VISIBLE_DEVICES is not set, returns
+    the device ID unchanged.
+
+    Args:
+        pytorch_device_id: The logical device ID from PyTorch (e.g., torch.cuda.current_device())
+
+    Returns:
+        The physical device ID
     """
+    device_idx = int(pytorch_device_id)
     cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
-    assert (
-        cuda_visible_devices is not None
-    ), "CUDA_VISIBLE_DEVICES should be set in a scheduler"
-    device_list = cuda_visible_devices.split(",")
-    assert (
-        len(device_list) == 1
-    ), "CUDA_VISIBLE_DEVICES should be set to a single device in a scheduler"
-    return int(device_list[0])
+    if cuda_visible_devices:
+        device_list = cuda_visible_devices.split(",")
+        return int(device_list[device_idx])
+    else:
+        return device_idx
 
 
 def get_device_sm_nvidia_smi():
@@ -3508,7 +3516,7 @@ def numa_bind_to_node(node: int):
         raise SystemError("numa not available on this system")
 
     libnuma.numa_run_on_node(ctypes.c_int(node))
-    libnuma.numa_set_localalloc()
+    libnuma.numa_set_preferred(ctypes.c_int(node))
 
 
 def json_list_type(value):
@@ -3808,3 +3816,107 @@ def get_or_create_event_loop():
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
         return loop
+
+
+def get_numa_node_count() -> int:
+    """
+    Get the number of NUMA nodes available on the system.
+    Must be called after is_numa_available() is True.
+    Returns:
+        int: The number of NUMA nodes.
+    """
+    libnuma = ctypes.CDLL("libnuma.so")
+    return libnuma.numa_max_node() + 1
+
+
+def is_numa_available() -> bool:
+    try:
+        libnuma = ctypes.CDLL("libnuma.so")
+        return libnuma.numa_available() >= 0
+    except Exception:
+        return False
+
+
+def get_system_gpu_count() -> int:
+    """
+    Get the total number of GPUs in the system (not affected by CUDA_VISIBLE_DEVICES).
+
+    Returns:
+        int: The total number of physical GPUs.
+    """
+    result = subprocess.run(
+        ["nvidia-smi", "--list-gpus"],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+    gpu_lines = [
+        line
+        for line in result.stdout.strip().split("\n")
+        if line.strip().startswith("GPU")
+    ]
+    return len(gpu_lines)
+
+
+@lru_cache(maxsize=1)
+def get_current_device_numa_node() -> int:
+    """
+    Retrieve the NUMA node ID of the CPU socket closest to the currently active CUDA device.
+
+    First tries to query nvidia-smi topology. If it returns a single NUMA ID, uses that directly.
+    If it returns multiple NUMA IDs (comma/dash separated), falls back to distributing GPUs
+    evenly across NUMA nodes based on GPU ID intervals.
+
+    For example, with 8 GPUs and 2 NUMA nodes: GPUs 0-3 -> node 0, GPUs 4-7 -> node 1.
+
+    Returns:
+        int: The NUMA node ID (e.g., 0, 1).
+
+    Raises:
+        RuntimeError: If device information cannot be retrieved.
+    """
+    import torch
+
+    logical_device_id = torch.cuda.current_device()
+    physical_device_id = get_physical_device_id(logical_device_id)
+
+    # Query NUMA topology from nvidia-smi
+    result = subprocess.run(
+        ["nvidia-smi", "topo", "-C", "-i", str(physical_device_id)],
+        capture_output=True,
+        text=True,
+        check=True,
+    )
+
+    output_line = result.stdout.strip()
+    prefix = "NUMA IDs of closest CPU:"
+
+    if output_line.startswith(prefix):
+        numa_id_str = output_line[len(prefix) :].strip()
+        if numa_id_str.isdigit():
+            return int(numa_id_str)
+
+    # Fall back: distribute GPUs evenly across NUMA nodes
+    numa_count = get_numa_node_count()
+    gpu_count = get_system_gpu_count()
+
+    if gpu_count >= numa_count:
+        gpus_per_numa = gpu_count // numa_count  # >= 1
+        numa_node = physical_device_id // gpus_per_numa  # 0 ~ numa_count - 1
+    else:
+        logger.warning(
+            f"GPU count {gpu_count} is less than NUMA count {numa_count}. Using first NUMA node."
+        )
+        numa_node = 0
+
+    return numa_node
+
+
+def bind_to_closest_numa_node():
+    """
+    Bind the current process to the NUMA node closest to the active CUDA device.
+
+    Uses `numa` library calls via ctypes to set the CPU affinity of the process.
+    """
+    node_id = get_current_device_numa_node()
+    numa_bind_to_node(node_id)