From 95fb722dd23ba921e2f00a0c8b5801c49bc014eb Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Thu, 7 May 2026 16:11:45 -0700 Subject: [PATCH] Add registry for custom speculative algorithms (#23991) --- python/sglang/srt/server_args.py | 23 +- python/sglang/srt/speculative/spec_info.py | 57 ++++- .../sglang/srt/speculative/spec_registry.py | 123 +++++++++ .../unit/spec/test_spec_registry.py | 234 ++++++++++++++++++ 4 files changed, 430 insertions(+), 7 deletions(-) create mode 100644 python/sglang/srt/speculative/spec_registry.py create mode 100644 test/registered/unit/spec/test_spec_registry.py diff --git a/python/sglang/srt/server_args.py b/python/sglang/srt/server_args.py index de7f06468..bd167be0d 100644 --- a/python/sglang/srt/server_args.py +++ b/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", diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index a0e6c9c24..8a2588c0c 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py new file mode 100644 index 000000000..c7e062596 --- /dev/null +++ b/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()) diff --git a/test/registered/unit/spec/test_spec_registry.py b/test/registered/unit/spec/test_spec_registry.py new file mode 100644 index 000000000..ee41e16e0 --- /dev/null +++ b/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)