[Intel][XPU] Add NUMA node binding support for Intel XPU (#31113)

This commit is contained in:
Siju Samuel
2026-09-07 10:39:21 +08:00
committed by GitHub
parent 1c992bbd94
commit c8207e32b6
3 changed files with 518 additions and 9 deletions
+159 -8
View File
@@ -16,9 +16,10 @@ import torch
from sglang.srt.environ import envs
from sglang.srt.server_args import ServerArgs
from sglang.srt.utils import get_cpu_ids_by_node, is_cuda
from sglang.srt.utils import get_cpu_ids_by_node, is_cuda, is_xpu
_is_cuda = is_cuda()
_is_xpu = is_xpu()
logger = logging.getLogger(__name__)
@@ -59,11 +60,17 @@ def configure_subprocess(server_args: ServerArgs, gpu_id: int):
executable, debug_str = _create_numactl_executable(
numactl_args=numactl_args
)
debug_str += (
f", logical_gpu_id={gpu_id}, "
f"physical_gpu_id={_get_nvml_device_index(gpu_id)}, "
f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}"
)
if _is_xpu:
debug_str += (
f", logical_gpu_id={gpu_id}, "
f"ZE_AFFINITY_MASK={os.environ.get('ZE_AFFINITY_MASK', '')}"
)
else:
debug_str += (
f", logical_gpu_id={gpu_id}, "
f"physical_gpu_id={_get_nvml_device_index(gpu_id)}, "
f"CUDA_VISIBLE_DEVICES={os.environ.get('CUDA_VISIBLE_DEVICES', '')}"
)
with _mp_set_executable(executable=executable, debug_str=debug_str):
yield
return
@@ -356,7 +363,7 @@ def _is_numa_available() -> bool:
"""
Check if NUMA is available and not already configured externally.
"""
if not _is_cuda:
if not (_is_cuda or _is_xpu):
return False
# Check if this is a numa system.
@@ -383,10 +390,13 @@ def _query_numa_node_for_gpu(device_id: int):
Get the NUMA node affinity list for a GPU device.
Args:
device_id: CUDA logical device index (post-CUDA_VISIBLE_DEVICES).
device_id: Logical device index (post-CUDA_VISIBLE_DEVICES / ZE_AFFINITY_MASK).
Returns:
List of NUMA node IDs that have affinity with the device.
"""
if _is_xpu:
return _query_numa_node_for_xpu(device_id)
try:
import pynvml
except ModuleNotFoundError:
@@ -486,3 +496,144 @@ def init_threads_binding(
f"local_omp_cpuid={local_omp_cpuid}"
)
return local_omp_cpuid
def _list_xpu_pci_addresses():
"""Intel GPU PCI addresses in Level Zero's physical device order."""
# sysfs, not the XPU runtime: this also runs in the launcher parent, where a
# torch XPU init pins ~170MB of device memory for the life of the server.
addresses = []
# Walk /dev/dri, which is what Level Zero opens and what a container narrows
# to the devices it was given; the sysfs PCI tree always shows every host GPU.
for node in glob.glob("/dev/dri/renderD*"):
if not os.access(node, os.R_OK | os.W_OK):
continue
device_link = os.path.join("/sys/class/drm", os.path.basename(node), "device")
try:
with open(os.path.join(device_link, "vendor")) as f:
vendor = f.read().strip()
except OSError:
continue
if vendor == "0x8086":
addresses.append(os.path.basename(os.path.realpath(device_link)))
# Level Zero sorts by PCI domain/bus/device/function, but places integrated
# GPUs last unless ZE_ENABLE_PCI_ID_DEVICE_ORDER is set
# (intel/compute-runtime, ExecutionEnvironment::comparePciIdBusNumber).
pci_id_order = os.environ.get("ZE_ENABLE_PCI_ID_DEVICE_ORDER", "").strip() not in (
"",
"0",
)
return sorted(
addresses,
key=lambda address: _xpu_pci_sort_key(
address, integrated_last=not pci_id_order
),
)
def _is_integrated_xpu(pci_address: str) -> bool:
# An Intel integrated GPU is always at domain 0000, bus 00.
domain, bus, _ = pci_address.split(":")
return int(domain, 16) == 0 and int(bus, 16) == 0
def _xpu_pci_sort_key(pci_address: str, *, integrated_last: bool):
domain, bus, device_function = pci_address.split(":")
device, function = device_function.split(".")
return (
integrated_last and _is_integrated_xpu(pci_address),
int(domain, 16),
int(bus, 16),
int(device, 16),
int(function),
)
def _xpu_visible_device_indices(num_devices: int):
"""Physical XPU indices ZE_AFFINITY_MASK exposes, in logical index order."""
affinity_mask = os.environ.get("ZE_AFFINITY_MASK", "").strip()
if not affinity_mask or affinity_mask == "default":
return list(range(num_devices))
# ZE_AFFINITY_MASK filters instead of permuting like CUDA_VISIBLE_DEVICES
# (intel/compute-runtime, ExecutionEnvironment::parseAffinityMask).
indices = set()
for entry in affinity_mask.split(","):
try:
index = int(entry.strip().split(".")[0])
except ValueError:
continue
if 0 <= index < num_devices:
indices.add(index)
return sorted(indices)
def _xpu_visible_pci_addresses(addresses: list, *, device_count: int):
"""The candidate device list whose length matches what the runtime reports."""
# Level Zero can leave an integrated GPU out of its device list altogether,
# and it orders them last anyway, so retry without them before giving up.
for candidate in (addresses, [a for a in addresses if not _is_integrated_xpu(a)]):
visible = [candidate[i] for i in _xpu_visible_device_indices(len(candidate))]
if len(visible) == device_count:
return visible
return None
def _xpu_pci_address(device_id: int) -> Optional[str]:
"""Logical XPU device index -> its PCI address, or None if unresolvable.
The ZE_AFFINITY_MASK-aware counterpart of _get_nvml_device_index.
"""
addresses = _list_xpu_pci_addresses()
if not addresses:
logger.warning(
"No Intel GPU render nodes found under /dev/dri, skipping NUMA node "
"configuration for XPU"
)
return None
# device_count() is the runtime's own count and needs no XPU init.
device_count = torch.xpu.device_count()
visible = _xpu_visible_pci_addresses(addresses, device_count=device_count)
if visible is None:
logger.warning(
f"Found {len(addresses)} Intel GPU(s) via /dev/dri but torch reports "
f"{device_count} XPU device(s) with ZE_AFFINITY_MASK="
f"{os.environ.get('ZE_AFFINITY_MASK', '')!r}, skipping NUMA node "
"configuration for XPU"
)
return None
if not 0 <= device_id < len(visible):
logger.warning(
f"XPU device {device_id} is out of range of the {len(visible)} "
"Intel GPU(s) found in sysfs, skipping NUMA node configuration for XPU"
)
return None
return visible[device_id]
def _read_pci_numa_node(pci_address: str):
numa_path = f"/sys/bus/pci/devices/{pci_address}/numa_node"
try:
with open(numa_path) as f:
node = int(f.read().strip())
except (OSError, ValueError) as e:
logger.warning(
f"Could not read {numa_path}: {e}, skipping NUMA node configuration for XPU"
)
return []
# The kernel reports -1 when the device has no NUMA affinity.
if node < 0:
return []
return [node]
def _query_numa_node_for_xpu(device_id: int):
"""NUMA node affinity list for an Intel XPU device, via sysfs."""
pci_address = _xpu_pci_address(device_id)
if pci_address is None:
return []
return _read_pci_numa_node(pci_address)
+5 -1
View File
@@ -27,8 +27,9 @@ register_cuda_ci(est_time=6, stage="base-c", runner_config="4-gpu-b200")
class TestIsNumaAvailable(unittest.TestCase):
"""Tests for _is_numa_available on both NUMA and non-NUMA systems."""
@patch("sglang.srt.utils.numa_utils._is_xpu", False)
@patch("sglang.srt.utils.numa_utils._is_cuda", False)
def test_returns_false_when_not_cuda(self):
def test_returns_false_when_not_cuda_or_xpu(self):
self.assertFalse(_is_numa_available())
@patch("sglang.srt.utils.numa_utils._is_cuda", True)
@@ -65,6 +66,9 @@ class TestIsNumaAvailable(unittest.TestCase):
mock_isdir.assert_called_with("/sys/devices/system/node/node1")
# Pin _is_xpu=False so these cases still reach the mocked pynvml on a real XPU
# host, where _query_numa_node_for_gpu short-circuits into the XPU sysfs branch.
@patch("sglang.srt.utils.numa_utils._is_xpu", False)
class TestQueryNumaNodeForGpu(unittest.TestCase):
"""Tests for _query_numa_node_for_gpu with mocked pynvml."""
+354
View File
@@ -0,0 +1,354 @@
"""
XPU NUMA-binding tests for numa_utils. The CUDA/CPU cases live in
test/registered/utils/test_numa_utils.py.
Usage:
python3 -m unittest test_numa_utils_xpu
"""
import os
import unittest
from unittest.mock import mock_open, patch
import torch
from sglang.srt.utils.numa_utils import (
_is_numa_available,
_list_xpu_pci_addresses,
_query_numa_node_for_gpu,
_query_numa_node_for_xpu,
_xpu_pci_address,
_xpu_visible_device_indices,
)
from sglang.test.ci.ci_register import register_cpu_ci, register_xpu_ci
# CPU too: only TestXpuOrderMatchesRuntime needs hardware.
register_cpu_ci(est_time=7, suite="base-a-test-cpu")
register_xpu_ci(est_time=10, suite="stage-a-test-1-gpu-xpu")
def _fake_sysfs_reader(pci_meta: dict):
"""``open`` replacement serving sysfs path -> contents from ``pci_meta``."""
def _open(path, *args, **kwargs):
if path in pci_meta:
return mock_open(read_data=pci_meta[path])(path, *args, **kwargs)
raise FileNotFoundError(path)
return _open
def _fake_dri_glob(render_nodes: list):
"""``glob.glob`` replacement listing ``render_nodes`` under /dev/dri."""
def _glob(pattern):
if pattern == "/dev/dri/renderD*":
return [f"/dev/dri/{node}" for node in render_nodes]
raise AssertionError(f"unexpected glob pattern {pattern!r}")
return _glob
def _bdf_from_xpu_uuid(uuid_bytes: bytes):
"""Intel device UUID -> PCI address, or None if not in the Intel layout."""
# vendor(2 LE) device(2 LE) revision(2 LE) domain(2 LE) bus(1) device(1)
# function(1) reserved(4) sub-device(1) -- intel/compute-runtime,
# Device::generateUuidFromPciBusInfo.
if len(uuid_bytes) != 16 or int.from_bytes(uuid_bytes[0:2], "little") != 0x8086:
return None
domain = int.from_bytes(uuid_bytes[6:8], "little")
bus, device, function = uuid_bytes[8], uuid_bytes[9], uuid_bytes[10]
return f"{domain:04x}:{bus:02x}:{device:02x}.{function}"
class TestIsNumaAvailableXpu(unittest.TestCase):
"""_is_numa_available must admit XPU just like CUDA."""
@patch("sglang.srt.utils.numa_utils._can_set_mempolicy", return_value=True)
@patch("sglang.srt.utils.numa_utils.shutil.which", return_value="/usr/bin/numactl")
@patch("sglang.srt.utils.numa_utils._is_xpu", True)
@patch("sglang.srt.utils.numa_utils._is_cuda", False)
@patch("os.path.isdir", return_value=True)
def test_returns_true_on_xpu_numa_system(
self, _mock_isdir, _mock_which, _mock_mempolicy
):
self.assertTrue(_is_numa_available())
class TestListXpuPciAddresses(unittest.TestCase):
"""_list_xpu_pci_addresses filters and orders like Level Zero does."""
def _run(self, devices: dict, env: dict = None, unopenable: tuple = ()):
"""``devices`` maps a render node to (PCI address, vendor or None)."""
pci_meta = {
f"/sys/class/drm/{node}/device/vendor": vendor + "\n"
for node, (_, vendor) in devices.items()
if vendor is not None
}
realpaths = {
f"/sys/class/drm/{node}/device": f"/sys/devices/pci0000:00/{address}"
for node, (address, _) in devices.items()
}
with (
patch(
"sglang.srt.utils.numa_utils.glob.glob", _fake_dri_glob(list(devices))
),
patch("builtins.open", _fake_sysfs_reader(pci_meta)),
patch(
"os.access", lambda path, mode: os.path.basename(path) not in unopenable
),
patch("os.path.realpath", lambda path: realpaths.get(path, path)),
patch.dict(os.environ, env or {}, clear=False),
):
if not (env or {}).get("ZE_ENABLE_PCI_ID_DEVICE_ORDER"):
os.environ.pop("ZE_ENABLE_PCI_ID_DEVICE_ORDER", None)
return _list_xpu_pci_addresses()
def test_keeps_only_intel_render_nodes_sorted(self):
devices = {
"renderD129": ("0000:29:00.0", "0x8086"), # Intel GPU -> keep
"renderD128": ("0000:18:00.0", "0x8086"), # Intel GPU -> keep
"renderD130": ("0000:3a:00.0", "0x10de"), # NVIDIA GPU -> drop (vendor)
}
self.assertEqual(self._run(devices), ["0000:18:00.0", "0000:29:00.0"])
def test_render_node_this_process_cannot_open_is_skipped(self):
# A container is handed a subset of /dev/dri; a node it cannot open is not
# an XPU here, and counting it would shift every logical device index.
devices = {
"renderD128": ("0000:18:00.0", "0x8086"),
"renderD129": ("0000:29:00.0", "0x8086"),
}
self.assertEqual(
self._run(devices, unopenable=("renderD128",)), ["0000:29:00.0"]
)
def test_integrated_gpu_is_ordered_last_by_default(self):
devices = {
"renderD128": ("0000:00:02.0", "0x8086"),
"renderD129": ("0000:03:00.0", "0x8086"),
}
self.assertEqual(self._run(devices), ["0000:03:00.0", "0000:00:02.0"])
def test_bus_00_in_another_domain_is_not_integrated(self):
# Only domain 0000 bus 00 is an iGPU; a second host bridge keeps BDF order.
devices = {
"renderD128": ("0001:00:00.0", "0x8086"),
"renderD129": ("0001:03:00.0", "0x8086"),
}
self.assertEqual(self._run(devices), ["0001:00:00.0", "0001:03:00.0"])
def test_pci_id_device_order_env_restores_plain_bdf_order(self):
devices = {
"renderD128": ("0000:00:02.0", "0x8086"),
"renderD129": ("0000:03:00.0", "0x8086"),
}
result = self._run(devices, env={"ZE_ENABLE_PCI_ID_DEVICE_ORDER": "1"})
self.assertEqual(result, ["0000:00:02.0", "0000:03:00.0"])
def test_sorts_by_domain_bus_device_function(self):
# Only case with hex digits in the BDF; catches decimal bus/device parsing.
devices = {
"renderD128": ("0001:0a:00.0", "0x8086"),
"renderD129": ("0000:0a:00.1", "0x8086"),
"renderD130": ("0000:0a:00.0", "0x8086"),
"renderD131": ("0000:09:1f.0", "0x8086"),
}
self.assertEqual(
self._run(devices),
["0000:09:1f.0", "0000:0a:00.0", "0000:0a:00.1", "0001:0a:00.0"],
)
def test_skips_nodes_without_sysfs_attrs(self):
self.assertEqual(self._run({"renderD128": ("0000:18:00.0", None)}), [])
class TestXpuVisibleDeviceIndices(unittest.TestCase):
"""ZE_AFFINITY_MASK filters the device list; it does not permute it."""
def _indices(self, mask, num_devices=4):
env = {} if mask is None else {"ZE_AFFINITY_MASK": mask}
with patch.dict(os.environ, env, clear=False):
if mask is None:
os.environ.pop("ZE_AFFINITY_MASK", None)
return _xpu_visible_device_indices(num_devices)
def test_unset_mask_exposes_every_device(self):
self.assertEqual(self._indices(None), [0, 1, 2, 3])
self.assertEqual(self._indices(""), [0, 1, 2, 3])
self.assertEqual(self._indices("default"), [0, 1, 2, 3])
def test_mask_order_is_ignored(self):
# Verified against Level Zero on an 8-GPU host: "3,0" exposes physical
# 0 as xpu:0 and physical 3 as xpu:1.
self.assertEqual(self._indices("3,0"), [0, 3])
self.assertEqual(self._indices("0,3"), [0, 3])
self.assertEqual(self._indices("3,2,1,0"), [0, 1, 2, 3])
def test_duplicate_and_out_of_range_entries_are_dropped(self):
self.assertEqual(self._indices("3,3,0"), [0, 3])
self.assertEqual(self._indices("1,99"), [1])
self.assertEqual(self._indices("99"), [])
self.assertEqual(self._indices("-1"), [])
def test_composite_entry_selects_its_root_device(self):
self.assertEqual(self._indices("1.0,3.0"), [1, 3])
self.assertEqual(self._indices("2.1"), [2])
def test_surrounding_whitespace_is_tolerated(self):
self.assertEqual(self._indices(" 2 , 1 "), [1, 2])
def test_unparsable_entries_are_dropped(self):
self.assertEqual(self._indices("1,abc,2"), [1, 2])
class TestQueryNumaNodeForXpu(unittest.TestCase):
"""_query_numa_node_for_xpu: mask handling, runtime cross-check, node read."""
ADDRS = [
"0000:18:00.0",
"0000:29:00.0",
"0000:3a:00.0",
"0000:5c:00.0",
]
NODES = {
"0000:18:00.0": "0",
"0000:29:00.0": "0",
"0000:3a:00.0": "1",
"0000:5c:00.0": "-1", # no NUMA affinity
}
def _query(
self, device_id, mask=None, device_count=None, nodes=None, addresses=None
):
addresses = self.ADDRS if addresses is None else addresses
pci_meta = {
f"/sys/bus/pci/devices/{a}/numa_node": node + "\n"
for a, node in (self.NODES if nodes is None else nodes).items()
}
env = {} if mask is None else {"ZE_AFFINITY_MASK": mask}
if device_count is None:
device_count = len(addresses)
with (
patch(
"sglang.srt.utils.numa_utils._list_xpu_pci_addresses",
return_value=list(addresses),
),
patch("builtins.open", _fake_sysfs_reader(pci_meta)),
patch.dict(os.environ, env, clear=False),
patch(
"torch.xpu.device_count",
return_value=device_count,
),
):
if mask is None:
os.environ.pop("ZE_AFFINITY_MASK", None)
return _query_numa_node_for_xpu(device_id)
def test_plain_device_index(self):
self.assertEqual(self._query(0), [0])
self.assertEqual(self._query(2), [1])
def test_negative_node_returns_empty(self):
self.assertEqual(self._query(3), [])
def test_affinity_mask_selects_by_ascending_physical_index(self):
# Treating "3,0" as a permutation would swap these two results.
self.assertEqual(self._query(0, mask="3,0", device_count=2), [0])
self.assertEqual(self._query(1, mask="3,0", device_count=2), [])
self.assertEqual(self._query(0, mask="2,1", device_count=2), [0])
self.assertEqual(self._query(1, mask="2,1", device_count=2), [1])
def test_device_id_out_of_range_returns_empty(self):
self.assertEqual(self._query(9, device_count=4), [])
# Core + iGPU + 1x Arc Pro B60: sysfs lists both Intel GPUs but Level Zero
# enumerates only the discrete one, so xpu:0 is the B60, not the iGPU.
IGPU_ADDRS = ["0000:03:00.0", "0000:00:02.0"] # discrete first, iGPU last
IGPU_NODES = {"0000:03:00.0": "1", "0000:00:02.0": "0"}
def test_integrated_gpu_the_runtime_does_not_enumerate_is_dropped(self):
self.assertEqual(
self._query(
0, addresses=self.IGPU_ADDRS, nodes=self.IGPU_NODES, device_count=1
),
[1],
)
def test_integrated_gpu_the_runtime_does_enumerate_stays_last(self):
for device_id, node in ((0, [1]), (1, [0])):
self.assertEqual(
self._query(
device_id,
addresses=self.IGPU_ADDRS,
nodes=self.IGPU_NODES,
device_count=2,
),
node,
)
def test_runtime_device_count_mismatch_returns_empty(self):
self.assertEqual(self._query(0, device_count=8), [])
self.assertEqual(self._query(0, device_count=0), [])
def test_unreadable_numa_node_returns_empty(self):
self.assertEqual(self._query(0, nodes={}), [])
self.assertEqual(self._query(0, nodes={self.ADDRS[0]: "not-a-node"}), [])
def test_no_intel_gpus_returns_empty(self):
with patch(
"sglang.srt.utils.numa_utils._list_xpu_pci_addresses", return_value=[]
):
self.assertEqual(_query_numa_node_for_xpu(0), [])
@patch("sglang.srt.utils.numa_utils._is_xpu", True)
def test_query_numa_node_for_gpu_delegates_to_xpu(self):
with patch(
"sglang.srt.utils.numa_utils._query_numa_node_for_xpu",
return_value=[2],
) as mock_xpu:
self.assertEqual(_query_numa_node_for_gpu(5), [2])
mock_xpu.assert_called_once_with(5)
@unittest.skipUnless(torch.xpu.is_available(), "requires XPU hardware")
class TestXpuOrderMatchesRuntime(unittest.TestCase):
"""The /dev/dri reconstruction must match the device order a worker sees."""
def setUp(self):
# A host whose /dev/dri and runtime disagree is one where declining to
# bind is the correct answer, so there is no mapping to assert; the
# mocked cases above pin that reconciliation logic.
if _xpu_pci_address(0) is None:
self.skipTest(
f"/dev/dri lists {len(_list_xpu_pci_addresses())} Intel GPU(s) but "
f"torch reports {torch.xpu.device_count()}; no mapping to check"
)
def test_logical_index_resolves_to_the_same_pci_address_as_torch(self):
for device_id in range(torch.xpu.device_count()):
device_uuid = torch.xpu.get_device_properties(device_id).uuid
runtime_address = _bdf_from_xpu_uuid(device_uuid.bytes)
if runtime_address is None:
self.skipTest(
f"xpu:{device_id} UUID {device_uuid} is not in the Intel layout"
)
self.assertEqual(
_xpu_pci_address(device_id),
runtime_address,
f"xpu:{device_id} mismatch",
)
def test_every_device_resolves_to_its_sysfs_numa_node(self):
for device_id in range(torch.xpu.device_count()):
address = _xpu_pci_address(device_id)
with open(f"/sys/bus/pci/devices/{address}/numa_node") as f:
sysfs_node = int(f.read().strip())
expected = [] if sysfs_node < 0 else [sysfs_node]
self.assertEqual(
_query_numa_node_for_xpu(device_id), expected, f"xpu:{device_id}"
)
if __name__ == "__main__":
unittest.main()