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

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

benchmark/hicache/bench_multiturn.py

@@ -475,6 +475,20 @@ class WorkloadGenerator:
         self.pbar.close()
 
         duration = self.finished_time - self.start_time
+        sorted_ttft = sorted(self.performance_metrics["ttft"])
+        sorted_latency = sorted(self.performance_metrics["latency"])
+
+        def percentile(sorted_vals, q):
+            if not sorted_vals:
+                return 0.0
+            idx = int(q * len(sorted_vals))
+            if idx >= len(sorted_vals):
+                idx = len(sorted_vals) - 1
+            return sorted_vals[idx]
+
+        def max_or_zero(sorted_vals):
+            return sorted_vals[-1] if sorted_vals else 0.0
+
         performance_data = {
             "summary": {
                 "total_requests": len(self.performance_metrics["ttft"]),
@@ -493,20 +507,16 @@ class WorkloadGenerator:
                 ),
                 "average_ttft": sum(self.performance_metrics["ttft"])
                 / len(self.performance_metrics["ttft"]),
-                "p90_ttft": sorted(self.performance_metrics["ttft"])[
+                "p90_ttft": percentile(sorted_ttft, 0.9),
-                    int(0.9 * len(self.performance_metrics["ttft"]))
+                "p99_ttft": percentile(sorted_ttft, 0.99),
-                ],
+                "median_ttft": percentile(sorted_ttft, 0.5),
-                "median_ttft": sorted(self.performance_metrics["ttft"])[
+                "max_ttft": max_or_zero(sorted_ttft),
-                    len(self.performance_metrics["ttft"]) // 2
-                ],
                 "average_latency": sum(self.performance_metrics["latency"])
                 / len(self.performance_metrics["latency"]),
-                "p90_latency": sorted(self.performance_metrics["latency"])[
+                "p90_latency": percentile(sorted_latency, 0.9),
-                    int(0.9 * len(self.performance_metrics["latency"]))
+                "p99_latency": percentile(sorted_latency, 0.99),
-                ],
+                "median_latency": percentile(sorted_latency, 0.5),
-                "median_latency": sorted(self.performance_metrics["latency"])[
+                "max_latency": max_or_zero(sorted_latency),
-                    len(self.performance_metrics["latency"]) // 2
-                ],
                 "input_token_throughput": sum(self.performance_metrics["prompt_len"])
                 / duration,
                 "output_token_throughput": sum(
@@ -554,12 +564,16 @@ class WorkloadGenerator:
         )
         print(f"  Average TTFT: {performance_data['summary']['average_ttft']:.2f}")
         print(f"  P90 TTFT: {performance_data['summary']['p90_ttft']:.2f}")
+        print(f"  P99 TTFT: {performance_data['summary']['p99_ttft']:.2f}")
         print(f"  Median TTFT: {performance_data['summary']['median_ttft']:.2f}")
+        print(f"  Max TTFT: {performance_data['summary']['max_ttft']:.2f}")
         print(
             f"  Average latency: {performance_data['summary']['average_latency']:.2f}"
         )
         print(f"  P90 latency: {performance_data['summary']['p90_latency']:.2f}")
+        print(f"  P99 latency: {performance_data['summary']['p99_latency']:.2f}")
         print(f"  Median latency: {performance_data['summary']['median_latency']:.2f}")
+        print(f"  Max latency: {performance_data['summary']['max_latency']:.2f}")
         print(
             f"  Input token throughput: {performance_data['summary']['input_token_throughput']:.2f} tokens per second"
         )

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 (
+    if cuda_visible_devices:
-        cuda_visible_devices is not None
-    ), "CUDA_VISIBLE_DEVICES should be set in a scheduler"
         device_list = cuda_visible_devices.split(",")
-    assert (
+        return int(device_list[device_idx])
-        len(device_list) == 1
+    else:
-    ), "CUDA_VISIBLE_DEVICES should be set to a single device in a scheduler"
+        return device_idx
-    return int(device_list[0])
 
 
 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)