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=(
-            help="Speculative algorithm.",
+                "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())

test/registered/unit/spec/test_spec_registry.py

@@ -0,0 +1,234 @@
+"""Unit tests for the speculative algorithm plugin registry."""
+
+import unittest
+from unittest.mock import MagicMock
+
+from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
+from sglang.srt.speculative.spec_registry import (
+    _REGISTRY,
+    _RESERVED_NAMES,
+    CustomSpecAlgo,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
+
+
+class _RegistryIsolated(CustomTestCase):
+    """Snapshot and restore the global registry so tests don't leak."""
+
+    def setUp(self):
+        self._snapshot = _REGISTRY.copy()
+        _REGISTRY.clear()
+
+    def tearDown(self):
+        _REGISTRY.clear()
+        _REGISTRY.update(self._snapshot)
+
+
+class TestFromString(_RegistryIsolated):
+    def test_none_input_returns_none_member(self):
+        self.assertIs(SpeculativeAlgorithm.from_string(None), SpeculativeAlgorithm.NONE)
+
+    def test_builtin_name_returns_enum(self):
+        self.assertIs(
+            SpeculativeAlgorithm.from_string("EAGLE"), SpeculativeAlgorithm.EAGLE
+        )
+        self.assertIs(
+            SpeculativeAlgorithm.from_string("NGRAM"), SpeculativeAlgorithm.NGRAM
+        )
+
+    def test_builtin_name_is_case_insensitive(self):
+        self.assertIs(
+            SpeculativeAlgorithm.from_string("eagle"), SpeculativeAlgorithm.EAGLE
+        )
+
+    def test_unknown_name_raises(self):
+        with self.assertRaisesRegex(ValueError, "Unknown speculative algorithm"):
+            SpeculativeAlgorithm.from_string("NOT_REGISTERED")
+
+    def test_registered_plugin_returns_custom_spec(self):
+        @SpeculativeAlgorithm.register("MY_FOO")
+        def _factory(server_args):
+            return MagicMock
+
+        algo = SpeculativeAlgorithm.from_string("MY_FOO")
+        self.assertIsInstance(algo, CustomSpecAlgo)
+        self.assertEqual(algo.name, "MY_FOO")
+
+    def test_registered_plugin_lookup_is_case_insensitive(self):
+        @SpeculativeAlgorithm.register("MY_FOO")
+        def _factory(server_args):
+            return MagicMock
+
+        self.assertIs(
+            SpeculativeAlgorithm.from_string("my_foo"),
+            SpeculativeAlgorithm.from_string("MY_FOO"),
+        )
+
+
+class TestRegister(_RegistryIsolated):
+    def test_register_returns_factory_unchanged(self):
+        def _factory(server_args):
+            return MagicMock
+
+        decorated = SpeculativeAlgorithm.register("MY_FOO")(_factory)
+        self.assertIs(decorated, _factory)
+
+    def test_two_distinct_registrations_are_independent(self):
+        @SpeculativeAlgorithm.register("FOO")
+        def _foo_factory(server_args):
+            return MagicMock
+
+        @SpeculativeAlgorithm.register("BAR")
+        def _bar_factory(server_args):
+            return MagicMock
+
+        foo = SpeculativeAlgorithm.from_string("FOO")
+        bar = SpeculativeAlgorithm.from_string("BAR")
+        self.assertIsNot(foo, bar)
+        self.assertNotEqual(foo, bar)
+        self.assertEqual(foo.name, "FOO")
+        self.assertEqual(bar.name, "BAR")
+
+    def test_duplicate_name_raises(self):
+        @SpeculativeAlgorithm.register("MY_FOO")
+        def _factory(server_args):
+            return MagicMock
+
+        with self.assertRaisesRegex(ValueError, "already registered"):
+
+            @SpeculativeAlgorithm.register("MY_FOO")
+            def _factory2(server_args):
+                return MagicMock
+
+    def test_reserved_name_raises(self):
+        for reserved in _RESERVED_NAMES:
+            with self.assertRaisesRegex(ValueError, "reserved"):
+                SpeculativeAlgorithm.register(reserved)
+
+    def test_register_is_case_insensitive_on_collision(self):
+        @SpeculativeAlgorithm.register("MY_FOO")
+        def _factory(server_args):
+            return MagicMock
+
+        with self.assertRaisesRegex(ValueError, "already registered"):
+
+            @SpeculativeAlgorithm.register("my_foo")
+            def _factory2(server_args):
+                return MagicMock
+
+
+class TestCustomSpecAlgoInterface(_RegistryIsolated):
+    """CustomSpecAlgo must duck-type SpeculativeAlgorithm enum values."""
+
+    def setUp(self):
+        super().setUp()
+
+        @SpeculativeAlgorithm.register("MY_FOO", supports_overlap=False)
+        def _factory(server_args):
+            return MagicMock
+
+        self.algo = SpeculativeAlgorithm.from_string("MY_FOO")
+
+    def test_is_predicates_all_false_except_speculative(self):
+        self.assertFalse(self.algo.is_none())
+        self.assertFalse(self.algo.is_eagle())
+        self.assertFalse(self.algo.is_eagle3())
+        self.assertFalse(self.algo.is_dflash())
+        self.assertFalse(self.algo.is_standalone())
+        self.assertFalse(self.algo.is_ngram())
+        self.assertTrue(self.algo.is_speculative())
+
+    def test_supports_spec_v2_follows_supports_overlap(self):
+        # Plugin registered with supports_overlap=False -> not spec_v2.
+        self.assertFalse(self.algo.supports_spec_v2())
+
+        @SpeculativeAlgorithm.register("MY_V2", supports_overlap=True)
+        def _factory(server_args):
+            return MagicMock
+
+        v2 = SpeculativeAlgorithm.from_string("MY_V2")
+        self.assertTrue(v2.supports_spec_v2())
+
+    def test_create_worker_calls_factory(self):
+        server_args = MagicMock()
+        server_args.disable_overlap_schedule = True
+        worker_cls = self.algo.create_worker(server_args)
+        self.assertIs(worker_cls, MagicMock)
+
+    def test_create_worker_raises_on_overlap_mismatch(self):
+        server_args = MagicMock()
+        server_args.disable_overlap_schedule = False
+        with self.assertRaisesRegex(ValueError, "does not support overlap"):
+            self.algo.create_worker(server_args)
+
+
+class TestValidatorHook(_RegistryIsolated):
+    def test_validator_invocation_is_caller_driven(self):
+        validator = MagicMock()
+
+        @SpeculativeAlgorithm.register("MY_FOO", validate_server_args=validator)
+        def _factory(server_args):
+            return MagicMock
+
+        algo = SpeculativeAlgorithm.from_string("MY_FOO")
+        self.assertIs(algo.validate_server_args, validator)
+        # Callers (e.g. ServerArgs.__post_init__) must invoke the hook themselves;
+        # CustomSpecAlgo does not call it from create_worker.
+        validator.assert_not_called()
+
+
+class TestSubclassOverride(_RegistryIsolated):
+    """Plugins can subclass CustomSpecAlgo to override is_*() / create_worker."""
+
+    def test_subclass_overrides_is_eagle(self):
+        class EagleLike(CustomSpecAlgo):
+            def is_eagle(self) -> bool:
+                return True
+
+        @SpeculativeAlgorithm.register(
+            "MY_LIKE_EAGLE", supports_overlap=True, spec_class=EagleLike
+        )
+        def _factory(server_args):
+            return MagicMock
+
+        algo = SpeculativeAlgorithm.from_string("MY_LIKE_EAGLE")
+        self.assertIsInstance(algo, EagleLike)
+        self.assertIsInstance(algo, CustomSpecAlgo)
+        self.assertTrue(algo.is_eagle())
+        # Other predicates default to False
+        self.assertFalse(algo.is_ngram())
+        self.assertFalse(algo.is_dflash())
+
+    def test_subclass_overrides_create_worker(self):
+        class CustomDispatch(CustomSpecAlgo):
+            def create_worker(self, server_args):
+                return "custom-dispatched"
+
+        @SpeculativeAlgorithm.register("MY_CUSTOM", spec_class=CustomDispatch)
+        def _factory(server_args):
+            return MagicMock
+
+        algo = SpeculativeAlgorithm.from_string("MY_CUSTOM")
+        # Custom dispatch bypasses default overlap check
+        self.assertEqual(algo.create_worker(MagicMock()), "custom-dispatched")
+
+
+class TestCrossTypeIdentity(_RegistryIsolated):
+    """A plugin algo and a builtin enum value must never compare equal."""
+
+    def test_plugin_not_equal_to_builtin(self):
+        @SpeculativeAlgorithm.register("MY_FOO")
+        def _factory(server_args):
+            return MagicMock
+
+        algo = SpeculativeAlgorithm.from_string("MY_FOO")
+        self.assertNotEqual(algo, SpeculativeAlgorithm.EAGLE)
+        self.assertNotEqual(algo, SpeculativeAlgorithm.NONE)
+        self.assertIsNot(algo, SpeculativeAlgorithm.EAGLE)
+
+
+if __name__ == "__main__":
+    unittest.main(verbosity=3)