Introduce CpuDeviceMixin and CpuSRTPlatform (#26385)

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)
 ```

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."

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

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):