Introduce CpuDeviceMixin and CpuSRTPlatform (#26385)

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