diff --git a/lightllm/platform/__init__.py b/lightllm/platform/__init__.py new file mode 100644 index 0000000000..b2a324cbfc --- /dev/null +++ b/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"] diff --git a/lightllm/platform/backends/__init__.py b/lightllm/platform/backends/__init__.py new file mode 100644 index 0000000000..968d3179ec --- /dev/null +++ b/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", +) diff --git a/lightllm/platform/backends/ascend/__init__.py b/lightllm/platform/backends/ascend/__init__.py new file mode 100644 index 0000000000..2a2a9099e3 --- /dev/null +++ b/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), + ) diff --git a/lightllm/platform/backends/ascend/graph.py b/lightllm/platform/backends/ascend/graph.py new file mode 100644 index 0000000000..9abf8c60fa --- /dev/null +++ b/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() diff --git a/lightllm/platform/backends/ascend/ops/__init__.py b/lightllm/platform/backends/ascend/ops/__init__.py new file mode 100644 index 0000000000..cc27207009 --- /dev/null +++ b/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"] diff --git a/lightllm/platform/backends/ascend/ops/act.py b/lightllm/platform/backends/ascend/ops/act.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/platform/backends/ascend/ops/embedding.py b/lightllm/platform/backends/ascend/ops/embedding.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/platform/backends/ascend/ops/norm.py b/lightllm/platform/backends/ascend/ops/norm.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/platform/backends/ascend/ops/rotary.py b/lightllm/platform/backends/ascend/ops/rotary.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/platform/backends/ascend/ops/sampling.py b/lightllm/platform/backends/ascend/ops/sampling.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/platform/backends/ascend/runtime.py b/lightllm/platform/backends/ascend/runtime.py new file mode 100644 index 0000000000..7bc5b96894 --- /dev/null +++ b/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) diff --git a/lightllm/platform/backends/cuda_like/__init__.py b/lightllm/platform/backends/cuda_like/__init__.py new file mode 100644 index 0000000000..d715da0dfa --- /dev/null +++ b/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" diff --git a/lightllm/platform/backends/cuda_like/graph.py b/lightllm/platform/backends/cuda_like/graph.py new file mode 100644 index 0000000000..987013a81e --- /dev/null +++ b/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() diff --git a/lightllm/platform/backends/cuda_like/ops/__init__.py b/lightllm/platform/backends/cuda_like/ops/__init__.py new file mode 100644 index 0000000000..14a0877ae1 --- /dev/null +++ b/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"] diff --git a/lightllm/platform/backends/cuda_like/ops/act.py b/lightllm/platform/backends/cuda_like/ops/act.py new file mode 100644 index 0000000000..680a630055 --- /dev/null +++ b/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, +) diff --git a/lightllm/platform/backends/cuda_like/ops/embedding.py b/lightllm/platform/backends/cuda_like/ops/embedding.py new file mode 100644 index 0000000000..8c7aa486ed --- /dev/null +++ b/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, +) diff --git a/lightllm/platform/backends/cuda_like/ops/norm.py b/lightllm/platform/backends/cuda_like/ops/norm.py new file mode 100644 index 0000000000..44b2f812a6 --- /dev/null +++ b/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, +) diff --git a/lightllm/platform/backends/cuda_like/ops/rotary.py b/lightllm/platform/backends/cuda_like/ops/rotary.py new file mode 100644 index 0000000000..b82c7f354d --- /dev/null +++ b/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, +) diff --git a/lightllm/platform/backends/cuda_like/ops/sampling.py b/lightllm/platform/backends/cuda_like/ops/sampling.py new file mode 100644 index 0000000000..200ec71025 --- /dev/null +++ b/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, +) diff --git a/lightllm/platform/backends/cuda_like/runtime.py b/lightllm/platform/backends/cuda_like/runtime.py new file mode 100644 index 0000000000..88fa595cf5 --- /dev/null +++ b/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) diff --git a/lightllm/platform/base/__init__.py b/lightllm/platform/base/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/lightllm/platform/base/backend.py b/lightllm/platform/base/backend.py new file mode 100644 index 0000000000..f5d6f0ecc7 --- /dev/null +++ b/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 diff --git a/lightllm/platform/base/graph.py b/lightllm/platform/base/graph.py new file mode 100644 index 0000000000..ea015cbd2a --- /dev/null +++ b/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 diff --git a/lightllm/platform/base/registry.py b/lightllm/platform/base/registry.py new file mode 100644 index 0000000000..acbb0c5956 --- /dev/null +++ b/lightllm/platform/base/registry.py @@ -0,0 +1,46 @@ +from __future__ import annotations + +import importlib +from dataclasses import dataclass +from typing import Type + +from lightllm.platform.base.backend import HardwareBackend + + +@dataclass(frozen=True) +class PlatformSpec: + name: str + backend: str + + def load_backend_cls(self) -> Type[HardwareBackend]: + module_name, class_name = self.backend.rsplit(":", 1) + module = importlib.import_module(module_name) + backend_cls = getattr(module, class_name) + if not issubclass(backend_cls, HardwareBackend): + raise TypeError(f"Platform {self.name!r} backend {self.backend!r} must be a HardwareBackend subclass.") + + platform_name = getattr(backend_cls, "platform_name", None) + if platform_name != self.name: + raise ValueError( + f"Registered platform {self.name!r} does not match " + f"{backend_cls.__name__}.platform_name {platform_name!r}." + ) + + return backend_cls + + +PLATFORMS: dict[str, PlatformSpec] = {} + + +def register_platform(name: str, backend: str) -> None: + if name in PLATFORMS: + raise ValueError(f"Platform {name!r} is already registered.") + + PLATFORMS[name] = PlatformSpec(name=name, backend=backend) + + +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 diff --git a/lightllm/platform/base/runtime.py b/lightllm/platform/base/runtime.py new file mode 100644 index 0000000000..d1f834ca3b --- /dev/null +++ b/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 diff --git a/lightllm/platform/ops/__init__.py b/lightllm/platform/ops/__init__.py new file mode 100644 index 0000000000..3eab918294 --- /dev/null +++ b/lightllm/platform/ops/__init__.py @@ -0,0 +1,19 @@ +from dataclasses import dataclass + +from lightllm.platform.ops.act import ActOps +from lightllm.platform.ops.embedding import EmbeddingOps +from lightllm.platform.ops.norm import NormOps +from lightllm.platform.ops.rotary import RotaryOps +from lightllm.platform.ops.sampling import SamplingOps + + +@dataclass(frozen=True) +class PlatformOps: + norm: NormOps + act: ActOps + sampling: SamplingOps + embedding: EmbeddingOps + rotary: RotaryOps + + +__all__ = ["ActOps", "EmbeddingOps", "NormOps", "PlatformOps", "RotaryOps", "SamplingOps"] diff --git a/lightllm/platform/ops/act.py b/lightllm/platform/ops/act.py new file mode 100644 index 0000000000..0b390f250a --- /dev/null +++ b/lightllm/platform/ops/act.py @@ -0,0 +1,18 @@ +from dataclasses import dataclass +from typing import Protocol + +import torch + + +class SiluAndMul(Protocol): + def __call__(self, input: torch.Tensor, output: torch.Tensor) -> None: ... + + +class GeluAndMul(Protocol): + def __call__(self, input: torch.Tensor, output: torch.Tensor) -> torch.Tensor: ... + + +@dataclass(frozen=True) +class ActOps: + silu_and_mul: SiluAndMul + gelu_and_mul: GeluAndMul diff --git a/lightllm/platform/ops/embedding.py b/lightllm/platform/ops/embedding.py new file mode 100644 index 0000000000..be673e65b0 --- /dev/null +++ b/lightllm/platform/ops/embedding.py @@ -0,0 +1,48 @@ +from dataclasses import dataclass +from typing import Protocol + +import torch + + +class Embedding(Protocol): + def __call__( + self, + input_ids: torch.Tensor, + weight: torch.Tensor, + vob_start_id: int, + vob_end_id: int, + out: torch.Tensor, + ) -> None: ... + + +class MultimodalEmb(Protocol): + def __call__( + self, + out: torch.Tensor, + prompt_ids: torch.Tensor, + text_weight_embs: torch.Tensor, + embed_cache: torch.Tensor, + img_token_lens: torch.Tensor, + img_start_token_ids: torch.Tensor, + img_start_locs_in_cache: torch.Tensor, + tp_text_start_token_id: int, + tp_text_end_token_id: int, + tp_world_size: int, + text_embed_scale: float = 1.0, + ) -> None: ... + + +class OffloadEmbedTensorToCache(Protocol): + def __call__( + self, + embed_tensor: torch.Tensor, + cache_tensor: torch.Tensor, + start_index_in_cache: int, + ) -> None: ... + + +@dataclass(frozen=True) +class EmbeddingOps: + embedding: Embedding + multimodal_emb: MultimodalEmb + offload_embed_tensor_to_cache: OffloadEmbedTensorToCache diff --git a/lightllm/platform/ops/norm.py b/lightllm/platform/ops/norm.py new file mode 100644 index 0000000000..f0e0025180 --- /dev/null +++ b/lightllm/platform/ops/norm.py @@ -0,0 +1,69 @@ +from dataclasses import dataclass +from typing import Any, Protocol + +import torch + + +class RmsNorm(Protocol): + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + out: torch.Tensor | None = None, + ) -> torch.Tensor: ... + + +class GatedRmsNorm(Protocol): + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor | None, + eps: float, + z: torch.Tensor, + out: torch.Tensor | None = None, + group_size: int | None = None, + norm_before_gate: bool = True, + run_config: dict[str, Any] | None = None, + ) -> torch.Tensor: ... + + +class LayerNorm(Protocol): + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + bias: torch.Tensor, + eps: float, + ) -> torch.Tensor: ... + + +class QkRmsNorm(Protocol): + def __call__( + self, + x: torch.Tensor, + weight: torch.Tensor, + eps: float, + ) -> torch.Tensor: ... + + +class FusedQkRmsNorm(Protocol): + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor, + w_q: torch.Tensor, + w_k: torch.Tensor, + eps: float = 1e-6, + fp32_multiply: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + +@dataclass(frozen=True) +class NormOps: + rms_norm: RmsNorm + gated_rms_norm: GatedRmsNorm + layer_norm: LayerNorm + qk_rms_norm: QkRmsNorm + fused_qk_rms_norm: FusedQkRmsNorm diff --git a/lightllm/platform/ops/rotary.py b/lightllm/platform/ops/rotary.py new file mode 100644 index 0000000000..fa04ffdd02 --- /dev/null +++ b/lightllm/platform/ops/rotary.py @@ -0,0 +1,20 @@ +from dataclasses import dataclass +from typing import Protocol + +import torch + + +class RotaryEmb(Protocol): + def __call__( + self, + q: torch.Tensor, + k: torch.Tensor | None, + cos: torch.Tensor, + sin: torch.Tensor, + partial_rotary_factor: float = 1.0, + ) -> None: ... + + +@dataclass(frozen=True) +class RotaryOps: + rotary_emb: RotaryEmb diff --git a/lightllm/platform/ops/sampling.py b/lightllm/platform/ops/sampling.py new file mode 100644 index 0000000000..73d205d2c2 --- /dev/null +++ b/lightllm/platform/ops/sampling.py @@ -0,0 +1,83 @@ +from dataclasses import dataclass +from typing import TYPE_CHECKING, Protocol + +import torch + +if TYPE_CHECKING: + from lightllm.common.req_manager.req_sampling_params import ReqSamplingParamsManager + from lightllm.server.router.model_infer.infer_batch import InferReq + + +class ApplyPenalty(Protocol): + def __call__( + self, + Logits: torch.Tensor, + b_req_idx: torch.Tensor, + b_length_penalty_param: torch.Tensor, + b_mask_eos_reqs: torch.Tensor, + p_token_ids: torch.Tensor, + p_token_counts: torch.Tensor, + p_cumsum_seq_len: torch.Tensor, + eos_ids: torch.Tensor, + sampling_params_manager: "ReqSamplingParamsManager", + ) -> None: ... + + +class ApplyPenaltyGpuCache(Protocol): + def __call__( + self, + Logits: torch.Tensor, + b_req_idx: torch.Tensor, + b_length_penalty_param: torch.Tensor, + b_mask_eos_reqs: torch.Tensor, + eos_ids: torch.Tensor, + sampling_params_manager: "ReqSamplingParamsManager", + ) -> None: ... + + +class ApplyInvalidTokenIds(Protocol): + def __call__( + self, + Logits: torch.Tensor, + invalid_token_ids: torch.Tensor, + cu_invalid_token_num: torch.Tensor, + ) -> None: ... + + +class TokenIdCounter(Protocol): + def __call__( + self, + prompt_ids: torch.Tensor, + out_token_id_counter: torch.Tensor, + ) -> None: ... + + +class UpdateReqToTokenIdCounter(Protocol): + def __call__( + self, + b_req_idx: torch.Tensor, + next_token_ids: torch.Tensor, + req_to_out_token_id_counter: torch.Tensor, + mask: torch.Tensor | None = None, + ) -> None: ... + + +class TopPTopKSample(Protocol): + def __call__( + self, + reqs: list["InferReq"], + probs: torch.Tensor, + b_top_ps: torch.Tensor, + b_top_ks: torch.Tensor, + exist_req_use_random_seed: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: ... + + +@dataclass(frozen=True) +class SamplingOps: + apply_penalty: ApplyPenalty + apply_penalty_gpu_cache: ApplyPenaltyGpuCache + apply_invalid_token_ids: ApplyInvalidTokenIds + token_id_counter: TokenIdCounter + update_req_to_token_id_counter: UpdateReqToTokenIdCounter + top_p_top_k_sample: TopPTopKSample diff --git a/lightllm/platform/plugin/__init__.py b/lightllm/platform/plugin/__init__.py new file mode 100644 index 0000000000..6ac3140dff --- /dev/null +++ b/lightllm/platform/plugin/__init__.py @@ -0,0 +1,13 @@ +from lightllm.platform.plugin.common import Plugin, plugin_registration + + +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: + with plugin_registration(): + for plugin in _PLUGINS: + plugin.load() diff --git a/lightllm/platform/plugin/att.py b/lightllm/platform/plugin/att.py new file mode 100644 index 0000000000..6426573e00 --- /dev/null +++ b/lightllm/platform/plugin/att.py @@ -0,0 +1,66 @@ +from dataclasses import dataclass + +from lightllm.platform.plugin.common import is_plugin_registering +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + +AttKey = tuple[str, str, str, str] + + +@dataclass(frozen=True) +class _AttImpl: + backend_cls: type + plugin: bool + + +ATT_IMPLS: dict[AttKey, _AttImpl] = {} + + +def register_att_backend( + name: str, + *, + category: str, + kv_type: str = "None", + platforms: tuple[str, ...] = ("cuda",), +): + def decorator(backend_cls): + plugin = is_plugin_registering() + staged: dict[AttKey, _AttImpl] = {} + for platform in platforms: + key = (platform, name, category, kv_type) + current = staged.get(key, ATT_IMPLS.get(key)) + if current is None: + staged[key] = _AttImpl(backend_cls, plugin) + continue + if plugin and not current.plugin: + logger.info( + "Att backend %s category %s kv type %s on platform %s: plugin replaced builtin implementation.", + name, + category, + kv_type, + platform, + ) + staged[key] = _AttImpl(backend_cls, plugin=True) + continue + raise ValueError(f"The att backend {name} is already registered for platform {platform!r}.") + ATT_IMPLS.update(staged) + return backend_cls + + return decorator + + +def get_att_backend_class( + name: str, + category: str, + kv_type: str, + platform: str, +) -> type: + key = (platform, name, category, kv_type) + entry = ATT_IMPLS.get(key) + if entry is None: + raise ValueError( + f"The att backend {name} is not registered for category {category}, " + f"kv type {kv_type}, and platform {platform}." + ) + return entry.backend_cls diff --git a/lightllm/platform/plugin/common.py b/lightllm/platform/plugin/common.py new file mode 100644 index 0000000000..4411a4dcfa --- /dev/null +++ b/lightllm/platform/plugin/common.py @@ -0,0 +1,145 @@ +import importlib +from contextlib import contextmanager +from dataclasses import dataclass +from importlib.metadata import entry_points +from typing import Any, Iterable, List, Mapping + +from lightllm.utils.envs_utils import get_env_start_args + +_registering_plugin = False + + +@contextmanager +def plugin_registration(): + global _registering_plugin + previous = _registering_plugin + _registering_plugin = True + try: + yield + finally: + _registering_plugin = previous + + +def is_plugin_registering() -> bool: + return _registering_plugin + + +@dataclass(frozen=True) +class PluginConfig: + modules: tuple[str, ...] = () + + +class Plugin: + def __init__(self, name: str, entry_point_group: str) -> None: + """加载 pip 安装的插件。 + + Args: + name: 插件种类。 + entry_point_group: entry point 组名，用 ``importlib.metadata`` 查找。 + + pyproject.toml: + [project.entry-points."lightllm.ops_plugin"] + example_ops = "lightllm_example_plugin.register:register_ops_plugin" + + [project.entry-points."lightllm.att_plugin"] + example_att = "lightllm_example_plugin.register:register_att_plugin" + + setup.py: + setup( + entry_points={ + "lightllm.ops_plugin": [ + "example_ops = lightllm_example_plugin.register:register_ops_plugin", + ], + "lightllm.att_plugin": [ + "example_att = lightllm_example_plugin.register:register_att_plugin", + ], + }, + ) + + ``pip install`` 之后，传入 ``--extra_ops example_ops`` 或 ``--extra_att example_att`` 加载插件。 + 被加载的模块调用 ``register_ops`` 或 ``register_att_backend``。同一个 key 已有内置实现，则替换。 + """ + self.name = name + self.entry_point_group = entry_point_group + + def load(self) -> None: + names = self._name_from_cli() + if not names: + return + + configs = self._load_entry_point_plugins(names) + config = PluginConfig(modules=merge_config_field(configs, "modules")) + + for module in config.modules: + importlib.import_module(module) + + def _name_from_cli(self) -> tuple[str, ...]: + start_args = get_env_start_args() + return _to_str_tuple(getattr(start_args, f"extra_{self.name}", None)) + + def _load_entry_point_plugins(self, plugin_names: tuple[str, ...]) -> List[PluginConfig]: + given_names = set(plugin_names) + configs: List[PluginConfig] = [] + loaded_names: set[str] = set() + available_entry_points = list(entry_points().select(group=self.entry_point_group)) + for entry_point in available_entry_points: + # 只从 CLI 加载 entry point + if entry_point.name not in given_names: + continue + + # example_ops = "lightllm_example_plugin.register:register_ops_plugin" + # -> register_fn = lightllm_example_plugin.register.register_ops_plugin + register_fn = entry_point.load() + config = parse_plugin_config(register_fn(), plugin_kind=self.name) + + configs.append(config) + loaded_names.add(entry_point.name) + + # 检查是否所有需要的插件都已加载 + missing = given_names - loaded_names + if missing: + available = tuple(sorted(ep.name for ep in available_entry_points)) + message = ( + f"{self.name} plugin(s) not found in entry point group " + f"{self.entry_point_group!r}: {sorted(missing)}" + ) + if available: + message += f". Installed plugins: {available}" + else: + message += ( + f". No {self.name} plugins installed; register entry points in group " + f"{self.entry_point_group!r} and pip install -e your plugin package." + ) + raise RuntimeError(message) + + return configs + + +def parse_plugin_config(value: Any, plugin_kind: str) -> PluginConfig: + if not isinstance(value, Mapping): + raise TypeError(f"{plugin_kind} plugin config must be a mapping, got {type(value)}") + + return PluginConfig(modules=_to_str_tuple(value.get("modules"))) + + +def merge_config_field(configs: Iterable[Any], field_name: str) -> tuple[str, ...]: + merged: list[str] = [] + seen: set[str] = set() + for config in configs: + for item in getattr(config, field_name): + if item in seen: + continue + seen.add(item) + merged.append(item) + return tuple(merged) + + +def _to_str_tuple(value: str | Iterable[str] | None) -> tuple[str, ...]: + if value is None: + return () + + if isinstance(value, str): + parts: Iterable[str] = value.split(",") + else: + parts = value + return tuple(item.strip() for item in parts if item and item.strip()) diff --git a/lightllm/platform/plugin/ops.py b/lightllm/platform/plugin/ops.py new file mode 100644 index 0000000000..462d0ea70f --- /dev/null +++ b/lightllm/platform/plugin/ops.py @@ -0,0 +1,84 @@ +from dataclasses import dataclass, fields +from typing import Any, Callable + +from lightllm.platform.ops import PlatformOps +from lightllm.platform.plugin.common import is_plugin_registering +from lightllm.utils.log_utils import init_logger + +logger = init_logger(__name__) + +OpKey = tuple[str, type, str] + + +@dataclass(frozen=True) +class _OpImpl: + fn: Callable[..., Any] + plugin: bool + + +OP_IMPLS: dict[OpKey, _OpImpl] = {} + + +def register_ops( + group: type, + *, + platforms: tuple[str, ...], + **ops: Callable[..., Any], +) -> None: + unknown = sorted(set(ops) - {field.name for field in fields(group)}) + if unknown: + raise ValueError(f"Unknown op(s) {unknown} on {group.__name__}.") + + plugin = is_plugin_registering() + group_name = next((field.name for field in fields(PlatformOps) if field.type is group), group.__name__) + staged: dict[OpKey, _OpImpl] = {} + for name, fn in ops.items(): + for platform in platforms: + key = (platform, group, name) + current = staged.get(key, OP_IMPLS.get(key)) + # 当前 key 不存在，则直接注册 + if current is None: + staged[key] = _OpImpl(fn, plugin) + continue + # 当前是插件且已有内置实现，则替换 + if plugin and not current.plugin: + logger.info( + "Op %s.%s on platform %s: plugin replaced builtin implementation.", + group_name, + name, + platform, + ) + staged[key] = _OpImpl(fn, plugin=True) + continue + # 按照同样 key 注册多次，报错 + raise ValueError(f"Op {group.__name__}.{name} is already registered for platform {platform!r}.") + OP_IMPLS.update(staged) + + +def build_platform_ops(platform: str) -> PlatformOps: + # 按平台把已登记的算子组装成 PlatformOps，无对应实现置为 unsupported op 在运行时报错 + groups = {} + for group_field in fields(PlatformOps): + group_type = group_field.type + impls = {} + for op_field in fields(group_type): + full_name = f"{group_field.name}.{op_field.name}" + impls[op_field.name] = _lookup(platform, group_type, op_field.name, full_name) + groups[group_field.name] = group_type(**impls) + return PlatformOps(**groups) + + +def _lookup(platform: str, group: type, name: str, full_name: str) -> Callable[..., Any]: + entry = OP_IMPLS.get((platform, group, name)) + if entry is None: + return _unsupported_op(platform, full_name) + return entry.fn + + +def _unsupported_op(platform: str, full_name: str) -> Callable[..., Any]: + + def op(*args: Any, **kwargs: Any) -> Any: + raise NotImplementedError(f"Op {full_name!r} is not implemented for platform {platform!r}.") + + op.__name__ = full_name.rsplit(".", 1)[-1] + return op diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index e6c077dbe4..a123c503ee 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -973,6 +973,20 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: (you should set it up by yourself). A NVTX range named 'LIGHTLLM_PROFILE' will be added within the profiling range.""", ) + parser.add_argument( + "--extra_ops", + type=str, + default=None, + help="""Extra ops plugin(s) to load at startup, comma-separated names. + Each name must match a pip-installed plugin.""", + ) + parser.add_argument( + "--extra_att", + type=str, + default=None, + help="""Extra att plugin(s) to load at startup, comma-separated names. + Each name must match a pip-installed plugin.""", + ) return parser diff --git a/lightllm/server/core/objs/start_args_type.py b/lightllm/server/core/objs/start_args_type.py index 9c89975de7..d063f41bd5 100644 --- a/lightllm/server/core/objs/start_args_type.py +++ b/lightllm/server/core/objs/start_args_type.py @@ -240,3 +240,6 @@ class StartArgs: disable_linear_att_small_page_cpu_cache: bool = field(default=False) linear_att_cache_size: Optional[int] = field(default=None) linear_att_ssm_data_type: Optional[str] = field(default="float32", metadata={"choices": ["bfloat16", "float32"]}) + + extra_ops: Optional[str] = field(default=None) + extra_att: Optional[str] = field(default=None)