sgl-project--sglang/python/sglang/srt/platforms/device_mixin.py

"""
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 ``<major><minor>`` (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})"