feat: add hardware platform and plugin support

lightllm/platform/__init__.py

@@ -0,0 +1,34 @@
+from __future__ import annotations
+
+from typing import Optional
+
+import lightllm.platform.backends  # noqa: F401
+from lightllm.platform.base.backend import HardwareBackend
+from lightllm.platform.base.registry import get_platform_spec
+from lightllm.platform.plugin import configure_plugins
+from lightllm.utils.envs_utils import get_env_start_args
+
+_backend: Optional[HardwareBackend] = None
+
+
+def get_hardware_backend() -> HardwareBackend:
+    global _backend
+
+    if _backend is not None:
+        return _backend
+
+    configure_plugins()
+
+    platform_name = get_env_start_args().hardware_platform
+    spec = get_platform_spec(platform_name)
+
+    backend_cls = spec.backend_cls
+    _backend = backend_cls()
+
+    if not _backend.runtime.is_available():
+        raise RuntimeError(f"Hardware backend {backend_cls.__name__} is not available.")
+
+    return _backend
+
+
+__all__ = ["get_hardware_backend"]

lightllm/platform/backends/__init__.py

@@ -0,0 +1 @@
+from lightllm.platform.backends import cuda_like  # noqa: F401

lightllm/platform/backends/cuda_like/__init__.py

@@ -0,0 +1,14 @@
+from lightllm.platform.base.backend import HardwareBackend
+from lightllm.platform.base.registry import register_platform
+from lightllm.platform.backends.cuda_like.runtime import CudaLikeRuntime
+from lightllm.platform.backends.cuda_like.graph import CudaLikeGraph
+
+
+class CudaLikeBackend(HardwareBackend):
+    def __init__(self) -> None:
+        super().__init__(CudaLikeRuntime(), CudaLikeGraph())
+
+
+@register_platform("cuda")
+class CudaBackend(CudaLikeBackend):
+    pass

lightllm/platform/backends/cuda_like/graph.py

@@ -0,0 +1,23 @@
+from typing import Any, ContextManager, Optional
+
+import torch
+from lightllm.platform.base.graph import HardwareBackendGraph
+
+
+class CudaLikeGraph(HardwareBackendGraph):
+    def create_graph(self) -> Any:
+        return torch.cuda.CUDAGraph()
+
+    def graph(
+        self,
+        graph_obj: Any,
+        pool: Optional[Any] = None,
+        stream: Optional[Any] = None,
+    ) -> ContextManager:
+        return torch.cuda.graph(graph_obj, pool=pool, stream=stream)
+
+    def graph_pool_handle(self) -> Any:
+        return torch.cuda.graph_pool_handle()
+
+    def is_capturing(self) -> bool:
+        return torch.cuda.is_current_stream_capturing()

lightllm/platform/backends/cuda_like/runtime.py

@@ -0,0 +1,63 @@
+from typing import Any, ContextManager, Optional, Tuple
+
+import torch
+from lightllm.platform.base.runtime import DeviceLike, HardwareBackendRuntime
+
+
+class CudaLikeRuntime(HardwareBackendRuntime):
+    @property
+    def device_type(self) -> str:
+        return "cuda"
+
+    @property
+    def dist_backend(self) -> str:
+        return "nccl"
+
+    def mem_get_info(self, device: DeviceLike) -> Tuple[int, int]:
+        return torch.cuda.mem_get_info(self._parse(device))
+
+    def get_device_properties(self, device: DeviceLike) -> Any:
+        return torch.cuda.get_device_properties(self._parse(device))
+
+    def device_count(self) -> int:
+        return torch.cuda.device_count()
+
+    def is_available(self) -> bool:
+        return torch.cuda.is_available()
+
+    def current_device(self) -> int:
+        return torch.cuda.current_device()
+
+    def get_device_name(self, device_id: Optional[int] = None) -> str:
+        if device_id is None:
+            device_id = self.current_device()
+        return torch.cuda.get_device_name(device_id)
+
+    def set_device(self, device: DeviceLike) -> None:
+        torch.cuda.set_device(self._parse(device))
+
+    def create_stream(self, **kwargs) -> Any:
+        return torch.cuda.Stream(**kwargs)
+
+    def stream(self, stream: Any) -> ContextManager[Any]:
+        return torch.cuda.stream(stream)
+
+    def current_stream(self, device_id: Optional[int] = None) -> Any:
+        if device_id is None:
+            device_id = self.current_device()
+        return torch.cuda.current_stream(device_id)
+
+    def create_event(self, **kwargs) -> torch.Event:
+        return torch.cuda.Event(**kwargs)
+
+    def synchronize(self, device: Optional[DeviceLike] = None) -> None:
+        if device is None:
+            torch.cuda.synchronize()
+            return
+        torch.cuda.synchronize(self._parse(device))
+
+    def empty_cache(self) -> None:
+        torch.cuda.empty_cache()
+
+    def manual_seed_all(self, seed: int) -> None:
+        torch.cuda.manual_seed_all(seed)

lightllm/platform/base/backend.py

@@ -0,0 +1,24 @@
+from abc import ABC
+
+from lightllm.platform.base.graph import HardwareBackendGraph
+from lightllm.platform.base.runtime import HardwareBackendRuntime
+
+
+class HardwareBackend(ABC):
+    platform_name: str
+
+    def __init__(self, runtime: HardwareBackendRuntime, graph: HardwareBackendGraph) -> None:
+        self._runtime = runtime
+        self._graph = graph
+
+    @property
+    def name(self) -> str:
+        return self.platform_name
+
+    @property
+    def runtime(self) -> HardwareBackendRuntime:
+        return self._runtime
+
+    @property
+    def graph(self) -> HardwareBackendGraph:
+        return self._graph

lightllm/platform/base/graph.py

@@ -0,0 +1,28 @@
+from abc import ABC, abstractmethod
+from typing import Any, ContextManager, Optional
+
+
+class HardwareBackendGraph(ABC):
+    @abstractmethod
+    def create_graph(self) -> Any:
+        pass
+
+    @abstractmethod
+    def graph(
+        self,
+        graph_obj: Any,
+        pool: Optional[Any] = None,
+        stream: Optional[Any] = None,
+    ) -> ContextManager:
+        pass
+
+    def replay_graph(self, graph_obj: Any) -> Any:
+        graph_obj.replay()
+
+    @abstractmethod
+    def graph_pool_handle(self) -> Any:
+        pass
+
+    @abstractmethod
+    def is_capturing(self) -> bool:
+        pass

lightllm/platform/base/registry.py

@@ -0,0 +1,37 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable, Type
+
+from lightllm.platform.base.backend import HardwareBackend
+
+
+@dataclass(frozen=True)
+class PlatformSpec:
+    name: str
+    backend_cls: Type[HardwareBackend]
+
+
+PLATFORMS: dict[str, PlatformSpec] = {}
+
+
+def register_platform(name: str) -> Callable[[Type[HardwareBackend]], Type[HardwareBackend]]:
+    def decorator(backend_cls: Type[HardwareBackend]) -> Type[HardwareBackend]:
+        if name in PLATFORMS:
+            raise ValueError(f"Platform {name!r} is already registered.")
+
+        backend_cls.platform_name = name
+        PLATFORMS[name] = PlatformSpec(
+            name=name,
+            backend_cls=backend_cls,
+        )
+        return backend_cls
+
+    return decorator
+
+
+def get_platform_spec(name: str) -> PlatformSpec:
+    spec = PLATFORMS.get(name)
+    if spec is None:
+        raise RuntimeError(f"Platform {name!r} is not registered, registered: {sorted(PLATFORMS.keys())}.")
+    return spec

lightllm/platform/base/runtime.py

@@ -0,0 +1,122 @@
+from abc import ABC, abstractmethod
+from typing import Any, ContextManager, Optional, Tuple, Union
+
+import torch
+import torch.distributed as dist
+
+DeviceLike = Union[int, str, torch.device]
+
+
+class HardwareBackendRuntime(ABC):
+    @property
+    @abstractmethod
+    def device_type(self) -> str:
+        pass
+
+    @property
+    @abstractmethod
+    def dist_backend(self) -> str:
+        pass
+
+    def dist_init_kwargs(self, target_device: torch.device) -> dict[str, Any]:
+        return {}
+
+    def init_process_group(
+        self,
+        *,
+        host: str,
+        port: int,
+        rank: int,
+        world_size: int,
+        device_id: int,
+    ) -> None:
+        target_device = self.target_device(device_id)
+        self.set_device(target_device)
+
+        kwargs: dict[str, Any] = {
+            "backend": self.dist_backend,
+            "init_method": f"tcp://{host}:{port}",
+            "rank": rank,
+            "world_size": world_size,
+        }
+        kwargs.update(self.dist_init_kwargs(target_device))
+        dist.init_process_group(**kwargs)
+
+    @abstractmethod
+    def mem_get_info(self, device: DeviceLike) -> Tuple[int, int]:
+        pass
+
+    @abstractmethod
+    def get_device_properties(self, device: DeviceLike) -> Any:
+        pass
+
+    def target_device(self, device_id: Optional[int] = None) -> torch.device:
+        if device_id is None:
+            device_id = self.current_device()
+        return torch.device(self.device_type, device_id)
+
+    @abstractmethod
+    def device_count(self) -> int:
+        pass
+
+    @abstractmethod
+    def is_available(self) -> bool:
+        pass
+
+    @abstractmethod
+    def current_device(self) -> int:
+        pass
+
+    @abstractmethod
+    def get_device_name(self, device_id: Optional[int] = None) -> str:
+        pass
+
+    def _parse(self, device: DeviceLike) -> torch.device:
+        if isinstance(device, torch.device):
+            _device = device
+        elif isinstance(device, int):
+            _device = torch.device(self.device_type, device)
+        elif isinstance(device, str):
+            _device = torch.device(device)
+        else:
+            raise TypeError(f"Invalid device: {device}")
+
+        if _device.type != self.device_type:
+            raise ValueError(f"Expected device type {self.device_type!r}, got {_device.type!r} ({_device})")
+
+        if _device.index is None:
+            _device = torch.device(self.device_type, self.current_device())
+
+        return _device
+
+    @abstractmethod
+    def set_device(self, device: DeviceLike) -> None:
+        pass
+
+    @abstractmethod
+    def create_stream(self, **kwargs) -> Any:
+        pass
+
+    @abstractmethod
+    def stream(self, stream: Any) -> ContextManager[Any]:
+        pass
+
+    @abstractmethod
+    def current_stream(self, device_id: Optional[int] = None) -> Any:
+        pass
+
+    @abstractmethod
+    def create_event(self, **kwargs) -> torch.Event:
+        pass
+
+    @abstractmethod
+    def synchronize(self, device: Optional[DeviceLike] = None) -> None:
+        pass
+
+    @abstractmethod
+    def empty_cache(self) -> None:
+        pass
+
+    @abstractmethod
+    def manual_seed_all(self, seed: int) -> None:
+        pass

lightllm/platform/plugin/__init__.py

@@ -0,0 +1,12 @@
+from lightllm.platform.plugin.common import Plugin
+
+
+OPS = Plugin(name="ops", entry_point_group="lightllm.ops_plugin")
+ATT = Plugin(name="att", entry_point_group="lightllm.att_plugin")
+
+_PLUGINS = (OPS, ATT)
+
+
+def configure_plugins() -> None:
+    for plugin in _PLUGINS:
+        plugin.load()