diff --git a/test/registered/utils/test_numa_utils.py b/test/registered/utils/test_numa_utils.py index bec6d687f..ecc29d465 100644 --- a/test/registered/utils/test_numa_utils.py +++ b/test/registered/utils/test_numa_utils.py @@ -4,6 +4,8 @@ import unittest from contextlib import ExitStack from unittest.mock import MagicMock, patch +import torch + from sglang.srt.environ import envs from sglang.srt.utils.numa_utils import ( _handle_numa_bind_failure, @@ -12,6 +14,7 @@ from sglang.srt.utils.numa_utils import ( _numactl_cpu_mem_args, _probe_numactl_args, _query_numa_node_for_gpu, + _read_pci_numa_node, _strip_memory_args, configure_subprocess, get_numa_node_if_available, @@ -306,18 +309,20 @@ class TestGraceBlackwellNumaTopology(unittest.TestCase): "Requires 4-GPU B200 hardware", ) class TestB200NumaTopology(unittest.TestCase): - """Hardware test validating expected NUMA topology on 4-GPU B200.""" + """Hardware test for the 4-GPU B200 NUMA mapping; layout is BIOS-dependent.""" def test_gpu_numa_mapping(self): self.assertEqual(_gpu_count, 4) - numa_nodes = { - _query_single_numa_node_for_gpu(gpu_id) for gpu_id in range(_gpu_count) - } - self.assertEqual( - len(numa_nodes), - 1, - f"Expected all visible 4-GPU B200 devices on one NUMA node, got {numa_nodes}", - ) + for gpu_id in range(_gpu_count): + # Not NVML: the index mapping under test must not validate itself. + props = torch.cuda.get_device_properties(gpu_id) + pci_address = f"{props.pci_domain_id:04x}:{props.pci_bus_id:02x}:{props.pci_device_id:02x}.0" + self.assertIn( + _read_pci_numa_node(pci_address)[0], + _query_numa_node_for_gpu(gpu_id), + f"GPU {gpu_id} ({pci_address}): NVML memory affinity omits the " + f"sysfs numa_node", + ) class TestNumaBindIntersection(unittest.TestCase):