feat: add hardware platform and plugin support

lightllm/platform/__init__.py

@@ -0,0 +1,36 @@
+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
+
+    platform_name = get_env_start_args().hardware_platform
+    spec = get_platform_spec(platform_name)
+
+    backend_cls = spec.load_backend_cls()
+
+    configure_plugins()
+
+    backend = backend_cls()
+
+    if not backend.runtime.is_available():
+        raise RuntimeError(f"Hardware backend {backend_cls.__name__} is not available.")
+
+    _backend = backend
+    return _backend
+
+
+__all__ = ["get_hardware_backend"]

lightllm/platform/backends/__init__.py

@@ -0,0 +1,16 @@
+from lightllm.platform.base.registry import register_platform
+
+register_platform(
+    "cuda",
+    backend="lightllm.platform.backends.cuda_like:CudaBackend",
+)
+
+register_platform(
+    "musa",
+    backend="lightllm.platform.backends.cuda_like:MusaBackend",
+)
+
+register_platform(
+    "ascend",
+    backend="lightllm.platform.backends.ascend:AscendBackend",
+)

lightllm/platform/backends/ascend/__init__.py

@@ -0,0 +1,16 @@
+import lightllm.platform.backends.ascend.ops  # noqa: F401
+from lightllm.platform.base.backend import HardwareBackend
+from lightllm.platform.backends.ascend.graph import AscendGraph
+from lightllm.platform.backends.ascend.runtime import AscendRuntime
+from lightllm.platform.plugin.ops import build_platform_ops
+
+
+class AscendBackend(HardwareBackend):
+    platform_name = "ascend"
+
+    def __init__(self) -> None:
+        super().__init__(
+            runtime=AscendRuntime(),
+            graph=AscendGraph(),
+            ops=build_platform_ops(self.platform_name),
+        )

lightllm/platform/backends/ascend/graph.py

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

lightllm/platform/backends/ascend/ops/__init__.py

@@ -0,0 +1,3 @@
+from lightllm.platform.backends.ascend.ops import act, embedding, norm, rotary, sampling
+
+__all__ = ["act", "embedding", "norm", "rotary", "sampling"]

lightllm/platform/backends/ascend/runtime.py

@@ -0,0 +1,69 @@
+import importlib
+from typing import Any, ContextManager, Optional, Tuple
+
+import torch
+
+from lightllm.platform.base.runtime import DeviceLike, HardwareBackendRuntime
+
+
+class AscendRuntime(HardwareBackendRuntime):
+    @property
+    def device_type(self) -> str:
+        return "npu"
+
+    @property
+    def dist_backend(self) -> str:
+        return "hccl"
+
+    def mem_get_info(self, device: DeviceLike) -> Tuple[int, int]:
+        return torch.npu.mem_get_info(self._parse(device))
+
+    def get_device_properties(self, device: DeviceLike) -> Any:
+        return torch.npu.get_device_properties(self._parse(device))
+
+    def device_count(self) -> int:
+        return torch.npu.device_count()
+
+    def is_available(self) -> bool:
+        try:
+            importlib.import_module("torch_npu")
+        except ImportError:
+            return False
+        return torch.npu.is_available()
+
+    def current_device(self) -> int:
+        return torch.npu.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.npu.get_device_name(device_id)
+
+    def set_device(self, device: DeviceLike) -> None:
+        torch.npu.set_device(self._parse(device))
+
+    def create_stream(self, **kwargs) -> Any:
+        return torch.npu.Stream(**kwargs)
+
+    def stream(self, stream: Any) -> ContextManager[Any]:
+        return torch.npu.stream(stream)
+
+    def current_stream(self, device_id: Optional[int] = None) -> Any:
+        if device_id is None:
+            device_id = self.current_device()
+        return torch.npu.current_stream(device_id)
+
+    def create_event(self, **kwargs) -> Any:
+        return torch.npu.Event(**kwargs)
+
+    def synchronize(self, device: Optional[DeviceLike] = None) -> None:
+        if device is None:
+            torch.npu.synchronize()
+            return
+        torch.npu.synchronize(self._parse(device))
+
+    def empty_cache(self) -> None:
+        torch.npu.empty_cache()
+
+    def manual_seed_all(self, seed: int) -> None:
+        torch.npu.manual_seed_all(seed)

lightllm/platform/backends/cuda_like/__init__.py

@@ -0,0 +1,22 @@
+import lightllm.platform.backends.cuda_like.ops  # noqa: F401
+from lightllm.platform.base.backend import HardwareBackend
+from lightllm.platform.backends.cuda_like.graph import CudaLikeGraph
+from lightllm.platform.backends.cuda_like.runtime import CudaLikeRuntime
+from lightllm.platform.plugin.ops import build_platform_ops
+
+
+class CudaLikeBackend(HardwareBackend):
+    def __init__(self) -> None:
+        super().__init__(
+            runtime=CudaLikeRuntime(),
+            graph=CudaLikeGraph(),
+            ops=build_platform_ops(self.platform_name),
+        )
+
+
+class CudaBackend(CudaLikeBackend):
+    platform_name = "cuda"
+
+
+class MusaBackend(CudaLikeBackend):
+    platform_name = "musa"

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/ops/__init__.py

@@ -0,0 +1,3 @@
+from lightllm.platform.backends.cuda_like.ops import act, embedding, norm, rotary, sampling
+
+__all__ = ["act", "embedding", "norm", "rotary", "sampling"]

lightllm/platform/backends/cuda_like/ops/act.py

@@ -0,0 +1,11 @@
+from lightllm.models.gemma_2b.triton_kernel.gelu_and_mul import gelu_and_mul_fwd
+from lightllm.models.llama.triton_kernel.silu_and_mul import silu_and_mul_fwd
+from lightllm.platform.ops.act import ActOps
+from lightllm.platform.plugin.ops import register_ops
+
+register_ops(
+    ActOps,
+    platforms=("cuda", "musa"),
+    silu_and_mul=silu_and_mul_fwd,
+    gelu_and_mul=gelu_and_mul_fwd,
+)

lightllm/platform/backends/cuda_like/ops/embedding.py

@@ -0,0 +1,13 @@
+from lightllm.common.basemodel.triton_kernel.embedding import embedding
+from lightllm.common.basemodel.triton_kernel.multimodal_emb import multimodal_emb
+from lightllm.platform.ops.embedding import EmbeddingOps
+from lightllm.platform.plugin.ops import register_ops
+from lightllm.server.embed_cache.copy_to_cache import offload_embed_tensor_to_cache
+
+register_ops(
+    EmbeddingOps,
+    platforms=("cuda", "musa"),
+    embedding=embedding,
+    multimodal_emb=multimodal_emb,
+    offload_embed_tensor_to_cache=offload_embed_tensor_to_cache,
+)

lightllm/platform/backends/cuda_like/ops/norm.py

@@ -0,0 +1,16 @@
+from lightllm.common.basemodel.triton_kernel.norm.gated_rmsnorm import gated_rmsnorm_forward
+from lightllm.common.basemodel.triton_kernel.norm.layernorm import layernorm_forward
+from lightllm.common.basemodel.triton_kernel.norm.qk_norm import qk_rmsnorm_forward, qk_rmsnorm_fused_forward
+from lightllm.common.basemodel.triton_kernel.norm.rmsnorm import rmsnorm_forward
+from lightllm.platform.ops.norm import NormOps
+from lightllm.platform.plugin.ops import register_ops
+
+register_ops(
+    NormOps,
+    platforms=("cuda", "musa"),
+    rms_norm=rmsnorm_forward,
+    gated_rms_norm=gated_rmsnorm_forward,
+    layer_norm=layernorm_forward,
+    qk_rms_norm=qk_rmsnorm_forward,
+    fused_qk_rms_norm=qk_rmsnorm_fused_forward,
+)

lightllm/platform/backends/cuda_like/ops/rotary.py

@@ -0,0 +1,9 @@
+from lightllm.models.llama.triton_kernel.rotary_emb import rotary_emb_fwd
+from lightllm.platform.ops.rotary import RotaryOps
+from lightllm.platform.plugin.ops import register_ops
+
+register_ops(
+    RotaryOps,
+    platforms=("cuda", "musa"),
+    rotary_emb=rotary_emb_fwd,
+)

lightllm/platform/backends/cuda_like/ops/sampling.py

@@ -0,0 +1,21 @@
+from lightllm.common.basemodel.triton_kernel.gen_sampling_params import (
+    token_id_counter,
+    update_req_to_token_id_counter,
+)
+from lightllm.common.basemodel.triton_kernel.post_process.apply_invalid_token import apply_invalid_token_ids
+from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty import apply_penalty
+from lightllm.common.basemodel.triton_kernel.post_process.apply_penalty_gpu_cache import apply_penalty_gpu_cache
+from lightllm.platform.ops.sampling import SamplingOps
+from lightllm.platform.plugin.ops import register_ops
+from lightllm.server.router.model_infer.mode_backend.generic_post_process import _top_p_top_k_sample
+
+register_ops(
+    SamplingOps,
+    platforms=("cuda", "musa"),
+    apply_penalty=apply_penalty,
+    apply_penalty_gpu_cache=apply_penalty_gpu_cache,
+    apply_invalid_token_ids=apply_invalid_token_ids,
+    token_id_counter=token_id_counter,
+    update_req_to_token_id_counter=update_req_to_token_id_counter,
+    top_p_top_k_sample=_top_p_top_k_sample,
+)

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,33 @@
+from lightllm.platform.base.graph import HardwareBackendGraph
+from lightllm.platform.base.runtime import HardwareBackendRuntime
+from lightllm.platform.ops import PlatformOps
+
+
+class HardwareBackend:
+    platform_name: str
+
+    def __init__(
+        self,
+        runtime: HardwareBackendRuntime,
+        graph: HardwareBackendGraph,
+        ops: PlatformOps,
+    ) -> None:
+        self._runtime = runtime
+        self._graph = graph
+        self._ops = ops
+
+    @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
+
+    @property
+    def ops(self) -> PlatformOps:
+        return self._ops