From fd655128f59d9714910beb808ad7651967af4c81 Mon Sep 17 00:00:00 2001 From: zhangtaoshan Date: Fri, 11 Sep 2026 09:17:40 +0000 Subject: [PATCH 1/3] feat: add hardware platform and plugin support --- lightllm/platform/__init__.py | 34 +++++ lightllm/platform/backends/__init__.py | 1 + .../platform/backends/cuda_like/__init__.py | 14 ++ lightllm/platform/backends/cuda_like/graph.py | 23 ++++ .../platform/backends/cuda_like/runtime.py | 63 +++++++++ lightllm/platform/base/__init__.py | 0 lightllm/platform/base/backend.py | 24 ++++ lightllm/platform/base/graph.py | 28 ++++ lightllm/platform/base/registry.py | 37 +++++ lightllm/platform/base/runtime.py | 122 +++++++++++++++++ lightllm/platform/plugin/__init__.py | 12 ++ lightllm/platform/plugin/att.py | 63 +++++++++ lightllm/platform/plugin/common.py | 128 ++++++++++++++++++ lightllm/platform/plugin/ops.py | 43 ++++++ lightllm/server/api_cli.py | 14 ++ lightllm/server/core/objs/start_args_type.py | 3 + 16 files changed, 609 insertions(+) create mode 100644 lightllm/platform/__init__.py create mode 100644 lightllm/platform/backends/__init__.py create mode 100644 lightllm/platform/backends/cuda_like/__init__.py create mode 100644 lightllm/platform/backends/cuda_like/graph.py create mode 100644 lightllm/platform/backends/cuda_like/runtime.py create mode 100644 lightllm/platform/base/__init__.py create mode 100644 lightllm/platform/base/backend.py create mode 100644 lightllm/platform/base/graph.py create mode 100644 lightllm/platform/base/registry.py create mode 100644 lightllm/platform/base/runtime.py create mode 100644 lightllm/platform/plugin/__init__.py create mode 100644 lightllm/platform/plugin/att.py create mode 100644 lightllm/platform/plugin/common.py create mode 100644 lightllm/platform/plugin/ops.py diff --git a/lightllm/platform/__init__.py b/lightllm/platform/__init__.py new file mode 100644 index 0000000000..70fcd51c65 --- /dev/null +++ b/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"] diff --git a/lightllm/platform/backends/__init__.py b/lightllm/platform/backends/__init__.py new file mode 100644 index 0000000000..25cb41dae8 --- /dev/null +++ b/lightllm/platform/backends/__init__.py @@ -0,0 +1 @@ +from lightllm.platform.backends import cuda_like # noqa: F401 diff --git a/lightllm/platform/backends/cuda_like/__init__.py b/lightllm/platform/backends/cuda_like/__init__.py new file mode 100644 index 0000000000..7ea2d3a81a --- /dev/null +++ b/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 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/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..3669cc0e8a --- /dev/null +++ b/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 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..61f4bbef21 --- /dev/null +++ b/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 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/plugin/__init__.py b/lightllm/platform/plugin/__init__.py new file mode 100644 index 0000000000..21c2ab8e6d --- /dev/null +++ b/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() diff --git a/lightllm/platform/plugin/att.py b/lightllm/platform/plugin/att.py new file mode 100644 index 0000000000..d113853b31 --- /dev/null +++ b/lightllm/platform/plugin/att.py @@ -0,0 +1,63 @@ +from dataclasses import dataclass + +ATT_TABLE: dict[tuple[str, str, str], list["AttBackendEntry"]] = {} + + +@dataclass(frozen=True) +class AttBackendEntry: + # The name of the att backend. e.g., "triton", "flashinfer", "fa3". + name: str + # The category of the att backend. e.g., "standard", "mla", "nsa". + category: str + # The type of the kv. e.g., "int8kv", "int4kv", "fp8kv_sph". + kv_type: str + # The class of the att backend. + backend_cls: type + # The platforms of the att backend. e.g., "cuda", "ascend". + platforms: tuple[str, ...] + + +def register_att_backend( + name: str, + *, + category: str, + kv_type: str = "None", + platforms: tuple[str, ...] = ("cuda",), +): + def decorator(backend_cls): + if not platforms: + raise ValueError("The value of platforms must be specified for att backend registration.") + + key = (name, category, kv_type) + # Check if the att backend is already registered. + for existing in ATT_TABLE.get(key, []): + if set(existing.platforms) & set(platforms): + raise ValueError(f"The att backend {name} is already registered for platforms {existing.platforms}.") + + ATT_TABLE.setdefault(key, []).append( + AttBackendEntry( + name=name, + category=category, + kv_type=kv_type, + backend_cls=backend_cls, + platforms=platforms, + ) + ) + return backend_cls + + return decorator + + +def get_att_backend_class( + name: str, + category: str, + kv_type: str, + platform: str, +) -> type: + key = (name, category, kv_type) + for entry in ATT_TABLE.get(key, []): + if platform in entry.platforms: + return entry.backend_cls + raise ValueError( + f"The att backend {name} is not registered for category {category}, kv type {kv_type}, and platform {platform}." + ) diff --git a/lightllm/platform/plugin/common.py b/lightllm/platform/plugin/common.py new file mode 100644 index 0000000000..c44b674565 --- /dev/null +++ b/lightllm/platform/plugin/common.py @@ -0,0 +1,128 @@ +import importlib +from dataclasses import dataclass +from importlib.metadata import entry_points, EntryPoint +from typing import Any, Iterable, List, Mapping + +from lightllm.utils.envs_utils import get_env_start_args + + +@dataclass(frozen=True) +class PluginConfig: + modules: tuple[str, ...] = () + + +class Plugin: + def __init__(self, name: str, entry_point_group: str) -> None: + """Loader for pip-installed plugins. + + Args: + name: Plugin kind. + entry_point_group: Entry point group name. Looked up with ``importlib.metadata``. + + pyproject.toml: + [project.entry-points."lightllm.ops_plugin"] + example_ops = "lightllm_example_plugin.register:register_ops" + + setup.py: + setup( + entry_points={ + "lightllm.ops_plugin": [ + "example_ops = lightllm_example_plugin.register:register_ops", + ], + }, + ) + + After ``pip install``, pass ``--extra-ops example_ops`` to load and register. + """ + 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(_iter_entry_points(self.entry_point_group)) + for entry_point in available_entry_points: + # Only load entry points from CLI. + if entry_point.name not in given_names: + continue + + # example_ops = "lightllm_example_plugin.register:register_ops" + # -> register_fn = lightllm_example_plugin.register.register_ops + register_fn = entry_point.load() + config = parse_plugin_config(register_fn(), plugin_kind=self.name) + + configs.append(config) + loaded_names.add(entry_point.name) + + # Check if all required plugins are loaded. + 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 _iter_entry_points(entry_point_group: str) -> Iterable[EntryPoint]: + eps = entry_points() + if hasattr(eps, "select"): + yield from eps.select(group=entry_point_group) + else: + yield from eps.get(entry_point_group, []) + + +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..771b8f925f --- /dev/null +++ b/lightllm/platform/plugin/ops.py @@ -0,0 +1,43 @@ +from dataclasses import dataclass, field +from typing import Callable + +# The default order of fallback implementations. +DEFAULT_ORDER = ("triton", "torch") + +OPS_TABLE: dict[str, "OpEntry"] = {} + + +@dataclass +class OpEntry: + fns: dict[str, Callable] = field(default_factory=dict) + extras: list[str] = field(default_factory=list) + + @property + def impls(self) -> tuple[str, ...]: + extras = tuple(name for name in self.extras if name in self.fns) + defaults = tuple(kind for kind in DEFAULT_ORDER if kind in self.fns) + return extras + defaults + + +def register_op(op_name: str, *, impl: str): + def decorator(fn: Callable) -> Callable: + entry = OPS_TABLE.setdefault(op_name, OpEntry()) + if impl in entry.fns: + raise RuntimeError(f"The {op_name} op has already been registered with the {impl!r} implementation.") + entry.fns[impl] = fn + if impl not in DEFAULT_ORDER: + entry.extras.insert(0, impl) + return fn + + return decorator + + +def get_op(op_name: str) -> Callable: + entry = OPS_TABLE.get(op_name) + if entry is None: + raise NotImplementedError(f"The {op_name!r} op is not registered.") + for impl in entry.impls: + fn = entry.fns.get(impl) + if fn is not None: + return fn + raise NotImplementedError(f"The {op_name!r} op has no usable implementation in {entry.impls}.") diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index e6c077dbe4..d278d6329e 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) From 9c2ba571d1eea2ff2b17d4a99b87c70b69be3b71 Mon Sep 17 00:00:00 2001 From: zhangtaoshan Date: Tue, 29 Sep 2026 11:08:42 +0800 Subject: [PATCH 2/3] feat: add grouped platform ops --- lightllm/platform/__init__.py | 7 +- lightllm/platform/backends/__init__.py | 17 ++++- lightllm/platform/backends/ascend/__init__.py | 16 +++++ lightllm/platform/backends/ascend/graph.py | 24 +++++++ .../platform/backends/ascend/ops/__init__.py | 13 ++++ lightllm/platform/backends/ascend/ops/act.py | 8 +++ .../platform/backends/ascend/ops/common.py | 10 +++ lightllm/platform/backends/ascend/ops/norm.py | 11 +++ .../platform/backends/ascend/ops/sampling.py | 12 ++++ lightllm/platform/backends/ascend/runtime.py | 69 +++++++++++++++++++ .../platform/backends/cuda_like/__init__.py | 16 +++-- .../backends/cuda_like/ops/__init__.py | 13 ++++ .../platform/backends/cuda_like/ops/act.py | 9 +++ .../platform/backends/cuda_like/ops/norm.py | 14 ++++ .../backends/cuda_like/ops/sampling.py | 17 +++++ lightllm/platform/base/backend.py | 17 +++-- lightllm/platform/base/registry.py | 37 ++++++---- lightllm/platform/ops/__init__.py | 15 ++++ lightllm/platform/ops/act.py | 10 +++ lightllm/platform/ops/norm.py | 17 +++++ lightllm/platform/ops/sampling.py | 13 ++++ lightllm/platform/plugin/ops.py | 62 ++++++++--------- lightllm/server/api_cli.py | 4 +- 23 files changed, 370 insertions(+), 61 deletions(-) create mode 100644 lightllm/platform/backends/ascend/__init__.py create mode 100644 lightllm/platform/backends/ascend/graph.py create mode 100644 lightllm/platform/backends/ascend/ops/__init__.py create mode 100644 lightllm/platform/backends/ascend/ops/act.py create mode 100644 lightllm/platform/backends/ascend/ops/common.py create mode 100644 lightllm/platform/backends/ascend/ops/norm.py create mode 100644 lightllm/platform/backends/ascend/ops/sampling.py create mode 100644 lightllm/platform/backends/ascend/runtime.py create mode 100644 lightllm/platform/backends/cuda_like/ops/__init__.py create mode 100644 lightllm/platform/backends/cuda_like/ops/act.py create mode 100644 lightllm/platform/backends/cuda_like/ops/norm.py create mode 100644 lightllm/platform/backends/cuda_like/ops/sampling.py create mode 100644 lightllm/platform/ops/__init__.py create mode 100644 lightllm/platform/ops/act.py create mode 100644 lightllm/platform/ops/norm.py create mode 100644 lightllm/platform/ops/sampling.py diff --git a/lightllm/platform/__init__.py b/lightllm/platform/__init__.py index 70fcd51c65..d1c8f5bc7e 100644 --- a/lightllm/platform/__init__.py +++ b/lightllm/platform/__init__.py @@ -22,12 +22,13 @@ def get_hardware_backend() -> HardwareBackend: platform_name = get_env_start_args().hardware_platform spec = get_platform_spec(platform_name) - backend_cls = spec.backend_cls - _backend = backend_cls() + backend_cls = spec.load_backend_cls() + backend = backend_cls() - if not _backend.runtime.is_available(): + if not backend.runtime.is_available(): raise RuntimeError(f"Hardware backend {backend_cls.__name__} is not available.") + _backend = backend return _backend diff --git a/lightllm/platform/backends/__init__.py b/lightllm/platform/backends/__init__.py index 25cb41dae8..968d3179ec 100644 --- a/lightllm/platform/backends/__init__.py +++ b/lightllm/platform/backends/__init__.py @@ -1 +1,16 @@ -from lightllm.platform.backends import cuda_like # noqa: F401 +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..56ab5af9aa --- /dev/null +++ b/lightllm/platform/backends/ascend/__init__.py @@ -0,0 +1,16 @@ +from lightllm.platform.base.backend import HardwareBackend +from lightllm.platform.backends.ascend.graph import AscendGraph +from lightllm.platform.backends.ascend.ops import ASCEND_OPS +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, ASCEND_OPS), + ) 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..e8ae9514f5 --- /dev/null +++ b/lightllm/platform/backends/ascend/ops/__init__.py @@ -0,0 +1,13 @@ +from .act import ASCEND_ACT_OPS +from .norm import ASCEND_NORM_OPS +from .sampling import ASCEND_SAMPLING_OPS +from lightllm.platform.ops import PlatformOps + + +ASCEND_OPS = PlatformOps( + norm=ASCEND_NORM_OPS, + act=ASCEND_ACT_OPS, + sampling=ASCEND_SAMPLING_OPS, +) + +__all__ = ["ASCEND_OPS"] diff --git a/lightllm/platform/backends/ascend/ops/act.py b/lightllm/platform/backends/ascend/ops/act.py new file mode 100644 index 0000000000..2fe9627527 --- /dev/null +++ b/lightllm/platform/backends/ascend/ops/act.py @@ -0,0 +1,8 @@ +from lightllm.platform.ops import ActOps + +from .common import unsupported_op + +ASCEND_ACT_OPS = ActOps( + silu_and_mul=unsupported_op("act.silu_and_mul"), + gelu_and_mul=unsupported_op("act.gelu_and_mul"), +) diff --git a/lightllm/platform/backends/ascend/ops/common.py b/lightllm/platform/backends/ascend/ops/common.py new file mode 100644 index 0000000000..de334d4a09 --- /dev/null +++ b/lightllm/platform/backends/ascend/ops/common.py @@ -0,0 +1,10 @@ +from collections.abc import Callable +from typing import Any + + +def unsupported_op(name: str) -> Callable[..., Any]: + def op(*args, **kwargs): + raise NotImplementedError(f"Ascend op {name!r} is not implemented.") + + op.__name__ = name.rsplit(".", 1)[-1] + return op diff --git a/lightllm/platform/backends/ascend/ops/norm.py b/lightllm/platform/backends/ascend/ops/norm.py new file mode 100644 index 0000000000..ce1671f304 --- /dev/null +++ b/lightllm/platform/backends/ascend/ops/norm.py @@ -0,0 +1,11 @@ +from lightllm.platform.ops import NormOps + +from .common import unsupported_op + +ASCEND_NORM_OPS = NormOps( + rms_norm=unsupported_op("norm.rms_norm"), + gated_rms_norm=unsupported_op("norm.gated_rms_norm"), + layer_norm=unsupported_op("norm.layer_norm"), + qk_rms_norm=unsupported_op("norm.qk_rms_norm"), + fused_qk_rms_norm=unsupported_op("norm.fused_qk_rms_norm"), +) diff --git a/lightllm/platform/backends/ascend/ops/sampling.py b/lightllm/platform/backends/ascend/ops/sampling.py new file mode 100644 index 0000000000..3758d72e84 --- /dev/null +++ b/lightllm/platform/backends/ascend/ops/sampling.py @@ -0,0 +1,12 @@ +from lightllm.platform.ops import SamplingOps + +from .common import unsupported_op + + +ASCEND_SAMPLING_OPS = SamplingOps( + apply_penalty=unsupported_op("sampling.apply_penalty"), + apply_penalty_gpu_cache=unsupported_op("sampling.apply_penalty_gpu_cache"), + apply_invalid_token_ids=unsupported_op("sampling.apply_invalid_token_ids"), + token_id_counter=unsupported_op("sampling.token_id_counter"), + update_req_to_token_id_counter=unsupported_op("sampling.update_req_to_token_id_counter"), +) 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 index 7ea2d3a81a..efd85d7271 100644 --- a/lightllm/platform/backends/cuda_like/__init__.py +++ b/lightllm/platform/backends/cuda_like/__init__.py @@ -1,14 +1,22 @@ from lightllm.platform.base.backend import HardwareBackend -from lightllm.platform.base.registry import register_platform +from lightllm.platform.backends.cuda_like.ops import CUDA_LIKE_OPS from lightllm.platform.backends.cuda_like.runtime import CudaLikeRuntime from lightllm.platform.backends.cuda_like.graph import CudaLikeGraph +from lightllm.platform.plugin.ops import build_platform_ops class CudaLikeBackend(HardwareBackend): def __init__(self) -> None: - super().__init__(CudaLikeRuntime(), CudaLikeGraph()) + super().__init__( + runtime=CudaLikeRuntime(), + graph=CudaLikeGraph(), + ops=build_platform_ops(self.platform_name, CUDA_LIKE_OPS), + ) -@register_platform("cuda") class CudaBackend(CudaLikeBackend): - pass + platform_name = "cuda" + + +class MusaBackend(CudaLikeBackend): + platform_name = "musa" 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..2ee453559c --- /dev/null +++ b/lightllm/platform/backends/cuda_like/ops/__init__.py @@ -0,0 +1,13 @@ +from lightllm.platform.backends.cuda_like.ops.act import CUDA_LIKE_ACT_OPS +from lightllm.platform.backends.cuda_like.ops.norm import CUDA_LIKE_NORM_OPS +from lightllm.platform.backends.cuda_like.ops.sampling import CUDA_LIKE_SAMPLING_OPS +from lightllm.platform.ops import PlatformOps + + +CUDA_LIKE_OPS = PlatformOps( + norm=CUDA_LIKE_NORM_OPS, + act=CUDA_LIKE_ACT_OPS, + sampling=CUDA_LIKE_SAMPLING_OPS, +) + +__all__ = ["CUDA_LIKE_OPS"] 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..c45853d1da --- /dev/null +++ b/lightllm/platform/backends/cuda_like/ops/act.py @@ -0,0 +1,9 @@ +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 + + +CUDA_LIKE_ACT_OPS = ActOps( + silu_and_mul=silu_and_mul_fwd, + gelu_and_mul=gelu_and_mul_fwd, +) 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..8a6cc9a7db --- /dev/null +++ b/lightllm/platform/backends/cuda_like/ops/norm.py @@ -0,0 +1,14 @@ +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 + + +CUDA_LIKE_NORM_OPS = NormOps( + 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/sampling.py b/lightllm/platform/backends/cuda_like/ops/sampling.py new file mode 100644 index 0000000000..a1bb71f63b --- /dev/null +++ b/lightllm/platform/backends/cuda_like/ops/sampling.py @@ -0,0 +1,17 @@ +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 + + +CUDA_LIKE_SAMPLING_OPS = SamplingOps( + 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, +) diff --git a/lightllm/platform/base/backend.py b/lightllm/platform/base/backend.py index 3669cc0e8a..f5d6f0ecc7 100644 --- a/lightllm/platform/base/backend.py +++ b/lightllm/platform/base/backend.py @@ -1,15 +1,20 @@ -from abc import ABC - from lightllm.platform.base.graph import HardwareBackendGraph from lightllm.platform.base.runtime import HardwareBackendRuntime +from lightllm.platform.ops import PlatformOps -class HardwareBackend(ABC): +class HardwareBackend: platform_name: str - def __init__(self, runtime: HardwareBackendRuntime, graph: HardwareBackendGraph) -> None: + def __init__( + self, + runtime: HardwareBackendRuntime, + graph: HardwareBackendGraph, + ops: PlatformOps, + ) -> None: self._runtime = runtime self._graph = graph + self._ops = ops @property def name(self) -> str: @@ -22,3 +27,7 @@ def runtime(self) -> HardwareBackendRuntime: @property def graph(self) -> HardwareBackendGraph: return self._graph + + @property + def ops(self) -> PlatformOps: + return self._ops diff --git a/lightllm/platform/base/registry.py b/lightllm/platform/base/registry.py index 61f4bbef21..acbb0c5956 100644 --- a/lightllm/platform/base/registry.py +++ b/lightllm/platform/base/registry.py @@ -1,7 +1,8 @@ from __future__ import annotations +import importlib from dataclasses import dataclass -from typing import Callable, Type +from typing import Type from lightllm.platform.base.backend import HardwareBackend @@ -9,25 +10,33 @@ @dataclass(frozen=True) class PlatformSpec: name: str - backend_cls: Type[HardwareBackend] + 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] = {} +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 +def register_platform(name: str, backend: str) -> None: + if name in PLATFORMS: + raise ValueError(f"Platform {name!r} is already registered.") - return decorator + PLATFORMS[name] = PlatformSpec(name=name, backend=backend) def get_platform_spec(name: str) -> PlatformSpec: diff --git a/lightllm/platform/ops/__init__.py b/lightllm/platform/ops/__init__.py new file mode 100644 index 0000000000..6b5cec1630 --- /dev/null +++ b/lightllm/platform/ops/__init__.py @@ -0,0 +1,15 @@ +from dataclasses import dataclass + +from lightllm.platform.ops.act import ActOps +from lightllm.platform.ops.norm import NormOps +from lightllm.platform.ops.sampling import SamplingOps + + +@dataclass(frozen=True) +class PlatformOps: + norm: NormOps + act: ActOps + sampling: SamplingOps + + +__all__ = ["ActOps", "NormOps", "PlatformOps", "SamplingOps"] diff --git a/lightllm/platform/ops/act.py b/lightllm/platform/ops/act.py new file mode 100644 index 0000000000..54c990c6db --- /dev/null +++ b/lightllm/platform/ops/act.py @@ -0,0 +1,10 @@ +from dataclasses import dataclass +from typing import Any, Callable + +ActOp = Callable[..., Any] + + +@dataclass(frozen=True) +class ActOps: + silu_and_mul: ActOp + gelu_and_mul: ActOp diff --git a/lightllm/platform/ops/norm.py b/lightllm/platform/ops/norm.py new file mode 100644 index 0000000000..8c4de84ccd --- /dev/null +++ b/lightllm/platform/ops/norm.py @@ -0,0 +1,17 @@ +from dataclasses import dataclass +from typing import Any, Callable + +import torch + + +TensorOp = Callable[..., torch.Tensor] +AnyOp = Callable[..., Any] + + +@dataclass(frozen=True) +class NormOps: + rms_norm: TensorOp + gated_rms_norm: TensorOp + layer_norm: TensorOp + qk_rms_norm: TensorOp + fused_qk_rms_norm: AnyOp diff --git a/lightllm/platform/ops/sampling.py b/lightllm/platform/ops/sampling.py new file mode 100644 index 0000000000..ee01cf3f3d --- /dev/null +++ b/lightllm/platform/ops/sampling.py @@ -0,0 +1,13 @@ +from dataclasses import dataclass +from typing import Any, Callable + +SamplingOp = Callable[..., Any] + + +@dataclass(frozen=True) +class SamplingOps: + apply_penalty: SamplingOp + apply_penalty_gpu_cache: SamplingOp + apply_invalid_token_ids: SamplingOp + token_id_counter: SamplingOp + update_req_to_token_id_counter: SamplingOp diff --git a/lightllm/platform/plugin/ops.py b/lightllm/platform/plugin/ops.py index 771b8f925f..1255769b56 100644 --- a/lightllm/platform/plugin/ops.py +++ b/lightllm/platform/plugin/ops.py @@ -1,43 +1,39 @@ -from dataclasses import dataclass, field -from typing import Callable +from dataclasses import fields, replace +from typing import Callable, TypeVar -# The default order of fallback implementations. -DEFAULT_ORDER = ("triton", "torch") +from lightllm.platform.ops import PlatformOps -OPS_TABLE: dict[str, "OpEntry"] = {} +OpsGroupT = TypeVar("OpsGroupT") -@dataclass -class OpEntry: - fns: dict[str, Callable] = field(default_factory=dict) - extras: list[str] = field(default_factory=list) +OPS_PLUGINS: dict[tuple[str, str], dict[str, Callable]] = {} - @property - def impls(self) -> tuple[str, ...]: - extras = tuple(name for name in self.extras if name in self.fns) - defaults = tuple(kind for kind in DEFAULT_ORDER if kind in self.fns) - return extras + defaults +def register_ops(name: str, platforms: tuple[str, ...] = ("cuda",), **ops: Callable) -> None: + for platform in platforms: + current = OPS_PLUGINS.setdefault((platform, name), {}) + duplicates = current.keys() & ops.keys() + if duplicates: + raise ValueError( + f"Ops {sorted(duplicates)} are already registered for platform {platform!r}, group {name!r}." + ) + current.update(ops) -def register_op(op_name: str, *, impl: str): - def decorator(fn: Callable) -> Callable: - entry = OPS_TABLE.setdefault(op_name, OpEntry()) - if impl in entry.fns: - raise RuntimeError(f"The {op_name} op has already been registered with the {impl!r} implementation.") - entry.fns[impl] = fn - if impl not in DEFAULT_ORDER: - entry.extras.insert(0, impl) - return fn - return decorator +def build_ops_group(name: str, platform: str, default: OpsGroupT) -> OpsGroupT: + updates = OPS_PLUGINS.get((platform, name), {}) + if not updates: + return default + return replace(default, **updates) -def get_op(op_name: str) -> Callable: - entry = OPS_TABLE.get(op_name) - if entry is None: - raise NotImplementedError(f"The {op_name!r} op is not registered.") - for impl in entry.impls: - fn = entry.fns.get(impl) - if fn is not None: - return fn - raise NotImplementedError(f"The {op_name!r} op has no usable implementation in {entry.impls}.") +def build_platform_ops(platform: str, default: PlatformOps) -> PlatformOps: + ops = { + field.name: build_ops_group( + field.name, + platform, + getattr(default, field.name), + ) + for field in fields(default) + } + return replace(default, **ops) diff --git a/lightllm/server/api_cli.py b/lightllm/server/api_cli.py index d278d6329e..a123c503ee 100644 --- a/lightllm/server/api_cli.py +++ b/lightllm/server/api_cli.py @@ -978,14 +978,14 @@ def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: type=str, default=None, help="""Extra ops plugin(s) to load at startup, comma-separated names. - Each name must match a pip-installed plugin.""" + 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.""" + Each name must match a pip-installed plugin.""", ) return parser From f00497198c7bd5bdaf931028c88237b3f4e27195 Mon Sep 17 00:00:00 2001 From: zhangtaoshan Date: Fri, 23 Oct 2026 17:13:33 +0800 Subject: [PATCH 3/3] feat: load builtin platform ops before pip plugins --- lightllm/platform/__init__.py | 5 +- lightllm/platform/backends/ascend/__init__.py | 4 +- .../platform/backends/ascend/ops/__init__.py | 14 +-- lightllm/platform/backends/ascend/ops/act.py | 8 -- .../platform/backends/ascend/ops/common.py | 10 -- .../platform/backends/ascend/ops/embedding.py | 0 lightllm/platform/backends/ascend/ops/norm.py | 11 -- .../platform/backends/ascend/ops/rotary.py | 0 .../platform/backends/ascend/ops/sampling.py | 12 --- .../platform/backends/cuda_like/__init__.py | 6 +- .../backends/cuda_like/ops/__init__.py | 14 +-- .../platform/backends/cuda_like/ops/act.py | 6 +- .../backends/cuda_like/ops/embedding.py | 13 +++ .../platform/backends/cuda_like/ops/norm.py | 6 +- .../platform/backends/cuda_like/ops/rotary.py | 9 ++ .../backends/cuda_like/ops/sampling.py | 8 +- lightllm/platform/ops/__init__.py | 6 +- lightllm/platform/ops/act.py | 16 ++- lightllm/platform/ops/embedding.py | 48 +++++++++ lightllm/platform/ops/norm.py | 68 ++++++++++-- lightllm/platform/ops/rotary.py | 20 ++++ lightllm/platform/ops/sampling.py | 84 +++++++++++++-- lightllm/platform/plugin/__init__.py | 7 +- lightllm/platform/plugin/att.py | 75 ++++++------- lightllm/platform/plugin/common.py | 57 ++++++---- lightllm/platform/plugin/ops.py | 101 +++++++++++++----- 26 files changed, 423 insertions(+), 185 deletions(-) delete mode 100644 lightllm/platform/backends/ascend/ops/common.py create mode 100644 lightllm/platform/backends/ascend/ops/embedding.py create mode 100644 lightllm/platform/backends/ascend/ops/rotary.py create mode 100644 lightllm/platform/backends/cuda_like/ops/embedding.py create mode 100644 lightllm/platform/backends/cuda_like/ops/rotary.py create mode 100644 lightllm/platform/ops/embedding.py create mode 100644 lightllm/platform/ops/rotary.py diff --git a/lightllm/platform/__init__.py b/lightllm/platform/__init__.py index d1c8f5bc7e..b2a324cbfc 100644 --- a/lightllm/platform/__init__.py +++ b/lightllm/platform/__init__.py @@ -17,12 +17,13 @@ def get_hardware_backend() -> HardwareBackend: 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.load_backend_cls() + + configure_plugins() + backend = backend_cls() if not backend.runtime.is_available(): diff --git a/lightllm/platform/backends/ascend/__init__.py b/lightllm/platform/backends/ascend/__init__.py index 56ab5af9aa..2a2a9099e3 100644 --- a/lightllm/platform/backends/ascend/__init__.py +++ b/lightllm/platform/backends/ascend/__init__.py @@ -1,6 +1,6 @@ +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.ops import ASCEND_OPS from lightllm.platform.backends.ascend.runtime import AscendRuntime from lightllm.platform.plugin.ops import build_platform_ops @@ -12,5 +12,5 @@ def __init__(self) -> None: super().__init__( runtime=AscendRuntime(), graph=AscendGraph(), - ops=build_platform_ops(self.platform_name, ASCEND_OPS), + ops=build_platform_ops(self.platform_name), ) diff --git a/lightllm/platform/backends/ascend/ops/__init__.py b/lightllm/platform/backends/ascend/ops/__init__.py index e8ae9514f5..cc27207009 100644 --- a/lightllm/platform/backends/ascend/ops/__init__.py +++ b/lightllm/platform/backends/ascend/ops/__init__.py @@ -1,13 +1,3 @@ -from .act import ASCEND_ACT_OPS -from .norm import ASCEND_NORM_OPS -from .sampling import ASCEND_SAMPLING_OPS -from lightllm.platform.ops import PlatformOps +from lightllm.platform.backends.ascend.ops import act, embedding, norm, rotary, sampling - -ASCEND_OPS = PlatformOps( - norm=ASCEND_NORM_OPS, - act=ASCEND_ACT_OPS, - sampling=ASCEND_SAMPLING_OPS, -) - -__all__ = ["ASCEND_OPS"] +__all__ = ["act", "embedding", "norm", "rotary", "sampling"] diff --git a/lightllm/platform/backends/ascend/ops/act.py b/lightllm/platform/backends/ascend/ops/act.py index 2fe9627527..e69de29bb2 100644 --- a/lightllm/platform/backends/ascend/ops/act.py +++ b/lightllm/platform/backends/ascend/ops/act.py @@ -1,8 +0,0 @@ -from lightllm.platform.ops import ActOps - -from .common import unsupported_op - -ASCEND_ACT_OPS = ActOps( - silu_and_mul=unsupported_op("act.silu_and_mul"), - gelu_and_mul=unsupported_op("act.gelu_and_mul"), -) diff --git a/lightllm/platform/backends/ascend/ops/common.py b/lightllm/platform/backends/ascend/ops/common.py deleted file mode 100644 index de334d4a09..0000000000 --- a/lightllm/platform/backends/ascend/ops/common.py +++ /dev/null @@ -1,10 +0,0 @@ -from collections.abc import Callable -from typing import Any - - -def unsupported_op(name: str) -> Callable[..., Any]: - def op(*args, **kwargs): - raise NotImplementedError(f"Ascend op {name!r} is not implemented.") - - op.__name__ = name.rsplit(".", 1)[-1] - return op 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 index ce1671f304..e69de29bb2 100644 --- a/lightllm/platform/backends/ascend/ops/norm.py +++ b/lightllm/platform/backends/ascend/ops/norm.py @@ -1,11 +0,0 @@ -from lightllm.platform.ops import NormOps - -from .common import unsupported_op - -ASCEND_NORM_OPS = NormOps( - rms_norm=unsupported_op("norm.rms_norm"), - gated_rms_norm=unsupported_op("norm.gated_rms_norm"), - layer_norm=unsupported_op("norm.layer_norm"), - qk_rms_norm=unsupported_op("norm.qk_rms_norm"), - fused_qk_rms_norm=unsupported_op("norm.fused_qk_rms_norm"), -) 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 index 3758d72e84..e69de29bb2 100644 --- a/lightllm/platform/backends/ascend/ops/sampling.py +++ b/lightllm/platform/backends/ascend/ops/sampling.py @@ -1,12 +0,0 @@ -from lightllm.platform.ops import SamplingOps - -from .common import unsupported_op - - -ASCEND_SAMPLING_OPS = SamplingOps( - apply_penalty=unsupported_op("sampling.apply_penalty"), - apply_penalty_gpu_cache=unsupported_op("sampling.apply_penalty_gpu_cache"), - apply_invalid_token_ids=unsupported_op("sampling.apply_invalid_token_ids"), - token_id_counter=unsupported_op("sampling.token_id_counter"), - update_req_to_token_id_counter=unsupported_op("sampling.update_req_to_token_id_counter"), -) diff --git a/lightllm/platform/backends/cuda_like/__init__.py b/lightllm/platform/backends/cuda_like/__init__.py index efd85d7271..d715da0dfa 100644 --- a/lightllm/platform/backends/cuda_like/__init__.py +++ b/lightllm/platform/backends/cuda_like/__init__.py @@ -1,7 +1,7 @@ +import lightllm.platform.backends.cuda_like.ops # noqa: F401 from lightllm.platform.base.backend import HardwareBackend -from lightllm.platform.backends.cuda_like.ops import CUDA_LIKE_OPS -from lightllm.platform.backends.cuda_like.runtime import CudaLikeRuntime 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 @@ -10,7 +10,7 @@ def __init__(self) -> None: super().__init__( runtime=CudaLikeRuntime(), graph=CudaLikeGraph(), - ops=build_platform_ops(self.platform_name, CUDA_LIKE_OPS), + ops=build_platform_ops(self.platform_name), ) diff --git a/lightllm/platform/backends/cuda_like/ops/__init__.py b/lightllm/platform/backends/cuda_like/ops/__init__.py index 2ee453559c..14a0877ae1 100644 --- a/lightllm/platform/backends/cuda_like/ops/__init__.py +++ b/lightllm/platform/backends/cuda_like/ops/__init__.py @@ -1,13 +1,3 @@ -from lightllm.platform.backends.cuda_like.ops.act import CUDA_LIKE_ACT_OPS -from lightllm.platform.backends.cuda_like.ops.norm import CUDA_LIKE_NORM_OPS -from lightllm.platform.backends.cuda_like.ops.sampling import CUDA_LIKE_SAMPLING_OPS -from lightllm.platform.ops import PlatformOps +from lightllm.platform.backends.cuda_like.ops import act, embedding, norm, rotary, sampling - -CUDA_LIKE_OPS = PlatformOps( - norm=CUDA_LIKE_NORM_OPS, - act=CUDA_LIKE_ACT_OPS, - sampling=CUDA_LIKE_SAMPLING_OPS, -) - -__all__ = ["CUDA_LIKE_OPS"] +__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 index c45853d1da..680a630055 100644 --- a/lightllm/platform/backends/cuda_like/ops/act.py +++ b/lightllm/platform/backends/cuda_like/ops/act.py @@ -1,9 +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 - -CUDA_LIKE_ACT_OPS = ActOps( +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 index 8a6cc9a7db..44b2f812a6 100644 --- a/lightllm/platform/backends/cuda_like/ops/norm.py +++ b/lightllm/platform/backends/cuda_like/ops/norm.py @@ -3,9 +3,11 @@ 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 - -CUDA_LIKE_NORM_OPS = NormOps( +register_ops( + NormOps, + platforms=("cuda", "musa"), rms_norm=rmsnorm_forward, gated_rms_norm=gated_rmsnorm_forward, layer_norm=layernorm_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 index a1bb71f63b..200ec71025 100644 --- a/lightllm/platform/backends/cuda_like/ops/sampling.py +++ b/lightllm/platform/backends/cuda_like/ops/sampling.py @@ -6,12 +6,16 @@ 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 - -CUDA_LIKE_SAMPLING_OPS = SamplingOps( +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/ops/__init__.py b/lightllm/platform/ops/__init__.py index 6b5cec1630..3eab918294 100644 --- a/lightllm/platform/ops/__init__.py +++ b/lightllm/platform/ops/__init__.py @@ -1,7 +1,9 @@ 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 @@ -10,6 +12,8 @@ class PlatformOps: norm: NormOps act: ActOps sampling: SamplingOps + embedding: EmbeddingOps + rotary: RotaryOps -__all__ = ["ActOps", "NormOps", "PlatformOps", "SamplingOps"] +__all__ = ["ActOps", "EmbeddingOps", "NormOps", "PlatformOps", "RotaryOps", "SamplingOps"] diff --git a/lightllm/platform/ops/act.py b/lightllm/platform/ops/act.py index 54c990c6db..0b390f250a 100644 --- a/lightllm/platform/ops/act.py +++ b/lightllm/platform/ops/act.py @@ -1,10 +1,18 @@ from dataclasses import dataclass -from typing import Any, Callable +from typing import Protocol -ActOp = Callable[..., Any] +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: ActOp - gelu_and_mul: ActOp + 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 index 8c4de84ccd..f0e0025180 100644 --- a/lightllm/platform/ops/norm.py +++ b/lightllm/platform/ops/norm.py @@ -1,17 +1,69 @@ from dataclasses import dataclass -from typing import Any, Callable +from typing import Any, Protocol import torch -TensorOp = Callable[..., torch.Tensor] -AnyOp = Callable[..., Any] +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: TensorOp - gated_rms_norm: TensorOp - layer_norm: TensorOp - qk_rms_norm: TensorOp - fused_qk_rms_norm: AnyOp + 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 index ee01cf3f3d..73d205d2c2 100644 --- a/lightllm/platform/ops/sampling.py +++ b/lightllm/platform/ops/sampling.py @@ -1,13 +1,83 @@ from dataclasses import dataclass -from typing import Any, Callable +from typing import TYPE_CHECKING, Protocol -SamplingOp = Callable[..., Any] +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: SamplingOp - apply_penalty_gpu_cache: SamplingOp - apply_invalid_token_ids: SamplingOp - token_id_counter: SamplingOp - update_req_to_token_id_counter: SamplingOp + 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 index 21c2ab8e6d..6ac3140dff 100644 --- a/lightllm/platform/plugin/__init__.py +++ b/lightllm/platform/plugin/__init__.py @@ -1,4 +1,4 @@ -from lightllm.platform.plugin.common import Plugin +from lightllm.platform.plugin.common import Plugin, plugin_registration OPS = Plugin(name="ops", entry_point_group="lightllm.ops_plugin") @@ -8,5 +8,6 @@ def configure_plugins() -> None: - for plugin in _PLUGINS: - plugin.load() + with plugin_registration(): + for plugin in _PLUGINS: + plugin.load() diff --git a/lightllm/platform/plugin/att.py b/lightllm/platform/plugin/att.py index d113853b31..6426573e00 100644 --- a/lightllm/platform/plugin/att.py +++ b/lightllm/platform/plugin/att.py @@ -1,20 +1,20 @@ from dataclasses import dataclass -ATT_TABLE: dict[tuple[str, str, str], list["AttBackendEntry"]] = {} +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 AttBackendEntry: - # The name of the att backend. e.g., "triton", "flashinfer", "fa3". - name: str - # The category of the att backend. e.g., "standard", "mla", "nsa". - category: str - # The type of the kv. e.g., "int8kv", "int4kv", "fp8kv_sph". - kv_type: str - # The class of the att backend. +class _AttImpl: backend_cls: type - # The platforms of the att backend. e.g., "cuda", "ascend". - platforms: tuple[str, ...] + plugin: bool + + +ATT_IMPLS: dict[AttKey, _AttImpl] = {} def register_att_backend( @@ -25,24 +25,26 @@ def register_att_backend( platforms: tuple[str, ...] = ("cuda",), ): def decorator(backend_cls): - if not platforms: - raise ValueError("The value of platforms must be specified for att backend registration.") - - key = (name, category, kv_type) - # Check if the att backend is already registered. - for existing in ATT_TABLE.get(key, []): - if set(existing.platforms) & set(platforms): - raise ValueError(f"The att backend {name} is already registered for platforms {existing.platforms}.") - - ATT_TABLE.setdefault(key, []).append( - AttBackendEntry( - name=name, - category=category, - kv_type=kv_type, - backend_cls=backend_cls, - platforms=platforms, - ) - ) + 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 @@ -54,10 +56,11 @@ def get_att_backend_class( kv_type: str, platform: str, ) -> type: - key = (name, category, kv_type) - for entry in ATT_TABLE.get(key, []): - if platform in entry.platforms: - return entry.backend_cls - raise ValueError( - f"The att backend {name} is not registered for category {category}, kv type {kv_type}, and platform {platform}." - ) + 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 index c44b674565..4411a4dcfa 100644 --- a/lightllm/platform/plugin/common.py +++ b/lightllm/platform/plugin/common.py @@ -1,10 +1,28 @@ import importlib +from contextlib import contextmanager from dataclasses import dataclass -from importlib.metadata import entry_points, EntryPoint +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: @@ -13,26 +31,33 @@ class PluginConfig: class Plugin: def __init__(self, name: str, entry_point_group: str) -> None: - """Loader for pip-installed plugins. + """加载 pip 安装的插件。 Args: - name: Plugin kind. - entry_point_group: Entry point group name. Looked up with ``importlib.metadata``. + name: 插件种类。 + entry_point_group: entry point 组名，用 ``importlib.metadata`` 查找。 pyproject.toml: [project.entry-points."lightllm.ops_plugin"] - example_ops = "lightllm_example_plugin.register:register_ops" + 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", + "example_ops = lightllm_example_plugin.register:register_ops_plugin", + ], + "lightllm.att_plugin": [ + "example_att = lightllm_example_plugin.register:register_att_plugin", ], }, ) - After ``pip install``, pass ``--extra-ops example_ops`` to load and register. + ``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 @@ -56,21 +81,21 @@ def _load_entry_point_plugins(self, plugin_names: tuple[str, ...]) -> List[Plugi given_names = set(plugin_names) configs: List[PluginConfig] = [] loaded_names: set[str] = set() - available_entry_points = list(_iter_entry_points(self.entry_point_group)) + available_entry_points = list(entry_points().select(group=self.entry_point_group)) for entry_point in available_entry_points: - # Only load entry points from CLI. + # 只从 CLI 加载 entry point if entry_point.name not in given_names: continue - # example_ops = "lightllm_example_plugin.register:register_ops" - # -> register_fn = lightllm_example_plugin.register.register_ops + # 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) - # Check if all required plugins are loaded. + # 检查是否所有需要的插件都已加载 missing = given_names - loaded_names if missing: available = tuple(sorted(ep.name for ep in available_entry_points)) @@ -109,14 +134,6 @@ def merge_config_field(configs: Iterable[Any], field_name: str) -> tuple[str, .. return tuple(merged) -def _iter_entry_points(entry_point_group: str) -> Iterable[EntryPoint]: - eps = entry_points() - if hasattr(eps, "select"): - yield from eps.select(group=entry_point_group) - else: - yield from eps.get(entry_point_group, []) - - def _to_str_tuple(value: str | Iterable[str] | None) -> tuple[str, ...]: if value is None: return () diff --git a/lightllm/platform/plugin/ops.py b/lightllm/platform/plugin/ops.py index 1255769b56..462d0ea70f 100644 --- a/lightllm/platform/plugin/ops.py +++ b/lightllm/platform/plugin/ops.py @@ -1,39 +1,84 @@ -from dataclasses import fields, replace -from typing import Callable, TypeVar +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 -OpsGroupT = TypeVar("OpsGroupT") +logger = init_logger(__name__) +OpKey = tuple[str, type, str] -OPS_PLUGINS: dict[tuple[str, str], dict[str, Callable]] = {} +@dataclass(frozen=True) +class _OpImpl: + fn: Callable[..., Any] + plugin: bool -def register_ops(name: str, platforms: tuple[str, ...] = ("cuda",), **ops: Callable) -> None: - for platform in platforms: - current = OPS_PLUGINS.setdefault((platform, name), {}) - duplicates = current.keys() & ops.keys() - if duplicates: - raise ValueError( - f"Ops {sorted(duplicates)} are already registered for platform {platform!r}, group {name!r}." - ) - current.update(ops) +OP_IMPLS: dict[OpKey, _OpImpl] = {} -def build_ops_group(name: str, platform: str, default: OpsGroupT) -> OpsGroupT: - updates = OPS_PLUGINS.get((platform, name), {}) - if not updates: - return default - return replace(default, **updates) +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__}.") -def build_platform_ops(platform: str, default: PlatformOps) -> PlatformOps: - ops = { - field.name: build_ops_group( - field.name, - platform, - getattr(default, field.name), - ) - for field in fields(default) - } - return replace(default, **ops) + 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