From 74e2e48c823c57d065e7ef7f0cd5c6e32fdef6b7 Mon Sep 17 00:00:00 2001 From: zijiexia <37504505+zijiexia@users.noreply.github.com> Date: Wed, 17 Jun 2026 17:41:17 -0700 Subject: [PATCH] Introduce CpuDeviceMixin and CpuSRTPlatform (#26385) --- docs_new/docs/hardware-platforms/plugin.mdx | 5 +- python/sglang/srt/platforms/__init__.py | 14 +- python/sglang/srt/platforms/cpu.py | 133 ++++++++++++++ .../unit/platforms/test_platform_interface.py | 163 ++++++++++++++++++ 4 files changed, 313 insertions(+), 2 deletions(-) create mode 100644 python/sglang/srt/platforms/cpu.py diff --git a/docs_new/docs/hardware-platforms/plugin.mdx b/docs_new/docs/hardware-platforms/plugin.mdx index 91686ce18..0debdc29b 100644 --- a/docs_new/docs/hardware-platforms/plugin.mdx +++ b/docs_new/docs/hardware-platforms/plugin.mdx @@ -85,7 +85,10 @@ entry_points("sglang.srt.platforms") → Enumerate ALL plugins by name (metadat │ └─ activate() returns None → RuntimeError (hardware unavailable) │ └─ SGLANG_PLATFORM unset (auto-discover, activate all): - ├─ 0 activated → fallback base SRTPlatform + ├─ 0 activated + SGLANG_USE_CPU_ENGINE=1 → fallback CpuSRTPlatform + ├─ 0 activated + CUDA available → fallback CudaSRTPlatform + ├─ 0 activated + ROCm available → fallback RocmSRTPlatform + ├─ 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 849f0c312..860936549 100644 --- a/python/sglang/srt/platforms/__init__.py +++ b/python/sglang/srt/platforms/__init__.py @@ -11,12 +11,14 @@ Usage: """ import logging +import os import pkgutil from importlib.metadata import entry_points import torch from sglang.srt.environ import envs +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 @@ -35,6 +37,10 @@ def _is_rocm_available() -> bool: return bool(torch.cuda.is_available() and torch.version.hip is not None) +def _is_cpu_available() -> bool: + return os.getenv("SGLANG_USE_CPU_ENGINE", "0") == "1" + + def _resolve_platform() -> SRTPlatform: """ Discover and instantiate the active platform. @@ -51,9 +57,12 @@ def _resolve_platform() -> SRTPlatform: SGLANG_PLATFORM unset (auto-discover): - Import and activate all discovered plugins + - 0 activated + SGLANG_USE_CPU_ENGINE=1 → fallback CpuSRTPlatform + (checked first; an explicit opt-in wins over CUDA/ROCm availability, + 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 + neither → fallback base SRTPlatform + - 0 activated + none of the above → fallback base SRTPlatform - 1 activated → use it - N activated → RuntimeError (must set SGLANG_PLATFORM) @@ -104,6 +113,9 @@ def _resolve_platform() -> SRTPlatform: logger.exception("Failed to activate platform plugin: %s", name) if len(activated) == 0: + if _is_cpu_available(): + logger.debug("SGLANG_USE_CPU_ENGINE=1. Using CPU SRTPlatform defaults.") + return CpuSRTPlatform() if _is_cuda_available(): logger.debug( "No platform plugin detected. Using CUDA SRTPlatform defaults." diff --git a/python/sglang/srt/platforms/cpu.py b/python/sglang/srt/platforms/cpu.py new file mode 100644 index 000000000..46fd0251a --- /dev/null +++ b/python/sglang/srt/platforms/cpu.py @@ -0,0 +1,133 @@ +"""CPU device operations for the SRT platform layer.""" + +import gc +import platform as _platform +from functools import cached_property +from typing import Optional + +import psutil +import torch + +from sglang.srt.platforms.device_mixin import ( + CpuArchEnum, + DeviceCapability, + DeviceMixin, + PlatformEnum, +) +from sglang.srt.platforms.interface import SRTPlatform + + +class CpuDeviceMixin(DeviceMixin): + """CPU implementation of the shared device operations.""" + + _enum: PlatformEnum = PlatformEnum.CPU + device_name: str = "cpu" + device_type: str = "cpu" + + @cached_property + def cpu_arch(self) -> CpuArchEnum: + """Host CPU architecture (X86 / ARM / UNSPECIFIED), resolved once. + + First-class identity attribute parallel to ``_enum`` — callers branch + on CPU arch through this instead of recomputing ``platform.machine()``. + ``get_cpu_architecture()`` is process-stable, so caching is safe. + """ + return self.get_cpu_architecture() + + def get_device_total_memory(self, device_id: int = 0) -> int: + return int(psutil.virtual_memory().total) + + def get_current_memory_usage( + self, device: Optional["torch.device"] = None + ) -> float: + """Whole-machine used memory (``total - available``) in bytes. + + Chosen so the [Active] contract + ``free = get_device_total_memory() - get_current_memory_usage()`` + yields ``psutil.available`` — the real free RAM on a machine shared + with the OS and other processes. Per-process RSS would wrongly ignore + their usage. There is no per-device allocator peak on CPU (unlike + ``torch.cuda.max_memory_allocated``), so this is current usage, not a + peak. Returns whole-machine bytes; per-rank NUMA division for CPU TP + is the caller's concern (kept in ``get_available_gpu_memory``'s CPU + branch), not here. + """ + vm = psutil.virtual_memory() + return float(vm.total - vm.available) + + def get_device(self, local_rank: int) -> "torch.device": + # local_rank is ignored: all CPU ranks share the one CPU device, so + # there is nothing rank-specific to return. PyTorch enforces this — + # Device::validate() asserts a CPU index must be -1 or 0 (c10/core/ + # Device.h). Per-rank isolation is done via OpenMP/numactl binding + # (ModelRunner.init_threads_binding), not the device object. + # TODO(zijiexia): make per-rank placement NUMA-affinity aware + # (rank -> NUMA node) when the platform layer takes this over. + return torch.device("cpu") + + def set_device(self, device: "torch.device") -> None: + # Documented no-op on CPU — torch.cpu.set_device is "in CPU we do + # nothing". Called (rather than left as ``pass``) for symmetry with + # CudaDeviceMixin.set_device. Note this is deliberately NOT + # torch.set_default_device("cpu"), which would flip the process-wide + # default tensor device; per-rank CPU isolation is via OpenMP/numactl + # binding (see get_device), not here. + torch.cpu.set_device(device) + + def get_device_name(self, device_id: int = 0) -> str: + # Arch-only label. We deliberately avoid platform.processor(): it + # spawns a subprocess (~ms) on some platforms (e.g. macOS) and on Linux + # is usually empty or redundant with the arch (e.g. "x86_64: x86_64"). + if self.cpu_arch == CpuArchEnum.ARM: + return "cpu (aarch64)" + if self.cpu_arch == CpuArchEnum.X86: + return "cpu (x86_64)" + return "cpu" + + def get_device_uuid(self, device_id: int = 0) -> str: + # CPU has no per-device UUID; return the arch string as a stable + # host-level identifier (matches the multimodal CpuPlatform). + return _platform.machine() + + def get_device_capability(self, device_id: int = 0) -> Optional[DeviceCapability]: + return None + + def empty_cache(self) -> None: + # No torch.cpu.empty_cache() exists; do a GC pass at the teardown + # points where this is called (flush_cache, idle sleep, weight reload). + # + # gc.collect() caveats: + # - the pause grows with heap size (full walk of tracked objects); + # - it only reclaims reference cycles — refcounting already frees + # everything else, so it may do little; + # - freed memory returns to the allocator, not the OS, so RSS may not + # drop. glibc malloc_trim would not help: it is a no-op under the + # tcmalloc / TBB malloc the CPU guide preloads via LD_PRELOAD. Real + # RSS reclaim belongs in a separate allocator-aware, benchmarked + # change. + gc.collect() + + def synchronize(self) -> None: + # Documented no-op on CPU (no async streams to drain). Called for + # symmetry with CudaDeviceMixin's torch.cuda.synchronize(). + torch.cpu.synchronize() + + def get_available_memory(self, device_id: int = 0) -> tuple[int, int]: + vm = psutil.virtual_memory() + return (vm.available, vm.total) + + def get_torch_distributed_backend_str(self) -> str: + return "gloo" + + +class CpuSRTPlatform(CpuDeviceMixin, SRTPlatform): + """Default in-tree CPU SRT platform. + + supports_fp8 / support_cuda_graph / support_piecewise_cuda_graph keep the + conservative SRTPlatform defaults (all False), so they are not repeated + here. Only is_pin_memory_available is overridden: the base defaults to + True, but CPU has no GPU to pin host memory to. + """ + + def is_pin_memory_available(self) -> bool: + return False diff --git a/test/registered/unit/platforms/test_platform_interface.py b/test/registered/unit/platforms/test_platform_interface.py index 8b8ef1454..2dcc2efbf 100644 --- a/test/registered/unit/platforms/test_platform_interface.py +++ b/test/registered/unit/platforms/test_platform_interface.py @@ -10,6 +10,7 @@ from unittest.mock import MagicMock, patch import torch from sglang.srt.platforms import _load_platform_class, _resolve_platform +from sglang.srt.platforms.cpu import CpuDeviceMixin, CpuSRTPlatform from sglang.srt.platforms.cuda import CudaDeviceMixin, CudaSRTPlatform from sglang.srt.platforms.device_mixin import ( CpuArchEnum, @@ -331,6 +332,123 @@ class TestCudaDeviceMixin(CustomTestCase): self.assertTrue(base.support_piecewise_cuda_graph()) +class TestCpuDeviceMixin(CustomTestCase): + """Tests for CPU device operation defaults (covers both x86 and ARM).""" + + def test_cpu_platform_identity(self): + base = CpuSRTPlatform() + self.assertTrue(base.is_cpu()) + self.assertFalse(base.is_cuda()) + self.assertFalse(base.is_cuda_alike()) + self.assertIsInstance(base, CpuDeviceMixin) + + def test_default_get_device_returns_cpu_device(self): + base = CpuSRTPlatform() + # ``local_rank`` is ignored — CPU has no per-rank device. + self.assertEqual(base.get_device(0), torch.device("cpu")) + self.assertEqual(base.get_device(7), torch.device("cpu")) + + @patch("sglang.srt.platforms.cpu.psutil.virtual_memory") + def test_default_get_device_total_memory_uses_psutil(self, mock_vm): + mock_vm.return_value.total = 12345 + base = CpuSRTPlatform() + self.assertEqual(base.get_device_total_memory(), 12345) + + @patch("sglang.srt.platforms.cpu.psutil.virtual_memory") + def test_default_get_available_memory_uses_psutil(self, mock_vm): + mock_vm.return_value.available = 100 + mock_vm.return_value.total = 200 + base = CpuSRTPlatform() + self.assertEqual(base.get_available_memory(), (100, 200)) + + @patch("sglang.srt.platforms.cpu.psutil.virtual_memory") + def test_default_get_current_memory_usage_is_system_used(self, mock_vm): + mock_vm.return_value.total = 1000 + mock_vm.return_value.available = 300 + base = CpuSRTPlatform() + # system-used == total - available (not per-process RSS) + self.assertEqual(base.get_current_memory_usage(), 700.0) + + @patch("sglang.srt.platforms.cpu.psutil.virtual_memory") + def test_memory_free_contract_yields_available(self, mock_vm): + # The [Active] contract free = total - used must yield psutil.available. + mock_vm.return_value.total = 1000 + mock_vm.return_value.available = 300 + base = CpuSRTPlatform() + free = base.get_device_total_memory() - base.get_current_memory_usage() + self.assertEqual(free, 300) + + @patch("torch.cpu.set_device") + def test_default_set_device_uses_torch_cpu(self, mock_set_device): + base = CpuSRTPlatform() + device = torch.device("cpu") + base.set_device(device) + # Documented CPU no-op, but called for symmetry with CudaDeviceMixin. + mock_set_device.assert_called_once_with(device) + + def test_default_set_device_does_not_flip_default(self): + base = CpuSRTPlatform() + # Must not call torch.set_default_device — process-wide default stays put. + before = torch.empty(0).device + base.set_device(torch.device("cpu")) + after = torch.empty(0).device + self.assertEqual(before, after) + + @patch("sglang.srt.platforms.cpu.gc.collect") + def test_default_empty_cache_calls_gc_collect(self, mock_collect): + base = CpuSRTPlatform() + base.empty_cache() + mock_collect.assert_called_once_with() + + @patch("torch.cpu.synchronize") + def test_default_synchronize_uses_torch_cpu(self, mock_synchronize): + base = CpuSRTPlatform() + base.synchronize() + mock_synchronize.assert_called_once_with() + + def test_default_distributed_backend_is_gloo(self): + base = CpuSRTPlatform() + self.assertEqual(base.get_torch_distributed_backend_str(), "gloo") + + @patch("platform.machine", return_value="aarch64") + def test_cpu_arch_property_resolves_and_caches(self, mock_machine): + base = CpuSRTPlatform() + self.assertEqual(base.cpu_arch, CpuArchEnum.ARM) + # cached_property: second access must not re-query platform.machine + call_count = mock_machine.call_count + self.assertEqual(base.cpu_arch, CpuArchEnum.ARM) + self.assertEqual(mock_machine.call_count, call_count) + + @patch("platform.machine", return_value="aarch64") + def test_get_device_name_arm_branch(self, _mock_machine): + base = CpuSRTPlatform() + name = base.get_device_name() + self.assertIn("aarch64", name) + + @patch("platform.machine", return_value="x86_64") + def test_get_device_name_x86_branch(self, _mock_machine): + base = CpuSRTPlatform() + name = base.get_device_name() + self.assertIn("x86_64", name) + + @patch("platform.machine", return_value="aarch64") + def test_get_device_uuid_returns_machine(self, _mock_machine): + base = CpuSRTPlatform() + self.assertEqual(base.get_device_uuid(), "aarch64") + + def test_get_device_capability_returns_none(self): + base = CpuSRTPlatform() + self.assertIsNone(base.get_device_capability()) + + def test_cpu_srt_platform_capabilities(self): + base = CpuSRTPlatform() + self.assertFalse(base.supports_fp8()) + self.assertFalse(base.support_cuda_graph()) + self.assertFalse(base.support_piecewise_cuda_graph()) + # Override of the SRTPlatform default (True) — no GPU to pin to. + self.assertFalse(base.is_pin_memory_available()) + + class TestSRTPlatformOverrides(CustomTestCase): """Tests for SRTPlatform method overrides via plugins.""" @@ -496,6 +614,51 @@ class TestResolvePlatformAutoDiscover(CustomTestCase): self.assertIsInstance(result, SRTPlatform) self.assertNotIsInstance(result, CudaSRTPlatform) + @patch("sglang.srt.platforms.load_plugins_by_group") + @patch("sglang.srt.platforms._is_cuda_available") + @patch("sglang.srt.platforms._is_cpu_available") + @patch("sglang.srt.platforms.envs") + def test_no_plugin_cpu_engine_enabled_activates_cpu_fallback( + self, mock_envs, mock_is_cpu, mock_is_cuda, mock_load + ): + """SGLANG_USE_CPU_ENGINE=1 + no plugins → CpuSRTPlatform.""" + mock_envs.SGLANG_PLATFORM.get.return_value = "" + mock_is_cpu.return_value = True + mock_is_cuda.return_value = False + mock_load.return_value = {} + result = _resolve_platform() + self.assertIsInstance(result, CpuSRTPlatform) + + @patch("sglang.srt.platforms.load_plugins_by_group") + @patch("sglang.srt.platforms._is_cuda_available") + @patch("sglang.srt.platforms._is_cpu_available") + @patch("sglang.srt.platforms.envs") + def test_cpu_engine_wins_over_cuda( + self, mock_envs, mock_is_cpu, mock_is_cuda, mock_load + ): + """When both CPU engine and CUDA are available, explicit opt-in wins.""" + mock_envs.SGLANG_PLATFORM.get.return_value = "" + mock_is_cpu.return_value = True + mock_is_cuda.return_value = True + mock_load.return_value = {} + result = _resolve_platform() + self.assertIsInstance(result, CpuSRTPlatform) + + @patch("sglang.srt.platforms.load_plugins_by_group") + @patch("sglang.srt.platforms._is_cuda_available") + @patch("sglang.srt.platforms._is_cpu_available") + @patch("sglang.srt.platforms.envs") + def test_no_plugin_cpu_engine_disabled_prefers_cuda( + self, mock_envs, mock_is_cpu, mock_is_cuda, mock_load + ): + """Regression: CPU opt-out leaves the existing CUDA fallback path intact.""" + mock_envs.SGLANG_PLATFORM.get.return_value = "" + mock_is_cpu.return_value = False + mock_is_cuda.return_value = True + mock_load.return_value = {} + result = _resolve_platform() + self.assertIsInstance(result, CudaSRTPlatform) + @patch("sglang.srt.platforms.load_plugins_by_group") @patch("sglang.srt.platforms.envs") def test_multiple_plugins_activate_raises(self, mock_envs, mock_load):