Add registry for custom speculative algorithms (#23991)

python/sglang/srt/server_args.py

@@ -3321,12 +3321,28 @@ class ServerArgs:
                 self.speculative_moe_runner_backend
             ).is_flashinfer_trtllm(), "Currently speculative MoE runner backend doesn't support flashinfer_trtllm, please use triton or auto backend for speculative moe runner instead."
 
+        if self.speculative_algorithm is not None:
+            self.speculative_algorithm = self.speculative_algorithm.upper()
+
         self.speculative_algorithm = _resolve_speculative_algorithm_alias(
             self.speculative_algorithm,
             self.speculative_draft_model_path,
             trust_remote_code=self.trust_remote_code,
         )
 
+        if self.speculative_algorithm is not None:
+            from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+            from sglang.srt.speculative.spec_registry import CustomSpecAlgo
+
+            algo = SpeculativeAlgorithm.from_string(self.speculative_algorithm)
+
+            # TODO: move the per-algorithm validation below into spec module hooks.
+            if (
+                isinstance(algo, CustomSpecAlgo)
+                and algo.validate_server_args is not None
+            ):
+                algo.validate_server_args(self)
+
         if self.speculative_skip_dp_mlp_sync:
             assert self.speculative_algorithm == "EAGLE", (
                 "--speculative-skip-dp-mlp-sync is only supported with "
@@ -5421,8 +5437,11 @@ class ServerArgs:
         parser.add_argument(
             "--speculative-algorithm",
             type=str,
-            choices=["DFLASH", "EAGLE", "EAGLE3", "NEXTN", "STANDALONE", "NGRAM"],
-            help="Speculative algorithm.",
+            help=(
+                "Speculative algorithm. Builtins: EAGLE, EAGLE3, NEXTN, STANDALONE, "
+                "NGRAM, DFLASH. Or any name registered via "
+                "`SpeculativeAlgorithm.register`."
+            ),
         )
         parser.add_argument(
             "--speculative-draft-model-path",

python/sglang/srt/speculative/spec_info.py

@@ -2,7 +2,17 @@ from __future__ import annotations
 
 from abc import ABC, abstractmethod
 from enum import Enum, IntEnum, auto
-from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union
+from typing import TYPE_CHECKING, Callable, List, Optional, Tuple, Type, Union
+
+from sglang.srt.speculative.spec_registry import (
+    CustomSpecAlgo,
+    ServerArgsValidator,
+    WorkerFactory,
+)
+from sglang.srt.speculative.spec_registry import get_spec as _get_registered_spec
+from sglang.srt.speculative.spec_registry import (
+    register_algorithm as _register_algorithm,
+)
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import ModelWorkerBatch
@@ -13,7 +23,11 @@ if TYPE_CHECKING:
 
 
 class SpeculativeAlgorithm(Enum):
-    """Speculative decoding algorithms."""
+    """Builtin speculative decoding algorithms. Plugin-registered ones are
+    ``CustomSpecAlgo`` instances; ``from_string`` returns either type, and
+    both expose the same ``is_*()`` / ``create_worker`` interface so callers
+    dispatch uniformly without isinstance checks.
+    """
 
     DFLASH = auto()
     EAGLE = auto()
@@ -24,13 +38,46 @@ class SpeculativeAlgorithm(Enum):
     NONE = auto()
 
     @classmethod
-    def from_string(cls, name: Optional[str]) -> SpeculativeAlgorithm:
+    def from_string(
+        cls, name: Optional[str]
+    ) -> Union[SpeculativeAlgorithm, CustomSpecAlgo]:
         if name is None:
             return cls.NONE
+        upper = name.upper()
         try:
-            return cls[name.upper()]
+            return cls[upper]
         except KeyError:
-            raise ValueError(f"Unknown speculative algorithm name: {name}")
+            pass
+        spec = _get_registered_spec(upper)
+        if spec is not None:
+            return spec
+        raise ValueError(f"Unknown speculative algorithm name: {name}")
+
+    @classmethod
+    def register(
+        cls,
+        name: str,
+        *,
+        supports_overlap: bool = False,
+        validate_server_args: Optional[ServerArgsValidator] = None,
+        spec_class: Type[CustomSpecAlgo] = CustomSpecAlgo,
+    ) -> Callable[[WorkerFactory], WorkerFactory]:
+        """Decorator to register a plugin speculative algorithm. The factory
+        takes ``server_args`` and returns the worker class. Pass a
+        ``CustomSpecAlgo`` subclass via ``spec_class`` to override any
+        ``is_*()`` / ``create_worker`` method.
+
+        Example:
+            @SpeculativeAlgorithm.register("MY_SPEC", supports_overlap=False)
+            def _factory(server_args):
+                return MySpecWorker
+        """
+        return _register_algorithm(
+            name,
+            supports_overlap=supports_overlap,
+            validate_server_args=validate_server_args,
+            spec_class=spec_class,
+        )
 
     def is_none(self) -> bool:
         return self == SpeculativeAlgorithm.NONE

python/sglang/srt/speculative/spec_registry.py

@@ -0,0 +1,123 @@
+"""Internal storage backing ``SpeculativeAlgorithm.register``. Plugins
+should use that classmethod API; do not import from this module directly.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Callable, Dict, Optional, Type
+
+if TYPE_CHECKING:
+    from sglang.srt.server_args import ServerArgs
+
+WorkerFactory = Callable[["ServerArgs"], Type]
+ServerArgsValidator = Callable[["ServerArgs"], None]
+
+
+class CustomSpecAlgo:
+    """A plugin-registered speculative algorithm. Duck-types
+    ``SpeculativeAlgorithm`` enum values (same ``is_*()`` / ``create_worker``
+    interface).
+
+    Plugins may subclass this to override any ``is_*()`` / ``supports_*()`` /
+    ``create_worker`` method (e.g. to integrate with builtin-specific
+    branches like ``if spec_algorithm.is_eagle():`` in scheduler /
+    model_runner). Pass the subclass via ``spec_class=...`` at registration.
+
+    Defaults: all ``is_*()`` return ``False`` except ``is_speculative``;
+    ``supports_spec_v2`` follows ``supports_overlap``.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        factory: WorkerFactory,
+        *,
+        supports_overlap: bool = False,
+        validate_server_args: Optional[ServerArgsValidator] = None,
+    ):
+        self.name = name
+        self.factory = factory
+        self.supports_overlap = supports_overlap
+        self.validate_server_args = validate_server_args
+
+    def __repr__(self) -> str:
+        return f"CustomSpecAlgo({self.name!r})"
+
+    def is_none(self) -> bool:
+        return False
+
+    def is_speculative(self) -> bool:
+        return True
+
+    def is_eagle(self) -> bool:
+        return False
+
+    def is_eagle3(self) -> bool:
+        return False
+
+    def is_dflash(self) -> bool:
+        return False
+
+    def is_standalone(self) -> bool:
+        return False
+
+    def is_ngram(self) -> bool:
+        return False
+
+    def supports_spec_v2(self) -> bool:
+        return self.supports_overlap
+
+    def create_worker(self, server_args: "ServerArgs") -> Type:
+        if not server_args.disable_overlap_schedule and not self.supports_overlap:
+            raise ValueError(
+                f"Speculative algorithm {self.name} does not support overlap scheduling."
+            )
+        return self.factory(server_args)
+
+
+_REGISTRY: Dict[str, CustomSpecAlgo] = {}
+
+# Builtin enum members + the NEXTN alias; plugins cannot shadow these.
+_RESERVED_NAMES = frozenset(
+    {"DFLASH", "EAGLE", "EAGLE3", "NEXTN", "STANDALONE", "NGRAM", "NONE"}
+)
+
+
+def register_algorithm(
+    name: str,
+    *,
+    supports_overlap: bool = False,
+    validate_server_args: Optional[ServerArgsValidator] = None,
+    spec_class: Type[CustomSpecAlgo] = CustomSpecAlgo,
+) -> Callable[[WorkerFactory], WorkerFactory]:
+    """Return a decorator that registers a plugin algorithm under ``name``.
+
+    Pass a ``spec_class`` subclass of ``CustomSpecAlgo`` to override any
+    ``is_*()`` / ``supports_*()`` / ``create_worker`` method.
+    """
+    upper = name.upper()
+    if upper in _RESERVED_NAMES:
+        raise ValueError(
+            f"'{upper}' is a reserved speculative algorithm name; cannot be re-registered."
+        )
+    if upper in _REGISTRY:
+        raise ValueError(f"Speculative algorithm '{upper}' already registered.")
+
+    def decorator(factory: WorkerFactory) -> WorkerFactory:
+        _REGISTRY[upper] = spec_class(
+            name=upper,
+            factory=factory,
+            supports_overlap=supports_overlap,
+            validate_server_args=validate_server_args,
+        )
+        return factory
+
+    return decorator
+
+
+def get_spec(name: Optional[str]) -> Optional[CustomSpecAlgo]:
+    """Return the registered spec for ``name``, or ``None`` for builtin /
+    unknown names."""
+    if name is None:
+        return None
+    return _REGISTRY.get(name.upper())