diff --git a/docs_new/docs/hardware-platforms/plugin.mdx b/docs_new/docs/hardware-platforms/plugin.mdx index cc63cee75..64de53ef9 100644 --- a/docs_new/docs/hardware-platforms/plugin.mdx +++ b/docs_new/docs/hardware-platforms/plugin.mdx @@ -88,6 +88,7 @@ entry_points("sglang.srt.platforms") → Enumerate ALL plugins by name (metadat ├─ 0 activated + SGLANG_USE_CPU_ENGINE=1 → fallback CpuSRTPlatform ├─ 0 activated + CUDA available → fallback CudaSRTPlatform ├─ 0 activated + ROCm available → fallback RocmSRTPlatform + ├─ 0 activated + XPU available → fallback XpuSRTPlatform ├─ 0 activated + none of the above → fallback base SRTPlatform ├─ 1 activated → use it └─ N activated → RuntimeError (must set SGLANG_PLATFORM) diff --git a/python/sglang/srt/platforms/__init__.py b/python/sglang/srt/platforms/__init__.py index 860936549..cd740f253 100644 --- a/python/sglang/srt/platforms/__init__.py +++ b/python/sglang/srt/platforms/__init__.py @@ -22,6 +22,7 @@ from sglang.srt.platforms.cpu import CpuSRTPlatform from sglang.srt.platforms.cuda import CudaSRTPlatform from sglang.srt.platforms.interface import SRTPlatform from sglang.srt.platforms.rocm import RocmSRTPlatform +from sglang.srt.platforms.xpu import XpuSRTPlatform from sglang.srt.plugins import PLATFORM_PLUGINS_GROUP, load_plugins_by_group logger = logging.getLogger(__name__) @@ -41,6 +42,10 @@ def _is_cpu_available() -> bool: return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" +def _is_xpu_available() -> bool: + return torch.xpu.is_available() + + def _resolve_platform() -> SRTPlatform: """ Discover and instantiate the active platform. @@ -62,6 +67,7 @@ def _resolve_platform() -> SRTPlatform: so developers on GPU hosts can intentionally exercise the CPU path) - 0 activated + CUDA available → fallback CudaSRTPlatform - 0 activated + ROCm available → fallback RocmSRTPlatform + - 0 activated + XPU available → fallback XpuSRTPlatform - 0 activated + none of the above → fallback base SRTPlatform - 1 activated → use it - N activated → RuntimeError (must set SGLANG_PLATFORM) @@ -126,6 +132,9 @@ def _resolve_platform() -> SRTPlatform: "No platform plugin detected. Using ROCm SRTPlatform defaults." ) return RocmSRTPlatform() + if _is_xpu_available(): + logger.debug("No platform plugin detected. Using XPU SRTPlatform defaults.") + return XpuSRTPlatform() logger.debug("No platform detected. Using base SRTPlatform.") return SRTPlatform() diff --git a/python/sglang/srt/platforms/xpu.py b/python/sglang/srt/platforms/xpu.py new file mode 100644 index 000000000..2995d7250 --- /dev/null +++ b/python/sglang/srt/platforms/xpu.py @@ -0,0 +1,104 @@ +"""XPU device operations for the SRT platform layer.""" + +import logging +from typing import Optional + +import torch + +from sglang.srt.platforms.device_mixin import ( + DeviceCapability, + DeviceMixin, + PlatformEnum, +) +from sglang.srt.platforms.interface import SRTPlatform + +logger = logging.getLogger(__name__) + + +class XpuDeviceMixin(DeviceMixin): + """XPU implementation of the shared device operations.""" + + _enum: PlatformEnum = PlatformEnum.XPU + device_name: str = "xpu" + device_type: str = "xpu" + + def get_device_total_memory(self, device_id: int = 0) -> int: + return int(torch.xpu.get_device_properties(device_id).total_memory) + + def get_current_memory_usage( + self, device: Optional["torch.device"] = None + ) -> float: + return float(torch.xpu.max_memory_allocated(device)) + + def get_device(self, local_rank: int) -> "torch.device": + return torch.device("xpu", local_rank) + + def set_device(self, device: "torch.device") -> None: + torch.xpu.set_device(device) + + def get_device_name(self, device_id: int = 0) -> str: + return str(torch.xpu.get_device_name(device_id)) + + def get_device_uuid(self, device_id: int = 0) -> str: + return str(torch.xpu.get_device_properties(device_id).uuid) + + def get_device_capability(self, device_id: int = 0) -> DeviceCapability: + # TODO: torch.xpu.get_device_capability + device = torch.xpu.current_device() + major, minor = torch.ops.sgl_kernel.query_device.default(device) + return DeviceCapability(major, minor) + + def empty_cache(self) -> None: + torch.xpu.empty_cache() + + def synchronize(self) -> None: + torch.xpu.synchronize() + + def get_available_memory(self, device_id: int = 0) -> tuple[int, int]: + # TODO: simple return of torch.xpu.mem_get_info + """Return the available and total device memory in Bytes.""" + + if not (hasattr(torch, "xpu") and torch.xpu.is_available()): + return 0, 0 + + num_gpus = torch.xpu.device_count() + if device_id < 0 or device_id >= num_gpus: + raise ValueError(f"Invalid XPU device_id={device_id}. num_gpus={num_gpus}") + + current = torch.xpu.current_device() + if current != device_id: + logger.warning( + "current device is not %s, but %s; this may cause useless memory allocation for torch XPU context.", + device_id, + current, + ) + + used_memory = torch.xpu.memory_allocated(device_id) + total_gpu_memory = torch.xpu.get_device_properties(device_id).total_memory + free_gpu_memory = total_gpu_memory - used_memory + + return free_gpu_memory, total_gpu_memory + + def is_pin_memory_available(self, device=None) -> bool: + if device is not None and str(device) == "cpu": + return False + return True + + @classmethod + def seed_everything(cls, seed: int | None = None) -> None: + if seed is not None: + super().seed_everything(seed) + torch.xpu.manual_seed_all(seed) + + +class XpuSRTPlatform(XpuDeviceMixin, SRTPlatform): + """Default in-tree XPU SRT platform.""" + + def supports_fp8(self) -> bool: + return False + + def support_cuda_graph(self) -> bool: + return True + + def support_piecewise_cuda_graph(self) -> bool: + return True diff --git a/test/registered/unit/platforms/test_platform_interface.py b/test/registered/unit/platforms/test_platform_interface.py index 50809184f..40871c926 100644 --- a/test/registered/unit/platforms/test_platform_interface.py +++ b/test/registered/unit/platforms/test_platform_interface.py @@ -20,6 +20,7 @@ from sglang.srt.platforms.device_mixin import ( ) from sglang.srt.platforms.interface import SRTPlatform from sglang.srt.platforms.rocm import RocmSRTPlatform +from sglang.srt.platforms.xpu import XpuSRTPlatform from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.test_utils import CustomTestCase @@ -255,6 +256,47 @@ class TestCudaDeviceMixin(CustomTestCase): self.assertTrue(base.support_piecewise_cuda_graph()) +class TestXpuDeviceMixin(CustomTestCase): + """Tests for XPU device operation defaults.""" + + def test_default_get_device_returns_xpu_device(self): + base = XpuSRTPlatform() + self.assertEqual(base.get_device(2), torch.device("xpu", 2)) + + # TODO: @patch("torch.xpu.get_device_capability", return_value=(9, 0)) + @patch("torch.ops.sgl_kernel.query_device.default", return_value=(9, 0)) + def test_default_get_device_capability_uses_xpu(self, mock_get_device_capability): + base = XpuSRTPlatform() + self.assertEqual(base.get_device_capability(0), DeviceCapability(9, 0)) + mock_get_device_capability.assert_called_once_with(0) + + def test_pin_memory_available_for_xpu_targets(self): + base = XpuSRTPlatform() + self.assertTrue(base.is_pin_memory_available()) + self.assertTrue(base.is_pin_memory_available(device="xpu")) + self.assertTrue(base.is_pin_memory_available(device=torch.device("xpu", 0))) + self.assertFalse(base.is_pin_memory_available(device="cpu")) + + @patch("torch.xpu.manual_seed_all") + @patch("torch.manual_seed") + @patch("sglang.srt.platforms.device_mixin.np.random.seed") + @patch("sglang.srt.platforms.device_mixin.random.seed") + def test_default_seed_everything_seeds_xpu( + self, mock_random_seed, mock_np_seed, mock_torch_seed, mock_xpu_seed + ): + XpuSRTPlatform.seed_everything(123) + mock_random_seed.assert_called_once_with(123) + mock_np_seed.assert_called_once_with(123) + mock_torch_seed.assert_called_once_with(123) + mock_xpu_seed.assert_called_once_with(123) + + def test_xpu_srt_platform_capabilities(self): + base = XpuSRTPlatform() + self.assertFalse(base.supports_fp8()) + self.assertTrue(base.support_cuda_graph()) + self.assertTrue(base.support_piecewise_cuda_graph()) + + class TestCpuDeviceMixin(CustomTestCase): """Tests for CPU device operation defaults (covers both x86 and ARM)."""