minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[CPU] Fix NUMA/core binding for DP ranks (#32856)

Browse Source
This commit is contained in:
Chunyuan WU
2026-08-24 12:15:38 +08:00
committed by GitHub
parent 3e30649064
commit 1daa94a069
3 changed files with 81 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/srt/model_executor/model_runner.py
+7 -1
View File
@@ -1415,8 +1415,14 @@ class ModelRunner:
)
def init_threads_binding(self):
# With --enable-dp-attention, dp partitions the existing TP group
# rather than spawning additional processes, so dp_size must not be
# multiplied into the process count here (unlike regular DP, where
# dp_size * tp_size * pp_size is the true worker count).
dp_size = 1 if get_parallel().enable_dp_attention else self.ps.dp_size
self.local_omp_cpuid = numa_utils.init_threads_binding(
tp_rank=self.ps.tp_rank, tp_size=self.ps.tp_size
numa_index=self.gpu_id,
world_size=dp_size * self.ps.tp_size * self.ps.pp_size,
)
def apply_torch_tp(self):
python/sglang/srt/utils/numa_utils.py
+29 -16
View File
@@ -433,38 +433,51 @@ def _query_numa_node_for_gpu(device_id: int):
def init_threads_binding(
*,
tp_rank: int,
tp_size: int,
numa_index: int,
world_size: int,
):
omp_cpuids = os.environ.get("SGLANG_CPU_OMP_THREADS_BIND", "all")
cpu_ids_by_node = get_cpu_ids_by_node()
n_numa_node = len(cpu_ids_by_node)
if omp_cpuids == "all":
assert tp_size <= n_numa_node, (
assert world_size <= n_numa_node, (
f"SGLANG_CPU_OMP_THREADS_BIND is not set, in this case, "
f"tp_size {tp_size} should be smaller than or equal to number of numa node on the machine {n_numa_node}. "
f"If you need tp_size to be larger than number of numa node, please set the CPU cores for each tp rank via SGLANG_CPU_OMP_THREADS_BIND explicitly. "
f"the total number of ranks (dp_size * tp_size * pp_size = {world_size}) should be smaller than or equal to number of numa node on the machine {n_numa_node}. "
f"If you need more ranks than the number of numa nodes, please set the CPU cores for each rank via SGLANG_CPU_OMP_THREADS_BIND explicitly. "
f"For example, on a machine with 2 numa nodes, where core 0-31 are on numa node 0 and core 32-63 are on numa node 1, "
f"it is suggested to use -tp 2 and bind tp rank 0 to core 0-31 and tp rank 1 to core 32-63. "
f"This is the default behavior if SGLANG_CPU_OMP_THREADS_BIND is not set and it is the same as setting SGLANG_CPU_OMP_THREADS_BIND=0-31|32-63. "
f"If you do need tp_size to be larger than the number of numa nodes, you could set SGLANG_CPU_OMP_THREADS_BIND explicitly for example SGLANG_CPU_OMP_THREADS_BIND=0-15|16-31|32-47|48-63 and run with -tp 4. "
f"If you don't want each tp rank to use all the cores on one numa node, you could set for example SGLANG_CPU_OMP_THREADS_BIND=0-15|32-47 and run with -tp 2."
f"If you do need more ranks than the number of numa nodes, you could set SGLANG_CPU_OMP_THREADS_BIND explicitly for example SGLANG_CPU_OMP_THREADS_BIND=0-15|16-31|32-47|48-63 and run with -tp 4. "
f"If you don't want each rank to use all the cores on one numa node, you could set for example SGLANG_CPU_OMP_THREADS_BIND=0-15|32-47 and run with -tp 2."
)
if tp_size < n_numa_node:
if world_size < n_numa_node:
logger.warning(
f"Detected the current machine has {n_numa_node} numa nodes available, but tp_size is set to {tp_size}, so only {tp_size} numa nodes are used."
f"Detected the current machine has {n_numa_node} numa nodes available, but the total number of ranks (dp_size * tp_size * pp_size) is {world_size}, so only {world_size} numa nodes are used."
)
local_omp_cpuid = cpu_ids_by_node[tp_rank]
assert 0 <= numa_index < n_numa_node, (
f"NUMA index {numa_index} (derived from the worker's global device id / gpu_id) "
f"is out of range for {n_numa_node} numa nodes. This usually means dp_size * tp_size "
f"exceeds the number of numa nodes; reduce it, or set SGLANG_CPU_OMP_THREADS_BIND explicitly."
)
local_omp_cpuid = cpu_ids_by_node[numa_index]
else:
threads_bind_list = omp_cpuids.split("|")
assert tp_size == len(threads_bind_list), (
f"SGLANG_CPU_OMP_THREADS_BIND setting must be aligned with TP size parameter ({tp_size}). "
f"Please double check your settings."
# Bound-check numa_index against the bind list rather than asserting
# world_size == len(...): in router mode each worker is an independent
# dp_size=1 server, so world_size is locally 1 and can't equal a
# multi-group bind string. numa_index (== the global gpu_id) is the
# correct frame here, so this per-rank bound both prevents IndexError
# and catches an under-sized bind list, without needing the global
# rank count.
assert 0 <= numa_index < len(threads_bind_list), (
f"NUMA index {numa_index} (derived from the worker's global device id / gpu_id) "
f"is out of range for the {len(threads_bind_list)} SGLANG_CPU_OMP_THREADS_BIND entries. "
f"Ensure the number of '|'-separated bind groups matches dp_size * tp_size * pp_size (across all DP workers)."
)
local_omp_cpuid = threads_bind_list[tp_rank]
if tp_size > n_numa_node:
local_omp_cpuid = threads_bind_list[numa_index]
if world_size > n_numa_node:
logger.warning(
f"TP size ({tp_size})is larger than numa node number ({n_numa_node}), "
f"The total number of ranks ({world_size}) is larger than numa node number ({n_numa_node}), "
f"in this case the available memory amount of each rank cannot be determined in prior. "
f"Please set proper `--max-total-tokens` to avoid the out-of-memory error."
)
test/registered/cpu/test_binding.py
+45
View File
@@ -1,10 +1,13 @@
import os
import re
import unittest
from unittest.mock import patch
import torch
kernel = torch.ops.sgl_kernel
from sglang.srt.utils.numa_utils import init_threads_binding
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -26,5 +29,47 @@ class TestBinding(CustomTestCase):
self.assertEqual(bindings, expected_cores)
class TestInitThreadsBinding(unittest.TestCase):
"""Tests for init_threads_binding: NUMA/core selection by global rank.
numa_index is the worker's global device id (gpu_id) across all DP
replicas; world_size is dp_size * tp_size * pp_size.
"""
@patch(
"sglang.srt.utils.numa_utils.get_cpu_ids_by_node",
return_value=["0,1,2,3", "4,5,6,7", "8,9,10,11", "12,13,14,15"],
)
@patch.dict(os.environ, {"SGLANG_CPU_OMP_THREADS_BIND": "all"})
def test_dp_ranks_get_distinct_numa_nodes(self, _mock_nodes):
# dp_size=2, tp_size=2 -> world_size=4, numa_index 0..3 global rank.
results = [init_threads_binding(numa_index=i, world_size=4) for i in range(4)]
self.assertEqual(results, ["0,1,2,3", "4,5,6,7", "8,9,10,11", "12,13,14,15"])
@patch(
"sglang.srt.utils.numa_utils.get_cpu_ids_by_node",
return_value=["0,1,2,3", "4,5,6,7"],
)
@patch.dict(os.environ, {"SGLANG_CPU_OMP_THREADS_BIND": "0-1|4-5"})
def test_explicit_bind_list_indexed_by_numa_index(self, _mock_nodes):
self.assertEqual(init_threads_binding(numa_index=0, world_size=2), "0-1")
self.assertEqual(init_threads_binding(numa_index=1, world_size=2), "4-5")
@patch(
"sglang.srt.utils.numa_utils.get_cpu_ids_by_node",
return_value=["0,1", "2,3", "4,5"],
)
@patch.dict(os.environ, {"SGLANG_CPU_OMP_THREADS_BIND": "0-1|2-3|4-5"})
def test_router_worker_uses_global_numa_index(self, _mock_nodes):
# Router mode: each worker is an independent dp_size=1 server, so
# world_size is locally 1 even though the bind string has multiple
# groups. numa_index is still the global gpu_id and must be able to
# select any group, not just index 0.
self.assertEqual(
init_threads_binding(numa_index=2, world_size=1),
"4-5",
)
if __name__ == "__main__":
unittest.main()