""" Shared device abstraction for SGLang platforms. DeviceMixin provides the common device identity queries and operations shared between the SRT (LLM inference) and Multimodal (diffusion) platform hierarchies. Concrete per-device mixins (e.g. MyDeviceMixin) implement the abstract operations; subsystem-specific platforms (SRTPlatform, MMPlatform) inherit DeviceMixin and add their own methods. Hierarchy example (OOT plugin):: DeviceMixin ├── MyDeviceMixin(DeviceMixin) # vendor-specific device operations ├── SRTPlatform(DeviceMixin) # + graph runner, KV pool, … │ └── MySRTPlatform(SRTPlatform, MyDeviceMixin) └── MMPlatform(DeviceMixin) # + attention backend, VAE, … └── MyMMPlatform(MMPlatform, MyDeviceMixin) Method status annotations: - ``[Active]`` — SGLang core calls this method through ``current_platform``. OOT implementations take effect immediately. - ``[Planned]`` — Reserved interface. SGLang core still uses hardcoded calls (e.g. ``torch.cuda.empty_cache()``). OOT implementations will NOT take effect until the core is migrated in a future PR. """ import enum import random from typing import NamedTuple, Optional import numpy as np import torch from sglang.srt.environ import envs class PlatformEnum(enum.Enum): """Enumeration of known platform types. Superset of both SRT and MM enums so that a single PlatformEnum can be shared across subsystems. """ CUDA = enum.auto() ROCM = enum.auto() CPU = enum.auto() XPU = enum.auto() MUSA = enum.auto() NPU = enum.auto() TPU = enum.auto() MPS = enum.auto() OOT = enum.auto() # Out-of-tree (external plugin) UNSPECIFIED = enum.auto() class CpuArchEnum(enum.Enum): """CPU architecture enumeration.""" X86 = enum.auto() ARM = enum.auto() UNSPECIFIED = enum.auto() class DeviceCapability(NamedTuple): """Device compute capability (major, minor). Uses NamedTuple for built-in comparison support: ``DeviceCapability(9, 0) >= DeviceCapability(8, 9)`` works naturally. """ major: int minor: int def as_version_str(self) -> str: return f"{self.major}.{self.minor}" def to_int(self) -> int: """Express capability as ```` (minor is single digit).""" assert 0 <= self.minor < 10 return self.major * 10 + self.minor _DEVICE_TO_DISTRIBUTED_BACKEND: dict[str, str] = { "cuda": "nccl", "xpu": "xccl", "hpu": "hccl", "cpu": "gloo", "npu": "hccl" if not envs.SGLANG_ZBAL_LOCAL_MEM_SIZE.get() > 0 else "zbal", "musa": "mccl", } class DeviceMixin: """Mixin providing device identity queries and basic device operations. Class-level attributes (override in subclasses): _enum: PlatformEnum identifying this platform. device_name: Human-readable short name (e.g. "cuda", "npu"). device_type: ``torch.device`` type string (e.g. "cuda", "npu"). """ _enum: PlatformEnum = PlatformEnum.UNSPECIFIED device_name: str = "unknown" device_type: str = "cpu" # ------------------------------------------------------------------ # Platform identity queries # ------------------------------------------------------------------ def is_cuda(self) -> bool: return self._enum == PlatformEnum.CUDA def is_rocm(self) -> bool: return self._enum == PlatformEnum.ROCM def is_cpu(self) -> bool: return self._enum == PlatformEnum.CPU def is_xpu(self) -> bool: return self._enum == PlatformEnum.XPU def is_musa(self) -> bool: return self._enum == PlatformEnum.MUSA def is_npu(self) -> bool: return self._enum == PlatformEnum.NPU def is_tpu(self) -> bool: return self._enum == PlatformEnum.TPU def is_mps(self) -> bool: return self._enum == PlatformEnum.MPS def is_cuda_alike(self) -> bool: """True for CUDA, ROCm, or MUSA (all expose CUDA-like APIs).""" return self._enum in ( PlatformEnum.CUDA, PlatformEnum.ROCM, PlatformEnum.MUSA, ) def is_out_of_tree(self) -> bool: """True for externally-registered OOT platforms.""" return self._enum == PlatformEnum.OOT # ------------------------------------------------------------------ # Active methods — core calls these through current_platform. # OOT implementations take effect immediately. # ------------------------------------------------------------------ def get_device_total_memory(self, device_id: int = 0) -> int: """[Active] Get total device memory in bytes.""" raise NotImplementedError def get_current_memory_usage( self, device: Optional["torch.device"] = None ) -> float: """[Active] Get current peak memory usage in bytes.""" raise NotImplementedError # ------------------------------------------------------------------ # Planned methods — reserved interface. Core still uses hardcoded # calls (e.g. torch.cuda.*). OOT implementations will NOT take # effect until the core is migrated in a future PR. # ------------------------------------------------------------------ # ---- Device management ---- def get_device(self, device_id: int = 0) -> str: """[Planned] Return ``torch.device`` for the given device id.""" raise NotImplementedError def set_device(self, device: "torch.device") -> None: """[Planned] Set the current device.""" raise NotImplementedError def get_device_name(self, device_id: int = 0) -> str: """[Planned] Get human-readable device name.""" raise NotImplementedError def get_device_uuid(self, device_id: int = 0) -> str: """[Planned] Get unique device identifier string.""" raise NotImplementedError def get_device_capability(self, device_id: int = 0) -> Optional["DeviceCapability"]: """[Planned] Get device compute capability. None if N/A.""" raise NotImplementedError def empty_cache(self) -> None: """[Planned] Release cached device memory. No-op for CPU-like platforms.""" pass def synchronize(self) -> None: """[Planned] Synchronize device operations. No-op for CPU-like platforms.""" pass # ---- Memory ---- def get_available_memory(self, device_id: int = 0) -> tuple[int, int]: """[Planned] Return ``(free_bytes, total_bytes)``.""" raise NotImplementedError # ---- Distributed ---- def get_torch_distributed_backend_str(self) -> str: """Return the torch.distributed backend string (e.g. "nccl", "hccl"). Default: lookup ``self.device_type`` in ``_DEVICE_TO_DISTRIBUTED_BACKEND``, falling back to ``"gloo"``. Subclasses override only when they need a non-default backend (e.g. mooncake, or a brand-new device). """ return _DEVICE_TO_DISTRIBUTED_BACKEND.get(self.device_type, "gloo") def get_communicator_class(self) -> type | None: """[Planned] Return platform-specific communicator class, or None for default.""" return None # ---- Misc ---- @classmethod def inference_mode(cls): """[Planned] Return inference mode context manager.""" return torch.inference_mode(mode=True) @classmethod def seed_everything(cls, seed: int | None = None) -> None: """[Planned] Set random seeds for reproducibility across all libraries.""" if seed is not None: random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) def verify_quantization(self, quant: str) -> None: """[Planned] Validate that a quantization method is supported. No-op by default.""" pass @classmethod def get_cpu_architecture(cls) -> "CpuArchEnum": """[Planned] Detect CPU architecture.""" import platform as _platform machine = _platform.machine().lower() if machine in ("x86_64", "amd64", "i386", "i686"): return CpuArchEnum.X86 elif machine in ("arm64", "aarch64"): return CpuArchEnum.ARM return CpuArchEnum.UNSPECIFIED # ------------------------------------------------------------------ # Dunder helpers # ------------------------------------------------------------------ def __repr__(self) -> str: return f"{self.__class__.__name__}(device={self.device_name})"