diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index e11bfbf8d..e69d4b9ad 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -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): diff --git a/python/sglang/srt/utils/numa_utils.py b/python/sglang/srt/utils/numa_utils.py index d9f765c69..af50e6d9b 100644 --- a/python/sglang/srt/utils/numa_utils.py +++ b/python/sglang/srt/utils/numa_utils.py @@ -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." ) diff --git a/test/registered/cpu/test_binding.py b/test/registered/cpu/test_binding.py index ff9195280..26e529a37 100644 --- a/test/registered/cpu/test_binding.py +++ b/test/registered/cpu/test_binding.py @@ -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()